Add test for SDE problem with de.jit

diffeqpy/tests/test_sde.py

@@ -35,3 +35,29 @@ def g(du,u,p,t):
     p = [10.0,28.0,2.66]
     prob = de.SDEProblem(numba_f, numba_g, u0, tspan, p)
     sol = de.solve(prob)
+
+
+def test_jit():
+
+    def f(du, u, p, t):
+        x, y, z = u
+        sigma, rho, beta = p
+        du[0] = sigma * (y - x)
+        du[1] = x * (rho - z) - y
+        du[2] = x * y - beta * z
+
+    def g(du, u, p, t):
+        du[0] = 0.3 * u[0]
+        du[1] = 0.3 * u[1]
+        du[2] = 0.3 * u[2]
+
+    u0 = [1.0, 0.0, 0.0]
+    tspan = (0.0, 100.0)
+    p = [10.0, 28.0, 2.66]
+    prob = de.jit(de.SDEProblem(f, g, u0, tspan, p))
+    sol = de.solve(prob)
+    assert sol.t[-1] == tspan[-1], f"Solver did not reach the final time. Last time: {sol.t[-1]}"
+    assert len(sol.u) > 0, "Solution is empty."
+    assert all(
+        abs(sol.u[i][j]) < float("inf") for j in range(len(u0)) for i in range(len(sol.t))
+    ), "Solution contains non-finite values."