Fix catastrophic cancellation in VanGenuchten near theta_range[0]

src/frontx/models.py

@@ -4,6 +4,7 @@
 
 import equinox as eqx
 import jax
+import jax.numpy as jnp
 import numpy as np
 
 from . import Param
@@ -275,7 +276,7 @@ def _kr(
         /,
     ) -> float | jax.Array | np.ndarray[tuple[int], np.dtype[np.floating | np.integer]]:
         Se = self._Se(theta)
-        return Se**self.l * (1 - (1 - Se ** (1 / self._m)) ** self._m) ** 2
+        return Se**self.l * jnp.expm1(self._m * jnp.log1p(-(Se ** (1 / self._m)))) ** 2
 
 
 class LETxs(_RichardsModel):

tests/test_models.py

@@ -0,0 +1,37 @@
+import jax
+import jax.numpy as jnp
+import pytest
+
+from frontx.models import VanGenuchten
+
+jax.config.update("jax_enable_x64", True)
+
+
+def test_van_genuchten_near_residual() -> None:
+    D = VanGenuchten(n=1.1)
+
+    # Reference values computed at high precision.
+    theta = 1e-3
+    expected = (
+        2.613452611709472e-36,
+        3.0054705034658917e-32,
+        3.1557440286391852e-28,
+    )
+
+    assert D(theta) == pytest.approx(expected[0], rel=1e-10, abs=0.0)
+    assert jax.grad(D)(theta) == pytest.approx(expected[1], rel=1e-10, abs=0.0)
+    assert jax.grad(jax.grad(D))(theta) == pytest.approx(
+        expected[2], rel=1e-10, abs=0.0
+    )
+
+
+def test_van_genuchten_kr_near_residual_float32() -> None:
+    model = VanGenuchten(n=1.5)
+
+    theta = jnp.asarray(1e-3, dtype=jnp.float32)
+
+    assert model._kr(theta) == pytest.approx(
+        3.5136418469739605e-21,
+        rel=1e-6,
+        abs=0.0,
+    )