Fix JAX test failures related to diffrax/optimistix release (#3129)

* adjust test_jax tols

* increase max steps for jax benchmarks

* pin optax for notebook test

* increase max steps for petab tests

* increase petab max steps again

* even higher max steps for Weber benchmark

* pin optax in docs build

Author: BSnelling

doc/rtd_requirements.txt

@@ -9,7 +9,7 @@ setuptools>=67.7.2
 git+https://github.com/jmuhlich/pysb@22d69a350b472f33d85ba64ffb10b190483c1c98
 # For forward type definition in generate_equinox
 matplotlib>=3.7.1
-optax
+optax==0.2.6
 nbsphinx
 nbformat
 myst-parser

python/sdist/amici/jax/petab.py

@@ -1550,7 +1550,7 @@ def run_simulations(
     steady_state_event: Callable[
         ..., diffrax._custom_types.BoolScalarLike
     ] = diffrax.steady_state_event(),
-    max_steps: int = 2**10,
+    max_steps: int = 2**13,
     ret: ReturnValue | str = ReturnValue.llh,
 ):
     """
@@ -1653,7 +1653,7 @@ def petab_simulate(
     steady_state_event: Callable[
         ..., diffrax._custom_types.BoolScalarLike
     ] = diffrax.steady_state_event(),
-    max_steps: int = 2**10,
+    max_steps: int = 2**13,
 ):
     """
     Run simulations for a problem and return the results as a petab simulation dataframe.

python/tests/test_jax.py

@@ -204,7 +204,7 @@ def check_fields_jax(
         "iy_trafos": jnp.array(iy_trafos),
         "x_preeq": jnp.array([]),
         "solver": diffrax.Kvaerno5(),
-        "controller": diffrax.PIDController(atol=ATOL_SIM, rtol=RTOL_SIM),
+        "controller": diffrax.PIDController(atol=1e-8, rtol=1e-8),
         "root_finder": optimistix.Newton(atol=ATOL_SIM, rtol=RTOL_SIM),
         "adjoint": diffrax.RecursiveCheckpointAdjoint(),
         "steady_state_event": diffrax.steady_state_event(),

scripts/installAmiciSource.sh

@@ -39,7 +39,7 @@ python -m pip install --upgrade pip wheel
 python -m pip install --upgrade pip setuptools cmake_build_extension==0.6.0 numpy petab swig
 python -m pip install git+https://github.com/pysb/pysb@master # for SPM with compartments
 python -m pip install git+https://github.com/patrick-kidger/diffrax@main # for events with direction
-python -m pip install optax # for jax petab notebook
+python -m pip install 'optax<0.2.7' # for jax petab notebook
 AMICI_BUILD_TEMP="${AMICI_PATH}/python/sdist/build/temp" \
   python -m pip install --verbose -e "${AMICI_PATH}/python/sdist[petab,test,vis,jax]" --no-build-isolation
 deactivate

tests/benchmark_models/test_petab_benchmark_jax.py

@@ -10,7 +10,7 @@
 from amici.importers.petab.v1 import (
     import_petab_problem,
 )
-from amici.jax.petab import run_simulations
+from amici.jax.petab import run_simulations, DEFAULT_CONTROLLER_SETTINGS
 from amici.sim.sundials import SensitivityMethod, SensitivityOrder
 from amici.sim.sundials.petab.v1 import (
     LLH,
@@ -113,11 +113,11 @@ def test_jax_llh(benchmark_problem):
         if problem_id == "Weber_BMC2015":
             atol = cur_settings.atol_sim
             rtol = cur_settings.rtol_sim
-            max_steps = 2 * 10**5
+            max_steps = 4 * 10**7
         else:
             atol = 1e-8
             rtol = 1e-8
-            max_steps = 1024
+            max_steps = 2 * 10**5
         beartype(run_simulations)(jax_problem)
         (llh_jax, _), sllh_jax = eqx.filter_value_and_grad(
             run_simulations, has_aux=True
@@ -130,7 +130,10 @@ def test_jax_llh(benchmark_problem):
             ),
         )
     else:
-        llh_jax, _ = beartype(run_simulations)(jax_problem)
+        llh_jax, _ = beartype(run_simulations)(
+            jax_problem,
+            max_steps=2 * 10**5,
+        )
 
     np.testing.assert_allclose(
         llh_jax,