Jaxley export [WIP]

.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

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

src/neat/actions/install.py

@@ -254,13 +254,61 @@ 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
+import jax.numpy as jnp
+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,
+)
+
+
+"""
+    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()
+        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 +346,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
+        )

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():

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/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"
+        "Only used when the [action] is 'install'.",
+    )
     parser.add_argument(
         "-p",
         "--path",
@@ -118,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)

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.")

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):
@@ -383,7 +413,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):
@@ -1258,3 +1287,158 @@ def _replaceConc(expr_str, prefix="", suffix=""):
 
         fh.close()
         fcc.close()
+
+
+    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, 
+        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}}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(
+        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
+            ])
+        }
+"""     
+
+        # 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 = 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}
+        {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}}g{cname}']
+        erev = params[f'{{self.prefix}}e{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._create_jaxley_funcstr(sp.printing.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