proposed new api; not tested (#731)

* proposed new api; not tested

---------

Co-authored-by: Matthijspals <matthijs-pals@hotmail.com>

Author: michaeldeistler

docs/jaxley.rst

@@ -23,7 +23,7 @@ Simulation
 
    jaxley.integrate
    jaxley.integrate.build_init_and_step_fn
-   jaxley.utils.dynamics.build_step_dynamics_fn
+   jaxley.utils.dynamics.build_dynamic_state_utils
 
 
 Morphologies

jaxley/integrate.py

@@ -23,17 +23,48 @@ def build_init_and_step_fn(
     """Return ``init_fn`` and ``step_fn`` which initialize modules and run update steps.
 
     This method can be used to gain additional control over the simulation workflow.
-    It exposes the ``step`` function, which can be used to perform step-by-step updates
+    It exposes a step function, which can be used to perform step-by-step updates
     of the differential equations.
 
     Args:
-        module: A `Module` object that e.g. a cell.
-        voltage_solver: Voltage solver used in step. Defaults to "jaxley.stone".
+        module: A ``jx.Module`` object that, for example a ``jx.Cell``.
+        voltage_solver: Voltage solver used in step. Defaults to "jaxley.dhs".
         solver: ODE solver. Defaults to "bwd_euler".
 
     Returns:
-        init_fn, step_fn: Functions that initialize the state and parameters, and
-            perform a single integration step, respectively.
+
+        * ``init_fn(params, all_states=None, param_state=None, delta_t=0.025)``
+
+          Callable which initializes the states and parameters.
+
+          * Args:
+
+            * ``params`` (list[dict]): returned by `.get_parameters()`.
+            * ``all_states`` (dict | None = None): typically `None`.
+            * ``param_state`` (list[dict] | None = None): returned by `.data_set()`.
+            * ``delta_t`` (float = 0.025): the time step.
+
+          * Returns:
+
+            * ``all_states`` (dict).
+            * ``all_params`` (dict), which can be passed to the `step_fn`.
+
+        * ``step_fn(all_states, all_params, external_inds, externals, delta_t=0.025)``
+
+          Callable which performs a single integration step.
+
+          * Args:
+
+            * ``all_states`` (dict): returned by `init_fn()`.
+            * ``all_params`` (dict): returned by `init_fn()`.
+            * ``externals`` (dict): obtained with `module.externals.copy()` but using
+              only the external input at the current time step (see examples below).
+            * ``external_inds`` (dict): obtained with `module.external_inds.copy()`.
+            * ``delta_t`` (float): the time step.
+
+          * Returns:
+
+            * Updated ``all_states``  (dict).
 
     Example usage
     ^^^^^^^^^^^^^

jaxley/modules/base.py

@@ -175,6 +175,7 @@ def __init__(self):
         # List of types of all `jx.Channel`s.
         self.channels: List[Channel] = []
         self.membrane_current_names: List[str] = []
+        self.synapse_current_names: List[str] = []
 
         # List of all pumps.
         self.pumped_ions: List[str] = []

jaxley/modules/network.py

@@ -336,6 +336,10 @@ def _synapse_currents(
     ) -> tuple[dict, tuple[Array, Array]]:
         voltages = states["v"]
 
+        current_states = {}
+        for name in self.synapse_current_names:
+            current_states[name] = jnp.zeros_like(voltages)
+
         grouped_syns = edges.groupby("type", sort=False, group_keys=False)
         pre_syn_inds = grouped_syns["pre_index"].apply(list)
         post_syn_inds = grouped_syns["post_index"].apply(list)
@@ -617,6 +621,10 @@ def _append_multiple_synapses(self, pre_nodes, post_nodes, synapse_type):
         self.base.edges["controlled_by_param"] = 0
         self._edges_in_view = self.edges.index.to_numpy()
 
+        current_name = f"i_{synapse_type._name}"
+        if current_name not in self.base.synapse_current_names:
+            self.base.synapse_current_names.append(current_name)
+
     def _add_params_to_edges(self, synapse_type, indices):
         # Add parameters and states to the `.edges` table.
         for key, param_val in synapse_type.synapse_params.items():