Working directly with ODE dynamics function

[bantin]

Hi jaxley folks!

I frequently find myself needing to work directly with the jaxley ODE dynamics function. In order to do various mathematical operations (e.g kalman smoothing) I need a dynamics function that looks something like this:

next_state_vector = dynamics_fn(state_vector, trainable_params, input_current)

Since the channel currents can be computed deterministically from the voltage + channel states, I need the state vector to exclude the currents. While I can create this function using Jaxley, it is currently very cumbersome to do so. I think there are a few things that are making it difficult:

1. the step_fn returned by build_init_and_step_fn requires all parameters, rather than just the trainables. Depending on how I implement things, this means that I sometimes end up calling module.get_all_parameters() every time I need to evolve the state, which seems suboptimal and potentially slow (though I haven't timed it carefully)
2. jaxley's state is stored as a dictionary, and I need to work directly with jax arrays.
3. jaxley's state includes the currents, where I need to work only with channel states and voltages.

I'm wondering if the jaxley team would be interested in adding functionality to jaxley that makes it easier to get a dynamics_fn that works directly with vectors? I suspect that I'm not the only one interested in this, and I'd be happy to work on implementing it.

[michaeldeistler]

Hi Ben,

thanks for reaching out---I agree that some utilities for coupling Jaxley with other libraries that are based on vectors would be useful. I am not yet sure what the best place for such utilities would be though. A few options:

• We could have them as a small module as part of Jaxley/utils/
• We could open a new repo called Jaxely-utilities that implements such things?
• Or we simply write a clean example on how to do this, and expect people to copy relevant code pieces.

To get started, I tend towards the first option, at least as long as it is rather little code. As the code base grows bigger, we might want to migrate to option 2.

Regarding point 2: You might aware of this anyways, but I think ravel_pytree is quite useful for this:

import jax
import jax.numpy as jnp

# Example pytree: dict of arrays
pytree = {
    "a": jnp.array([1., 2.]),
    "b": jnp.array([3., 4., 5.])
}

# Flatten into 1D array
flat, unravel = jax.flatten_util.ravel_pytree(pytree)

print(flat)        # 1D array: [1. 2. 3. 4. 5.]
print(flat.shape)  # (5,)

# Recover original pytree structure
restored = unravel(flat)
print(restored)

Anyways...if you would be up for getting started with this, that would be awesome. If you need input, let us know (but I will be on vacation next week :) ).

All the best
Michael

[bantin]

Thanks for the quick reply @michaeldeistler ! I agree that putting it inside of the utils would make sense. Let me clean up my example code and send it over in this thread, then we can decide what to do with it.

I do have one question for you though. I've struggled to settle on the best design for handling the conversion between trainable_parameters and all parameters. If possible, I'd like to avoid calling module.get_all_parameters inside of the dynamics_fn. The way that I've avoided this so far is by having my dynamics_fn take in all of the parameters (instead of just the trainables). Then, for a single call to our kalman smoother, we only have to call module.get_all_parameters once. This does add a bit of overhead in terms of code complexity though. Curious if you have any thoughts here?

[michaeldeistler]

Hi Ben! Great!

I am not really sure I understand why this adds code complexity. For now, let's just go with the solution that you already have and then we can have a look :)

Michael

[bantin]

Thanks Michael! I have a working version that includes the channel currents as part of the state vector, but I'm struggling to figure out how to accomplish this when I need my state vector to exclude the channel currents. I'll paste what I have so far below. The issue I'm having is how to write the add_channel_currents_to_state_dict function. The goal of having this function is so that we can recompute the channel currents before passing the state dictionary to jaxley's step_fn. While what I have below works outside of JIT, I haven't managed to find a way to write this in a JIT-compatible way. Curious if you have any ideas? The problem with what I have below is that query_channel_states_and_params is indexing into a dictionary, which doesn't seem to be allowed on traced values.

def get_param_update_fn(cell: jx.Cell) -> callable:
    """Create a function to update all jaxley parameters given trainables."""

    def update_jaxley_params(trainable_params):
        pstate = params_to_pstate(trainable_params, cell.indices_set_by_trainables)
        return cell.get_all_parameters(pstate)

    return update_jaxley_params

def add_channel_currents_to_state_dict(reduced_state_dict, cell, 
        jx_params_dict):
            """
            Add channel currents back to a reduced state dictionary.
            
            Args:
                reduced_state_dict: State dict without channel currents
                cell: The jaxley cell object containing the channels
                jx_params_dict: Parameters dictionary from init_fn
            
            Returns:
                Complete state dict with channel currents added
            """
            # Create a copy to avoid modifying the original
            complete_state_dict = reduced_state_dict.copy()

            # Get voltage from the state dict
            voltages = reduced_state_dict["v"]

            # Get all channel indices (assuming all compartments have all channels)
            # TODO: fix this
            all_indices = jnp.arange(len(voltages))

            # Initialize current states for each channel
            for current_name in cell.membrane_current_names:
                complete_state_dict[current_name] = jnp.zeros_like(voltages)

            # Compute current for each channel
            for channel in cell.channels:
                # Get the channel-specific parameter and state names
                channel_param_names = list(channel.channel_params.keys())
                channel_state_names = list(channel.channel_states.keys())

                # Extract channel-specific parameters and states
                channel_params = query_channel_states_and_params(
                    jx_params_dict, channel_param_names, all_indices
                )
                channel_states = query_channel_states_and_params(
                    reduced_state_dict, channel_state_names, all_indices
                )

                # Compute current using the channel's compute_current method
                current = channel.compute_current(channel_states, voltages,
        channel_params)

                # Add the computed current to the state dict
                complete_state_dict[channel.current_name] = current

            return complete_state_dict

def build_dynamics_function(
    cell: jx.Cell,
):
    
    reduce_states = lambda states: {
        k: v for k, v in states.items() if k not in cell.membrane_current_names
    }

    cell.to_jax()
    init_fn, step_fn = build_init_and_step_fn(cell)
    trainable_params = cell.get_parameters()
    jx_state_dict, _ = init_fn(trainable_params)
    jx_state_dict_reduced = reduce_states(jx_state_dict)

    _, unflatten_state_fn = jax.flatten_util.ravel_pytree(jx_state_dict_reduced)

    def step_state(
        state_array: jnp.array,  # state_array excluding channel currents
        jx_params_dict: dict,
        input_current: float,
        input_current_idx: int,
    ):
        reduced_state_dict = unflatten_state_fn(state_array)
        state_dict = add_channel_currents_to_state_dict(reduced_state_dict, cell, jx_params_dict)

        # apply forward transforms: unconstrained -> constrained
        state_dict = apply_forward_state_transforms(state_dict)

        # apply dynamics
        externals = {"i": jnp.array([input_current])}
        external_inds = {"i": jnp.array([input_current_idx])}
        next_state_dict = step_fn(state_dict, jx_params_dict, externals, external_inds, delta_t=0.025)
        next_reduced_state_dict = reduce_states(next_state_dict)

        # apply inverse transforms: constrained -> unconstrained
        next_reduced_state_dict = apply_inverse_state_transforms(next_reduced_state_dict)

        # convert back to array
        next_state_array, _ = jax.flatten_util.ravel_pytree(next_reduced_state_dict)
        return next_state_array
    
    return step_state

[Matthijspals]

Hi @bantin!

I agree that this would be a super useful thing to have in Jaxley! Both for filtering algorithms and e.g., computing Jacobians (see also discussion 699), among others!

Coincidentally, I also just started working on this, here is what I got so far - it seems to be jit compilable:

def add_currents_to_voltages(module, states, delta_t, all_params):
    """Add the currents through channels and synapses to the voltages in states."""
    # Add to the states the initial current through every channel.
    states, _ = module.base._channel_currents(
        states, delta_t, module.channels + module.pumps, module.nodes, all_params
    )

    # Add to the states the initial current through every synapse.
    states, _ = module.base._synapse_currents(
        states,
        module.synapses,
        all_params,
        delta_t,
        module.edges,
    )
    return states

def remove_currents_from_voltages(states, current_keys):
    """Remove the currents through channels and synapses from the voltages in states."""
    # Remove from the states the initial current through every channel.
    for key in current_keys:
        del states[key]
    return states

def build_init_and_step_fn_no_currents(
    module: Module,
    voltage_solver: str = "jaxley.dhs",
    solver: str = "bwd_euler",
) -> Tuple[Callable, Callable]:
    """Return ``init_fn`` and ``step_fn`` which initialize modules and run update steps.

    Args:
        module: A `Module` object that e.g. a cell.
        voltage_solver: Voltage solver used in step. Defaults to "jaxley.stone".
        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.
    """

    # Initialize the external inputs and their indices.
    external_inds = module.external_inds.copy()

    def init_fn_no_currents(
        params: list[dict[str, Array]],
        all_states: dict | None = None,
        param_state: list[dict] | None = None,
        delta_t: float = 0.025,
    ) -> Tuple[Dict, Dict]:
        """Initializes the parameters and states of the neuron model.

        Args:
            params: List of trainable parameters.
            all_states: State if already initialized. Defaults to None.
            param_state: Parameters returned by `data_set`.. Defaults to None.
            delta_t: Step size. Defaults to 0.025.

        Returns:
            All states and parameters.
        """
        # Make the `trainable_params` of the same shape as the `param_state`, such that
        # they can be processed together by `get_all_parameters`.
        pstate = params_to_pstate(params, module.indices_set_by_trainables)
        if param_state is not None:
            pstate += param_state

        all_params = module.get_all_parameters(pstate)

        # todo: might be missing something, compare to:
        # https://github.com/jaxleyverse/jaxley/blob/659e332dea0138ef21d9e28495729eef1c1d6ee3/jaxley/modules/base.py#L1762
        all_states = (
            module._get_states_from_nodes_and_edges()
            if all_states is None
            else all_states
        )
        return all_states, all_params
    
    def step_fn_no_currents(
        all_states_no_currents: Dict,
        all_params: Dict,
        externals: Dict,
        external_inds: Dict = external_inds,
        delta_t: float = 0.025,
    ) -> Dict:
        """Performs a single integration step with step size delta_t.

        Args:
            all_states_no_currents: Current state of the neuron model.
            all_params: Current parameters of the neuron model.
            externals: External inputs.
            external_inds: External indices. Defaults to `module.external_inds`.
            delta_t: Time step. Defaults to 0.025.

        Returns:
            Updated states.
        """
        keys = list(all_states_no_currents.keys())
        state = add_currents_to_voltages(module, all_states_no_currents, delta_t, all_params)
        new_keys = list(state.keys())
        added_keys = [key for key in new_keys if key not in keys]
        state = module.step(
            state,
            delta_t,
            external_inds,
            externals,
            params=all_params,
            solver=solver,
            voltage_solver=voltage_solver,
        )
        #state_no_currents = module.base._get_states_from_nodes_and_edges()
        state_no_currents = remove_currents_from_voltages(state, added_keys)
        return state_no_currents

    return init_fn_no_currents, step_fn_no_currents

Here is an example run using these functions:

# Simulation parameters
n_comps=3
delta_t = 0.025
n_steps = 10

# Load the cell model and insert channels
cell = jx.Cell()
cell = jx.read_swc("../jaxley/tests/swc_files/morph_ca1_n120_250_single_point_soma.swc", ncomp=n_comps)
#visualise cell
cell.vis()

cell.insert(Leak())
cell.branch(0).insert(Na())


# set random initial voltages to get dynamics
for i in range(10):
    for j in range(n_comps):
        cell.branch(i).comp(j).set("v", -70+i)


cell.record()
params = cell.get_parameters()


cell.to_jax()
rec_inds = cell.recordings.rec_index.to_numpy()
rec_states = cell.recordings.state.to_numpy()

# Initialize.
init_fn_no_currents, step_fn_no_currents = build_init_and_step_fn_no_currents(cell,solver="bwd_euler")
states, params = init_fn_no_currents(params)

# we need a step function that takes in a vector, we can use ravel_pytree for this.
# unravel_fn stays the same even if the states change, so we can use it to reconstruct the states later.
flat, unravel_fn = ravel_pytree(states)
print(len(flat)-np.sum(np.isnan(flat)))

# Loop over the ODE. The `step_fn` can be jitted for improving speed.
@jit
def step_fn_vec_to_vec(flat):
    states = unravel_fn(flat)
    states = step_fn_no_currents(states, params, {}, delta_t)
    return ravel_pytree(states)[0]

flats = []
for step in range(n_steps):
    
    flat = step_fn_vec_to_vec(flat)
    flats.append(flat)

I haven't had time to compare our implementations in detail, but will have a look again later in the coming week! In any case posting this already in case it helps you in some way.

[bantin]

Awesome, thanks @Matthijspals ! I was able to plug your code into our filtering setup quite easily. Would be awesome to have this as part of jaxley.

[Matthijspals]

Hi, I opened an Issue for this Issue #717, and will start working on a pull-request. Please let me know if you run into anything!