more fixes

pyproject.toml

@@ -9,15 +9,19 @@ authors = [
 requires-python = ">=3.12"
 dependencies = [
     "beartype>=0.21.0",
+    "dask>=2024.8.0",
     "equinox>=0.13.0",
-    "jax>=0.7.0",
+    "jax>=0.7.2",
     "jaxtyping>=0.3.2",
     "loguru>=0.7.3",
+    "optimistix>=0.0.10",
+    "sgkit>=0.10.0",
+    "sparse>=0.17.0",
 ]
 
 [project.optional-dependencies]
 gpu = [
-    "jax[cuda12]>=0.7.0",
+    "jax[cuda12]>=0.7.2",
     "cuda-bindings>=12.3,<13.0",
     "cuda-core>=0.3.1",
     "cuda-python>=12.2,<13",
@@ -30,3 +34,13 @@ dev = [
 [build-system]
 requires = ["uv_build>=0.8.4,<0.9.0"]
 build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "bio2zarr>=0.1.6",
+    "jupyter>=1.1.1",
+    "matplotlib>=3.10.6",
+    "notebook>=7.4.7",
+    "pytest>=8.4.2",
+    "stdpopsim>=0.3.0",
+]

src/phlashlib/iicr.py

@@ -3,6 +3,7 @@
 import equinox as eqx
 import jax
 import jax.numpy as jnp
+import optimistix as optx
 from jaxtyping import ArrayLike, Float, Scalar, ScalarLike
 
 

src/phlashlib/transition.py

@@ -50,7 +50,7 @@ def transition_matrix(eta: PiecewiseConstant, rho: ScalarLike, n: int = 2) -> ja
     dt0 = jnp.isclose(dt_aug, 0.0)
     dt_safe = jnp.where(dt0, 1.0, dt_aug)
     cr = jnp.repeat(eta.c, 2, axis=0)[:-1]
-    P = jax.vmap(_expQ, (0, 0, None))(2 * dt_safe * rho, dt_aug * cr, n)
+    P = jax.vmap(_expQ, (0, 0, None))(dt_safe * rho, dt_aug * cr, n)
     P = jnp.where(dt0[:, None, None], jnp.eye(3)[None], P)
     Pinf = jnp.array([[0.0, 0.0, 1.0]] * 3)
     P = jnp.concatenate([jnp.eye(3)[None], P, Pinf[None]], 0)