Simulation

Navigation

Simulation#

integrate(module, params=[], *, param_state=None, data_stimuli=None, data_clamps=None, t_max=None, delta_t=0.025, solver='bwd_euler', voltage_solver='jaxley.stone', checkpoint_lengths=None, all_states=None, return_states=False)[source]#

Solves ODE and simulates neuron model.

Parameters:

  • params (List[Dict[str, Array]]) – Trainable parameters returned by get_parameters().

  • param_state (List[Dict] | None) – Parameters returned by data_set.

  • data_stimuli (Tuple[Array, DataFrame] | None) – Outputs of .data_stimulate(), only needed if stimuli change across function calls.

  • data_clamps (Tuple[str, Array, DataFrame] | None) – Outputs of .data_clamp(), only needed if clamps change across function calls.

  • t_max (float | None) – Duration of the simulation in milliseconds. If t_max is greater than the length of the stimulus input, the stimulus will be padded at the end with zeros. If t_max is smaller, then the stimulus with be truncated.

  • delta_t (float) – Time step of the solver in milliseconds.

  • solver (str) – Which ODE solver to use. Either of [“fwd_euler”, “bwd_euler”, “crank_nicolson”].

  • tridiag_solver – Algorithm to solve tridiagonal systems. The different options only affect bwd_euler and crank_nicolson solvers. Either of [“stone”, “thomas”], where stone is much faster on GPU for long branches with many compartments and thomas is slightly faster on CPU (thomas is used in NEURON).

  • checkpoint_lengths (List[int] | None) – Number of timesteps at every level of checkpointing. The prod(checkpoint_lengths) must be larger or equal to the desired number of simulated timesteps. Warning: the simulation is run for prod(checkpoint_lengths) timesteps, and the result is posthoc truncated to the desired simulation length. Therefore, a poor choice of checkpoint_lengths can lead to longer simulation time. If None, no checkpointing is applied.

  • all_states (Dict | None) – An optional initial state that was returned by a previous jx.integrate(…, return_states=True) run. Overrides potentially trainable initial states.

  • return_states (bool) – If True, it returns all states such that the current state of the Module can be set with set_states.

  • module (Module)

  • voltage_solver (str)

Return type:

Array

Navigation