fix(maths): eliminate cancellation error in exact parameter info-gain weight

pymdp/maths.py

@@ -784,6 +784,11 @@ def _exact_wnorm(A: ArrayLike) -> ArrayLike:
     """
     Implements (-1) * eq. (D.15) in Da Costa et al. ‘Active inference on discrete state-spaces: A synthesis’, Journal of Mathematical Psychology, 2020.
 
+    Note: the direct expansion of eq. (D.15) sums `1/A - 1/sumA` with `digamma(A) - digamma(sumA)`; both terms
+    diverge as A shrinks and are supposed to cancel, which is numerically unstable in finite precision.
+    Using the digamma recurrence `digamma(x+1) = digamma(x) + 1/x` lets those terms cancel algebraically
+    instead, avoiding the cancellation entirely.
+
     Note: Like the legacy SPM implementation this function clips A for numerical stability. However note that if some values of Aare set to zero e.g. by Bayesian model reduction, these are non-zeroed in this calculation, and thus contribute a large amount to the information gain unless these are zeroed when multiplying by beliefs about states and expected observations. In principle, this should be the case.
 
     Parameters
@@ -802,8 +807,7 @@ def _exact_wnorm(A: ArrayLike) -> ArrayLike:
 
     wA = (
         jnp.log(safe_sumA) - jnp.log(safe_A)
-        + 1. / safe_A - 1. / safe_sumA
-        + digamma(safe_A) - digamma(safe_sumA)
+        + digamma(safe_A + 1.) - digamma(safe_sumA + 1.)
     )
 
     return -wA # TODO: minus sign here gives negative info gain for backward compatibility with spm implementation. Later will need to remove minus sign here to get positive info gain and adjust function documentation accordingly.

test/test_param_info_gain_jax.py

@@ -32,24 +32,43 @@ def test_exact_wnorm_mathematical_correctness():
     # Simple 2x2 case with known values
     A = jnp.array([[1.0, 2.0],
                    [3.0, 4.0]])
-    
+
     result = _exact_wnorm(A)
-
-    # Manually compute expected result using the formula
-    # wA = log(A.sum(0)) - log(A) + 1/A - 1/A.sum(0) + digamma(A) - digamma(A.sum(0))
+
+    # Reference: direct expansion of eq. (D.15), stable at this magnitude since A is
+    # far from MINVAL. Mathematically equivalent to the implementation (see the digamma
+    # recurrence note in _exact_wnorm's docstring), so both should agree up to float32
+    # rounding -- the two forms don't round identically since they go through different
+    # sequences of operations.
     A_sum = A.sum(axis=0)  # [4.0, 6.0]
-    
+
     expected = (
         jnp.log(A_sum) - jnp.log(A) +
         1.0 / A - 1.0 / A_sum +
         jnp.array(digamma(A)) - jnp.array(digamma(A_sum))
     )
     expected = -expected  # minus sign in the implementation
-    
-    np.testing.assert_allclose(result, expected, rtol=1e-6, atol=1e-8)
-    
+
+    np.testing.assert_allclose(result, expected, rtol=1e-4, atol=1e-4)
+
     # Verify output shape matches input
-    assert result.shape == A.shape 
+    assert result.shape == A.shape
+
+
+@pytest.mark.parametrize("K", [2, 5, 10, 50])
+@pytest.mark.parametrize("a", [1e-8, 1e-12, 1e-15])
+def test_exact_wnorm_symmetric_tiny_concentrations_match_log_k(K, a):
+    """For a symmetric Dirichlet (all K entries equal to `a`), the exact value of eq. (D.15)
+    as `a -> 0` is `-log(K)`: both `digamma(a+1)` and `digamma(K*a+1)` tend to `digamma(1)`
+    and cancel, leaving only `log(K)`. This is an independent closed-form reference (no
+    digamma evaluation needed at the limit), so it directly catches the cancellation bug:
+    the pre-fix direct-expansion formula drifts further from `-log(K)` as `a` shrinks
+    instead of converging to it.
+    """
+    A = jnp.full((K, 1), a)
+    result = _exact_wnorm(A)
+    expected = -jnp.log(float(K))
+    np.testing.assert_allclose(result, expected, rtol=1e-4, atol=1e-4)
 
 
 # -----------------------------------------------------------------------------