Move amici.jax to amici.exporters.jax and amici.sim.jax

Related to #3041.

Author: dweindl

doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb

@@ -65,7 +65,7 @@
    "source": [
     "## Simulation\n",
     "\n",
-    "We can now run efficient simulation using [amici.jax.run_simulations]((https://amici.readthedocs.io/en/latest/generated/amici.jax.html#amici.jax.run_simulations)."
+    "We can now run efficient simulation using [amici.sim.jax.run_simulations]((https://amici.readthedocs.io/en/latest/generated/amici.jax.html#amici.jax.run_simulations)."
    ]
   },
   {
@@ -75,7 +75,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from amici.jax import run_simulations\n",
+    "from amici.sim.jax import run_simulations\n",
     "\n",
     "# Run simulations and compute the log-likelihood\n",
     "llh, results = run_simulations(jax_problem)"
@@ -386,7 +386,7 @@
     "import diffrax\n",
     "import jax.numpy as jnp\n",
     "import optimistix\n",
-    "from amici.jax import ReturnValue\n",
+    "from amici.sim.jax import ReturnValue\n",
     "\n",
     "# Define the simulation condition\n",
     "experiment_condition = \"_petab_experiment_condition___default__\"\n",

doc/python_modules.rst

@@ -15,7 +15,7 @@ AMICI Python API
    amici.importers.petab
    amici.importers.petab.v1
    amici.importers.utils
-   amici.jax
+   amici.sim.jax
    amici.sim.sundials
    amici.sim.sundials.plotting
    amici.sim.sundials.gradient_check

python/sdist/amici/_symbolic/de_model.py

@@ -2499,7 +2499,8 @@ def _process_hybridization(self, hybridization: dict) -> None:
         orig_obs = tuple([s.get_sym() for s in self._observables])
         for net_id, net in hybridization.items():
             if net["static"]:
-                continue  # do not integrate into ODEs, handle in amici.jax.petab
+                # do not integrate into ODEs, handle in amici.sim.jax.petab
+                continue
             inputs = [
                 comp
                 for comp in self._components

python/sdist/amici/exporters/jax/__init__.py

@@ -0,0 +1,24 @@
+"""
+Code generation for JAX models for simulation with diffrax solvers.
+
+This module provides an interface to generate and use AMICI models with JAX.
+Please note that this module is experimental, the API may substantially change
+in the future. Use at your own risk and do not expect backward compatibility.
+"""
+
+from warnings import warn
+
+from .nn import Flatten, cat, generate_equinox, tanhshrink
+
+warn(
+    "The JAX module is experimental and the API may change in the future.",
+    ImportWarning,
+    stacklevel=2,
+)
+
+__all__ = [
+    "Flatten",
+    "generate_equinox",
+    "tanhshrink",
+    "cat",
+]

python/sdist/amici/jax/jax.template.py …dist/amici/exporters/jax/jax.template.pypython/sdist/amici/jax/jax.template.py renamed to python/sdist/amici/exporters/jax/jax.template.py python/sdist/amici/jax/jax.template.py renamed to python/sdist/amici/exporters/jax/jax.template.py

@@ -1,16 +1,16 @@
 # ruff: noqa: F401, F821, F841
 from pathlib import Path
 
-import equinox as eqx
+import equinox as eqx  # noqa: F401
 import jax.numpy as jnp
-import jax.random as jr
-import jaxtyping as jt
-from interpax import interp1d
-from jax.numpy import inf as oo
-from jax.numpy import nan as nan
-
-from amici import _module_from_path
-from amici.jax.model import JAXModel, safe_div, safe_log
+import jax.random as jr  # noqa: F401
+import jaxtyping as jt  # noqa: F401
+from interpax import interp1d  # noqa: F401
+from jax.numpy import inf as oo  # noqa: F401
+from jax.numpy import nan as nan  # noqa: F401
+
+from amici import _module_from_path  # noqa: F401
+from amici.sim.jax.model import JAXModel, safe_div, safe_log  # noqa: F401
 
 TPL_NET_IMPORTS
 

python/sdist/amici/jax/jaxcodeprinter.py …st/amici/exporters/jax/jaxcodeprinter.pypython/sdist/amici/jax/jaxcodeprinter.py renamed to python/sdist/amici/exporters/jax/jaxcodeprinter.py python/sdist/amici/jax/jaxcodeprinter.py renamed to python/sdist/amici/exporters/jax/jaxcodeprinter.py

@@ -4,8 +4,7 @@
 import jax.numpy as jnp
 
 from amici import amiciModulePath
-
-from ..exporters.template import apply_template
+from amici.exporters.template import apply_template
 
 
 class Flatten(eqx.Module):
@@ -185,7 +184,7 @@ def _generate_layer(layer: "Layer", indent: int, ilayer: int) -> str:  # noqa: F
     layer_map = {
         "Dropout1d": "eqx.nn.Dropout",
         "Dropout2d": "eqx.nn.Dropout",
-        "Flatten": "amici.jax.Flatten",
+        "Flatten": "amici.export.jax.Flatten",
     }
 
     # mapping of keyword argument names in sciml yaml format to equinox/custom amici implementations
@@ -321,9 +320,9 @@ def _process_activation_call(node: "Node") -> str:  # noqa: F821
         "hardtanh": "jax.nn.hard_tanh",
         "hardsigmoid": "jax.nn.hard_sigmoid",
         "hardswish": "jax.nn.hard_swish",
-        "tanhshrink": "amici.jax.tanhshrink",
+        "tanhshrink": "amici.export.jax.tanhshrink",
         "softsign": "jax.nn.soft_sign",
-        "cat": "amici.jax.cat",
+        "cat": "amici.export.jax.cat",
     }
 
     # Validate hardtanh parameters

python/sdist/amici/jax/nn.py python/sdist/amici/exporters/jax/nn.pypython/sdist/amici/jax/nn.py renamed to python/sdist/amici/exporters/jax/nn.py

@@ -1,12 +1,7 @@
 # ruff: noqa: F401, F821, F841
 import equinox as eqx
-import jax
-import jax.nn
-import jax.numpy as jnp
 import jax.random as jr
 
-import amici.jax.nn
-
 
 class TPL_MODEL_ID(eqx.Module):
     layers: dict

python/sdist/amici/jax/nn.template.py …sdist/amici/exporters/jax/nn.template.pypython/sdist/amici/jax/nn.template.py renamed to python/sdist/amici/exporters/jax/nn.template.py python/sdist/amici/jax/nn.template.py renamed to python/sdist/amici/exporters/jax/nn.template.py

@@ -18,19 +18,17 @@
 
 import sympy as sp
 
-from amici import (
-    amiciModulePath,
-)
 from amici._symbolic.de_model import DEModel
 from amici._symbolic.sympy_utils import (
     _monkeypatch_sympy,
 )
+from amici.exporters.jax.nn import generate_equinox
 from amici.exporters.sundials.de_export import is_valid_identifier
 from amici.exporters.template import apply_template
-from amici.jax.jaxcodeprinter import AmiciJaxCodePrinter, _jnp_array_str
-from amici.jax.model import JAXModel
-from amici.jax.nn import generate_equinox
 from amici.logging import get_logger, log_execution_time, set_log_level
+from amici.sim.jax.model import JAXModel
+
+from .jaxcodeprinter import AmiciJaxCodePrinter, _jnp_array_str
 
 #: python log manager
 logger = get_logger(__name__, logging.ERROR)
@@ -303,7 +301,7 @@ def _generate_jax_code(self) -> None:
         }
 
         apply_template(
-            Path(amiciModulePath) / "jax" / "jax.template.py",
+            Path(__file__).parent / "jax.template.py",
             self.model_path / "__init__.py",
             tpl_data,
         )

python/sdist/amici/jax/ode_export.py …/sdist/amici/exporters/jax/ode_export.pypython/sdist/amici/jax/ode_export.py renamed to python/sdist/amici/exporters/jax/ode_export.py python/sdist/amici/jax/ode_export.py renamed to python/sdist/amici/exporters/jax/ode_export.py

@@ -0,0 +1 @@
+"""Code generation for models for simulation with SUNDIALS solvers."""