From d1c8250b37ac75ab184146232f3846cbc7a6a1bf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= Date: Thu, 5 Dec 2024 21:33:46 +0000 Subject: [PATCH 1/2] Mark JAX experimental (#2618) * mark jax experimental * add JAX import warning to pyest.ini --- pytest.ini | 1 + python/sdist/amici/jax/__init__.py | 17 ++++++++++++++++- 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/pytest.ini b/pytest.ini index adbf313922..8cc45e0fd9 100644 --- a/pytest.ini +++ b/pytest.ini @@ -12,6 +12,7 @@ filterwarnings = ignore:Conservation laws for non-constant species in models with Species-AssignmentRules are currently not supported and will be turned off.:UserWarning ignore:Conservation laws for non-constant species in combination with parameterized stoichiometric coefficients are not currently supported and will be turned off.:UserWarning ignore:Support for PEtab2.0 is experimental!:UserWarning + ignore:The JAX module is experimental and the API may change in the future.:ImportWarning # hundreds of SBML <=5.17 warnings ignore:.*inspect.getargspec\(\) is deprecated.*:DeprecationWarning # pysb warnings diff --git a/python/sdist/amici/jax/__init__.py b/python/sdist/amici/jax/__init__.py index e14d231e1e..8b67abda27 100644 --- a/python/sdist/amici/jax/__init__.py +++ b/python/sdist/amici/jax/__init__.py @@ -1,6 +1,21 @@ -"""Interface to facilitate AMICI generated models using JAX""" +""" +JAX +--- + +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 amici.jax.petab import JAXProblem, run_simulations from amici.jax.model import JAXModel +warn( + "The JAX module is experimental and the API may change in the future.", + ImportWarning, + stacklevel=2, +) + __all__ = ["JAXModel", "JAXProblem", "run_simulations"] From 449041d81dc23f44a6e30bda73d20d0f0586e00b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= Date: Thu, 5 Dec 2024 23:40:34 +0000 Subject: [PATCH 2/2] Separate pre-equilibration and dynamic simulation in jax (#2617) * disentangle sim & preeq * disentangle sim & preeq * run preequilibration once * fix symlink * separate default dirs for jax/cpp, honour model dir/name * fix notebook * fix path SNAFU * fix models without preequilibration * fix tests * fixup * fix doc typehints * fix notebook * fix output dict construction * fix notebook --- .../example_jax_petab/ExampleJaxPEtab.ipynb | 11 ++- python/sdist/amici/jax/model.py | 64 ++++++++++++------ python/sdist/amici/jax/ode_export.py | 5 +- python/sdist/amici/jax/petab.py | 67 +++++++++++++++---- python/sdist/amici/petab/import_helpers.py | 24 ++++++- python/sdist/amici/petab/petab_import.py | 14 ++-- python/sdist/amici/petab/sbml_import.py | 9 ++- python/tests/test_jax.py | 46 +++++++------ .../benchmark-models/test_petab_benchmark.py | 4 +- 9 files changed, 167 insertions(+), 77 deletions(-) diff --git a/python/examples/example_jax_petab/ExampleJaxPEtab.ipynb b/python/examples/example_jax_petab/ExampleJaxPEtab.ipynb index 855860e242..f6a4f10e98 100644 --- a/python/examples/example_jax_petab/ExampleJaxPEtab.ipynb +++ b/python/examples/example_jax_petab/ExampleJaxPEtab.ipynb @@ -151,7 +151,7 @@ " results (dict): Simulation results from run_simulations.\n", " \"\"\"\n", " # Extract the simulation results for the specific condition\n", - " sim_results = results[simulation_condition][1]\n", + " sim_results = results[simulation_condition]\n", "\n", " # Create a new figure for the state trajectories\n", " plt.figure(figsize=(8, 6))\n", @@ -357,14 +357,12 @@ "simulation_condition = (\"model1_data1\",)\n", "\n", "# Load condition-specific data\n", - "ts_preeq, ts_dyn, ts_posteq, my, iys = jax_problem._measurements[\n", + "ts_init, ts_dyn, ts_posteq, my, iys = jax_problem._measurements[\n", " simulation_condition\n", "]\n", "\n", "# Load parameters for the specified condition\n", "p = jax_problem.load_parameters(simulation_condition[0])\n", - "# Disable preequilibration\n", - "p_preeq = jnp.array([])\n", "\n", "\n", "# Define a function to compute the gradient with respect to dynamic timepoints\n", @@ -372,12 +370,12 @@ "def grad_ts_dyn(tt):\n", " return jax_problem.model.simulate_condition(\n", " p=p,\n", - " p_preeq=p_preeq,\n", - " ts_preeq=ts_preeq,\n", + " ts_init=ts_init,\n", " ts_dyn=tt,\n", " ts_posteq=ts_posteq,\n", " my=jnp.array(my),\n", " iys=jnp.array(iys),\n", + " x_preeq=jnp.array([]),\n", " solver=diffrax.Kvaerno5(),\n", " controller=diffrax.PIDController(atol=1e-8, rtol=1e-8),\n", " max_steps=2**10,\n", @@ -489,7 +487,6 @@ "amici_model = import_petab_problem(\n", " petab_problem,\n", " verbose=False,\n", - " compile_=True,\n", " jax=False, # load the amici model this time\n", ")\n", "\n", diff --git a/python/sdist/amici/jax/model.py b/python/sdist/amici/jax/model.py index ac86b547a6..8f9650ef0f 100644 --- a/python/sdist/amici/jax/model.py +++ b/python/sdist/amici/jax/model.py @@ -427,12 +427,12 @@ def _sigmays( def simulate_condition( self, p: jt.Float[jt.Array, "np"], - p_preeq: jt.Float[jt.Array, "*np"], - ts_preeq: jt.Float[jt.Array, "nt_preeq"], + ts_init: jt.Float[jt.Array, "nt_preeq"], ts_dyn: jt.Float[jt.Array, "nt_dyn"], ts_posteq: jt.Float[jt.Array, "nt_posteq"], my: jt.Float[jt.Array, "nt"], iys: jt.Int[jt.Array, "nt"], + x_preeq: jt.Float[jt.Array, "nx"], solver: diffrax.AbstractSolver, controller: diffrax.AbstractStepSizeController, adjoint: diffrax.AbstractAdjoint, @@ -444,12 +444,9 @@ def simulate_condition( :param p: parameters for simulation ordered according to ids in :ivar parameter_ids: - :param p_preeq: - parameters for pre-equilibration ordered according to ids in :ivar parameter_ids:. May be empty to - disable pre-equilibration. - :param ts_preeq: - time points for pre-equilibration. Usually valued 0.0, but needs to be shaped according to - the number of observables that are evaluated after pre-equilibration. + :param ts_init: + time points that do not require simulation. Usually valued 0.0, but needs to be shaped according to + the number of observables that are evaluated before dynamic simulation. :param ts_dyn: time points for dynamic simulation. Usually valued > 0.0 and sorted in monotonically increasing order. Duplicate time points are allowed to facilitate the evaluation of multiple observables at specific time @@ -486,24 +483,16 @@ def simulate_condition( output according to `ret` and statistics """ # Pre-equilibration - if p_preeq.shape[0] > 0: - x0 = self._x0(p_preeq) - tcl = self._tcl(x0, p_preeq) - current_x = self._x_solver(x0) - current_x, stats_preeq = self._eq( - p_preeq, tcl, current_x, solver, controller, max_steps - ) + if x_preeq.shape[0] > 0: + current_x = self._x_solver(x_preeq) # update tcl with new parameters - tcl = self._tcl(self._x_rdata(current_x, tcl), p) + tcl = self._tcl(x_preeq, p) else: x0 = self._x0(p) current_x = self._x_solver(x0) - stats_preeq = None tcl = self._tcl(x0, p) - x_preq = jnp.repeat( - current_x.reshape(1, -1), ts_preeq.shape[0], axis=0 - ) + x_preq = jnp.repeat(current_x.reshape(1, -1), ts_init.shape[0], axis=0) # Dynamic simulation if ts_dyn.shape[0] > 0: @@ -536,7 +525,7 @@ def simulate_condition( current_x.reshape(1, -1), ts_posteq.shape[0], axis=0 ) - ts = jnp.concatenate((ts_preeq, ts_dyn, ts_posteq), axis=0) + ts = jnp.concatenate((ts_init, ts_dyn, ts_posteq), axis=0) x = jnp.concatenate((x_preq, x_dyn, x_posteq), axis=0) nllhs = self._nllhs(ts, x, p, tcl, my, iys) @@ -555,11 +544,42 @@ def simulate_condition( }[ret], dict( ts=ts, x=x, - stats_preeq=stats_preeq, stats_dyn=stats_dyn, stats_posteq=stats_posteq, ) + @eqx.filter_jit + def preequilibrate_condition( + self, + p: jt.Float[jt.Array, "np"], + solver: diffrax.AbstractSolver, + controller: diffrax.AbstractStepSizeController, + max_steps: int | jnp.int_, + ) -> tuple[jt.Float[jt.Array, "nx"], dict]: + r""" + Simulate a condition. + + :param p: + parameters for simulation ordered according to ids in :ivar parameter_ids: + :param solver: + ODE solver + :param controller: + step size controller + :param max_steps: + maximum number of solver steps + :return: + pre-equilibrated state variables and statistics + """ + # Pre-equilibration + x0 = self._x0(p) + tcl = self._tcl(x0, p) + current_x = self._x_solver(x0) + current_x, stats_preeq = self._eq( + p, tcl, current_x, solver, controller, max_steps + ) + + return self._x_rdata(current_x, tcl), dict(stats_preeq=stats_preeq) + def safe_log(x: jnp.float_) -> jnp.float_: """ diff --git a/python/sdist/amici/jax/ode_export.py b/python/sdist/amici/jax/ode_export.py index 7ea4a29d8a..cec5104ded 100644 --- a/python/sdist/amici/jax/ode_export.py +++ b/python/sdist/amici/jax/ode_export.py @@ -234,12 +234,10 @@ def _generate_jax_code(self) -> None: "MODEL_API_VERSION": f"'{JAXModel.MODEL_API_VERSION}'", }, } - outdir = self.model_path / (self.model_name + "_jax") - outdir.mkdir(parents=True, exist_ok=True) apply_template( Path(amiciModulePath) / "jax" / "jax.template.py", - outdir / "__init__.py", + self.model_path / "__init__.py", tpl_data, ) @@ -258,6 +256,7 @@ def set_paths(self, output_dir: str | Path | None = None) -> None: output_dir = Path(os.getcwd()) / f"amici-{self.model_name}" self.model_path = Path(output_dir).resolve() + self.model_path.mkdir(parents=True, exist_ok=True) def set_name(self, model_name: str) -> None: """ diff --git a/python/sdist/amici/jax/petab.py b/python/sdist/amici/jax/petab.py index 2c823259fe..0411e5e2df 100644 --- a/python/sdist/amici/jax/petab.py +++ b/python/sdist/amici/jax/petab.py @@ -154,7 +154,7 @@ def _get_parameter_mappings( def _get_measurements( self, simulation_conditions: pd.DataFrame ) -> dict[ - tuple[str], + tuple[str, ...], tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray], ]: """ @@ -307,19 +307,21 @@ def run_simulation( solver: diffrax.AbstractSolver, controller: diffrax.AbstractStepSizeController, max_steps: jnp.int_, + x_preeq: jt.Float[jt.Array, "*nx"] = jnp.array([]), # noqa: F821, F722 ) -> tuple[jnp.float_, dict]: """ Run a simulation for a given simulation condition. :param simulation_condition: - Tuple of simulation conditions to run the simulation for. can be a single string (simulation only) or a - tuple of strings (pre-equilibration followed by simulation). + Simulation condition to run simulation for. :param solver: ODE solver to use for simulation :param controller: Step size controller to use for simulation :param max_steps: Maximum number of steps to take during simulation + :param x_preeq: + Pre-equilibration state if available :return: Tuple of log-likelihood and simulation statistics """ @@ -327,29 +329,53 @@ def run_simulation( simulation_condition ] p = self.load_parameters(simulation_condition[0]) - p_preeq = ( - self.load_parameters(simulation_condition[1]) - if len(simulation_condition) > 1 - else jnp.array([]) - ) return self.model.simulate_condition( p=p, - p_preeq=p_preeq, - ts_preeq=jax.lax.stop_gradient(jnp.array(ts_preeq)), + ts_init=jax.lax.stop_gradient(jnp.array(ts_preeq)), ts_dyn=jax.lax.stop_gradient(jnp.array(ts_dyn)), ts_posteq=jax.lax.stop_gradient(jnp.array(ts_posteq)), my=jax.lax.stop_gradient(jnp.array(my)), iys=jax.lax.stop_gradient(jnp.array(iys)), + x_preeq=x_preeq, solver=solver, controller=controller, max_steps=max_steps, adjoint=diffrax.RecursiveCheckpointAdjoint(), ) + def run_preequilibration( + self, + simulation_condition: str, + solver: diffrax.AbstractSolver, + controller: diffrax.AbstractStepSizeController, + max_steps: jnp.int_, + ) -> tuple[jt.Float[jt.Array, "nx"], dict]: # noqa: F821 + """ + Run a pre-equilibration simulation for a given simulation condition. + + :param simulation_condition: + Simulation condition to run simulation for. + :param solver: + ODE solver to use for simulation + :param controller: + Step size controller to use for simulation + :param max_steps: + Maximum number of steps to take during simulation + :return: + Pre-equilibration state + """ + p = self.load_parameters(simulation_condition) + return self.model.preequilibrate_condition( + p=p, + solver=solver, + controller=controller, + max_steps=max_steps, + ) + def run_simulations( problem: JAXProblem, - simulation_conditions: Iterable[tuple] | None = None, + simulation_conditions: Iterable[tuple[str, ...]] | None = None, solver: diffrax.AbstractSolver = diffrax.Kvaerno5(), controller: diffrax.AbstractStepSizeController = diffrax.PIDController( rtol=1e-8, @@ -379,8 +405,23 @@ def run_simulations( if simulation_conditions is None: simulation_conditions = problem.get_all_simulation_conditions() + preeqs = { + sc: problem.run_preequilibration(sc, solver, controller, max_steps) + # only run preequilibration once per condition + for sc in {sc[1] for sc in simulation_conditions if len(sc) > 1} + } + results = { - sc: problem.run_simulation(sc, solver, controller, max_steps) + sc: problem.run_simulation( + sc, + solver, + controller, + max_steps, + preeqs.get(sc[1])[0] if len(sc) > 1 else jnp.array([]), + ) for sc in simulation_conditions } - return sum(llh for llh, _ in results.values()), results + return sum(llh for llh, _ in results.values()), { + sc: res[1] | preeqs[sc[1]][1] if len(sc) > 1 else res[1] + for sc, res in results.items() + } diff --git a/python/sdist/amici/petab/import_helpers.py b/python/sdist/amici/petab/import_helpers.py index daa902efb0..d42e99b1e3 100644 --- a/python/sdist/amici/petab/import_helpers.py +++ b/python/sdist/amici/petab/import_helpers.py @@ -138,10 +138,9 @@ def _can_import_model( Check whether a module of that name can already be imported. """ # try to import (in particular checks version) - suffix = "_jax" if jax else "" try: model_module = amici.import_model_module( - model_name + suffix, model_output_dir + *_get_package_name_and_path(model_name, model_output_dir, jax) ) except ModuleNotFoundError: return False @@ -271,3 +270,24 @@ def check_model( "the current model might also resolve this. Parameters: " f"{amici_ids_free_required.difference(amici_ids_free)}" ) + + +def _get_package_name_and_path( + model_name: str, model_output_dir: str | Path, jax: bool = False +) -> tuple[str, Path]: + """ + Get the package name and path for the generated model module. + + :param model_name: + Name of the model + :param model_output_dir: + Target directory for the generated model module + :param jax: + Whether to generate the paths for a JAX or CPP model + :return: + """ + if jax: + outdir = Path(model_output_dir) + return outdir.stem, outdir.parent + else: + return model_name, Path(model_output_dir) diff --git a/python/sdist/amici/petab/petab_import.py b/python/sdist/amici/petab/petab_import.py index 63bade9bb8..b7fccca241 100644 --- a/python/sdist/amici/petab/petab_import.py +++ b/python/sdist/amici/petab/petab_import.py @@ -16,7 +16,12 @@ from petab.v1.models import MODEL_TYPE_PYSB, MODEL_TYPE_SBML from ..logging import get_logger -from .import_helpers import _can_import_model, _create_model_name, check_model +from .import_helpers import ( + _can_import_model, + _create_model_name, + check_model, + _get_package_name_and_path, +) from .sbml_import import import_model_sbml try: @@ -114,7 +119,7 @@ def import_petab_problem( from .sbml_import import _create_model_output_dir_name model_output_dir = _create_model_output_dir_name( - petab_problem.sbml_model, model_name + petab_problem.sbml_model, model_name, jax=jax ) else: model_output_dir = os.path.abspath(model_output_dir) @@ -136,7 +141,7 @@ def import_petab_problem( ) # remove folder if exists - if os.path.exists(model_output_dir): + if not jax and os.path.exists(model_output_dir): shutil.rmtree(model_output_dir) logger.info(f"Compiling model {model_name} to {model_output_dir}.") @@ -160,9 +165,8 @@ def import_petab_problem( ) # import model - suffix = "_jax" if jax else "" model_module = amici.import_model_module( - model_name + suffix, model_output_dir + *_get_package_name_and_path(model_name, model_output_dir, jax=jax) ) if jax: diff --git a/python/sdist/amici/petab/sbml_import.py b/python/sdist/amici/petab/sbml_import.py index 02a2c4e12c..e4e7efd7fc 100644 --- a/python/sdist/amici/petab/sbml_import.py +++ b/python/sdist/amici/petab/sbml_import.py @@ -588,7 +588,9 @@ def _get_fixed_parameters_sbml( def _create_model_output_dir_name( - sbml_model: "libsbml.Model", model_name: str | None = None + sbml_model: "libsbml.Model", + model_name: str | None = None, + jax: bool = False, ) -> Path: """ Find a folder for storing the compiled amici model. @@ -599,12 +601,13 @@ def _create_model_output_dir_name( BASE_DIR = Path("amici_models").absolute() BASE_DIR.mkdir(exist_ok=True) # try model_name + suffix = "_jax" if jax else "" if model_name: - return BASE_DIR / model_name + return BASE_DIR / (model_name + suffix) # try sbml model id if sbml_model_id := sbml_model.getId(): - return BASE_DIR / sbml_model_id + return BASE_DIR / (sbml_model_id + suffix) # create random folder name return Path(tempfile.mkdtemp(dir=BASE_DIR)) diff --git a/python/tests/test_jax.py b/python/tests/test_jax.py index 8f4c68510b..ce7018e078 100644 --- a/python/tests/test_jax.py +++ b/python/tests/test_jax.py @@ -47,7 +47,7 @@ def test_conversion(): module_name=model.name, module_path=outdir ) jax_module = amici.import_model_module( - module_name=model.name + "_jax", module_path=outdir + module_name=Path(outdir).stem, module_path=Path(outdir).parent ) ts = tuple(np.linspace(0, 1, 10)) @@ -108,7 +108,7 @@ def test_dimerization(): module_name=model.name, module_path=outdir ) jax_module = amici.import_model_module( - module_name=model.name + "_jax", module_path=outdir + module_name=Path(outdir).stem, module_path=Path(outdir).parent ) ts = tuple(np.linspace(0, 1, 10)) @@ -178,7 +178,7 @@ def check_fields_jax( ts = ts.flatten() iys = iys.flatten() - ts_preeq = ts[ts == 0] + ts_init = ts[ts == 0] ts_dyn = ts[ts > 0] ts_posteq = np.array([]) @@ -188,31 +188,37 @@ def check_fields_jax( } p = jnp.array([par_dict[par_id] for par_id in jax_model.parameter_ids]) - args = ( - jnp.array([]), # p_preeq - jnp.array(ts_preeq), # ts_preeq - jnp.array(ts_dyn), # ts_dyn - jnp.array(ts_posteq), # ts_posteq - jnp.array(my), # my - jnp.array(iys), # iys - diffrax.Kvaerno5(), # solver - diffrax.PIDController(atol=ATOL_SIM, rtol=RTOL_SIM), # controller - diffrax.RecursiveCheckpointAdjoint(), # adjoint - 2**8, # max_steps - ) + kwargs = { + "ts_init": jnp.array(ts_init), + "ts_dyn": jnp.array(ts_dyn), + "ts_posteq": jnp.array(ts_posteq), + "my": jnp.array(my), + "iys": jnp.array(iys), + "x_preeq": jnp.array([]), + "solver": diffrax.Kvaerno5(), + "controller": diffrax.PIDController(atol=ATOL_SIM, rtol=RTOL_SIM), + "adjoint": diffrax.RecursiveCheckpointAdjoint(), + "max_steps": 2**8, # max_steps + } fun = beartype(jax_model.simulate_condition) for output in ["llh", "x0", "x", "y", "res"]: - oargs = (*args[:-2], diffrax.DirectAdjoint(), 2**8, output) + okwargs = kwargs | { + "adjoint": diffrax.DirectAdjoint(), + "max_steps": 2**8, + "ret": output, + } if sensi_order == amici.SensitivityOrder.none: - r_jax[output] = fun(p, *oargs)[0] + r_jax[output] = fun(p, **okwargs)[0] if sensi_order == amici.SensitivityOrder.first: if output == "llh": - r_jax[f"s{output}"] = jax.grad(fun, has_aux=True)(p, *args)[0] - else: - r_jax[f"s{output}"] = jax.jacfwd(fun, has_aux=True)(p, *oargs)[ + r_jax[f"s{output}"] = jax.grad(fun, has_aux=True)(p, **kwargs)[ 0 ] + else: + r_jax[f"s{output}"] = jax.jacfwd(fun, has_aux=True)( + p, **okwargs + )[0] amici_par_idx = np.array( [jax_model.parameter_ids.index(par_id) for par_id in parameter_ids] diff --git a/tests/benchmark-models/test_petab_benchmark.py b/tests/benchmark-models/test_petab_benchmark.py index 74c84d37a9..6a388f7493 100644 --- a/tests/benchmark-models/test_petab_benchmark.py +++ b/tests/benchmark-models/test_petab_benchmark.py @@ -328,7 +328,7 @@ def test_jax_llh(benchmark_problem): jax_model = import_petab_problem( petab_problem, - model_output_dir=benchmark_outdir / problem_id, + model_output_dir=benchmark_outdir / (problem_id + "_jax"), jax=True, ) jax_problem = JAXProblem(jax_model, petab_problem) @@ -340,7 +340,7 @@ def test_jax_llh(benchmark_problem): [problem_parameters[pid] for pid in jax_problem.parameter_ids] ), ) - if problem_id in problems_for_gradient_check_jax: + if problem_id in problems_for_gradient_check: (llh_jax, _), sllh_jax = eqx.filter_jit( eqx.filter_value_and_grad(run_simulations, has_aux=True) )(jax_problem)