Return amici.Model from sbml2amici, pysb2amici

CHANGELOG.md

@@ -74,6 +74,9 @@ See also our [versioning policy](https://amici.readthedocs.io/en/latest/versioni
     This only works on shared file systems, as the solver state is stored in a
     temporary HDF5 file.
   * `amici.ExpData` is now picklable.
+* The import function `sbml2amici`, `pysb2amici`, and `antimony2amici` now
+  return an instance of the generated model class if called with `compile=True`
+  (default).
 
 ## v0.X Series
 

python/sdist/amici/pysb_import.py

@@ -22,6 +22,8 @@
 import pysb.pattern
 import sympy as sp
 
+import amici
+
 from .de_export import (
     Constant,
     DEExporter,
@@ -161,7 +163,7 @@ def pysb2amici(
     generate_sensitivity_code: bool = True,
     model_name: str | None = None,
     pysb_model_has_obs_and_noise: bool = False,
-):
+) -> amici.Model | None:
     r"""
     Generate AMICI C++ files for the provided model.
 
@@ -238,6 +240,10 @@ def pysb2amici(
     :param pysb_model_has_obs_and_noise:
         if set to ``True``, the pysb model is expected to have extra
         observables and noise variables added
+
+    :return:
+        If `compile` is `True` and compilation was successful, an instance
+        of the generated model class, otherwise `None`.
     """
     if observation_model is None:
         observation_model = []
@@ -275,6 +281,14 @@ def pysb2amici(
     if compile:
         exporter.compile_model()
 
+        from . import import_model_module
+
+        return import_model_module(
+            module_name=model_name, module_path=output_dir
+        ).get_model()
+
+    return None
+
 
 @log_execution_time("creating ODE model", logger)
 def ode_model_from_pysb_importer(

python/sdist/amici/sbml_import.py

@@ -26,6 +26,8 @@
 from sympy.logic.boolalg import Boolean, BooleanFalse, BooleanTrue
 from sympy.matrices.dense import MutableDenseMatrix
 
+import amici
+
 from . import has_clibs
 from .constants import SymbolId
 from .de_export import (
@@ -290,7 +292,7 @@ def sbml2amici(
         cache_simplify: bool = False,
         generate_sensitivity_code: bool = True,
         hardcode_symbols: Sequence[str] = None,
-    ) -> None:
+    ) -> amici.Model | None:
         """
         Generate and compile AMICI C++ files for the model provided to the
         constructor.
@@ -379,6 +381,10 @@ def sbml2amici(
             Their values cannot be changed anymore after model import.
             Currently, only parameters that are not targets of rules or
             initial assignments are supported.
+
+        :return:
+            If `compile` is `True` and compilation was successful, an instance
+            of the generated model class, otherwise `None`.
         """
         set_log_level(logger, verbose)
 
@@ -413,6 +419,14 @@ def sbml2amici(
                 )
             exporter.compile_model()
 
+            from . import import_model_module
+
+            return import_model_module(
+                module_name=model_name, module_path=output_dir
+            ).get_model()
+
+        return None
+
     def sbml2jax(
         self,
         model_name: str,

python/tests/test_pysb.py

@@ -343,17 +343,13 @@ def test_heavyside_and_special_symbols():
     )
 
     with TemporaryDirectoryWinSafe(prefix=model.name) as outdir:
-        pysb2amici(
+        amici_model = pysb2amici(
             model,
             outdir,
             verbose=True,
             observation_model=[amici.MeasurementChannel("a")],
         )
 
-        model_module = amici.import_model_module(
-            module_name=model.name, module_path=outdir
-        )
-        amici_model = model_module.get_model()
         assert amici_model.ne
 
 

python/tests/test_sbml_import.py

@@ -377,7 +377,7 @@ def test_presimulation_events_and_sensitivities(tempdir):
     from amici.antimony_import import antimony2amici
 
     model_name = "test_presim_events2"
-    antimony2amici(
+    model = antimony2amici(
         """
     some_time = time
     some_time' = 1
@@ -394,9 +394,6 @@ def test_presimulation_events_and_sensitivities(tempdir):
         output_dir=tempdir,
     )
 
-    model_module = import_model_module(model_name, tempdir)
-
-    model = model_module.get_model()
     model.set_timepoints([0, 1, 2])
     edata = amici.ExpData(model)
     edata.t_presim = 2
@@ -992,11 +989,12 @@ def test_import_same_model_name(tempdir):
     if sys.platform == "win32":
         return
 
-    antimony2amici(
-        ant_model_3,
-        model_name=module_name,
-        output_dir=outdir_2,
-    )
+    with pytest.raises(RuntimeError, match="in the same location"):
+        antimony2amici(
+            ant_model_3,
+            model_name=module_name,
+            output_dir=outdir_2,
+        )
 
     with pytest.raises(RuntimeError, match="in the same location"):
         import_model_module(module_name=module_name, module_path=outdir_2)