separate default dirs for jax/cpp, honour model dir/name

AMICI-dev · Dec 5, 2024 · 1613aa2 · 1613aa2
1 parent 019efd6
commit 1613aa2
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 13 deletions.
diff --git a/python/examples/example_jax_petab/ExampleJaxPEtab.ipynb b/python/examples/example_jax_petab/ExampleJaxPEtab.ipynb
@@ -489,7 +489,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/petab/import_helpers.py b/python/sdist/amici/petab/import_helpers.py
@@ -138,11 +138,8 @@ 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
-        )
+        model_module = amici.import_model_module(model_name, model_output_dir)
     except ModuleNotFoundError:
         return False
 

diff --git a/python/sdist/amici/petab/petab_import.py b/python/sdist/amici/petab/petab_import.py
@@ -114,7 +114,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)
@@ -160,10 +160,7 @@ def import_petab_problem(
             )
 
     # import model
-    suffix = "_jax" if jax else ""
-    model_module = amici.import_model_module(
-        model_name + suffix, model_output_dir
-    )
+    model_module = amici.import_model_module(model_name, model_output_dir)
 
     if jax:
         model = model_module.Model()

diff --git 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))