Fix CI test failures from jax/numpy deprecation warnings

conda-env.yml

@@ -12,7 +12,7 @@ dependencies:
   - hdf5!=1.12.2 # https://github.com/SciTools/iris/issues/5187 and https://github.com/pydata/xarray/issues/7549
   - jax>=0.4.19
   - numba>=0.56
-  - numpy>=1.23
+  - numpy>=1.23,<2.5  # gdal_calc.py uses a numpy.shape= assignment deprecated in 2.5; fixed upstream (OSGeo/gdal#15127) but not yet in a GDAL release. Remove cap once GDAL ships that fix.
   - pydantic>=2.1
   - pyproj>=3.3
   - rasterio>=1.3

requirements.txt

@@ -2,7 +2,7 @@
 h5py>=3.6
 jax>=0.4.19
 numba>=0.56
-numpy>=1.23
+numpy>=1.23,<2.5  # gdal_calc.py uses a numpy.shape= assignment deprecated in 2.5; fixed upstream (OSGeo/gdal#15127) but not yet in a GDAL release. Remove cap once GDAL ships that fix.
 opera-utils>=0.25.9
 pydantic>=2.1
 pyproj>=3.3

src/dolphin/phase_link/crlb.py

@@ -184,7 +184,10 @@ def _crlb_from_x(
     # Σ = inverse of FIM
     if fim_jitter != 0.0:
         FIM = FIM + fim_jitter * eyeN1
-    Sigma = solve(FIM, eyeN1, assume_a="pos")
+    # Use `inv` directly (rather than `solve(FIM, eyeN1, ...)`) since passing an
+    # identity matrix as the RHS to request a full inverse hits an ambiguous
+    # batched-solve code path in jax's `cho_solve` that triggers a FutureWarning.
+    Sigma = jnp.linalg.inv(FIM)
     sig = jnp.sqrt(jnp.diagonal(Sigma, axis1=-2, axis2=-1))
     return jnp.insert(sig, reference_idx, 0.0, axis=-1)
 
@@ -238,7 +241,8 @@ def compute_crlb_jax(
     abs_G_safe = abs_G + gamma_jitter * eyeNb
     abs_G_safe = jnp.where(is_zero_block, eyeNb, abs_G_safe)
 
-    abs_G_inv = solve(abs_G_safe, eyeNb, assume_a="pos")
+    # `inv` instead of `solve(abs_G_safe, eyeNb, ...)`; see note in `_crlb_from_x`.
+    abs_G_inv = jnp.linalg.inv(abs_G_safe)
 
     # Build X once and do the inverse-free CRLB from X
     X = _build_fisher_from_abs_gamma(abs_G, abs_G_inv, num_looks)