Fix JAXProblem.save()/load() round-trip against current petab.v2 API (#3204)

petab.v2.Problem.to_files() was simplified from a many-keyword-argument
signature to to_files(base_path), which relies on each table/model
already having rel_path/base_path set (normally done by
Problem.from_yaml). JAXProblem.save() still called the old signature
(prefix_path=..., model_file=..., condition_file=..., ...), raising
TypeError. Fixed to set config.filepath = "problem.yaml" (so the
written YAML has the name JAXProblem.load() expects) and call
to_files(base_path=directory).

This was masked by a broader skip-guard in test_serialisation; getting
past that guard also required fixing a prerequisite bug in
import_petab_problem(..., jax=True): it passed the original PEtab v1
problem straight into JAXProblem(), which rejects v1 problems outright.
Upgrade v1->v2 via to_files_generic/from_yaml first, mirroring the
conversion used elsewhere in the codebase.

Removed the now-unneeded try/except skip-guard from test_serialisation
so it actually asserts the round trip. Left the equivalent guard on
test_preequilibration_failure alone, since that covers a separate,
already-known bug in the preequilibration-condition heuristic.

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

Author: FFroehlich

python/sdist/amici/importers/petab/v1/_petab_import.py

@@ -255,6 +255,10 @@ def import_petab_problem(
     )
 
     if jax:
+        import tempfile
+
+        import petab.v2 as petabv2
+
         from amici.sim.jax import JAXProblem
 
         model = model_module.Model()
@@ -263,9 +267,17 @@ def import_petab_problem(
             f"Successfully loaded jax model {model_name} from {output_dir}."
         )
 
+        # JAXProblem requires a PEtab v2 problem; upgrade the v1 problem by
+        # serializing it to a temporary PEtab v1 problem on disk and letting
+        # petab auto-upgrade it (``petab.v2.Problem.from_yaml`` upgrades v1
+        # YAML files via ``petab1to2``).
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            yaml_path = petab_problem.to_files_generic(prefix_path=tmp_dir)
+            petab_problem_v2 = petabv2.Problem.from_yaml(yaml_path)
+
         # Create and return JAXProblem
         logger.info(f"Successfully created JAXProblem for {model_name}.")
-        return JAXProblem(model, petab_problem)
+        return JAXProblem(model, petab_problem_v2)
 
     model = model_module.get_model()
     check_model(amici_model=model, petab_problem=petab_problem)

python/sdist/amici/sim/jax/petab.py

@@ -183,15 +183,12 @@ def save(self, directory: Path):
         :param directory:
             Directory to save the problem to.
         """
-        self._petab_problem.to_files(
-            prefix_path=directory,
-            model_file="model",
-            condition_file="conditions.tsv",
-            measurement_file="measurements.tsv",
-            parameter_file="parameters.tsv",
-            observable_file="observables.tsv",
-            yaml_file="problem.yaml",
-        )
+        if self._petab_problem.config is None:
+            self._petab_problem.config = petabv2.ProblemConfig(
+                format_version="2.0.0"
+            )
+        self._petab_problem.config.filepath = "problem.yaml"
+        self._petab_problem.to_files(base_path=directory)
         shutil.copy(self.model.jax_py_file, directory / "jax_py_file.py")
         with open(directory / "parameters.pkl", "wb") as f:
             eqx.tree_serialise_leaves(f, self)

python/tests/test_jax.py

@@ -313,27 +313,22 @@ def test_serialisation(lotka_volterra):  # noqa: F811
     with TemporaryDirectoryWinSafe(
         prefix=petab_problem.model.model_id
     ) as model_dir:
-        try:
-            jax_problem = import_petab_problem(
-                petab_problem, jax=True, output_dir=model_dir
-            )
-            # change parameters to random values to test serialisation
-            jax_problem.update_parameters(
-                jax_problem.parameters
-                + jr.normal(jr.PRNGKey(0), jax_problem.parameters.shape)
-            )
+        jax_problem = import_petab_problem(
+            petab_problem, jax=True, output_dir=model_dir
+        )
+        # change parameters to random values to test serialisation
+        jax_problem.update_parameters(
+            jax_problem.parameters
+            + jr.normal(jr.PRNGKey(0), jax_problem.parameters.shape)
+        )
 
-            with TemporaryDirectoryWinSafe() as outdir:
-                outdir = Path(outdir)
-                jax_problem.save(outdir)
-                jax_problem_loaded = JAXProblem.load(outdir)
-                assert_allclose(
-                    jax_problem.parameters, jax_problem_loaded.parameters
-                )
-        except (TypeError, NotImplementedError) as err:
-            if "JAXProblem does not support PEtab v1 problems" in str(err):
-                pytest.skip(str(err))
-            raise err
+        with TemporaryDirectoryWinSafe() as outdir:
+            outdir = Path(outdir)
+            jax_problem.save(outdir)
+            jax_problem_loaded = JAXProblem.load(outdir)
+            assert_allclose(
+                jax_problem.parameters, jax_problem_loaded.parameters
+            )
 
 
 @skip_on_valgrind