Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Set different nseg #427

Closed
wants to merge 6 commits into from
Closed

Set different nseg #427

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
83ae3d0
Enable different nseg per branch for custom jaxley sparse solvers
michaeldeistler Sep 19, 2024
68a7756
fix forward euler
michaeldeistler Sep 20, 2024
2c71bf0
modify type
michaeldeistler Sep 20, 2024
9e782e3
Use the new types to get rid of `init_conds_custom`
michaeldeistler Sep 20, 2024
54b45a5
Move some utilities to `solver_utils.py`
michaeldeistler Sep 20, 2024
091faf4
first pseudocode draft for setting nseg after init
michaeldeistler Sep 21, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 0 additions & 63 deletions jaxley/build_branched_tridiag.py
View file
Open in desktop

This file was deleted.

123 changes: 51 additions & 72 deletions jaxley/modules/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
from typing import Callable, Dict, List, Optional, Tuple, Union

import jax.numpy as jnp
import networkx as nx
import numpy as np
import pandas as pd
from jax import jit, vmap
Expand All @@ -17,16 +16,16 @@
from jaxley.channels import Channel
from jaxley.solver_voltage import (
step_voltage_explicit,
step_voltage_implicit_with_custom_spsolve,
step_voltage_implicit_with_jax_spsolve,
step_voltage_implicit_with_jaxley_spsolve,
)
from jaxley.synapses import Synapse
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,
interpolate_xyz,
loc_of_index,
query_channel_states_and_params,
Expand All @@ -35,6 +34,7 @@
from jaxley.utils.debug_solver import compute_morphology_indices
from jaxley.utils.misc_utils import childview, concat_and_ignore_empty
from jaxley.utils.plot_utils import plot_morph
from jaxley.utils.solver_utils import convert_to_csc


class Module(ABC):
Expand Down Expand Up @@ -258,7 +258,7 @@ def _show(

def init_morph(self):
"""Initialize the morphology such that it can be processed by the solvers."""
self._init_morph_custom_spsolve()
self._init_morph_jaxley_spsolve()
self._init_morph_jax_spsolve()
self.initialized_morph = True

Expand All @@ -268,36 +268,13 @@ def _init_morph_jax_spsolve(self):
raise NotImplementedError

@abstractmethod
def _init_morph_custom_spsolve(self):
def _init_morph_jaxley_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"):
return self._init_conds_custom_spsolve(params)
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 @@ -525,9 +502,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 @@ -576,11 +553,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 @@ -963,19 +937,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 @@ -985,29 +966,27 @@ 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,
"coupling_conds_upper": params["branch_uppers"],
"coupling_conds_lower": params["branch_lowers"],
"summed_coupling_conds": params["branch_diags"],
"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"],
"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
step_voltage_implicit = step_voltage_implicit_with_jaxley_spsolve

if solver == "bwd_euler":
u["v"] = step_voltage_implicit(**solver_kwargs, delta_t=delta_t)
Expand Down
Loading
Loading