Merge pull request #5 from cgiovanetti/fix-diffrax-deprecation-warning

Fix diffrax deprecation warning

Author: smsharma

.github/workflows/test.yml

@@ -26,4 +26,4 @@ jobs:
       - name: Run tests
         run: |
           cd pytest
-          pytest .
+          pytest -m "not slow"

linx/background.py

@@ -4,7 +4,7 @@
 
 import equinox as eqx
 
-from diffrax import diffeqsolve, ODETerm, Tsit5, PIDController, SaveAt, DiscreteTerminatingEvent 
+from diffrax import diffeqsolve, ODETerm, Tsit5, PIDController, SaveAt, Event
 
 import linx.thermo as thermo
 import linx.const as const 
@@ -109,14 +109,13 @@ def __call__(
 
         Y0 = (lna_init, T_EM_init, T_nu_init)
 
-        def T_EM_check(state, **kwargs): 
-
-            return state.y[1] < T_end
+        def T_EM_check(t, y, args, **kwargs): 
+            return y[1] < T_end
 
         sol = diffeqsolve(
             ODETerm(self.dY), solver, args=(lna_init, rho_extra_init),
             t0=0., t1=jnp.inf, dt0=None, y0=Y0, 
-            saveat=SaveAt(steps=True), discrete_terminating_event = DiscreteTerminatingEvent(T_EM_check),
+            saveat=SaveAt(steps=True), event=Event(T_EM_check),
             stepsize_controller = PIDController(
                 rtol=rtol, atol=atol
             ), 

pytest.ini

@@ -0,0 +1,3 @@
+[pytest]
+markers =
+    slow: marks tests as slow (deselect with '-m "not slow"')

pytest/test_numpyro.py

@@ -1,9 +1,13 @@
 import sys
-sys.path.append('../scripts')
+import os
+import pytest
+# Add absolute path to the scripts directory
+sys.path.append(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'scripts'))
 
 from run_numpyro import run
 
 
+@pytest.mark.slow
 def test_run_numpyro():
     try:
         run(bbn_only=True, n_steps_svi=5, n_warmup_mcmc=5, n_samples_mcmc=5, n_chains=1)

requirements.txt

@@ -5,12 +5,12 @@ cloudpickle
 corner
 cosmopower-jax
 Cython
-diffrax==0.4.1
+diffrax>=0.6.2
 dynesty
 emcee
 equinox
-jax==0.4.28
-jaxlib==0.4.28
+jax>=0.4.38
+jaxlib>=0.4.38
 jaxopt
 jaxtyping
 joblib