aquakin.integrate_ensemble#

aquakin.integrate_ensemble(model, tracks, C0_fn, params, *, rtol=1e-06, atol=None, n_save=None, integrator=IntegratorConfig(order=3, factormax=3.0, colored_jacobian='auto', dtmax=None, max_steps=100000, solver=None), diff=DifferentiationConfig(mode='reverse', method='stable', check_finite=True, adjoint_max_steps=100000, adjoint_low_memory=False))[source]#

Integrate the model along an ensemble of particle tracks.

Parameters:

  • model (CompiledModel)

  • tracks (mapping int -> Track) – particle_id -> Track.

  • C0_fn (callable) – Maps particle_id to its inlet concentration vector. Often lambda pid: model.default_concentrations().

  • params (jnp.ndarray) – Flat parameter vector shared across all particles.

  • rtol (passed through to each) – ParticleTrackReactor. integrator / diff let an ensemble of stiff tracks be differentiated the same way a single track can.

  • atol (passed through to each) – ParticleTrackReactor. integrator / diff let an ensemble of stiff tracks be differentiated the same way a single track can.

  • n_save (passed through to each) – ParticleTrackReactor. integrator / diff let an ensemble of stiff tracks be differentiated the same way a single track can.

  • integrator (passed through to each) – ParticleTrackReactor. integrator / diff let an ensemble of stiff tracks be differentiated the same way a single track can.

  • diff (passed through to each) – ParticleTrackReactor. integrator / diff let an ensemble of stiff tracks be differentiated the same way a single track can.

Returns:

One solution per particle, keyed by id.

Return type:

dict[int, TrackSolution]