From 1613aa245d437a40a21ac9bdb66ead43be6bdf39 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= <fabian@schaluck.com>
Date: Thu, 5 Dec 2024 11:11:27 +0000
Subject: [PATCH] separate default dirs for jax/cpp, honour model dir/name

---
 python/examples/example_jax_petab/ExampleJaxPEtab.ipynb | 1 -
 python/sdist/amici/petab/import_helpers.py              | 5 +----
 python/sdist/amici/petab/petab_import.py                | 7 ++-----
 python/sdist/amici/petab/sbml_import.py                 | 9 ++++++---
 4 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/python/examples/example_jax_petab/ExampleJaxPEtab.ipynb b/python/examples/example_jax_petab/ExampleJaxPEtab.ipynb
index fe7b5d8830..a77cf64c69 100644
--- 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
index daa902efb0..57bc551205 100644
--- 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
index 63bade9bb8..24cb21a466 100644
--- 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
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))