fixup merge

python/sdist/amici/jax/jax.template.py

@@ -20,7 +20,7 @@ def __init__(self):
         super().__init__()
 
     def _xdot(self, t, x, args):
-        pk, tcl = args
+        p, tcl = args
 
         TPL_X_SYMS = x
         TPL_P_SYMS = p
@@ -31,7 +31,6 @@ def _xdot(self, t, x, args):
 
         return TPL_XDOT_RET
 
-
     def _w(self, t, x, p, tcl):
         TPL_X_SYMS = x
         TPL_P_SYMS = p

python/sdist/amici/jax/ode_export.py

@@ -24,6 +24,7 @@
 from amici._codegen.template import apply_template
 from amici.jax.jaxcodeprinter import AmiciJaxCodePrinter
 from amici.jax.model import JAXModel
+from amici.jax.nn import generate_equinox
 from amici.de_model import DEModel
 from amici.de_export import is_valid_identifier
 from amici.import_utils import (
@@ -129,6 +130,7 @@ def __init__(
         outdir: Path | str | None = None,
         verbose: bool | int | None = False,
         model_name: str | None = "model",
+        hybridisation: dict[str, str] = {},
     ):
         """
         Generate AMICI jax files for the ODE provided to the constructor.
@@ -157,6 +159,8 @@ def __init__(
 
         self.model: DEModel = ode_model
 
+        self.hybridisation = hybridisation
+
         self._code_printer = AmiciJaxCodePrinter()
 
     @log_execution_time("generating jax code", logger)
@@ -169,6 +173,7 @@ def generate_model_code(self) -> None:
         ):
             self._prepare_model_folder()
             self._generate_jax_code()
+            self._generate_nn_code()
 
     def _prepare_model_folder(self) -> None:
         """
@@ -233,6 +238,14 @@ def _generate_jax_code(self) -> None:
                 # can flag conflicts in the future
                 "MODEL_API_VERSION": f"'{JAXModel.MODEL_API_VERSION}'",
             },
+            "NET_IMPORTS": "\n".join(
+                f"{net} = _module_from_path('{net}', Path(__file__).parent / '{net}.py')"
+                for net in self.hybridisation.keys()
+            ),
+            "NETS": ",\n".join(
+                f'"{net}": {net}.net(jr.PRNGKey(0))'
+                for net in self.hybridisation.keys()
+            ),
         }
         outdir = self.model_path / (self.model_name + "_jax")
         outdir.mkdir(parents=True, exist_ok=True)
@@ -243,6 +256,15 @@ def _generate_jax_code(self) -> None:
             tpl_data,
         )
 
+    def _generate_nn_code(self) -> None:
+        for net_name, net in self.hybridisation.items():
+            generate_equinox(
+                net["model"],
+                os.path.join(
+                    self.model_path, self.model_name + "_jax", f"{net_name}.py"
+                ),
+            )
+
     def set_paths(self, output_dir: str | Path | None = None) -> None:
         """
         Set output paths for the model and create if necessary

python/sdist/amici/sbml_import.py

@@ -460,6 +460,7 @@ def sbml2jax(
         simplify: Callable | None = _default_simplify,
         cache_simplify: bool = False,
         log_as_log10: bool = True,
+        hybridisation: dict = None,
     ) -> None:
         """
         Generate and compile AMICI jax files for the model provided to the
@@ -549,6 +550,7 @@ def sbml2jax(
             model_name=model_name,
             outdir=output_dir,
             verbose=verbose,
+            hybridisation=hybridisation,
         )
         exporter.generate_model_code()