From efbf4a77beca9877a5295766f690c46ef619400a Mon Sep 17 00:00:00 2001
From: Eric Giguere <eric.giguere@gmail.com>
Date: Thu, 11 Jan 2024 13:58:48 -0500
Subject: [PATCH] filter in tests

---
 .github/workflows/tests.yml  |  2 +-
 .github/workflows/weekly.yml |  2 +-
 src/qutip_jax/ode.py         | 27 ++++++++++-----------------
 3 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index bf5426a..dc5c5ea 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -57,7 +57,7 @@ jobs:
         # to see if something hung.
         timeout-minutes: 60
         run: |
-          pytest --durations=0 --durations-min=1.0 --verbosity=1 --cov=qutip_jax --cov-report= --color=yes -W ignore::UserWarning:qutip
+          pytest --durations=0 --durations-min=1.0 --verbosity=1 --cov=qutip_jax --cov-report= --color=yes -W ignore::UserWarning:qutip -W "ignore:Complex dtype:UserWarning"
           # Above flags are:
           #  --durations=0 --durations-min=1.0
           #     at the end, show a list of all the tests that took longer than a
diff --git a/.github/workflows/weekly.yml b/.github/workflows/weekly.yml
index 3cd7ade..260ec2c 100644
--- a/.github/workflows/weekly.yml
+++ b/.github/workflows/weekly.yml
@@ -64,7 +64,7 @@ jobs:
         # to see if something hung.
         timeout-minutes: 60
         run: |
-          pytest --durations=0 --durations-min=1.0 --verbosity=1 --color=yes -W ignore::UserWarning:qutip
+          pytest --durations=0 --durations-min=1.0 --verbosity=1 --color=yes -W ignore::UserWarning:qutip -W "ignore:Complex dtype:UserWarning"
           # Above flags are:
           #  --durations=0 --durations-min=1.0
           #     at the end, show a list of all the tests that took longer than a
diff --git a/src/qutip_jax/ode.py b/src/qutip_jax/ode.py
index 0b9ccdf..10a7448 100644
--- a/src/qutip_jax/ode.py
+++ b/src/qutip_jax/ode.py
@@ -1,5 +1,4 @@
 import diffrax
-import warnings
 from qutip.solver.integrator import Integrator
 import jax
 import jax.numpy as jnp
@@ -65,22 +64,16 @@ def get_state(self, copy=False):
         return self.t, JaxArray(_float2cplx(self.state))
 
     def integrate(self, t, copy=False, **kwargs):
-        with warnings.catch_warnings():
-            # Diffrax added partial support for complex number, but raise a
-            # warning when it find a complex anywhere in the tree.
-            warnings.filterwarnings("ignore",
-                message="Complex dtype support is work in progress,"
-            )
-            sol = diffrax.diffeqsolve(
-                self.ODEsystem,
-                t0=self.t,
-                t1=t,
-                y0=self.state,
-                saveat=diffrax.SaveAt(t1=True, solver_state=True),
-                solver_state=self.solver_state,
-                args=(self.system, kwargs),
-                **self._options,
-            )
+        sol = diffrax.diffeqsolve(
+            self.ODEsystem,
+            t0=self.t,
+            t1=t,
+            y0=self.state,
+            saveat=diffrax.SaveAt(t1=True, solver_state=True),
+            solver_state=self.solver_state,
+            args=(self.system, kwargs),
+            **self._options,
+        )
         self.t = t
         self.state = sol.ys[0, :]
         self.solver_state = sol.solver_state