From b28e193f1c320d3c2e203696b1c31bc89c9d7637 Mon Sep 17 00:00:00 2001 From: "w.wybo" Date: Mon, 22 Sep 2025 13:46:20 +0200 Subject: [PATCH 1/9] initial fix --- src/neat/modelreduction/compartmentfitter.py | 15 +++++++++++- src/neat/trees/compartmenttree.py | 24 ++++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/src/neat/modelreduction/compartmentfitter.py b/src/neat/modelreduction/compartmentfitter.py index 9b3b6aed..0a4ae6b2 100755 --- a/src/neat/modelreduction/compartmentfitter.py +++ b/src/neat/modelreduction/compartmentfitter.py @@ -1197,6 +1197,7 @@ def fit_e_eq(self, fit_arg): ctree, locs = self.convert_fit_arg(fit_arg) # compute the equilibirum potentials at fit locations + # v_eqs_fit, conc_eqs_fit = self.calc_e_eq(locs) v_eqs_fit, conc_eqs_fit = self.calc_e_eq(locs) # set the equilibria @@ -1205,7 +1206,19 @@ def fit_e_eq(self, fit_arg): ctree.set_conc_eq(ion, conc_eqs_fit[ion]) # fit the leak - ctree.fit_e_leak() + # print("---") + ctree.fit_e_leak_() + # lll = [] + # for n in ctree: lll.append(n.currents['L']) + # lll = np.array(lll) + # ctree.fit_e_leak() + # lll_ = [] + # for n in ctree: lll_.append(n.currents['L']) + # lll_ = np.array(lll_) + + # for e1, e2 in zip(lll, lll_): + # print(e1[1], e2[1]) + # breakpoint() return ctree, locs diff --git a/src/neat/trees/compartmenttree.py b/src/neat/trees/compartmenttree.py index eedf16f4..f770e57b 100755 --- a/src/neat/trees/compartmenttree.py +++ b/src/neat/trees/compartmenttree.py @@ -845,6 +845,30 @@ def _jac_e_leak_fit(self, e_l): node.currents["L"][1] = e_l[ii] jac_vals = np.array([-node.currents["L"][0] for node in self]) return np.diag(jac_vals) + + def fit_e_leak_(self): + """ + Fit the leak reversal potential to obtain the stored equilibirum potentials + as resting membrane potential + """ + e_l_0 = self.get_e_eq(indexing="tree") + + fun = self._fun_e_leak_fit(e_l_0) + jac = self._jac_e_leak_fit_(e_l_0) + e_l = (-fun + jac * e_l_0) / jac + # set the leak reversals + for ii, node in enumerate(self): + if not e_l.mask[ii]: + node.currents["L"] = [node.currents["L"][0], e_l[ii]] + + def _jac_e_leak_fit_(self, e_l): + for ii, node in enumerate(self): + node.currents["L"][1] = e_l[ii] + return np.ma.masked_array( + [-node.currents["L"][0] for node in self], + mask=[node.currents["L"][0] < 1e-16 for node in self], + ) + def add_channel_current(self, channel, e_rev): """ From e7f9ead461675f5b03f84fff9906e76448cc723e Mon Sep 17 00:00:00 2001 From: "w.wybo" Date: Tue, 9 Dec 2025 10:45:15 +0100 Subject: [PATCH 2/9] initial neuron model and channel generation --- src/neat/channels/ionchannels.py | 112 ++++++++++ tests/test_jaxleytree.py | 351 +++++++++++++++++++++++++++++++ 2 files changed, 463 insertions(+) create mode 100644 tests/test_jaxleytree.py diff --git a/src/neat/channels/ionchannels.py b/src/neat/channels/ionchannels.py index a55e9da9..dbeef3df 100755 --- a/src/neat/channels/ionchannels.py +++ b/src/neat/channels/ionchannels.py @@ -1258,3 +1258,115 @@ def _replaceConc(expr_str, prefix="", suffix=""): fh.close() fcc.close() + + + def _print_jaxley_pycode(self, expr): + pstr = sp.printing.pycode(expr) + pstr = pstr.replace("match.exp", "save_exp") + pstr = pstr.replace("math.", "jnp.") + return pstr + + def write_jaxley_code(self, + v_comp=-75.0, + g=0.0, + e=None, + ): + cname = self.__class__.__name__ + sv = [str(svar) for svar in self.ordered_statevars] + cs = [str(conc) for conc in self.conc] + e = self._get_reversal(e) + + channelstr = f""" +class {cname}(Channel): + # Automatically generated using neat + + def __init__(self, name=None): + self.current_is_in_mA_per_cm2 = True + super().__init__(name) + self.channel_params = {{ + f'{{self.prefix}}gbar_{cname}': 1e-4, + f'erev_{cname}': {e}, + }} + self.channel_states = {{ + {",\n ".join( + f"f'{{self.prefix}}{var}': 0.0" for var in sv + )} + }} + + + @property + def prefix(self): + return f'{{self._name}}_' + +""" + channelstr += f""" + def update_states( + self, + u: Dict[str, jnp.ndarray], + dt: float, + voltages: float, + params: Dict[str, jnp.ndarray], + ): + v = voltages + { + "\n ".join([ + f"{var} = u[f'{{self.prefix}}_{var}']" for var in sv + ]) + } +""" + # substitution for common neuron names + repl_pairs = [(str(c), str(c) + "i") for c in self.conc] + for var, svar in zip(sv, self.ordered_statevars): + vi = self._print_jaxley_pycode(self.varinf[svar])#, assign_to=f"{var}_inf") + ti = self._print_jaxley_pycode(self.tauinf[svar])#, assign_to=f"tau_{var}") + for repl_pair in repl_pairs: + vi = vi.replace(*repl_pair) + ti = ti.replace(*repl_pair) + + channelstr += f""" + # advance state variable {var} + {var}_new = solve_inf_gate_exponential( + {svar}, + dt, + {vi}, + {ti} + ) +""" + channelstr += f""" + return {{ + {",\n ".join( + f"f'{{self.prefix}}{var}': {var}_new" for var in sv + )} + }} +""" + + channelstr += f""" + def compute_current( + self, u: Dict[str, jnp.ndarray], voltages, params: Dict[str, jnp.ndarray] + ): + gbar = params[f'{{self.prefix}}gbar_{cname}'] + erev = params[f'{{self.prefix}}erev_{cname}'] + v = voltages + { + "\n ".join([ + f"{var} = u[f'{{self.prefix}}_{var}']" for var in sv + ]) + } + return gbar * {self._print_jaxley_pycode(self.p_open)} * (v - erev) +""" + channelstr += f""" + def init_state(self, states, voltages, params, delta_t): + v = voltages + { '\n '.join([ + f'{var} = {self._print_jaxley_pycode(self.varinf[svar])}' \ + for var, svar in zip(sv, self.ordered_statevars) + ]) + } + return {{ + {",\n ".join( + f"f'{{self.prefix}}{var}': {var}" for var in sv + )} + }} +""" + + return channelstr \ No newline at end of file diff --git a/tests/test_jaxleytree.py b/tests/test_jaxleytree.py new file mode 100644 index 00000000..0e24fb6c --- /dev/null +++ b/tests/test_jaxleytree.py @@ -0,0 +1,351 @@ +import numpy as np +import jaxley as jx + +import os + +from neat import PhysTree, PhysNode, MorphTree, CompartmentFitter, CompartmentTree + +import channelcollection_for_tests as channelcollection +import channel_installer + +channel_installer.load_or_install_neuron_test_channels() + + +MORPHOLOGIES_PATH_PREFIX = os.path.abspath( + os.path.join(os.path.dirname(__file__), "test_morphologies") +) + +class JaxleySimNode(PhysNode): + """ + Subclass of `neat.PhysNode` that implements functionality to instantiate a + cylindrical `neuron.h.Section` based on its physiological and geometrical + parameters. + """ + + def __init__(self, index, p3d=None): + super().__init__(index, p3d) + + def _make_branch(self, dx_max=15, factorlambda=1.0, pprint=False): + if dx_max is None: + n_comp = factorlambda + + if self.index == 1: + n_comp = 1 + l_comp = 2.0 * self.R + + else: + n_comp = self.L / dx_max + 1 + l_comp = self.L / n_comp + + compartments = [] + for _ in range(n_comp): + jx_comp = jx.Compartment() + jx_comp.set("length", l_comp) + jx_comp.set("radius", self.R) + jx_comp.set("axial_resistivity", self.r_a * 1e6) # MOhm*cm --> Ohm*cm + jx_comp.set("capacitance", self.c_m) # uF/cm^2 + compartments.append(jx_comp) + + branch = jx.Branch(compartments) + return branch + + +class JaxleySimTree(PhysTree): + """ + Tree class to define NEURON (Carnevale & Hines, 2004) based on `neat.PhysTree`. + + Attributes + ---------- + sections: dict of hoc sections + Storage for hoc sections. Keys are node indices. + shunts: list of hoc mechanisms + Storage container for shunts + syns: list of hoc mechanisms + Storage container for synapses + iclamps: list of hoc mechanisms + Storage container for current clamps + vclamps: lis of hoc mechanisms + Storage container for voltage clamps + vecstims: list of hoc mechanisms + Storage container for vecstim objects + netcons: list of hoc mechanisms + Storage container for netcon objects + vecs: list of hoc vectors + Storage container for hoc spike vectors + dt: float + timestep of the simulator ``[ms]`` + t_calibrate: float + Time for the model to equilibrate``[ms]``. Not counted as part of the + simulation. + factor_lambda : int or float + If int, the number of segments per section. If float, multiplies the + number of segments given by the standard lambda rule (Carnevale, 2004) + to give the number of compartments simulated (default value 1. gives + the number given by the lambda rule) + v_init: float + The initial voltage at which the model is initialized ``[mV]`` + + A `NeuronSimTree` can be extended easily with custom point process mechanisms. + Just make sure that you store the point process in an existing appropriate + storage container or in a custom storage container, since if all references + to the hocobject disappear, the object itself will be deleted as well. + + .. code-block:: python + class CustomSimTree(NeuronSimTree): + def addCustomPointProcessMech(self, loc, **kwargs): + loc = MorphLoc(loc, self) + + # create the point process + pp = h.custom_point_process(self.sections[loc['node']](loc['x'])) + pp.arg1 = kwargs['arg1'] + pp.arg2 = kwargs['arg2'] + ... + + self.storage_container_for_point_process.append(pp) + + If you define a custom storage container, make sure that you overwrite the + `__init__()` and `delete_model()` functions to make sure it is created and + deleted properly. + """ + + def __init__(self, arg=None, types=[1, 3, 4]): + # neuron storage + self.sections = {} + self.shunts = [] + self.syns = [] + self.iclamps = [] + self.vclamps = [] + self.vecstims = [] + self.netcons = [] + self.vecs = [] + # simulation parameters + self.dt = 0.1 # ms + self.t_calibrate = 0.0 # ms + self.factor_lambda = 1.0 + self.v_init = -75.0 # mV + self.indstart = 0 + # initialize the tree structure + super().__init__(arg=arg, types=types) + + def create_corresponding_node(self, node_index, p3d=None): + """ + Creates a node with the given index corresponding to the tree class. + + Parameters + ---------- + node_index: int + index of the new node + """ + return JaxleySimNode(node_index, p3d=p3d) + + + def init_model( + self, dt=0.025, t_calibrate=0.0, v_init=-75.0, factor_lambda=1.0, pprint=False + ): + """ + Initialize hoc-objects to simulate the neuron model implemented by this + tree. + + Parameters + ---------- + dt: float (default is ``.025`` ms) + Timestep of the simulation + t_calibrate: float (default ``0.`` ms) + The calibration time; time model runs without input to reach its + equilibrium state before the true simulation starts + v_init: float (default ``-75.`` mV) + The initial voltage at which the model is initialized + factor_lambda: float or int (default 1.) + If int, the number of segments per section. If float, multiplies the + number of segments given by the standard lambda rule (Carnevale, 2004) + to give the number of compartments simulated (default value 1. gives + the number given by the lambda rule) + pprint: bool (default ``False``) + Whether or not to print info on the NEURON model's creation + """ + # self.set_simulation_parameters( + # dt=dt, + # t_calibrate=t_calibrate, + # v_init=v_init, + # factor_lambda=factor_lambda, + # ) + # reset all storage + # self.delete_model() + # create the Jaxley + cell = self._create_jaxley_cell(pprint=pprint) + return cell + + def _create_jaxley_cell(self, pprint): + # simple tree copy + aux_tree = MorphTree(self) + aux_tree.reset_indices() + + branches = [node._make_branch() for node in self] + parents = [aux_node.parent_node.index if not aux_tree.is_root(aux_node) else -1 for aux_node in aux_tree] + + cell = jx.Cell(branches, parents) + return cell + + +class JaxleyCompartmentNode(JaxleySimNode): + """ + Subclass of `NeuronSimNode` that defines a cylinder with fake geometry in + NEURON to result in the effective simulation as a single compartment. + """ + + def __init__(self, index): + super().__init__(index) + + def get_child_nodes(self, skip_inds=[]): + return super().get_child_nodes(skip_inds=skip_inds) + + def _make_branch(self, pprint=False): + return super()._make_branch(dx_max=None, factorlambda=1, pprint=False) + +class JaxleyCompartmentTree(JaxleySimTree): + """ + Creates a `neat.NeuronCompartmentTree` to simulate reduced compartmentment + models from a `neat.CompartmentTree`. + + Parameters + ---------- + ctree: `neat.CompartmentTree` + The tree containing the parameters of the reduced compartmental model + to be simulated + fake_c_m: float + Fake value for the membrance capacitance density, rescales cylinder + surface + fake_r_a: float + Fake value for the axial resistance, rescales cylinder length + + Attributes + ---------- + equivalent_locs: list of tuples + 'Fake' locations corresponding to each compartment, which are + used to insert hoc point process at the compartments using + the same functions definitions as for as for a morphological + `neat.NeuronSimTree`. + + Notes + ----- + - Note that this class inherits from `neat.NeuronSimTree` and *not* from + `neat.CompartmentTree`. This is because NEAT defines a fake morphology to + implement the compartment model in NEURON, and also to reuse the functionality + implemented by `neat.NeuronSimTree`. Any function that is not explicitly + redefined from `neat.NeuronSimTree` can be called in the same way for this + compartment model. + - Locations to this class can be provided either as fake morphology locations + -- i.e. a tuple `(node.index, x-location in [0,1])` -- where the value of the + x-location is ignored since the nodes here are single compartments, as in + the `neat.CompartmentTree`, and not cylinders, as in `neat.MorphTree` or + subclasses, or as location indices, where the index corresponds to the location + in the original list of locations from which the `neat.CompartmentTree` was + derived. + """ + + def __init__(self, ctree, fake_c_m=1.0, fake_r_a=100.0 * 1e-6, method=2): + + try: + assert issubclass(ctree.__class__, CompartmentTree) + except AssertionError as e: + raise ValueError( + "`neat.NeuronCompartmentTree` can only be instantiated " + "from a `neat.CompartmentTree` or derived class" + ) + super().__init__(ctree, types=[1, 3, 4]) + self.equivalent_locs = ctree.get_equivalent_locs() + self._create_reduced_jaxley_model( + ctree, + fake_c_m=fake_c_m, + fake_r_a=fake_r_a, + ) + + def _create_reduced_jaxley_model(self, ctree, fake_c_m=1.0, fake_r_a=100.0 * 1e-6): + arg1, arg2 = ctree.compute_fake_geometry( + fake_c_m=fake_c_m, + fake_r_a=fake_r_a, + factor_r_a=1e-6, + delta=1e-10, + method=2, + ) + lengths = arg1 + radii = arg2 + surfaces = 2.0 * np.pi * radii * lengths + for ii, comp_node in enumerate(ctree): + sim_node = self.__getitem__(comp_node.index, skip_inds=[]) + if self.is_root(sim_node): + sim_node.set_p3d(np.array([0.0, 0.0, 0.0]), radii[ii] * 1e4, 1) + else: + sim_node.set_p3d( + np.array( + [sim_node.parent_node.xyz[0] + lengths[ii] * 1e4, 0.0, 0.0] + ), + radii[ii] * 1e4, + 3, + ) + + # # fill the tree with the currents + # for ii, sim_node in enumerate(self): + # comp_node = ctree[ii] + # sim_node.currents = { + # chan: [g / surfaces[comp_node.index], e] + # for chan, (g, e) in comp_node.currents.items() + # } + # sim_node.concmechs = copy.deepcopy(comp_node.concmechs) + # for concmech in sim_node.concmechs.values(): + # concmech.gamma *= surfaces[comp_node.index] * 1e6 + # sim_node.c_m = fake_c_m + # sim_node.r_a = fake_r_a + # sim_node.R = radii[comp_node.index] * 1e4 # convert to [um] + # sim_node.L = lengths[comp_node.index] * 1e4 # convert to [um] + + + +class TestJaxley: + def load_ball(self): + """ + Load point neuron model + """ + self.tree = PhysTree(os.path.join(MORPHOLOGIES_PATH_PREFIX, "ball.swc")) + # capacitance and axial resistance + self.tree.set_physiology(0.8, 100.0 / 1e6) + # ion channels + self.k_chan = channelcollection.Kv3_1() + self.tree.add_channel_current(self.k_chan, 0.766 * 1e6, -85.0) + self.na_chan = channelcollection.NaTa_t() + self.tree.add_channel_current(self.na_chan, 1.71 * 1e6, 50.0) + # fit leak current + self.tree.fit_leak_current(-75.0, 10.0) + # set equilibirum potententials + self.tree.set_v_ep(-75.0) + # set computational tree + self.tree.set_comp_tree() + + cfit = CompartmentFitter(self.tree, save_cache=False, recompute_cache=True) + self.ctree, _ = cfit.fit_model([(1, 0.5)]) + + + def create_jaxley_model(self): + self.load_ball() + + jt = JaxleySimTree(self.tree) + jcell = jt.init_model() + + jct = JaxleyCompartmentTree(self.ctree) + jcompcell = jt.init_model() + + def test_jaxley_channels(self): + self.load_ball() + + channelstr_k = self.k_chan.write_jaxley_code() + channelstr_na = self.na_chan.write_jaxley_code() + with open("jtest", "w") as file: + file.write(channelstr_k + "\n\n\n") + file.write(channelstr_na) + + + +if __name__ == "__main__": + tjx = TestJaxley() + # tjx.create_jaxley_model() + tjx.test_jaxley_channels() From 7d578fe2c0bf930a78ecf70728ab5a8d0f78a3a3 Mon Sep 17 00:00:00 2001 From: "w.wybo" Date: Tue, 9 Dec 2025 11:25:50 +0100 Subject: [PATCH 3/9] update actions to incorporate jaxley channel generation --- src/neat/actions/install.py | 47 +++++++++++++++++++++++++++++++++- src/neat/actions/list.py | 11 +++++--- src/neat/actions/neatmodels.py | 12 +++++++-- src/neat/actions/uninstall.py | 12 ++++++++- 4 files changed, 75 insertions(+), 7 deletions(-) diff --git a/src/neat/actions/install.py b/src/neat/actions/install.py index 04cfc2ac..0ef171f9 100644 --- a/src/neat/actions/install.py +++ b/src/neat/actions/install.py @@ -254,13 +254,53 @@ def _compile_nest(model_name, path_neat, channels, path_nestresource=None, ions= ) +def _compile_jaxley(model_name, path_neat, channels, path_jaxleyresource=None): + + # combine `model_name` with the neuron compilation path + path_for_jaxley_compilation = os.path.join( + path_neat, "simulations/jaxley/tmp/" + ) + + jaxley_file_path = os.path.join( + path_for_jaxley_compilation, f"{model_name}.py" + ) + + print(f"--- writing channels to \n" f" > {jaxley_file_path}") + os.makedirs(path_for_jaxley_compilation, exist_ok=True) + + + filestr = """ +from typing import Dict, Optional + +import jax.numpy as jnp +from jax.lax import select +from jaxley.channels import Channel +from jaxley.solver_gate import ( + exponential_euler, + save_exp, + solve_gate_exponential, + solve_inf_gate_exponential, +) + + +""" + for chan in channels: + print(" - writing .jaxley code for:", chan.__class__.__name__) + channelstr = chan.write_jaxley_code() + filestr += channelstr + + with open(jaxley_file_path, "w") as file: + file.write(filestr) + + def _install_models( model_name, path_neat, channel_path_arg, - simulators=["neuron", "nest"], + simulators=["neuron", "nest", "jaxley"], path_nestresource=None, path_neuronresource=None, + path_jaxleyresource=None, ): """ Compile a set of ion channels models specified by [channel_path_arg] @@ -298,3 +338,8 @@ def _install_models( _compile_nest( model_name, path_neat, channels, path_nestresource=path_nestresource ) + + if "jaxley" in simulators: + _compile_jaxley( + model_name, path_neat, channels, path_jaxleyresource=path_jaxleyresource + ) diff --git a/src/neat/actions/list.py b/src/neat/actions/list.py index 61e14c48..53d318c5 100644 --- a/src/neat/actions/list.py +++ b/src/neat/actions/list.py @@ -23,7 +23,7 @@ import glob -def _list_models(path_neat, simulators=["nest", "neuron"], pprint=True): +def _list_models(path_neat, simulators=["nest", "neuron", "jaxley"], pprint=True): """ Get (and print) a dictionary containing all installed model @@ -49,12 +49,17 @@ def _list_models(path_neat, simulators=["nest", "neuron"], pprint=True): if "neuron" in simulators: path_neuron = os.path.join(path_neat, "simulations/", "neuron/tmp/*/") - print(path_neuron) - for file_path in glob.glob(path_neuron): file_name = os.path.basename(os.path.normpath(file_path)) models["neuron"].append(file_name) + if "jaxley" in simulators: + path_jaxley = os.path.join(path_neat, "simulations/", "jaxley/tmp/*") + + for file_path in glob.glob(path_jaxley): + file_name = os.path.basename(os.path.normpath(file_path)).replace(".py", "") + models["jaxley"].append(file_name) + if pprint: print("\n------- installed models --------") for simulator, model_list in models.items(): diff --git a/src/neat/actions/neatmodels.py b/src/neat/actions/neatmodels.py index 145fcc24..c7577344 100755 --- a/src/neat/actions/neatmodels.py +++ b/src/neat/actions/neatmodels.py @@ -60,7 +60,7 @@ def parse_cmd_args(path_neat): "-s", "--simulator", nargs="*", - choices=["neuron", "nest"], + choices=["neuron", "nest", "jaxley"], default=["neuron"], help="The simulators to which the action is applied. \n" "Default is 'neuron'", ) @@ -75,11 +75,19 @@ def parse_cmd_args(path_neat): parser.add_argument( "--nestresource", default=os.path.join(path_neat, "simulations/nest/default_syns.nestml"), - help="Path to directory containing additional .mod-file mechanisms " + help="Path to directory containing additional .nestml-file mechanisms " "(e.g. synapses) that will be compiled together with the " "generated ion channel mod files.\n" "Only used when the [action] is 'install'.", ) + parser.add_argument( + "--jaxleyresource", + default=os.path.join(path_neat, "simulations/nest/default_syns.nestml"), + help="Path to directory containing additional jaxley mechanisms " + "(e.g. synapses) that will be combined together with the " + "generated ion channel mod files.\n" + "Only used when the [action] is 'install'.", + ) parser.add_argument( "-p", "--path", diff --git a/src/neat/actions/uninstall.py b/src/neat/actions/uninstall.py index 25ba8653..45bea9f8 100644 --- a/src/neat/actions/uninstall.py +++ b/src/neat/actions/uninstall.py @@ -39,7 +39,7 @@ def _check_model_name(model_name): def _uninstall_models( *model_names, path_neat, - simulators=["nest", "neuron"], + simulators=["nest", "neuron", "jaxley"], ): """ Uninstall the model with the given name from the provided simulators @@ -75,3 +75,13 @@ def _uninstall_models( print(f"> Uninstalled {model_name} from neuron") except FileNotFoundError as e: print(f"> {model_name} not found in neuron, nothing to uninstall.") + + if "jaxley" in simulators: + try: + path_jaxley = os.path.join( + path_neat, "simulations/", f"jaxley/tmp/{model_name}.py" + ) + os.remove(path_jaxley) + print(f"> Uninstalled {model_name} from jaxley") + except FileNotFoundError as e: + print(f"> {model_name} not found in jaxley, nothing to uninstall.") From c31d6a40a9bd7715055b034e879524fd4487c2e9 Mon Sep 17 00:00:00 2001 From: "w.wybo" Date: Tue, 9 Dec 2025 12:33:08 +0100 Subject: [PATCH 4/9] working jaxley tree creation --- .gitignore | 1 + src/neat/__init__.py | 8 + src/neat/channels/ionchannels.py | 25 +- src/neat/simulations/jaxley/jaxleymodel.py | 370 +++++++++++++++++++++ tests/channel_installer.py | 33 +- tests/test_jaxleytree.py | 310 +---------------- 6 files changed, 435 insertions(+), 312 deletions(-) create mode 100644 src/neat/simulations/jaxley/jaxleymodel.py diff --git a/.gitignore b/.gitignore index 047ea7e0..794f95c6 100755 --- a/.gitignore +++ b/.gitignore @@ -15,6 +15,7 @@ report/ # ignore installed models src/neat/simulations/nest/tmp/ src/neat/simulations/neuron/tmp/ +src/neat/simulations/jaxley/tmp/ # cython generated neat/simulations/netsim.cpp diff --git a/src/neat/__init__.py b/src/neat/__init__.py index 34c56994..793608e8 100755 --- a/src/neat/__init__.py +++ b/src/neat/__init__.py @@ -69,6 +69,14 @@ except ModuleNotFoundError: warnings.warn("NEST not available", UserWarning) +try: + from .simulations.jaxley.jaxleymodel import load_jaxley_model + from .simulations.jaxley.jaxleymodel import JaxleySimTree + from .simulations.jaxley.jaxleymodel import JaxleySimNode + from .simulations.jaxley.jaxleymodel import JaxleyCompartmentTree +except ModuleNotFoundError: + warnings.warn("Jaxley not available", UserWarning) + from .tools.kernelextraction import FourierQuadrature from .tools.kernelextraction import FourierTools from .tools.kernelextraction import Kernel diff --git a/src/neat/channels/ionchannels.py b/src/neat/channels/ionchannels.py index dbeef3df..50bc21c0 100755 --- a/src/neat/channels/ionchannels.py +++ b/src/neat/channels/ionchannels.py @@ -383,7 +383,6 @@ def __init__(self, **kwargs): except KeyError: warnings.warn("No default reversal potential defined.") - # self._lambdify_channel() self.set_default_params(**kwargs) def __getstate__(self): @@ -1262,7 +1261,7 @@ def _replaceConc(expr_str, prefix="", suffix=""): def _print_jaxley_pycode(self, expr): pstr = sp.printing.pycode(expr) - pstr = pstr.replace("match.exp", "save_exp") + pstr = pstr.replace("math.exp", "save_exp") pstr = pstr.replace("math.", "jnp.") return pstr @@ -1284,20 +1283,19 @@ def __init__(self, name=None): self.current_is_in_mA_per_cm2 = True super().__init__(name) self.channel_params = {{ - f'{{self.prefix}}gbar_{cname}': 1e-4, - f'erev_{cname}': {e}, + f'{{self.prefix}}g{cname}': 1e-4, + f'{{self.prefix}}e{cname}': {e}, }} self.channel_states = {{ {",\n ".join( f"f'{{self.prefix}}{var}': 0.0" for var in sv )} }} + self.current_name = "i_{self.ion}" - @property def prefix(self): return f'{{self._name}}_' - """ channelstr += f""" def update_states( @@ -1310,7 +1308,7 @@ def update_states( v = voltages { "\n ".join([ - f"{var} = u[f'{{self.prefix}}_{var}']" for var in sv + f"{var} = u[f'{{self.prefix}}{var}']" for var in sv ]) } """ @@ -1333,10 +1331,10 @@ def update_states( ) """ channelstr += f""" - return {{ + return {{ {",\n ".join( f"f'{{self.prefix}}{var}': {var}_new" for var in sv - )} + )} }} """ @@ -1344,15 +1342,15 @@ def update_states( def compute_current( self, u: Dict[str, jnp.ndarray], voltages, params: Dict[str, jnp.ndarray] ): - gbar = params[f'{{self.prefix}}gbar_{cname}'] - erev = params[f'{{self.prefix}}erev_{cname}'] + gbar = params[f'{{self.prefix}}g{cname}'] + erev = params[f'{{self.prefix}}e{cname}'] v = voltages { "\n ".join([ - f"{var} = u[f'{{self.prefix}}_{var}']" for var in sv + f"{var} = u[f'{{self.prefix}}{var}']" for var in sv ]) } - return gbar * {self._print_jaxley_pycode(self.p_open)} * (v - erev) + return gbar * ({self._print_jaxley_pycode(self.p_open)}) * (v - erev) """ channelstr += f""" def init_state(self, states, voltages, params, delta_t): @@ -1368,5 +1366,4 @@ def init_state(self, states, voltages, params, delta_t): )} }} """ - return channelstr \ No newline at end of file diff --git a/src/neat/simulations/jaxley/jaxleymodel.py b/src/neat/simulations/jaxley/jaxleymodel.py new file mode 100644 index 00000000..8b8260ef --- /dev/null +++ b/src/neat/simulations/jaxley/jaxleymodel.py @@ -0,0 +1,370 @@ +import numpy as np + +import os +import sys +import copy +import importlib + +from ...trees.phystree import PhysTree, PhysNode +from ...trees.morphtree import MorphTree +from ...trees.compartmenttree import CompartmentTree +from ...factorydefaults import DefaultPhysiology + +try: + import jaxley as jx + from jaxley.channels import Leak + +except ModuleNotFoundError: + warnings.warn( + "Jaxley not available, importing non-functional jx module only for doc generation", + UserWarning, + ) + + # universal iterable mock object + class JX(object): + def __init__(self): + pass + + def __getattr__(self, attr): + try: + return super(H, self).__getattr__(attr) + except AttributeError: + return self.__global_handler + + def __global_handler(self, *args, **kwargs): + return JX() + + def __iter__(self): # make iterable + return self + + def __next__(self): + raise StopIteration + + def __mul__(self, other): # make multipliable + return 1.0 + + def __rmul__(self, other): + return self * other + + def __call__(self, *args, **kwargs): # make callable + return JX() + + jx = JX() + np_array = np.array + + def array(*args, **kwargs): + if isinstance(args[0], JX): + print(args) + print(kwargs) + return np.eye(2) + else: + return np_array(*args, **kwargs) + + np.array = array + +JX_MECH = {} + +def load_jaxley_model(name): + jx_path = os.path.join( + os.path.dirname(__file__), + f"tmp/", + ) + if os.path.exists(jx_path + f"{name}.py"): + if jx_path not in sys.path: + sys.path.insert(0, jx_path) + JX_MECH[name] = importlib.import_module(name) + else: + raise FileNotFoundError( + f"The Jaxley model named '{name}' is not installed. " + f"Run 'neatmodels -h' in the terminal for help on " + f"installing new Jaxley models with NEAT. " + f"Installed models will be in '{path}'." + ) + + +class JaxleySimNode(PhysNode): + """ + Subclass of `neat.PhysNode` that implements functionality to instantiate a + cylindrical `neuron.h.Section` based on its physiological and geometrical + parameters. + """ + + def __init__(self, index, p3d=None): + super().__init__(index, p3d) + + def _make_branch(self, jx_channels, dx_max=15, factorlambda=1.0, pprint=False): + if dx_max is None: + n_comp = factorlambda + + if self.index == 1: + n_comp = 1 + l_comp = 2.0 * self.R + + else: + n_comp = self.L / dx_max + 1 + l_comp = self.L / n_comp + + compartments = [] + for _ in range(n_comp): + jx_comp = jx.Compartment() + jx_comp.set("length", l_comp) + jx_comp.set("radius", self.R) + jx_comp.set("axial_resistivity", self.r_a * 1e6) # MOhm*cm --> Ohm*cm + jx_comp.set("capacitance", self.c_m) # uF/cm^2 + + for key, current in self.currents.items(): + jx_channel = jx_channels[key] + if current[0] > 1e-10: + prefix = jx_channel.prefix if key != 'L' else 'Leak_' + suffix = key if key != 'L' else 'Leak' + jx_comp.insert(jx_channel) + jx_comp.set(f"{prefix}g{suffix}", current[0]*1e-6) # uS/cm^2 --> S/cm^2 + jx_comp.set(f"{prefix}e{suffix}", current[1]) # mV + + compartments.append(jx_comp) + + branch = jx.Branch(compartments) + return branch + + +class JaxleySimTree(PhysTree): + """ + Tree class to define NEURON (Carnevale & Hines, 2004) based on `neat.PhysTree`. + + Attributes + ---------- + sections: dict of hoc sections + Storage for hoc sections. Keys are node indices. + shunts: list of hoc mechanisms + Storage container for shunts + syns: list of hoc mechanisms + Storage container for synapses + iclamps: list of hoc mechanisms + Storage container for current clamps + vclamps: lis of hoc mechanisms + Storage container for voltage clamps + vecstims: list of hoc mechanisms + Storage container for vecstim objects + netcons: list of hoc mechanisms + Storage container for netcon objects + vecs: list of hoc vectors + Storage container for hoc spike vectors + dt: float + timestep of the simulator ``[ms]`` + t_calibrate: float + Time for the model to equilibrate``[ms]``. Not counted as part of the + simulation. + factor_lambda : int or float + If int, the number of segments per section. If float, multiplies the + number of segments given by the standard lambda rule (Carnevale, 2004) + to give the number of compartments simulated (default value 1. gives + the number given by the lambda rule) + v_init: float + The initial voltage at which the model is initialized ``[mV]`` + + A `NeuronSimTree` can be extended easily with custom point process mechanisms. + Just make sure that you store the point process in an existing appropriate + storage container or in a custom storage container, since if all references + to the hocobject disappear, the object itself will be deleted as well. + + .. code-block:: python + class CustomSimTree(NeuronSimTree): + def addCustomPointProcessMech(self, loc, **kwargs): + loc = MorphLoc(loc, self) + + # create the point process + pp = h.custom_point_process(self.sections[loc['node']](loc['x'])) + pp.arg1 = kwargs['arg1'] + pp.arg2 = kwargs['arg2'] + ... + + self.storage_container_for_point_process.append(pp) + + If you define a custom storage container, make sure that you overwrite the + `__init__()` and `delete_model()` functions to make sure it is created and + deleted properly. + """ + + def __init__(self, arg=None, types=[1, 3, 4]): + # neuron storage + self.sections = {} + self.shunts = [] + self.syns = [] + self.iclamps = [] + self.vclamps = [] + self.vecstims = [] + self.netcons = [] + self.vecs = [] + # simulation parameters + self.dt = 0.1 # ms + self.t_calibrate = 0.0 # ms + self.factor_lambda = 1.0 + self.v_init = -75.0 # mV + self.indstart = 0 + # initialize the tree structure + super().__init__(arg=arg, types=types) + + def create_corresponding_node(self, node_index, p3d=None): + """ + Creates a node with the given index corresponding to the tree class. + + Parameters + ---------- + node_index: int + index of the new node + """ + return JaxleySimNode(node_index, p3d=p3d) + + + def init_model( + self, model_name, dt=0.025, t_calibrate=0.0, v_init=-75.0, factor_lambda=1.0, pprint=False + ): + """ + Initialize hoc-objects to simulate the neuron model implemented by this + tree. + + Parameters + ---------- + dt: float (default is ``.025`` ms) + Timestep of the simulation + t_calibrate: float (default ``0.`` ms) + The calibration time; time model runs without input to reach its + equilibrium state before the true simulation starts + v_init: float (default ``-75.`` mV) + The initial voltage at which the model is initialized + factor_lambda: float or int (default 1.) + If int, the number of segments per section. If float, multiplies the + number of segments given by the standard lambda rule (Carnevale, 2004) + to give the number of compartments simulated (default value 1. gives + the number given by the lambda rule) + pprint: bool (default ``False``) + Whether or not to print info on the NEURON model's creation + """ + # simple tree copy + aux_tree = MorphTree(self) + aux_tree.reset_indices() + + jx_channels = { + chan: eval(f"JX_MECH['{model_name}'].{chan}()") for chan in self.channel_storage + } + jx_channels['L'] = Leak() + + branches = [node._make_branch(jx_channels) for node in self] + parents = [aux_node.parent_node.index if not aux_tree.is_root(aux_node) else -1 for aux_node in aux_tree] + + cell = jx.Cell(branches, parents) + + return cell + + +class JaxleyCompartmentNode(JaxleySimNode): + """ + Subclass of `NeuronSimNode` that defines a cylinder with fake geometry in + NEURON to result in the effective simulation as a single compartment. + """ + + def __init__(self, index): + super().__init__(index) + + def get_child_nodes(self, skip_inds=[]): + return super().get_child_nodes(skip_inds=skip_inds) + + def _make_branch(self, pprint=False): + return super()._make_branch(dx_max=None, factorlambda=1, pprint=False) + +class JaxleyCompartmentTree(JaxleySimTree): + """ + Creates a `neat.NeuronCompartmentTree` to simulate reduced compartmentment + models from a `neat.CompartmentTree`. + + Parameters + ---------- + ctree: `neat.CompartmentTree` + The tree containing the parameters of the reduced compartmental model + to be simulated + fake_c_m: float + Fake value for the membrance capacitance density, rescales cylinder + surface + fake_r_a: float + Fake value for the axial resistance, rescales cylinder length + + Attributes + ---------- + equivalent_locs: list of tuples + 'Fake' locations corresponding to each compartment, which are + used to insert hoc point process at the compartments using + the same functions definitions as for as for a morphological + `neat.NeuronSimTree`. + + Notes + ----- + - Note that this class inherits from `neat.NeuronSimTree` and *not* from + `neat.CompartmentTree`. This is because NEAT defines a fake morphology to + implement the compartment model in NEURON, and also to reuse the functionality + implemented by `neat.NeuronSimTree`. Any function that is not explicitly + redefined from `neat.NeuronSimTree` can be called in the same way for this + compartment model. + - Locations to this class can be provided either as fake morphology locations + -- i.e. a tuple `(node.index, x-location in [0,1])` -- where the value of the + x-location is ignored since the nodes here are single compartments, as in + the `neat.CompartmentTree`, and not cylinders, as in `neat.MorphTree` or + subclasses, or as location indices, where the index corresponds to the location + in the original list of locations from which the `neat.CompartmentTree` was + derived. + """ + + def __init__(self, ctree, fake_c_m=1.0, fake_r_a=100.0 * 1e-6, method=2): + + try: + assert issubclass(ctree.__class__, CompartmentTree) + except AssertionError as e: + raise ValueError( + "`neat.NeuronCompartmentTree` can only be instantiated " + "from a `neat.CompartmentTree` or derived class" + ) + super().__init__(ctree, types=[1, 3, 4]) + self.equivalent_locs = ctree.get_equivalent_locs() + self._create_reduced_jaxley_model( + ctree, + fake_c_m=fake_c_m, + fake_r_a=fake_r_a, + ) + + def _create_reduced_jaxley_model(self, ctree, fake_c_m=1.0, fake_r_a=100.0 * 1e-6): + arg1, arg2 = ctree.compute_fake_geometry( + fake_c_m=fake_c_m, + fake_r_a=fake_r_a, + factor_r_a=1e-6, + delta=1e-10, + method=2, + ) + lengths = arg1 + radii = arg2 + surfaces = 2.0 * np.pi * radii * lengths + for ii, comp_node in enumerate(ctree): + sim_node = self.__getitem__(comp_node.index, skip_inds=[]) + if self.is_root(sim_node): + sim_node.set_p3d(np.array([0.0, 0.0, 0.0]), radii[ii] * 1e4, 1) + else: + sim_node.set_p3d( + np.array( + [sim_node.parent_node.xyz[0] + lengths[ii] * 1e4, 0.0, 0.0] + ), + radii[ii] * 1e4, + 3, + ) + + # fill the tree with the currents + for ii, sim_node in enumerate(self): + comp_node = ctree[ii] + sim_node.currents = { + chan: [g / surfaces[comp_node.index], e] + for chan, (g, e) in comp_node.currents.items() + } + sim_node.concmechs = copy.deepcopy(comp_node.concmechs) + for concmech in sim_node.concmechs.values(): + concmech.gamma *= surfaces[comp_node.index] * 1e6 + sim_node.c_m = fake_c_m + sim_node.r_a = fake_r_a + sim_node.R = radii[comp_node.index] * 1e4 # convert to [um] + sim_node.L = lengths[comp_node.index] * 1e4 # convert to [um] diff --git a/tests/channel_installer.py b/tests/channel_installer.py index 5cd46ccc..e5bb516f 100644 --- a/tests/channel_installer.py +++ b/tests/channel_installer.py @@ -28,7 +28,7 @@ import os import subprocess -from neat import load_neuron_model, load_nest_model +from neat import load_neuron_model, load_nest_model, load_jaxley_model def load_or_install_neuron_test_channels(): @@ -85,3 +85,34 @@ def load_or_install_nest_test_channels(): ] ) load_nest_model("multichannel_test") + + +def load_or_install_jaxley_test_channels(): + """ + neatmodels install multichannel_test -s jaxley -p channelcollection_for_tests.py + """ + channel_file = os.path.abspath( + os.path.join(os.path.dirname(__file__), "channelcollection_for_tests.py") + ) + try: + # load_jaxley_model() calls will raise a RuntimeError is a compiled model + # is loaded multiple times + try: + # raises FileNotFoundError if not compiled + load_jaxley_model("multichannel_test") + except FileNotFoundError: + subprocess.call( + [ + "neatmodels", + "install", + "multichannel_test", + "-s", + "jaxley", + "-p", + channel_file, + ] + ) + load_jaxley_model("multichannel_test") + except RuntimeError as e: + # the neuron model "multichannel_test" has already been loaded + pass diff --git a/tests/test_jaxleytree.py b/tests/test_jaxleytree.py index 0e24fb6c..8c7974ca 100644 --- a/tests/test_jaxleytree.py +++ b/tests/test_jaxleytree.py @@ -3,303 +3,19 @@ import os -from neat import PhysTree, PhysNode, MorphTree, CompartmentFitter, CompartmentTree +from neat import PhysTree, JaxleySimTree, JaxleyCompartmentTree, CompartmentFitter import channelcollection_for_tests as channelcollection import channel_installer channel_installer.load_or_install_neuron_test_channels() +channel_installer.load_or_install_jaxley_test_channels() MORPHOLOGIES_PATH_PREFIX = os.path.abspath( os.path.join(os.path.dirname(__file__), "test_morphologies") ) -class JaxleySimNode(PhysNode): - """ - Subclass of `neat.PhysNode` that implements functionality to instantiate a - cylindrical `neuron.h.Section` based on its physiological and geometrical - parameters. - """ - - def __init__(self, index, p3d=None): - super().__init__(index, p3d) - - def _make_branch(self, dx_max=15, factorlambda=1.0, pprint=False): - if dx_max is None: - n_comp = factorlambda - - if self.index == 1: - n_comp = 1 - l_comp = 2.0 * self.R - - else: - n_comp = self.L / dx_max + 1 - l_comp = self.L / n_comp - - compartments = [] - for _ in range(n_comp): - jx_comp = jx.Compartment() - jx_comp.set("length", l_comp) - jx_comp.set("radius", self.R) - jx_comp.set("axial_resistivity", self.r_a * 1e6) # MOhm*cm --> Ohm*cm - jx_comp.set("capacitance", self.c_m) # uF/cm^2 - compartments.append(jx_comp) - - branch = jx.Branch(compartments) - return branch - - -class JaxleySimTree(PhysTree): - """ - Tree class to define NEURON (Carnevale & Hines, 2004) based on `neat.PhysTree`. - - Attributes - ---------- - sections: dict of hoc sections - Storage for hoc sections. Keys are node indices. - shunts: list of hoc mechanisms - Storage container for shunts - syns: list of hoc mechanisms - Storage container for synapses - iclamps: list of hoc mechanisms - Storage container for current clamps - vclamps: lis of hoc mechanisms - Storage container for voltage clamps - vecstims: list of hoc mechanisms - Storage container for vecstim objects - netcons: list of hoc mechanisms - Storage container for netcon objects - vecs: list of hoc vectors - Storage container for hoc spike vectors - dt: float - timestep of the simulator ``[ms]`` - t_calibrate: float - Time for the model to equilibrate``[ms]``. Not counted as part of the - simulation. - factor_lambda : int or float - If int, the number of segments per section. If float, multiplies the - number of segments given by the standard lambda rule (Carnevale, 2004) - to give the number of compartments simulated (default value 1. gives - the number given by the lambda rule) - v_init: float - The initial voltage at which the model is initialized ``[mV]`` - - A `NeuronSimTree` can be extended easily with custom point process mechanisms. - Just make sure that you store the point process in an existing appropriate - storage container or in a custom storage container, since if all references - to the hocobject disappear, the object itself will be deleted as well. - - .. code-block:: python - class CustomSimTree(NeuronSimTree): - def addCustomPointProcessMech(self, loc, **kwargs): - loc = MorphLoc(loc, self) - - # create the point process - pp = h.custom_point_process(self.sections[loc['node']](loc['x'])) - pp.arg1 = kwargs['arg1'] - pp.arg2 = kwargs['arg2'] - ... - - self.storage_container_for_point_process.append(pp) - - If you define a custom storage container, make sure that you overwrite the - `__init__()` and `delete_model()` functions to make sure it is created and - deleted properly. - """ - - def __init__(self, arg=None, types=[1, 3, 4]): - # neuron storage - self.sections = {} - self.shunts = [] - self.syns = [] - self.iclamps = [] - self.vclamps = [] - self.vecstims = [] - self.netcons = [] - self.vecs = [] - # simulation parameters - self.dt = 0.1 # ms - self.t_calibrate = 0.0 # ms - self.factor_lambda = 1.0 - self.v_init = -75.0 # mV - self.indstart = 0 - # initialize the tree structure - super().__init__(arg=arg, types=types) - - def create_corresponding_node(self, node_index, p3d=None): - """ - Creates a node with the given index corresponding to the tree class. - - Parameters - ---------- - node_index: int - index of the new node - """ - return JaxleySimNode(node_index, p3d=p3d) - - - def init_model( - self, dt=0.025, t_calibrate=0.0, v_init=-75.0, factor_lambda=1.0, pprint=False - ): - """ - Initialize hoc-objects to simulate the neuron model implemented by this - tree. - - Parameters - ---------- - dt: float (default is ``.025`` ms) - Timestep of the simulation - t_calibrate: float (default ``0.`` ms) - The calibration time; time model runs without input to reach its - equilibrium state before the true simulation starts - v_init: float (default ``-75.`` mV) - The initial voltage at which the model is initialized - factor_lambda: float or int (default 1.) - If int, the number of segments per section. If float, multiplies the - number of segments given by the standard lambda rule (Carnevale, 2004) - to give the number of compartments simulated (default value 1. gives - the number given by the lambda rule) - pprint: bool (default ``False``) - Whether or not to print info on the NEURON model's creation - """ - # self.set_simulation_parameters( - # dt=dt, - # t_calibrate=t_calibrate, - # v_init=v_init, - # factor_lambda=factor_lambda, - # ) - # reset all storage - # self.delete_model() - # create the Jaxley - cell = self._create_jaxley_cell(pprint=pprint) - return cell - - def _create_jaxley_cell(self, pprint): - # simple tree copy - aux_tree = MorphTree(self) - aux_tree.reset_indices() - - branches = [node._make_branch() for node in self] - parents = [aux_node.parent_node.index if not aux_tree.is_root(aux_node) else -1 for aux_node in aux_tree] - - cell = jx.Cell(branches, parents) - return cell - - -class JaxleyCompartmentNode(JaxleySimNode): - """ - Subclass of `NeuronSimNode` that defines a cylinder with fake geometry in - NEURON to result in the effective simulation as a single compartment. - """ - - def __init__(self, index): - super().__init__(index) - - def get_child_nodes(self, skip_inds=[]): - return super().get_child_nodes(skip_inds=skip_inds) - - def _make_branch(self, pprint=False): - return super()._make_branch(dx_max=None, factorlambda=1, pprint=False) - -class JaxleyCompartmentTree(JaxleySimTree): - """ - Creates a `neat.NeuronCompartmentTree` to simulate reduced compartmentment - models from a `neat.CompartmentTree`. - - Parameters - ---------- - ctree: `neat.CompartmentTree` - The tree containing the parameters of the reduced compartmental model - to be simulated - fake_c_m: float - Fake value for the membrance capacitance density, rescales cylinder - surface - fake_r_a: float - Fake value for the axial resistance, rescales cylinder length - - Attributes - ---------- - equivalent_locs: list of tuples - 'Fake' locations corresponding to each compartment, which are - used to insert hoc point process at the compartments using - the same functions definitions as for as for a morphological - `neat.NeuronSimTree`. - - Notes - ----- - - Note that this class inherits from `neat.NeuronSimTree` and *not* from - `neat.CompartmentTree`. This is because NEAT defines a fake morphology to - implement the compartment model in NEURON, and also to reuse the functionality - implemented by `neat.NeuronSimTree`. Any function that is not explicitly - redefined from `neat.NeuronSimTree` can be called in the same way for this - compartment model. - - Locations to this class can be provided either as fake morphology locations - -- i.e. a tuple `(node.index, x-location in [0,1])` -- where the value of the - x-location is ignored since the nodes here are single compartments, as in - the `neat.CompartmentTree`, and not cylinders, as in `neat.MorphTree` or - subclasses, or as location indices, where the index corresponds to the location - in the original list of locations from which the `neat.CompartmentTree` was - derived. - """ - - def __init__(self, ctree, fake_c_m=1.0, fake_r_a=100.0 * 1e-6, method=2): - - try: - assert issubclass(ctree.__class__, CompartmentTree) - except AssertionError as e: - raise ValueError( - "`neat.NeuronCompartmentTree` can only be instantiated " - "from a `neat.CompartmentTree` or derived class" - ) - super().__init__(ctree, types=[1, 3, 4]) - self.equivalent_locs = ctree.get_equivalent_locs() - self._create_reduced_jaxley_model( - ctree, - fake_c_m=fake_c_m, - fake_r_a=fake_r_a, - ) - - def _create_reduced_jaxley_model(self, ctree, fake_c_m=1.0, fake_r_a=100.0 * 1e-6): - arg1, arg2 = ctree.compute_fake_geometry( - fake_c_m=fake_c_m, - fake_r_a=fake_r_a, - factor_r_a=1e-6, - delta=1e-10, - method=2, - ) - lengths = arg1 - radii = arg2 - surfaces = 2.0 * np.pi * radii * lengths - for ii, comp_node in enumerate(ctree): - sim_node = self.__getitem__(comp_node.index, skip_inds=[]) - if self.is_root(sim_node): - sim_node.set_p3d(np.array([0.0, 0.0, 0.0]), radii[ii] * 1e4, 1) - else: - sim_node.set_p3d( - np.array( - [sim_node.parent_node.xyz[0] + lengths[ii] * 1e4, 0.0, 0.0] - ), - radii[ii] * 1e4, - 3, - ) - - # # fill the tree with the currents - # for ii, sim_node in enumerate(self): - # comp_node = ctree[ii] - # sim_node.currents = { - # chan: [g / surfaces[comp_node.index], e] - # for chan, (g, e) in comp_node.currents.items() - # } - # sim_node.concmechs = copy.deepcopy(comp_node.concmechs) - # for concmech in sim_node.concmechs.values(): - # concmech.gamma *= surfaces[comp_node.index] * 1e6 - # sim_node.c_m = fake_c_m - # sim_node.r_a = fake_r_a - # sim_node.R = radii[comp_node.index] * 1e4 # convert to [um] - # sim_node.L = lengths[comp_node.index] * 1e4 # convert to [um] - - class TestJaxley: def load_ball(self): @@ -329,23 +45,23 @@ def create_jaxley_model(self): self.load_ball() jt = JaxleySimTree(self.tree) - jcell = jt.init_model() + jcell = jt.init_model("multichannel_test") jct = JaxleyCompartmentTree(self.ctree) - jcompcell = jt.init_model() + jcompcell = jt.init_model("multichannel_test") - def test_jaxley_channels(self): - self.load_ball() + # def test_jaxley_channels(self): + # self.load_ball() - channelstr_k = self.k_chan.write_jaxley_code() - channelstr_na = self.na_chan.write_jaxley_code() - with open("jtest", "w") as file: - file.write(channelstr_k + "\n\n\n") - file.write(channelstr_na) + # channelstr_k = self.k_chan.write_jaxley_code() + # channelstr_na = self.na_chan.write_jaxley_code() + # with open("jtest", "w") as file: + # file.write(channelstr_k + "\n\n\n") + # file.write(channelstr_na) if __name__ == "__main__": tjx = TestJaxley() - # tjx.create_jaxley_model() - tjx.test_jaxley_channels() + tjx.create_jaxley_model() + # tjx.test_jaxley_channels() From 1e373301e6f4170022653896d9393fcb9c211057 Mon Sep 17 00:00:00 2001 From: "w.wybo" Date: Tue, 9 Dec 2025 17:04:33 +0100 Subject: [PATCH 5/9] almost running version --- src/neat/actions/install.py | 8 ++ src/neat/actions/neatmodels.py | 3 +- src/neat/simulations/jaxley/default_syns.py | 112 ++++++++++++++++++++ src/neat/simulations/jaxley/jaxleymodel.py | 50 +++++++-- tests/test_jaxleytree.py | 8 ++ 5 files changed, 173 insertions(+), 8 deletions(-) create mode 100644 src/neat/simulations/jaxley/default_syns.py diff --git a/src/neat/actions/install.py b/src/neat/actions/install.py index 0ef171f9..d3071f65 100644 --- a/src/neat/actions/install.py +++ b/src/neat/actions/install.py @@ -275,6 +275,7 @@ def _compile_jaxley(model_name, path_neat, channels, path_jaxleyresource=None): import jax.numpy as jnp from jax.lax import select from jaxley.channels import Channel +from jaxley.synapses.synapse import Synapse from jaxley.solver_gate import ( exponential_euler, save_exp, @@ -284,6 +285,13 @@ def _compile_jaxley(model_name, path_neat, channels, path_jaxleyresource=None): """ + if path_jaxleyresource is not None: + with open(path_jaxleyresource, 'r') as file: + resourcestr = file.read() + resourcestr = resourcestr.split("class", 1)[1] + + filestr += 'class' + resourcestr + for chan in channels: print(" - writing .jaxley code for:", chan.__class__.__name__) channelstr = chan.write_jaxley_code() diff --git a/src/neat/actions/neatmodels.py b/src/neat/actions/neatmodels.py index c7577344..2292851e 100755 --- a/src/neat/actions/neatmodels.py +++ b/src/neat/actions/neatmodels.py @@ -82,7 +82,7 @@ def parse_cmd_args(path_neat): ) parser.add_argument( "--jaxleyresource", - default=os.path.join(path_neat, "simulations/nest/default_syns.nestml"), + default=os.path.join(path_neat, "simulations/jaxley/default_syns.py"), help="Path to directory containing additional jaxley mechanisms " "(e.g. synapses) that will be combined together with the " "generated ion channel mod files.\n" @@ -126,6 +126,7 @@ def main(): simulators=cmd_args.simulator, path_neuronresource=cmd_args.neuronresource, path_nestresource=cmd_args.nestresource, + path_jaxleyresource=cmd_args.jaxleyresource, ) elif cmd_args.action == "list": _list_models(path_neat=path_neat, simulators=cmd_args.simulator) diff --git a/src/neat/simulations/jaxley/default_syns.py b/src/neat/simulations/jaxley/default_syns.py new file mode 100644 index 00000000..aea786bc --- /dev/null +++ b/src/neat/simulations/jaxley/default_syns.py @@ -0,0 +1,112 @@ +from typing import Dict, Optional + +import jax.numpy as jnp +from jax.lax import select +from jaxley.channels import Channel +from jaxley.synapses.synapse import Synapse +from jaxley.solver_gate import ( + exponential_euler, + save_exp, + solve_gate_exponential, + solve_inf_gate_exponential, +) + + + +class NEATJaxleySynapse(Synapse): + """ + Compute syanptic current and update synapse state. + """ + def __init__(self, name = None): + super().__init__(name) + self.channel_params = {} + self.channel_states = {} + + + def set_spiketrain(self, spktms, dt=0.1, t_calibrate=0.0, t_max=100., weight=1.0): + maxidx = jnp.round(t_max / dt).astype(int) + calidx = jnp.round(t_calibrate / dt).astype(int) + + spkidx = jnp.round(jnp.array(spktms) / dt).astype(int) + spkidx = spkidx[spkidx < maxidx] + spkidx += calidx + + self.cc = 0 # step index counter + self.spkidx = spkidx + self.channel_params["weight"] = weight + + +class DoubleExpSynapse(NEATJaxleySynapse): + def __init__(self, name = None, tau_r=.2, tau_d=3., e_r=0.): + super().__init__(name) + self.channel_params = { + "tau_r": tau_r, + "tau_d": tau_d, + "e_r": e_r, + "weight": 0.0, # meant to be reset + } + self.channel_states = { + "x_r": 0.0, + "x_d": 0.0 + } + + def compute_propagators(self, delta_t, params): + tau_r, tau_d = params['tau_r'], params['tau_d'] + tp = (tau_r * tau_d) / (tau_d - tau_r) * jnp.log( tau_d / tau_r ) + self.g_norm = 1. / ( -jnp.exp( -tp / tau_r ) + jnp.exp( -tp / tau_d ) ) + + p_r = jnp.exp(-delta_t / params["tau_r"]) + p_d = jnp.exp(-delta_t / params["tau_d"]) + return p_r, p_d + + def update_states(self, u, voltages, params): + """Return updated synapse state and current.""" + + # spike delivery + if self.cc == self.spktms[0]: + new_x_r = states["x_r"] - self.g_norm * params["weight"] + new_x_d = states["x_d"] + self.g_norm * params["weight"] + self.spktms.delete(0) + self.cc += 1 + + p_r, p_d = self.compute_propagators(delta_t, params) + new_x_r *= p_r + new_x_d *= p_d + + return {"x_r": new_x_r, "x_d": new_x_d} + + def compute_current(self, states, pre_voltage, post_voltage, params): + return params["weight"] * (states["x_d"] - states['x_r']) * (post_voltage - params["e_r"]) + + +class AMPASynapse(DoubleExpSynapse): + def __init__(self, name = None, tau_r_AMPA=.2, tau_d_AMPA=3., e_r_AMPA=0.): + super().__init__(tau_r=tau_r_AMPA, tau_d=tau_d_AMPA, e_r=e_r_AMPA) + self.current_name = "AMPA" + + +class GABASynapse(DoubleExpSynapse): + def __init__(self, name = None, tau_r_GABA=.2, tau_d_GABA=10., e_GABA=-80.): + super().__init__(tau_r=tau_r_GABA, tau_d=tau_d_GABA, e_r=e_r_GABA) + + +class AMPA_NMDASynapse(NEATJaxleySynapse): + def __init__(self, name=None, + tau_r_AMPA=.2, tau_AMPA=3., e_r_AMPA=0., + tau_r_NMDA=.2, tau_d_NMDA=10., e_r_NMDA=-80., + NMDA_ratio=2.0, + ): + self.channel_params = { + "tau_r_AMPA": tau_r_AMPA, + "tau_d_AMPA": tau_d_AMPA, + "e_r_AMPA": e_r_AMPA, + "tau_r_NMDA": tau_r_NMDA, + "tau_d_NMDA": tau_d_NMDA, + "e_r_NMDA": e_r_NMDA, + "NMDA_ratio": NMDA_ratio, + "weight": weight, # meant to be reset + } + self.channel_states = { + "x_r_AMPA": 0.0, + "x_d_NMDA": 0.0 + } \ No newline at end of file diff --git a/src/neat/simulations/jaxley/jaxleymodel.py b/src/neat/simulations/jaxley/jaxleymodel.py index 8b8260ef..4603daae 100644 --- a/src/neat/simulations/jaxley/jaxleymodel.py +++ b/src/neat/simulations/jaxley/jaxleymodel.py @@ -6,7 +6,7 @@ import importlib from ...trees.phystree import PhysTree, PhysNode -from ...trees.morphtree import MorphTree +from ...trees.morphtree import MorphTree, MorphLoc from ...trees.compartmenttree import CompartmentTree from ...factorydefaults import DefaultPhysiology @@ -78,7 +78,7 @@ def load_jaxley_model(name): f"The Jaxley model named '{name}' is not installed. " f"Run 'neatmodels -h' in the terminal for help on " f"installing new Jaxley models with NEAT. " - f"Installed models will be in '{path}'." + f"Installed models will be in '{jx_path}'." ) @@ -215,9 +215,8 @@ def create_corresponding_node(self, node_index, p3d=None): """ return JaxleySimNode(node_index, p3d=p3d) - def init_model( - self, model_name, dt=0.025, t_calibrate=0.0, v_init=-75.0, factor_lambda=1.0, pprint=False + self, model_name, dt=0.025, t_calibrate=0.0, t_max=100., v_init=-75.0, factor_lambda=1.0, pprint=False ): """ Initialize hoc-objects to simulate the neuron model implemented by this @@ -240,9 +239,20 @@ def init_model( pprint: bool (default ``False``) Whether or not to print info on the NEURON model's creation """ + # simulation control + self.sim_control = { + 'dt': dt, + 't_calibrate': t_calibrate, + 't_max': t_max, + } + self.model_name = model_name + self.syn_locs = [] + self.syn_models = [] + # simple tree copy aux_tree = MorphTree(self) aux_tree.reset_indices() + self.index_map = {n_orig.index: n_aux.index-1 for n_orig, n_aux in zip(self, aux_tree)} jx_channels = { chan: eval(f"JX_MECH['{model_name}'].{chan}()") for chan in self.channel_storage @@ -250,12 +260,37 @@ def init_model( jx_channels['L'] = Leak() branches = [node._make_branch(jx_channels) for node in self] - parents = [aux_node.parent_node.index if not aux_tree.is_root(aux_node) else -1 for aux_node in aux_tree] + parents = [aux_node.parent_node.index-1 if not aux_tree.is_root(aux_node) else -1 for aux_node in aux_tree] + + self.cell = jx.Cell(branches, parents) - cell = jx.Cell(branches, parents) + return self.cell + + def add_AMPA_synapse(self, loc): + synapse = JX_MECH[f'{self.model_name}'].AMPASynapse() + self.syn_locs.append(MorphLoc(loc, self)) + self.syn_models.append(synapse) - return cell + def set_spiketrain(self, syn_index, syn_weight, spike_times): + self.syn_models[syn_index].set_spiketrain(spike_times, weight=syn_weight, **self.sim_control) + def run(self): + for loc, syn in zip(self.syn_locs, self.syn_models): + breakpoint() + self.cell.branch( + self.index_map[loc['node']] + ).loc( + loc['x'] + ).insert(syn) + + self.cell.delete_recordings() + self.cell.branch(0).loc(0.0).record("v") + + current = jx.step_current(i_delay=1.0, i_dur=1.0, i_amp=0.0, delta_t=self.sim_control['dt'], t_max=self.sim_control['t_max']) + self.cell.branch(0).loc(0.0).stimulate(current) + + res = jx.integrate(self.cell, delta_t=self.sim_control['dt']) + return res class JaxleyCompartmentNode(JaxleySimNode): """ @@ -272,6 +307,7 @@ def get_child_nodes(self, skip_inds=[]): def _make_branch(self, pprint=False): return super()._make_branch(dx_max=None, factorlambda=1, pprint=False) + class JaxleyCompartmentTree(JaxleySimTree): """ Creates a `neat.NeuronCompartmentTree` to simulate reduced compartmentment diff --git a/tests/test_jaxleytree.py b/tests/test_jaxleytree.py index 8c7974ca..cea27a08 100644 --- a/tests/test_jaxleytree.py +++ b/tests/test_jaxleytree.py @@ -47,6 +47,14 @@ def create_jaxley_model(self): jt = JaxleySimTree(self.tree) jcell = jt.init_model("multichannel_test") + jt.add_AMPA_synapse((1,.5)) + jt.set_spiketrain(0, [10.], 1.) + res = jt.run() + + import matplotlib.pyplot as pl + pl.plot(res.T) + pl.show() + jct = JaxleyCompartmentTree(self.ctree) jcompcell = jt.init_model("multichannel_test") From 3531620c2f557d2987b92f8d820b12279aed1327 Mon Sep 17 00:00:00 2001 From: "w.wybo" Date: Wed, 17 Dec 2025 10:37:58 +0100 Subject: [PATCH 6/9] fix jax state variable computations --- src/neat/channels/ionchannels.py | 91 +++++++++++++++++++++++++++++--- 1 file changed, 83 insertions(+), 8 deletions(-) diff --git a/src/neat/channels/ionchannels.py b/src/neat/channels/ionchannels.py index 50bc21c0..6c371b20 100755 --- a/src/neat/channels/ionchannels.py +++ b/src/neat/channels/ionchannels.py @@ -25,9 +25,11 @@ import os import ast import warnings +from functools import reduce from ..factorydefaults import DefaultPhysiology + # CONC_DICT = { # 'na': 10., # mM # 'k': 54.4, # mM @@ -62,6 +64,34 @@ def find_IfExp_node(self, node): return_node = self.ifexp_node self.ifexp_node = None return return_node + + + +class BoolOpToJnpTransformer(ast.NodeTransformer): + def visit_BoolOp(self, node: ast.BoolOp): + # First, recursively transform child nodes + self.generic_visit(node) + + if isinstance(node.op, ast.And): + func_name = "logical_and" + elif isinstance(node.op, ast.Or): + func_name = "logical_or" + else: + return node + + # Reduce n-ary BoolOp into nested binary calls + def combine(left, right): + return ast.Call( + func=ast.Attribute( + value=ast.Name(id="jnp", ctx=ast.Load()), + attr=func_name, + ctx=ast.Load(), + ), + args=[left, right], + keywords=[], + ) + + return reduce(combine, node.values) class _func(object): @@ -1259,10 +1289,49 @@ def _replaceConc(expr_str, prefix="", suffix=""): fcc.close() - def _print_jaxley_pycode(self, expr): - pstr = sp.printing.pycode(expr) + def _create_jaxley_funcstr(self, code_str, n_spaces=0, indent=4): + """ + This function is used to recursively expand if... else... statements + across multiple lines, as by default the single line version is printed + by `sympy.pycode()` and `ast.unparse()` + """ + tree = ast.parse(code_str) + iev = IfExpVisitor() + ifexp = iev.find_IfExp_node(tree) + transformer = BoolOpToJnpTransformer() +# new_tree = transformer.visit(tree) + if ifexp is not None: + # sanity check + assert iev.find_IfExp_node(ifexp.test) is None + # if test is True + cond_1_str = self._create_jaxley_funcstr( + ast.unparse(ifexp.body), n_spaces=n_spaces, indent=n_spaces + indent + ) + # if test is False + cond_0_str = self._create_jaxley_funcstr( + ast.unparse(ifexp.orelse), n_spaces=n_spaces, indent=n_spaces + indent + ) + code_str = f"""jnp.where( + {ast.unparse(transformer.visit(ifexp.test))}, + {cond_1_str}, + {cond_0_str}, + )""" + else: + try: + code_str = \ + " " * indent + self._replace_jaxley_patterns(code_str) + except TypeError as e: + print(e) + return code_str + + def _replace_jaxley_patterns(self, pstr): pstr = pstr.replace("math.exp", "save_exp") pstr = pstr.replace("math.", "jnp.") + return pstr + + def _print_jaxley_pycode(self, expr): + pstr = sp.printing.pycode(expr) + pstr = self._replace_jaxley_patterns(pstr) return pstr def write_jaxley_code(self, @@ -1311,15 +1380,21 @@ def update_states( f"{var} = u[f'{{self.prefix}}{var}']" for var in sv ]) } -""" - # substitution for common neuron names +""" + + # if self.__class__.__name__ == "NaTa_t": + # breakpoint() + # substitution for common neuron names repl_pairs = [(str(c), str(c) + "i") for c in self.conc] for var, svar in zip(sv, self.ordered_statevars): - vi = self._print_jaxley_pycode(self.varinf[svar])#, assign_to=f"{var}_inf") - ti = self._print_jaxley_pycode(self.tauinf[svar])#, assign_to=f"tau_{var}") + vi = sp.printing.pycode(self.varinf[svar])#, assign_to=f"{var}_inf") + ti = sp.printing.pycode(self.tauinf[svar])#, assign_to=f"tau_{var}") for repl_pair in repl_pairs: vi = vi.replace(*repl_pair) ti = ti.replace(*repl_pair) + + vi = self._create_jaxley_funcstr(vi) + ti = self._create_jaxley_funcstr(ti) channelstr += f""" # advance state variable {var} @@ -1327,7 +1402,7 @@ def update_states( {svar}, dt, {vi}, - {ti} + {ti}, ) """ channelstr += f""" @@ -1356,7 +1431,7 @@ def compute_current( def init_state(self, states, voltages, params, delta_t): v = voltages { '\n '.join([ - f'{var} = {self._print_jaxley_pycode(self.varinf[svar])}' \ + f'{var} = {self._create_jaxley_funcstr(sp.printing.pycode(self.varinf[svar]))}' \ for var, svar in zip(sv, self.ordered_statevars) ]) } From bf463fa80621a11384b7c52d21f4fc708618dcba Mon Sep 17 00:00:00 2001 From: "w.wybo" Date: Fri, 19 Dec 2025 10:26:53 +0100 Subject: [PATCH 7/9] experiment with spiketrain generators --- src/neat/actions/install.py | 2 +- src/neat/simulations/jaxley/default_syns.py | 113 ++++++++++++++------ src/neat/simulations/jaxley/jaxleymodel.py | 12 ++- tests/test_jaxleytree.py | 4 +- 4 files changed, 90 insertions(+), 41 deletions(-) diff --git a/src/neat/actions/install.py b/src/neat/actions/install.py index d3071f65..5d15be68 100644 --- a/src/neat/actions/install.py +++ b/src/neat/actions/install.py @@ -272,8 +272,8 @@ def _compile_jaxley(model_name, path_neat, channels, path_jaxleyresource=None): filestr = """ from typing import Dict, Optional +import jax import jax.numpy as jnp -from jax.lax import select from jaxley.channels import Channel from jaxley.synapses.synapse import Synapse from jaxley.solver_gate import ( diff --git a/src/neat/simulations/jaxley/default_syns.py b/src/neat/simulations/jaxley/default_syns.py index aea786bc..eee97c71 100644 --- a/src/neat/simulations/jaxley/default_syns.py +++ b/src/neat/simulations/jaxley/default_syns.py @@ -1,7 +1,7 @@ from typing import Dict, Optional +import jax import jax.numpy as jnp -from jax.lax import select from jaxley.channels import Channel from jaxley.synapses.synapse import Synapse from jaxley.solver_gate import ( @@ -13,82 +13,124 @@ -class NEATJaxleySynapse(Synapse): +class NEATJaxleySynapse(Channel): """ Compute syanptic current and update synapse state. """ def __init__(self, name = None): + self.current_is_in_mA_per_cm2=True super().__init__(name) - self.channel_params = {} - self.channel_states = {} + # self.channel_params = {} + # self.channel_states = {} + @property + def prefix(self): + return f'{self._name}_' def set_spiketrain(self, spktms, dt=0.1, t_calibrate=0.0, t_max=100., weight=1.0): maxidx = jnp.round(t_max / dt).astype(int) calidx = jnp.round(t_calibrate / dt).astype(int) spkidx = jnp.round(jnp.array(spktms) / dt).astype(int) - spkidx = spkidx[spkidx < maxidx] + spkidx = spkidx[spkidx <= maxidx + 1] spkidx += calidx - self.cc = 0 # step index counter - self.spkidx = spkidx + # self.cc = 0 # step index counter + self.spkidx = spkidx[1:] + print("Setting spiketrain:", self.spkidx) self.channel_params["weight"] = weight + self.channel_states.update({ + f"{self.prefix}cc": 0, + f"{self.prefix}next_spk": spkidx[0], + }) class DoubleExpSynapse(NEATJaxleySynapse): def __init__(self, name = None, tau_r=.2, tau_d=3., e_r=0.): super().__init__(name) self.channel_params = { - "tau_r": tau_r, - "tau_d": tau_d, - "e_r": e_r, - "weight": 0.0, # meant to be reset + f"{self.prefix}tau_r": tau_r, + f"{self.prefix}tau_d": tau_d, + f"{self.prefix}e_r": e_r, + f"{self.prefix}weight": 0.0, # meant to be reset } self.channel_states = { - "x_r": 0.0, - "x_d": 0.0 + f"{self.prefix}x_r": 0.0, + f"{self.prefix}x_d": 0.0, + f"{self.prefix}cc": 0, + f"{self.prefix}next_spk": 0, } def compute_propagators(self, delta_t, params): - tau_r, tau_d = params['tau_r'], params['tau_d'] + tau_r, tau_d = params[f'{self.prefix}tau_r'], params[f'{self.prefix}tau_d'] tp = (tau_r * tau_d) / (tau_d - tau_r) * jnp.log( tau_d / tau_r ) - self.g_norm = 1. / ( -jnp.exp( -tp / tau_r ) + jnp.exp( -tp / tau_d ) ) + g_norm = 1. / ( -jnp.exp( -tp / tau_r ) + jnp.exp( -tp / tau_d ) ) - p_r = jnp.exp(-delta_t / params["tau_r"]) - p_d = jnp.exp(-delta_t / params["tau_d"]) - return p_r, p_d + p_r = jnp.exp(-delta_t / tau_r) + p_d = jnp.exp(-delta_t / tau_d) + return p_r, p_d, g_norm - def update_states(self, u, voltages, params): + def update_states(self, u, dt, voltages, params): """Return updated synapse state and current.""" + w = params[f'{self.prefix}weight'] + x_r = u[f'{self.prefix}x_r'] + x_d = u[f'{self.prefix}x_d'] + # spkidx = u[f'{self.prefix}spkidx'] + cc = u[f'{self.prefix}cc'] + next_spk = u[f'{self.prefix}next_spk'] + + p_r, p_d, g_norm = self.compute_propagators(dt, params) + + def branch1(): + print("!! in spike branch !!") + n_x_r = x_r - g_norm * w + n_x_d = x_d + g_norm * w + # self.spkidx.delete(0) + n_cc = cc + 1 + n_spk = next_spk + 1 + # n_spkidx = self.spkidx[n_cc:] + return n_x_r, n_x_d, n_cc, n_spk + + def branch2(): + n_cc = cc + 1 + return x_r, x_d, n_cc, next_spk # spike delivery - if self.cc == self.spktms[0]: - new_x_r = states["x_r"] - self.g_norm * params["weight"] - new_x_d = states["x_d"] + self.g_norm * params["weight"] - self.spktms.delete(0) - self.cc += 1 - - p_r, p_d = self.compute_propagators(delta_t, params) + new_x_r, new_x_d, new_cc, new_spk = jax.lax.cond( + (cc == next_spk)[0], + # True, + branch1, + branch2, + ) new_x_r *= p_r new_x_d *= p_d - return {"x_r": new_x_r, "x_d": new_x_d} + return { + f'{self.prefix}x_r': new_x_r, + f'{self.prefix}x_d': new_x_d, + # f'{self.prefix}spkidx': new_spkidx, + f'{self.prefix}cc': new_cc, + f'{self.prefix}next_spk': new_spk, + } - def compute_current(self, states, pre_voltage, post_voltage, params): - return params["weight"] * (states["x_d"] - states['x_r']) * (post_voltage - params["e_r"]) + def compute_current(self, u: Dict[str, jnp.ndarray], voltages, params: Dict[str, jnp.ndarray]): + print(self.prefix)#, params) + e_r = params[f'{self.prefix}e_r'] + v = voltages + x_r = u[f'{self.prefix}x_r'] + x_d = u[f'{self.prefix}x_d'] + return (x_d - x_r) * (v - e_r) class AMPASynapse(DoubleExpSynapse): def __init__(self, name = None, tau_r_AMPA=.2, tau_d_AMPA=3., e_r_AMPA=0.): - super().__init__(tau_r=tau_r_AMPA, tau_d=tau_d_AMPA, e_r=e_r_AMPA) - self.current_name = "AMPA" - + super().__init__(name=name, tau_r=tau_r_AMPA, tau_d=tau_d_AMPA, e_r=e_r_AMPA) + self.current_name = "i_AMPA" class GABASynapse(DoubleExpSynapse): - def __init__(self, name = None, tau_r_GABA=.2, tau_d_GABA=10., e_GABA=-80.): - super().__init__(tau_r=tau_r_GABA, tau_d=tau_d_GABA, e_r=e_r_GABA) - + def __init__(self, name = None, tau_r_GABA=.2, tau_d_GABA=10., e_r_GABA=-80.): + super().__init__(name=name, tau_r=tau_r_GABA, tau_d=tau_d_GABA, e_r=e_r_GABA) + self.current_name = "i_GABA" class AMPA_NMDASynapse(NEATJaxleySynapse): def __init__(self, name=None, @@ -96,6 +138,7 @@ def __init__(self, name=None, tau_r_NMDA=.2, tau_d_NMDA=10., e_r_NMDA=-80., NMDA_ratio=2.0, ): + self.current_is_in_mA_per_cm2=True self.channel_params = { "tau_r_AMPA": tau_r_AMPA, "tau_d_AMPA": tau_d_AMPA, diff --git a/src/neat/simulations/jaxley/jaxleymodel.py b/src/neat/simulations/jaxley/jaxleymodel.py index 4603daae..22e0153b 100644 --- a/src/neat/simulations/jaxley/jaxleymodel.py +++ b/src/neat/simulations/jaxley/jaxleymodel.py @@ -276,17 +276,23 @@ def set_spiketrain(self, syn_index, syn_weight, spike_times): def run(self): for loc, syn in zip(self.syn_locs, self.syn_models): - breakpoint() self.cell.branch( self.index_map[loc['node']] ).loc( loc['x'] ).insert(syn) + # for loc, syn in zip(self.syn_locs, self.syn_models): + # self.cell.branch( + # self.index_map[loc['node']] + # ).loc( + # loc['x'] + # ).insert(JX_MECH["multichannel_test"].Bla()) self.cell.delete_recordings() - self.cell.branch(0).loc(0.0).record("v") + self.cell.branch(0).loc(0.0).record() + self.cell.record("AMPASynapse_x_r") - current = jx.step_current(i_delay=1.0, i_dur=1.0, i_amp=0.0, delta_t=self.sim_control['dt'], t_max=self.sim_control['t_max']) + current = jx.step_current(i_delay=100., i_dur=10.0, i_amp=10.0, delta_t=self.sim_control['dt'], t_max=self.sim_control['t_max']) self.cell.branch(0).loc(0.0).stimulate(current) res = jx.integrate(self.cell, delta_t=self.sim_control['dt']) diff --git a/tests/test_jaxleytree.py b/tests/test_jaxleytree.py index cea27a08..dfe40c75 100644 --- a/tests/test_jaxleytree.py +++ b/tests/test_jaxleytree.py @@ -45,10 +45,10 @@ def create_jaxley_model(self): self.load_ball() jt = JaxleySimTree(self.tree) - jcell = jt.init_model("multichannel_test") + jcell = jt.init_model("multichannel_test", t_max=1000) jt.add_AMPA_synapse((1,.5)) - jt.set_spiketrain(0, [10.], 1.) + jt.set_spiketrain(0, 1., [200.]) res = jt.run() import matplotlib.pyplot as pl From 7f1c848444a336aca8f74f621d1b63e04577b1ac Mon Sep 17 00:00:00 2001 From: "w.wybo" Date: Wed, 24 Dec 2025 12:20:52 +0100 Subject: [PATCH 8/9] first working implementation --- src/neat/simulations/jaxley/default_syns.py | 87 ++++++------------- src/neat/simulations/jaxley/jaxleymodel.py | 95 +++++++++++++++------ tests/test_jaxleytree.py | 26 ++++-- 3 files changed, 113 insertions(+), 95 deletions(-) diff --git a/src/neat/simulations/jaxley/default_syns.py b/src/neat/simulations/jaxley/default_syns.py index eee97c71..7d0997a6 100644 --- a/src/neat/simulations/jaxley/default_syns.py +++ b/src/neat/simulations/jaxley/default_syns.py @@ -13,7 +13,7 @@ -class NEATJaxleySynapse(Channel): +class NEATJaxleySynapse(Synapse): """ Compute syanptic current and update synapse state. """ @@ -27,38 +27,19 @@ def __init__(self, name = None): def prefix(self): return f'{self._name}_' - def set_spiketrain(self, spktms, dt=0.1, t_calibrate=0.0, t_max=100., weight=1.0): - maxidx = jnp.round(t_max / dt).astype(int) - calidx = jnp.round(t_calibrate / dt).astype(int) - - spkidx = jnp.round(jnp.array(spktms) / dt).astype(int) - spkidx = spkidx[spkidx <= maxidx + 1] - spkidx += calidx - - # self.cc = 0 # step index counter - self.spkidx = spkidx[1:] - print("Setting spiketrain:", self.spkidx) - self.channel_params["weight"] = weight - self.channel_states.update({ - f"{self.prefix}cc": 0, - f"{self.prefix}next_spk": spkidx[0], - }) - class DoubleExpSynapse(NEATJaxleySynapse): def __init__(self, name = None, tau_r=.2, tau_d=3., e_r=0.): super().__init__(name) - self.channel_params = { + self.synapse_params = { f"{self.prefix}tau_r": tau_r, f"{self.prefix}tau_d": tau_d, f"{self.prefix}e_r": e_r, f"{self.prefix}weight": 0.0, # meant to be reset } - self.channel_states = { + self.synapse_states = { f"{self.prefix}x_r": 0.0, f"{self.prefix}x_d": 0.0, - f"{self.prefix}cc": 0, - f"{self.prefix}next_spk": 0, } def compute_propagators(self, delta_t, params): @@ -70,56 +51,44 @@ def compute_propagators(self, delta_t, params): p_d = jnp.exp(-delta_t / tau_d) return p_r, p_d, g_norm - def update_states(self, u, dt, voltages, params): + def update_states(self, states, delta_t, pre_voltage, post_voltage, params): """Return updated synapse state and current.""" w = params[f'{self.prefix}weight'] - x_r = u[f'{self.prefix}x_r'] - x_d = u[f'{self.prefix}x_d'] - # spkidx = u[f'{self.prefix}spkidx'] - cc = u[f'{self.prefix}cc'] - next_spk = u[f'{self.prefix}next_spk'] - - p_r, p_d, g_norm = self.compute_propagators(dt, params) - - def branch1(): - print("!! in spike branch !!") - n_x_r = x_r - g_norm * w - n_x_d = x_d + g_norm * w - # self.spkidx.delete(0) - n_cc = cc + 1 - n_spk = next_spk + 1 - # n_spkidx = self.spkidx[n_cc:] - return n_x_r, n_x_d, n_cc, n_spk - - def branch2(): - n_cc = cc + 1 - return x_r, x_d, n_cc, next_spk - - # spike delivery - new_x_r, new_x_d, new_cc, new_spk = jax.lax.cond( - (cc == next_spk)[0], - # True, - branch1, - branch2, + x_r = states[f'{self.prefix}x_r'] + x_d = states[f'{self.prefix}x_d'] + + tau_r, tau_d = params[f'{self.prefix}tau_r'], params[f'{self.prefix}tau_d'] + + p_r, p_d, g_norm = self.compute_propagators(delta_t, params) + + # add the spikes to the conductance window, weight is contained in pre_voltage + pred = pre_voltage >= 1e-15 + new_x_r = jnp.where( + pred, + x_r - g_norm * pre_voltage, + x_r + ) + new_x_d = jnp.where( + pred, + x_d + g_norm * pre_voltage, + x_d ) + # decay the conductance window new_x_r *= p_r new_x_d *= p_d return { f'{self.prefix}x_r': new_x_r, f'{self.prefix}x_d': new_x_d, - # f'{self.prefix}spkidx': new_spkidx, - f'{self.prefix}cc': new_cc, - f'{self.prefix}next_spk': new_spk, } - def compute_current(self, u: Dict[str, jnp.ndarray], voltages, params: Dict[str, jnp.ndarray]): + def compute_current(self, states, pre_voltage, post_voltage, params): print(self.prefix)#, params) e_r = params[f'{self.prefix}e_r'] - v = voltages - x_r = u[f'{self.prefix}x_r'] - x_d = u[f'{self.prefix}x_d'] - return (x_d - x_r) * (v - e_r) + v = post_voltage + x_r = states[f'{self.prefix}x_r'] + x_d = states[f'{self.prefix}x_d'] + return (x_d + x_r) * (v - e_r) class AMPASynapse(DoubleExpSynapse): diff --git a/src/neat/simulations/jaxley/jaxleymodel.py b/src/neat/simulations/jaxley/jaxleymodel.py index 22e0153b..0c6f761c 100644 --- a/src/neat/simulations/jaxley/jaxleymodel.py +++ b/src/neat/simulations/jaxley/jaxleymodel.py @@ -11,6 +11,7 @@ from ...factorydefaults import DefaultPhysiology try: + import jax.numpy as jnp import jaxley as jx from jaxley.channels import Leak @@ -186,15 +187,10 @@ def addCustomPointProcessMech(self, loc, **kwargs): """ def __init__(self, arg=None, types=[1, 3, 4]): - # neuron storage - self.sections = {} - self.shunts = [] - self.syns = [] - self.iclamps = [] - self.vclamps = [] - self.vecstims = [] - self.netcons = [] - self.vecs = [] + self.cell = None + self.pre_dummies = {} + self.syn_locs = [] + self.syn_models = [] # simulation parameters self.dt = 0.1 # ms self.t_calibrate = 0.0 # ms @@ -244,8 +240,10 @@ def init_model( 'dt': dt, 't_calibrate': t_calibrate, 't_max': t_max, + 't_sim': t_calibrate + t_max, } self.model_name = model_name + self.pre_dummies = {} self.syn_locs = [] self.syn_models = [] @@ -266,37 +264,78 @@ def init_model( return self.cell + def delete_model(self): + """ + Deletes all Jaxley objects created to simulate the model + implemented by this tree. + """ + self.cell = None + self.pre_dummies = [] + self.syn_locs = [] + self.syn_models = [] + def add_AMPA_synapse(self, loc): synapse = JX_MECH[f'{self.model_name}'].AMPASynapse() self.syn_locs.append(MorphLoc(loc, self)) self.syn_models.append(synapse) + def add_GABA_synapse(self, loc): + synapse = JX_MECH[f'{self.model_name}'].GABASynapse() + self.syn_locs.append(MorphLoc(loc, self)) + self.syn_models.append(synapse) + def set_spiketrain(self, syn_index, syn_weight, spike_times): - self.syn_models[syn_index].set_spiketrain(spike_times, weight=syn_weight, **self.sim_control) + spike_times = jnp.array(spike_times) + self.sim_control['t_calibrate'] + # create the pre synaptic dummy voltage representing weighted spikes + spks_v = jnp.zeros(int(self.sim_control['t_sim'] / self.sim_control['dt']) + 1) + spks_v = spks_v.at[(spike_times / self.sim_control['dt']).astype(int)].set(syn_weight) + # store the dummy cell and its 'spike' voltage + self.pre_dummies[syn_index] = (jx.Cell(), spks_v) def run(self): - for loc, syn in zip(self.syn_locs, self.syn_models): - self.cell.branch( + # create the network consisting of the pre dummies and the cell + net = jx.Network( + [ + self.pre_dummies[ii][0] for ii in range(len(self.pre_dummies)) + ] + [ + self.cell + ] + ) + # connect the synapses + cell_id = len(self.pre_dummies) + for ii, (loc, syn) in enumerate(zip(self.syn_locs, self.syn_models)): + # post cell synapse at correct location + post = net.cell(cell_id).branch( self.index_map[loc['node']] ).loc( loc['x'] - ).insert(syn) - # for loc, syn in zip(self.syn_locs, self.syn_models): - # self.cell.branch( - # self.index_map[loc['node']] - # ).loc( - # loc['x'] - # ).insert(JX_MECH["multichannel_test"].Bla()) - - self.cell.delete_recordings() - self.cell.branch(0).loc(0.0).record() - self.cell.record("AMPASynapse_x_r") - - current = jx.step_current(i_delay=100., i_dur=10.0, i_amp=10.0, delta_t=self.sim_control['dt'], t_max=self.sim_control['t_max']) - self.cell.branch(0).loc(0.0).stimulate(current) - - res = jx.integrate(self.cell, delta_t=self.sim_control['dt']) + ) + # clamp the pre dummy + net.cell(ii).clamp("v", self.pre_dummies[ii][1]) + # select correct pre dummy + pre = net.cell(ii).branch(0).loc(0.0) + jx.connect(pre, post, syn) + + net.delete_recordings() + net.cell(cell_id).branch(0).loc(0.0).record() + + current = jx.step_current( + i_delay=10.+self.sim_control['t_calibrate'], i_dur=10.0, i_amp=1.0, + delta_t=self.sim_control['dt'], + t_max=self.sim_control['t_sim'], + ) + net.cell(cell_id).branch(0).loc(0.5).stimulate(current) + + res_raw = jx.integrate(net, delta_t=self.sim_control['dt']) + + i0 = int(self.sim_control['t_calibrate'] / self.sim_control['dt']) + # breakpoint() + res = { + 't': jnp.arange(0, res_raw[0,i0:].shape[0]) * self.sim_control['dt'], + 'v_m': res_raw[:,i0:], + } return res + class JaxleyCompartmentNode(JaxleySimNode): """ diff --git a/tests/test_jaxleytree.py b/tests/test_jaxleytree.py index dfe40c75..5885c55c 100644 --- a/tests/test_jaxleytree.py +++ b/tests/test_jaxleytree.py @@ -3,7 +3,7 @@ import os -from neat import PhysTree, JaxleySimTree, JaxleyCompartmentTree, CompartmentFitter +from neat import PhysTree, JaxleySimTree, JaxleyCompartmentTree, NeuronSimTree, CompartmentFitter import channelcollection_for_tests as channelcollection import channel_installer @@ -45,18 +45,28 @@ def create_jaxley_model(self): self.load_ball() jt = JaxleySimTree(self.tree) - jcell = jt.init_model("multichannel_test", t_max=1000) + jcell = jt.init_model("multichannel_test", t_max=200., t_calibrate=100.) - jt.add_AMPA_synapse((1,.5)) - jt.set_spiketrain(0, 1., [200.]) - res = jt.run() + jt.add_AMPA_synapse((1,0.)) + jt.set_spiketrain(0, .001, [150.]) + jres = jt.run() + + # jct = JaxleyCompartmentTree(self.ctree) + # jcompcell = jt.init_model("multichannel_test") + + nt = NeuronSimTree(self.tree) + nt.init_model(t_calibrate=100.) + + nt.add_double_exp_synapse((1,.5), .2, 3., 0.) + nt.set_spiketrain(0, .001, [150.]) + nt.add_i_clamp((1,0.), 1., 10., 10.) + nres = nt.run(t_max=200.) import matplotlib.pyplot as pl - pl.plot(res.T) + pl.plot(nres['t'], nres['v_m'][0], c='r', label='neuron') + pl.plot(jres['t'], jres['v_m'][0], c='b', ls='--', label='jaxley') pl.show() - jct = JaxleyCompartmentTree(self.ctree) - jcompcell = jt.init_model("multichannel_test") # def test_jaxley_channels(self): # self.load_ball() From 934ead90f011a638c8a99c3fe0d1234323fffa5f Mon Sep 17 00:00:00 2001 From: "w.wybo" Date: Tue, 13 Jan 2026 11:02:41 +0100 Subject: [PATCH 9/9] ion channels issues with jaxley channels --- src/neat/simulations/jaxley/default_syns.py | 1 - src/neat/simulations/jaxley/jaxleymodel.py | 42 ++++++++----- tests/test_jaxleytree.py | 70 +++++++++++++++++++-- 3 files changed, 90 insertions(+), 23 deletions(-) diff --git a/src/neat/simulations/jaxley/default_syns.py b/src/neat/simulations/jaxley/default_syns.py index 7d0997a6..43d9830d 100644 --- a/src/neat/simulations/jaxley/default_syns.py +++ b/src/neat/simulations/jaxley/default_syns.py @@ -83,7 +83,6 @@ def update_states(self, states, delta_t, pre_voltage, post_voltage, params): } def compute_current(self, states, pre_voltage, post_voltage, params): - print(self.prefix)#, params) e_r = params[f'{self.prefix}e_r'] v = post_voltage x_r = states[f'{self.prefix}x_r'] diff --git a/src/neat/simulations/jaxley/jaxleymodel.py b/src/neat/simulations/jaxley/jaxleymodel.py index 0c6f761c..ba750669 100644 --- a/src/neat/simulations/jaxley/jaxleymodel.py +++ b/src/neat/simulations/jaxley/jaxleymodel.py @@ -102,7 +102,7 @@ def _make_branch(self, jx_channels, dx_max=15, factorlambda=1.0, pprint=False): l_comp = 2.0 * self.R else: - n_comp = self.L / dx_max + 1 + n_comp = int(self.L / dx_max) + 1 l_comp = self.L / n_comp compartments = [] @@ -243,22 +243,22 @@ def init_model( 't_sim': t_calibrate + t_max, } self.model_name = model_name - self.pre_dummies = {} + self.pre_dummies = [] self.syn_locs = [] self.syn_models = [] - - # simple tree copy - aux_tree = MorphTree(self) - aux_tree.reset_indices() - self.index_map = {n_orig.index: n_aux.index-1 for n_orig, n_aux in zip(self, aux_tree)} - + + # make the Jaxley channel current mechanisms in this model available jx_channels = { chan: eval(f"JX_MECH['{model_name}'].{chan}()") for chan in self.channel_storage } jx_channels['L'] = Leak() + # create Jaxley branches for all nodes branches = [node._make_branch(jx_channels) for node in self] - parents = [aux_node.parent_node.index-1 if not aux_tree.is_root(aux_node) else -1 for aux_node in aux_tree] + + # create index map for mapping branches to parents + self.index_map = {node.index: ii for ii, node in enumerate(self)} + parents = [self.index_map[node.parent_node.index] if not self.is_root(node) else -1 for node in self] self.cell = jx.Cell(branches, parents) @@ -273,16 +273,19 @@ def delete_model(self): self.pre_dummies = [] self.syn_locs = [] self.syn_models = [] - - def add_AMPA_synapse(self, loc): - synapse = JX_MECH[f'{self.model_name}'].AMPASynapse() - self.syn_locs.append(MorphLoc(loc, self)) - self.syn_models.append(synapse) - def add_GABA_synapse(self, loc): - synapse = JX_MECH[f'{self.model_name}'].GABASynapse() + def _append_synapse(self, loc, synapse): self.syn_locs.append(MorphLoc(loc, self)) self.syn_models.append(synapse) + self.pre_dummies.append(None) + + def add_AMPA_synapse(self, loc, tau_r_AMPA=.2, tau_d_AMPA=3., e_r_AMPA=0.): + synapse = JX_MECH[f'{self.model_name}'].AMPASynapse(tau_r_AMPA=tau_r_AMPA, tau_d_AMPA=tau_d_AMPA, e_r_AMPA=e_r_AMPA) + self._append_synapse(loc, synapse) + + def add_GABA_synapse(self, loc, tau_r_GABA=.2, tau_d_GABA=10., e_r_GABA=-80.): + synapse = JX_MECH[f'{self.model_name}'].GABASynapse(tau_r_GABA=tau_r_GABA, tau_d_GABA=tau_d_GABA, e_r_GABA=e_r_GABA) + self._append_synapse(loc, synapse) def set_spiketrain(self, syn_index, syn_weight, spike_times): spike_times = jnp.array(spike_times) + self.sim_control['t_calibrate'] @@ -293,10 +296,15 @@ def set_spiketrain(self, syn_index, syn_weight, spike_times): self.pre_dummies[syn_index] = (jx.Cell(), spks_v) def run(self): + # if there are synapses that don't receive a spiketrain, create empty ones + for ii, pd in enumerate(self.pre_dummies): + if pd is None: + self.set_spiketrain(ii, 0.0, []) + # create the network consisting of the pre dummies and the cell net = jx.Network( [ - self.pre_dummies[ii][0] for ii in range(len(self.pre_dummies)) + pre_dummy[0] for pre_dummy in self.pre_dummies ] + [ self.cell ] diff --git a/tests/test_jaxleytree.py b/tests/test_jaxleytree.py index 5885c55c..7e76a145 100644 --- a/tests/test_jaxleytree.py +++ b/tests/test_jaxleytree.py @@ -26,8 +26,8 @@ def load_ball(self): # capacitance and axial resistance self.tree.set_physiology(0.8, 100.0 / 1e6) # ion channels - self.k_chan = channelcollection.Kv3_1() - self.tree.add_channel_current(self.k_chan, 0.766 * 1e6, -85.0) + # self.k_chan = channelcollection.Kv3_1() + # self.tree.add_channel_current(self.k_chan, 0.766 * 1e6, -85.0) self.na_chan = channelcollection.NaTa_t() self.tree.add_channel_current(self.na_chan, 1.71 * 1e6, 50.0) # fit leak current @@ -40,8 +40,37 @@ def load_ball(self): cfit = CompartmentFitter(self.tree, save_cache=False, recompute_cache=True) self.ctree, _ = cfit.fit_model([(1, 0.5)]) - - def create_jaxley_model(self): + def load_axon_tree(self): + """ + Parameters taken from a BBP SST model for a subset of ion channels + """ + tree = PhysTree( + os.path.join(MORPHOLOGIES_PATH_PREFIX, "ball_and_axon.swc"), + types=[1, 2, 3, 4], + ) + # capacitance and axial resistance + tree.set_physiology(1.0, 100.0 / 1e6) + # ion channels + k_chan = channelcollection.SKv3_1() + tree.add_channel_current(k_chan, 0.653374 * 1e6, -85.0, node_arg=[tree[1]]) + # tree.add_channel_current(k_chan, 0.196957 * 1e6, -85.0, node_arg="axonal") + # na_chan = channelcollection.Na_Ta() + # tree.add_channel_current(na_chan, 1.0 * 1e6, 50.0, node_arg=[tree[1]]) + # tree.add_channel_current(na_chan, 3.418459 * 1e6, 50.0, node_arg="axonal") + # ca_chan = channelcollection.Ca_HVA() + # tree.add_channel_current( + # ca_chan, 0.000792 * 1e6, 132.4579341637009, node_arg=[tree[1]] + # ) + # tree.add_channel_current( + # ca_chan, 0.000138 * 1e6, 132.4579341637009, node_arg="axonal" + # ) + # passive leak current + tree.set_leak_current(0.000091 * 1e6, -62.442793, node_arg=[tree[1]]) + tree.set_leak_current(0.000094 * 1e6, -79.315740, node_arg="axonal") + + self.tree = tree + + def test_ball(self): self.load_ball() jt = JaxleySimTree(self.tree) @@ -67,6 +96,36 @@ def create_jaxley_model(self): pl.plot(jres['t'], jres['v_m'][0], c='b', ls='--', label='jaxley') pl.show() + def test_axon(self): + stim_locs = [(4,.9), (5,.9)] + + self.load_axon_tree() + + jt = JaxleySimTree(self.tree) + jcell = jt.init_model("multichannel_test", t_max=200., t_calibrate=100.) + + jt.add_AMPA_synapse(stim_locs[0]) + jt.set_spiketrain(0, .0001, [50.]) + jt.add_AMPA_synapse(stim_locs[1]) + jt.set_spiketrain(1, .0005, [150.]) + jres = jt.run() + + nt = NeuronSimTree(self.tree) + nt.init_model(t_calibrate=100.) + + nt.add_double_exp_synapse(stim_locs[0], .2, 3., 0.) + nt.set_spiketrain(0, .0001, [50.]) + nt.add_double_exp_synapse(stim_locs[1], .2, 3., 0.) + nt.set_spiketrain(1, .0005, [150.]) + nt.add_i_clamp((1,0.), 1., 10., 10.) + nres = nt.run(t_max=200.) + + import matplotlib.pyplot as pl + pl.plot(nres['t'], nres['v_m'][0], c='r', label='neuron') + pl.plot(jres['t'], jres['v_m'][0], c='b', ls='--', label='jaxley') + pl.legend() + pl.show() + # def test_jaxley_channels(self): # self.load_ball() @@ -81,5 +140,6 @@ def create_jaxley_model(self): if __name__ == "__main__": tjx = TestJaxley() - tjx.create_jaxley_model() + tjx.test_ball() + # tjx.test_axon() # tjx.test_jaxley_channels()