Skip to content

Commit

Permalink
Use the new types to get rid of init_conds_custom
Browse files Browse the repository at this point in the history
  • Loading branch information
@michaeldeistler
michaeldeistler committed Sep 20, 2024
1 parent 2c71bf0 commit c668807
Show file tree
Hide file tree
Showing 6 changed files with 88 additions and 334 deletions.
122 changes: 46 additions & 76 deletions jaxley/modules/base.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
from jaxley.utils.cell_utils import (
_compute_index_of_child,
_compute_num_children,
compute_axial_conductances,
compute_levels,
convert_point_process_to_distributed,
convert_to_csc,
Expand Down Expand Up @@ -272,36 +273,9 @@ def _init_morph_custom_spsolve(self):
"""Initialize the morphology for the custom Jaxley solver."""
raise NotImplementedError

def init_conds(self, params: Dict, voltage_solver: str):
def _compute_axial_conductances(self, params: Dict[str, jnp.ndarray]):
"""Given radius, length, r_a, compute the axial coupling conductances."""
if voltage_solver.startswith("jaxley"):
all_conds = self._init_conds_custom_spsolve(params)
conds = self._init_conds_jax_spsolve(params)
for key in conds.keys():
all_conds[key] = conds[key]
return all_conds
else:
return self._init_conds_jax_spsolve(params)

@abstractmethod
def _init_conds_jax_spsolve(self, params: Dict):
"""Initialize coupling conductances.
Args:
params: Conductances and morphology parameters, not yet including
coupling conductances.
"""
raise NotImplementedError

@abstractmethod
def _init_conds_custom_spsolve(self, params: Dict):
"""Initialize coupling conductances.
Args:
params: Conductances and morphology parameters, not yet including
coupling conductances.
"""
raise NotImplementedError
return compute_axial_conductances(self._comp_edges, params)

def _append_channel_to_nodes(self, view: pd.DataFrame, channel: "jx.Channel"):
"""Adds channel nodes from constituents to `self.channel_nodes`."""
Expand Down Expand Up @@ -529,9 +503,9 @@ def get_all_parameters(
) -> Dict[str, jnp.ndarray]:
"""Return all parameters (and coupling conductances) needed to simulate.
Runs `init_conds()` and return every parameter that is needed to solve the ODE.
This includes conductances, radiuses, lengths, axial_resistivities, but also
coupling conductances.
Runs `_compute_axial_conductances()` and return every parameter that is needed
to solve the ODE. This includes conductances, radiuses, lengths,
axial_resistivities, but also coupling conductances.
This is done by first obtaining the current value of every parameter (not only
the trainable ones) and then replacing the trainable ones with the value
Expand Down Expand Up @@ -580,11 +554,8 @@ def get_all_parameters(
# `.set()` to work. This is done with `[:, None]`.
params[key] = params[key].at[inds].set(set_param[:, None])

# Compute conductance params and append them.
cond_params = self.init_conds(params=params, voltage_solver=voltage_solver)
for key in cond_params:
params[key] = cond_params[key]

# Compute conductance params and add them to the params dictionary.
params["axial_conductances"] = self._compute_axial_conductances(params=params)
return params

def get_states_from_nodes_and_edges(self):
Expand Down Expand Up @@ -967,19 +938,26 @@ def step(
# Voltage steps.
cm = params["capacitance"] # Abbreviation.

# Arguments used by all solvers.
solver_kwargs = {
"voltages": voltages,
"voltage_terms": (v_terms + syn_v_terms) / cm,
"constant_terms": (const_terms + i_ext + syn_const_terms) / cm,
"axial_conductances": params["axial_conductances"],
"internal_node_inds": self._internal_node_inds,
}

# Add solver specific arguments.
if voltage_solver == "jax.sparse":
solver_kwargs = {
"voltages": voltages,
"voltage_terms": (v_terms + syn_v_terms) / cm,
"constant_terms": (const_terms + i_ext + syn_const_terms) / cm,
"axial_conductances": params["axial_conductances"],
"data_inds": self._data_inds,
"indices": self._indices_jax_spsolve,
"indptr": self._indptr_jax_spsolve,
"sinks": np.asarray(self._comp_edges["sink"].to_list()),
"n_nodes": self._n_nodes,
"internal_node_inds": self._internal_node_inds,
}
solver_kwargs.update(
{
"sinks": np.asarray(self._comp_edges["sink"].to_list()),
"data_inds": self._data_inds,
"indices": self._indices_jax_spsolve,
"indptr": self._indptr_jax_spsolve,
"n_nodes": self._n_nodes,
}
)
# Only for `bwd_euler` and `cranck-nicolson`.
step_voltage_implicit = step_voltage_implicit_with_jax_spsolve
else:
Expand All @@ -989,33 +967,25 @@ def step(
# Currently, the forward Euler solver also uses this format. However,
# this is only for historical reasons and we are planning to change this in
# the future.
solver_kwargs = {
"voltages": voltages,
"voltage_terms": (v_terms + syn_v_terms) / cm,
"constant_terms": (const_terms + i_ext + syn_const_terms) / cm,
"branchpoint_conds_children": params["branchpoint_conds_children"],
"branchpoint_conds_parents": params["branchpoint_conds_parents"],
"branchpoint_weights_children": params["branchpoint_weights_children"],
"branchpoint_weights_parents": params["branchpoint_weights_parents"],
"axial_conductances": params["axial_conductances"],
"sources": np.asarray(self._comp_edges["source"].to_list()),
"sinks": np.asarray(self._comp_edges["sink"].to_list()),
"types": np.asarray(self._comp_edges["type"].to_list()),
"internal_node_inds": self._internal_node_inds,
"masked_node_inds": self._remapped_node_indices,
"n_nodes": self._n_nodes,
"nseg_per_branch": self.nseg_per_branch,
"nseg": self.nseg,
"par_inds": self.par_inds,
"child_inds": self.child_inds,
"nbranches": self.total_nbranches,
"solver": voltage_solver,
"children_in_level": self.children_in_level,
"parents_in_level": self.parents_in_level,
"root_inds": self.root_inds,
"branchpoint_group_inds": self.branchpoint_group_inds,
"debug_states": self.debug_states,
}
solver_kwargs.update(
{
"sinks": np.asarray(self._comp_edges["sink"].to_list()),
"sources": np.asarray(self._comp_edges["source"].to_list()),
"types": np.asarray(self._comp_edges["type"].to_list()),
"masked_node_inds": self._remapped_node_indices,
"nseg_per_branch": self.nseg_per_branch,
"nseg": self.nseg,
"par_inds": self.par_inds,
"child_inds": self.child_inds,
"nbranches": self.total_nbranches,
"solver": voltage_solver,
"children_in_level": self.children_in_level,
"parents_in_level": self.parents_in_level,
"root_inds": self.root_inds,
"branchpoint_group_inds": self.branchpoint_group_inds,
"debug_states": self.debug_states,
}
)
# Only for `bwd_euler` and `cranck-nicolson`.
step_voltage_implicit = step_voltage_implicit_with_custom_spsolve

Expand Down
59 changes: 0 additions & 59 deletions jaxley/modules/branch.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -148,65 +148,6 @@ def _init_morph_jax_spsolve(self):
self._indices_jax_spsolve = indices
self._indptr_jax_spsolve = indptr

def _init_conds_custom_spsolve(self, params: Dict) -> Dict[str, jnp.ndarray]:
conds = self.init_branch_conds_custom_spsolve(
params["axial_resistivity"], params["radius"], params["length"], self.nseg
)
cond_params = {
"branchpoint_conds_children": jnp.asarray([]),
"branchpoint_conds_parents": jnp.asarray([]),
"branchpoint_weights_children": jnp.asarray([]),
"branchpoint_weights_parents": jnp.asarray([]),
}
cond_params["branch_lowers"] = conds[0]
cond_params["branch_uppers"] = conds[1]
cond_params["branch_diags"] = conds[2]

return cond_params

def _init_conds_jax_spsolve(self, params: Dict) -> Dict[str, jnp.ndarray]:
conds = compute_axial_conductances(self._comp_edges, params)
return {"axial_conductances": conds}

@staticmethod
def init_branch_conds_custom_spsolve(
axial_resistivity: jnp.ndarray,
radiuses: jnp.ndarray,
lengths: jnp.ndarray,
nseg: int,
) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
"""Given an axial resisitivity, set the coupling conductances.
Args:
axial_resistivity: Axial resistivity of each compartment.
radiuses: Radius of each compartment.
lengths: Length of each compartment.
nseg: Number of compartments in the branch.
Returns:
Tuple of forward coupling conductances, backward coupling conductances, and summed coupling conductances.
"""

# Compute coupling conductance for segments within a branch.
# `radius`: um
# `r_a`: ohm cm
# `length_single_compartment`: um
# `coupling_conds`: S * um / cm / um^2 = S / cm / um
r1 = radiuses[:-1]
r2 = radiuses[1:]
r_a1 = axial_resistivity[:-1]
r_a2 = axial_resistivity[1:]
l1 = lengths[:-1]
l2 = lengths[1:]
coupling_conds_bwd = compute_coupling_cond(r1, r2, r_a1, r_a2, l1, l2)
coupling_conds_fwd = compute_coupling_cond(r2, r1, r_a2, r_a1, l2, l1)

# Compute the summed coupling conductances of each compartment.
summed_coupling_conds = jnp.zeros((nseg))
summed_coupling_conds = summed_coupling_conds.at[1:].add(coupling_conds_fwd)
summed_coupling_conds = summed_coupling_conds.at[:-1].add(coupling_conds_bwd)
return coupling_conds_fwd, coupling_conds_bwd, summed_coupling_conds

def __len__(self) -> int:
return self.nseg

Expand Down
60 changes: 0 additions & 60 deletions jaxley/modules/cell.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -274,66 +274,6 @@ def _init_morph_jax_spsolve(self):
self._indices_jax_spsolve = indices
self._indptr_jax_spsolve = indptr

def _init_conds_custom_spsolve(self, params: Dict) -> Dict[str, jnp.ndarray]:
"""Given an axial resisitivity, set the coupling conductances."""
axial_resistivity = params["axial_resistivity"]
radiuses = params["radius"]
lengths = params["length"]

nsegs = jnp.concatenate(
[
jnp.asarray([0]),
jnp.cumsum(self.nseg_per_branch),
]
)

child_comp_ind = nsegs[self.child_inds]
par_comp_ind = nsegs[self.par_inds + 1] - 1

# The conductance from the children to the branch point.
branchpoint_conds_children = vmap(
compute_coupling_cond_branchpoint, in_axes=(0, 0, 0)
)(
radiuses[child_comp_ind],
axial_resistivity[child_comp_ind],
lengths[child_comp_ind],
)
# The conductance from the parents to the branch point.
branchpoint_conds_parents = vmap(
compute_coupling_cond_branchpoint, in_axes=(0, 0, 0)
)(
radiuses[par_comp_ind],
axial_resistivity[par_comp_ind],
lengths[par_comp_ind],
)

# Weights with which the compartments influence their nearby node.
# The impact of the children on the branch point.
branchpoint_weights_children = vmap(compute_impact_on_node, in_axes=(0, 0, 0))(
radiuses[child_comp_ind],
axial_resistivity[child_comp_ind],
lengths[child_comp_ind],
)
# The impact of parents on the branch point.
branchpoint_weights_parents = vmap(compute_impact_on_node, in_axes=(0, 0, 0))(
radiuses[par_comp_ind],
axial_resistivity[par_comp_ind],
lengths[par_comp_ind],
)

cond_params = {
"branchpoint_conds_children": branchpoint_conds_children,
"branchpoint_conds_parents": branchpoint_conds_parents,
"branchpoint_weights_children": branchpoint_weights_children,
"branchpoint_weights_parents": branchpoint_weights_parents,
}
return cond_params

def _init_conds_jax_spsolve(self, params: Dict) -> Dict[str, jnp.ndarray]:
"""Given length, radius, and r_a, set the coupling conductances."""
conds = compute_axial_conductances(self._comp_edges, params)
return {"axial_conductances": conds}

@staticmethod
def update_summed_coupling_conds_custom_spsolve(
summed_conds,
Expand Down
16 changes: 4 additions & 12 deletions jaxley/modules/compartment.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,18 +92,10 @@ def _init_morph_jax_spsolve(self):
self._indices_jax_spsolve = indices
self._indptr_jax_spsolve = indptr

def _init_conds_custom_spsolve(self, params):
return {
"branchpoint_conds_children": jnp.asarray([]),
"branchpoint_conds_parents": jnp.asarray([]),
"branchpoint_weights_children": jnp.asarray([]),
"branchpoint_weights_parents": jnp.asarray([]),
"branch_uppers": jnp.asarray([]),
"branch_lowers": jnp.asarray([]),
"branch_diags": jnp.asarray([0.0]),
}

def _init_conds_jax_spsolve(self, params):
def init_conds(self, params: Dict[str, jnp.ndarray]):
"""Override `Base.init_axial_conds()`.
This is because compartments do not have any axial conductances."""
return {"axial_conductances": jnp.asarray([])}


Expand Down
Loading

0 comments on commit c668807

Please sign in to comment.