Merge branch 'main' into jax_sciml

FFroehlich · web-flow · commit 94c407f985da · 2025-10-28T14:54:34.000Z
diff --git a/.github/workflows/deploy_branch.yml b/.github/workflows/deploy_branch.yml
@@ -31,7 +31,7 @@ jobs:
         scripts/buildSdist.sh
 
     - name: "Upload artifact: sdist"
-      uses: actions/upload-artifact@v4
+      uses: actions/upload-artifact@v5
       with:
         name: sdist
         path: python/sdist/dist/amici-*.gz
diff --git a/.github/workflows/test_benchmark_collection_models.yml b/.github/workflows/test_benchmark_collection_models.yml
@@ -84,7 +84,7 @@ jobs:
       run: |
           python3 tests/benchmark_models/evaluate_benchmark.py
 
-    - uses: actions/upload-artifact@v4
+    - uses: actions/upload-artifact@v5
       with:
         name: computation-times-${{ matrix.python-version }}-${{ matrix.extract_subexpressions }}
         path: |
diff --git a/.github/workflows/test_performance.yml b/.github/workflows/test_performance.yml
@@ -55,7 +55,7 @@ jobs:
         AMICI_IMPORT_NPROCS=2 check_time.sh petab_import python tests/performance/test.py import
 
     - name: "Upload artifact: CS_Signalling_ERBB_RAS_AKT_petab"
-      uses: actions/upload-artifact@v4
+      uses: actions/upload-artifact@v5
       with:
         name: model_performance_test
         path: model_performance_test
diff --git a/.github/workflows/test_sbml_semantic_test_suite.yml b/.github/workflows/test_sbml_semantic_test_suite.yml
@@ -47,7 +47,7 @@ jobs:
     - run: AMICI_PARALLEL_COMPILE="" ./scripts/run-SBMLTestsuite.sh ${{ matrix.cases }}
 
     - name: "Upload artifact: SBML semantic test suite results"
-      uses: actions/upload-artifact@v4
+      uses: actions/upload-artifact@v5
       with:
         name: amici-semantic-results-${{ matrix.cases }}
         path: tests/sbml/amici-semantic-results
diff --git a/.github/workflows/test_sbml_semantic_test_suite_jax.yml b/.github/workflows/test_sbml_semantic_test_suite_jax.yml
@@ -44,7 +44,7 @@ jobs:
     - run: ./scripts/run-SBMLTestsuite.sh --jax ${{ matrix.cases }}
 
     - name: "Upload artifact: SBML semantic test suite results"
-      uses: actions/upload-artifact@v4
+      uses: actions/upload-artifact@v5
       with:
         name: amici-semantic-results-jax-${{ matrix.cases }}
         path: tests/sbml/amici-semantic-results-jax
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -61,10 +61,19 @@ See also our [versioning policy](https://amici.readthedocs.io/en/latest/versioni
   This is a wrapper for both `amici.run_simulation` and
   `amici.run_simulations`, depending on the type of the `edata` argument.
   It also supports passing some `Solver` options as keyword arguments.
-* `amici.ModelPtr` now supports sufficient pickling for use in
-  multi-processing contexts. This works only if the amici-generated model
-  package exists in the same file system location and does not change until
-  unpickling.
+* Improved `pickle` support for `amici.{ModelPtr,Solver,ExpData`.
+  Note that AMICI's pickling support is only intended for short-term storage
+  or inter-process communication.
+  Reading pickled objects after updating AMICI or the model code will almost
+  certainly fail.
+  * `amici.ModelPtr` now supports sufficient pickling for use in
+    multi-processing contexts. This works only if the amici-generated model
+    package exists in the same file system location and does not change until
+    unpickling.
+  * `amici.Solver` is now picklable if amici was built with HDF5 support.
+    This only works on shared file systems, as the solver state is stored in a
+    temporary HDF5 file.
+  * `amici.ExpData` is now picklable.
 
 ## v0.X Series
 
diff --git a/include/amici/solver.h b/include/amici/solver.h
@@ -75,6 +75,12 @@ class Solver {
      */
     SUNContext get_sun_context() const;
 
+    /**
+     * @brief Get the name of this class.
+     * @return Class name.
+     */
+    virtual std::string get_class_name() const = 0;
+
     /**
      * @brief runs a forward simulation until the specified timepoint
      *
diff --git a/include/amici/solver_cvodes.h b/include/amici/solver_cvodes.h
@@ -38,6 +38,8 @@ class CVodeSolver : public Solver {
      */
     Solver* clone() const override;
 
+    std::string get_class_name() const override {return "CVodeSolver"; };
+
     void reinit(
         realtype t0, AmiVector const& yy0, AmiVector const& yp0
     ) const override;
diff --git a/include/amici/solver_idas.h b/include/amici/solver_idas.h
@@ -35,6 +35,8 @@ class IDASolver : public Solver {
      */
     Solver* clone() const override;
 
+    std::string get_class_name() const override {return "IDASolver"; };
+
     void reinit_post_process_f(realtype tnext) const override;
 
     void reinit_post_process_b(realtype tnext) const override;
diff --git a/python/sdist/amici/swig_wrappers.py b/python/sdist/amici/swig_wrappers.py
@@ -397,3 +397,52 @@ def file_checksum(
         for chunk in iter(lambda: f.read(chunk_size), b""):
             h.update(chunk)
     return h.hexdigest()
+
+
+def restore_edata(
+    init_args: Sequence,
+    simulation_parameter_dict: dict[str, Any],
+) -> amici_swig.ExpData:
+    """
+    Recreate an ExpData instance.
+
+    For use in ExpData.__reduce__.
+    """
+    edata = amici_swig.ExpData(*init_args)
+
+    edata.pscale = amici.parameter_scaling_from_int_vector(
+        simulation_parameter_dict.pop("pscale")
+    )
+    for key, value in simulation_parameter_dict.items():
+        if key == "timepoints":
+            # timepoints are set during ExpData construction
+            continue
+        assert hasattr(edata, key)
+        setattr(edata, key, value)
+    return edata
+
+
+def restore_solver(cls: type, cls_name: str, hdf5_file: str) -> Solver:
+    """
+    Recreate a Solver or SolverPtr instance from an HDF5 file.
+
+    For use in Solver.__reduce__.
+
+    :param cls:
+        Class of the original object ({CVode,IDA}Solver or SolverPtr).
+    :param cls_name:
+        Name of the (pointed to) solver class ("CVodeSolver" or "IDASolver").
+    :param hdf5_file:
+        HDF5 file from which to read the solver settings.
+    """
+    if cls_name == "CVodeSolver":
+        solver = amici.CVodeSolver()
+    elif cls_name == "IDASolver":
+        solver = amici.IDASolver()
+    else:
+        raise ValueError(f"Unknown solver class name: {cls_name}")
+
+    if not issubclass(cls, Solver):
+        solver = cls(solver)
+    read_solver_settings_from_hdf5(hdf5_file, solver)
+    return solver
diff --git a/python/tests/test_swig_interface.py b/python/tests/test_swig_interface.py
@@ -697,3 +697,40 @@ def test_pickle_model(sbml_example_presimulation_module):
         model.get_steady_state_sensitivity_mode()
         != model_pickled.get_steady_state_sensitivity_mode()
     )
+
+
+def test_pickle_edata():
+    ny = 2
+    nz = 3
+    ne = 4
+    nt = 5
+    edata = amici.ExpData(ny, nz, ne, range(nt))
+    edata.set_observed_data(list(np.arange(ny * nt, dtype=float)))
+    edata.pscale = amici.parameter_scaling_from_int_vector(
+        [amici.ParameterScaling.log10] * 5
+    )
+
+    edata_pickled = pickle.loads(pickle.dumps(edata))
+    assert edata == edata_pickled
+
+
+@pytest.mark.skipif(
+    not amici.hdf5_enabled,
+    reason="AMICI build without HDF5 support",
+)
+def test_pickle_solver():
+    for solver in (
+        amici.CVodeSolver(),
+        amici.IDASolver(),
+        amici.SolverPtr(amici.CVodeSolver()),
+        amici.SolverPtr(amici.IDASolver()),
+    ):
+        solver.set_max_steps(1234)
+        solver.set_sensitivity_order(amici.SensitivityOrder.first)
+        solver_pickled = pickle.loads(pickle.dumps(solver))
+        assert type(solver) is type(solver_pickled)
+        assert solver.get_max_steps() == solver_pickled.get_max_steps()
+        assert (
+            solver.get_sensitivity_order()
+            == solver_pickled.get_sensitivity_order()
+        )
diff --git a/python/tests/valgrind-python.supp b/python/tests/valgrind-python.supp
@@ -1136,6 +1136,17 @@
    ...
 }
 
+{
+   libsbml
+   Memcheck:Leak
+   match-leak-kinds: definite
+   fun:malloc
+   fun:_PyObject_New
+   fun:PyInit__libsbml
+   fun:_PyImport_RunModInitFunc
+   ...
+}
+
 {
    h5py_3.15
    Memcheck:Leak
diff --git a/swig/amici.i b/swig/amici.i
@@ -155,6 +155,23 @@ wrap_unique_ptr(ExpDataPtr, amici::ExpData)
 %naturalvar amici::SimulationParameters::reinitialization_state_idxs_sim;
 %naturalvar amici::SimulationParameters::reinitialization_state_idxs_presim;
 
+%extend amici::SimulationParameters {
+%pythoncode %{
+    def __iter__(self):
+        for attr_name in dir(self):
+            if (
+                not attr_name.startswith('_')
+                and attr_name not in ("this", "thisown")
+                and not callable(attr_val := getattr(self, attr_name))
+            ):
+                if isinstance(attr_val, (DoubleVector, ParameterScalingVector)):
+                    yield attr_name, tuple(attr_val)
+                else:
+                    yield attr_name, attr_val
+%}
+}
+
+
 // DO NOT IGNORE amici::SimulationParameters, amici::ModelDimensions, amici::CpuTimer
 %ignore amici::ModelContext;
 %ignore amici::ContextManager;
diff --git a/swig/edata.i b/swig/edata.i
@@ -102,6 +102,27 @@ def __deepcopy__(self, memo):
     # invoke copy constructor
     return type(self)(self)
 
+def __reduce__(self):
+    from amici.swig_wrappers import restore_edata
+
+    return (
+        restore_edata,
+        (
+            # ExpData ctor arguments
+            (
+                self.nytrue(),
+                self.nztrue(),
+                self.nmaxevent(),
+                self.get_timepoints(),
+                self.get_observed_data(),
+                self.get_observed_data_std_dev(),
+                self.get_observed_events(),
+                self.get_observed_events_std_dev(),
+            ),
+            dict(self)
+        ),
+        {}
+    )
 %}
 };
 %extend std::unique_ptr<amici::ExpData> {
diff --git a/swig/solver.i b/swig/solver.i
@@ -101,17 +101,46 @@ def _solver_repr(self: "Solver"):
         f"  state_ordering: {self.get_state_ordering()}",
         ">"
     ])
+
+def _solver_reduce(self: "Solver"):
+    """
+    For now, we just store solver settings in a temporary HDF5 file.
+    This is sufficient for multiprocessing use cases, but will not survive
+    reboots and will not work in distributed (MPI) settings.
+    This requires that amici was compiled with HDF5 support.
+    """
+    from amici.swig_wrappers import restore_solver
+    from tempfile import NamedTemporaryFile
+    from amici import write_solver_settings_to_hdf5
+    import os
+    with NamedTemporaryFile(suffix=".h5", delete=False) as tmpfile:
+        tmpfilename = tmpfile.name
+        write_solver_settings_to_hdf5(self, tmpfilename)
+
+    return (
+        restore_solver,
+        (self.__class__, self.get_class_name(), tmpfilename,),
+    )
+
 %}
 %extend amici::CVodeSolver {
 %pythoncode %{
 def __repr__(self):
     return _solver_repr(self)
+
+def __reduce__(self):
+    return _solver_reduce(self)
+
 %}
 };
 %extend amici::IDASolver {
 %pythoncode %{
 def __repr__(self):
     return _solver_repr(self)
+
+def __reduce__(self):
+    return _solver_reduce(self)
+
 %}
 };
 
@@ -122,13 +151,20 @@ def __repr__(self):
 
 def __deepcopy__(self, memo):
     return self.clone()
+
+def __reduce__(self):
+    return _solver_reduce(self)
 %}
 };
 
 %extend amici::Solver {
 %pythoncode %{
 def __deepcopy__(self, memo):
     return self.clone()
+
+def __reduce__(self):
+    return _solver_reduce(self)
+
 %}
 };
 
