Merge pull request #15 from kirkegaardlab/fix_jax_deprecation

fix: replace deprecated jax.numpy.trapz

.github/workflows/integration_test.yml

@@ -14,10 +14,10 @@ jobs:
     steps:
       - uses: actions/checkout@v2
 
-      - name: Set up Python 3.8
+      - name: Set up Python 3.10
         uses: actions/setup-python@v2
         with:
-          python-version: '3.8'
+          python-version: '3.10'
           architecture: 'x64'
 
       - name: apt-get

celegans/simulation.py

@@ -6,6 +6,7 @@
 
 import jax
 import jax.numpy as jnp
+from jax.scipy.integrate import trapezoid
 
 
 def _theta(t, s, params):
@@ -123,9 +124,9 @@ def solve(t, u, X, ds, alpha):
     fx = Ut * tx[jnp.newaxis] + alpha * Un * nx[jnp.newaxis]
     fy = Ut * ty[jnp.newaxis] + alpha * Un * ny[jnp.newaxis]
 
-    Fx = jnp.trapz(fx, dx=ds)
-    Fy = jnp.trapz(fy, dx=ds)
-    Tau = jnp.trapz(x * fy - y * fx, dx=ds)
+    Fx = trapezoid(fx, dx=ds)
+    Fy = trapezoid(fy, dx=ds)
+    Tau = trapezoid(x * fy - y * fx, dx=ds)
 
     b = -jnp.array([Fx[0], Fy[0], Tau[0]])
     A = jnp.array([Fx[1:], Fy[1:], Tau[1:]])

examples/detect.py

@@ -2,6 +2,12 @@
 import deeptangle as dt
 import matplotlib.pyplot as plt
 from skimage.exposure import equalize_adapthist
+import numpy
+
+# scikit-video uses deprecated numpy.float, numpy.int
+# hacky fix: https://github.com/scikit-video/scikit-video/issues/154
+numpy.float = numpy.float64
+numpy.int = numpy.int_
 import skvideo.io
 
 

examples/track.py

@@ -9,6 +9,12 @@
 import matplotlib.pyplot as plt
 import numpy as np
 from skimage.exposure import equalize_adapthist
+
+# scikit-video uses deprecated numpy.float, numpy.int
+# hacky fix: https://github.com/scikit-video/scikit-video/issues/154
+import numpy
+numpy.float = numpy.float64
+numpy.int = numpy.int_
 import skvideo.io
 
 import deeptangle as dt

requirements.txt

@@ -3,11 +3,12 @@ scikit-image
 scikit-video
 optax
 chex
-jax
+jax>=0.4.16
+jaxlib>=0.4.20
 dm-pix
 scikit-learn
-numpy==1.21.6
-numba==0.55
+numpy>=1.21.6
+numba>=0.55
 matplotlib
-https://github.com/alonfnt/dm-haiku/archive/refs/heads/avg_pool_perf.zip
+dm-haiku
 trackpy