Fix nan gradients in analytical likelihood

.pre-commit-config.yaml

@@ -14,4 +14,9 @@ repos:
     rev: v1.10.0 # Use the sha / tag you want to point at
     hooks:
       - id: mypy
-        args: [--no-strict-optional, --ignore-missing-imports]
+        args:
+          [
+            --no-strict-optional,
+            --ignore-missing-imports,
+            --config-file=pyproject.toml,
+          ]

pyproject.toml

@@ -159,6 +159,7 @@ convention = "numpy"
 
 [tool.mypy]
 ignore_missing_imports = true
+exclude = 'tests/*'
 
 [build-system]
 requires = ["poetry-core"]

src/hssm/likelihoods/analytical.py

@@ -9,7 +9,6 @@
 
 import numpy as np
 import pymc as pm
-import pytensor
 import pytensor.tensor as pt
 from numpy import inf
 from pymc.distributions.dist_math import check_parameters
@@ -25,7 +24,7 @@ def k_small(rt: np.ndarray, err: float) -> np.ndarray:
     Parameters
     ----------
     rt
-        A 1D numpy array of flipped R.... T.....s. (0, inf).
+        A 1D numpy array of flipped R.... pt.....s. (0, inf).
     err
         Error bound.
 
@@ -34,9 +33,11 @@ def k_small(rt: np.ndarray, err: float) -> np.ndarray:
     np.ndarray
         A 1D at array of k_small.
     """
-    ks = 2 + pt.sqrt(-2 * rt * pt.log(2 * np.sqrt(2 * np.pi * rt) * err))
-    ks = pt.max(pt.stack([ks, pt.sqrt(rt) + 1]), axis=0)
-    ks = pt.switch(2 * pt.sqrt(2 * np.pi * rt) * err < 1, ks, 2)
+    _a = 2 * pt.sqrt(2 * np.pi * rt) * err < 1
+    _b = 2 + pt.sqrt(-2 * rt * pt.log(2 * pt.sqrt(2 * np.pi * rt) * err))
+    _c = pt.sqrt(rt) + 1
+    _d = pt.max(pt.stack([_b, _c]), axis=0)
+    ks = _a * _d + (1 - _a) * 2
 
     return ks
 
@@ -56,9 +57,11 @@ def k_large(rt: np.ndarray, err: float) -> np.ndarray:
     np.ndarray
         A 1D at array of k_large.
     """
-    kl = pt.sqrt(-2 * pt.log(np.pi * rt * err) / (np.pi**2 * rt))
-    kl = pt.max(pt.stack([kl, 1.0 / (np.pi * pt.sqrt(rt))]), axis=0)
-    kl = pt.switch(np.pi * rt * err < 1, kl, 1.0 / (np.pi * pt.sqrt(rt)))
+    _a = np.pi * rt * err < 1
+    _b = 1.0 / (np.pi * pt.sqrt(rt))
+    _c = pt.sqrt(-2 * pt.log(np.pi * rt * err) / (np.pi**2 * rt))
+    _d = pt.max(pt.stack([_b, _c]), axis=0)
+    kl = _a * _b + (1 - _a) * _b
 
     return kl
 
@@ -81,34 +84,7 @@ def compare_k(rt: np.ndarray, err: float) -> np.ndarray:
     ks = k_small(rt, err)
     kl = k_large(rt, err)
 
-    return ks < kl
-
-
-def get_ks(k_terms: int, fast: bool) -> np.ndarray:
-    """Return an array of ks.
-
-    Returns an array of ks given the number of terms needed to approximate the sum of
-    the infinite series.
-
-    Parameters
-    ----------
-    k_terms
-        number of terms needed
-    fast
-        whether the function is used in the fast of slow expansion.
-
-    Returns
-    -------
-    np.ndarray
-        An array of ks.
-    """
-    ks = (
-        pt.arange(-pt.floor((k_terms - 1) / 2), pt.ceil((k_terms - 1) / 2) + 1)
-        if fast
-        else pt.arange(1, k_terms + 1).reshape((-1, 1))
-    )
-
-    return ks.astype(pytensor.config.floatX)
+    return pt.lt(ks, kl)
 
 
 def ftt01w_fast(tt: np.ndarray, w: float, k_terms: int) -> np.ndarray:
@@ -133,7 +109,10 @@ def ftt01w_fast(tt: np.ndarray, w: float, k_terms: int) -> np.ndarray:
     """
     # Slightly changed the original code to mimic the paper and
     # ensure correctness
-    k = get_ks(k_terms, fast=True)
+    k = pt.arange(
+        -pt.floor((k_terms - 1) / 2.0),
+        pt.ceil((k_terms - 1) / 2.0) + 1.0,
+    )
 
     # A log-sum-exp trick is used here
     y = w + 2 * k.reshape((-1, 1))
@@ -166,7 +145,7 @@ def ftt01w_slow(tt: np.ndarray, w: float, k_terms: int) -> np.ndarray:
     np.ndarray
         The approximated function f(tt|0, 1, w).
     """
-    k = get_ks(k_terms, fast=False)
+    k = pt.arange(1, k_terms + 1).reshape((-1, 1))
     y = k * pt.sin(k * np.pi * w)
     r = -pt.power(k, 2) * pt.power(np.pi, 2) * tt / 2
     p = pt.sum(y * pt.exp(r), axis=0) * np.pi
@@ -208,7 +187,7 @@ def ftt01w(
     p_fast = ftt01w_fast(tt, w, k_terms)
     p_slow = ftt01w_slow(tt, w, k_terms)
 
-    p = pt.switch(lambda_rt, p_fast, p_slow)
+    p = lambda_rt * p_fast + (1.0 - lambda_rt) * p_slow
 
     return p
 
@@ -220,7 +199,7 @@ def logp_ddm(
     z: float,
     t: float,
     err: float = 1e-15,
-    k_terms: int = 20,
+    k_terms: int = 7,
     epsilon: float = 1e-15,
 ) -> np.ndarray:
     """Compute analytical likelihood for the DDM model with `sv`.
@@ -262,15 +241,17 @@ def logp_ddm(
     z_flipped = pt.switch(flip, 1 - z, z)  # transform z if x is upper-bound response
     rt = rt - t
 
-    p = pt.maximum(ftt01w(rt, a, z_flipped, err, k_terms), pt.exp(LOGP_LB))
+    negative_rt = rt <= epsilon
 
-    logp = pt.where(
-        rt <= epsilon,
-        LOGP_LB,
+    tt = negative_rt * epsilon + (1 - negative_rt) * rt
+
+    p = pt.maximum(ftt01w(tt, a, z_flipped, err, k_terms), pt.exp(LOGP_LB))
+
+    logp = negative_rt * LOGP_LB + (1 - negative_rt) * (
         pt.log(p)
         - v_flipped * a * z_flipped
-        - (v_flipped**2 * rt / 2.0)
-        - 2.0 * pt.log(a),
+        - (v_flipped**2 * tt / 2.0)
+        - 2.0 * pt.log(pt.maximum(epsilon, a))
     )
 
     checked_logp = check_parameters(logp, a >= 0, msg="a >= 0")
@@ -333,7 +314,8 @@ def logp_ddm_sdv(
     z_flipped = pt.switch(flip, 1 - z, z)  # transform z if x is upper-bound response
     rt = rt - t
 
-    p = pt.maximum(ftt01w(rt, a, z_flipped, err, k_terms), pt.exp(LOGP_LB))
+    tt = pt.switch(rt <= epsilon, epsilon, rt)
+    p = pt.maximum(ftt01w(tt, a, z_flipped, err, k_terms), pt.exp(LOGP_LB))
 
     logp = pt.switch(
         rt <= epsilon,
@@ -342,11 +324,11 @@ def logp_ddm_sdv(
         + (
             (a * z_flipped * sv) ** 2
             - 2 * a * v_flipped * z_flipped
-            - (v_flipped**2) * rt
+            - (v_flipped**2) * tt
         )
-        / (2 * (sv**2) * rt + 2)
-        - 0.5 * pt.log(sv**2 * rt + 1)
-        - 2 * pt.log(a),
+        / (2 * (sv**2) * tt + 2)
+        - 0.5 * pt.log(sv**2 * tt + 1)
+        - 2 * pt.log(pt.maximum(epsilon, a)),
     )
 
     checked_logp = check_parameters(logp, a >= 0, msg="a >= 0")

tests/test_likelihoods.py

@@ -4,52 +4,28 @@
 old implementation of WFPT from (https://github.com/hddm-devs/hddm)
 """
 
-import math
 from pathlib import Path
 from itertools import product
 
 import numpy as np
+import pandas as pd
 import pymc as pm
+import pytensor
 import pytensor.tensor as pt
 import pytest
-from numpy.random import rand
+
+from pytensor.compile.nanguardmode import NanGuardMode
 
 import hssm
 
 # pylint: disable=C0413
-from hssm.likelihoods.analytical import compare_k, logp_ddm, logp_ddm_sdv
+from hssm.likelihoods.analytical import logp_ddm, logp_ddm_sdv
 from hssm.likelihoods.blackbox import logp_ddm_bbox, logp_ddm_sdv_bbox
 from hssm.distribution_utils import make_likelihood_callable
 
 hssm.set_floatX("float32")
 
 
-def test_kterm(data_ddm):
-    """This function defines a range of kterms and tests results to
-    makes sure they are not equal to infinity or unknown values.
-    """
-    for k_term in range(7, 12):
-        v = (rand() - 0.5) * 1.5
-        sv = 0
-        a = (1.5 + rand()) / 2
-        z = 0.5 * rand()
-        t = rand() * 0.5
-        err = 1e-7
-        logp = logp_ddm_sdv(data_ddm, v, a, z, t, sv, err, k_terms=k_term)
-        logp = sum(logp.eval())
-        assert not math.isinf(logp)
-        assert not math.isnan(logp)
-
-
-def test_compare_k(data_ddm):
-    """This function tests output of decision function."""
-    err = 1e-7
-    data = data_ddm["rt"] * data_ddm["response"]
-    lambda_rt = compare_k(np.abs(data.values), err)
-    assert all(not v for v in lambda_rt.eval())
-    assert data_ddm.shape[0] == lambda_rt.eval().shape[0]
-
-
 # def test_logp(data_fixture):
 #     """
 #     This function compares new and old implementation of logp calculation
@@ -128,13 +104,49 @@ def test_bbox(data_ddm):
     )
 
 
-cav_data = hssm.load_data("cavanagh_theta")
+cav_data: pd.DataFrame = hssm.load_data("cavanagh_theta")
 cav_data_numpy = cav_data[["rt", "response"]].values
 param_matrix = product(
     (0.0, 0.01, 0.05, 0.5), ("analytical", "approx_differentiable", "blackbox")
 )
 
 
+def test_analytical_gradient():
+    v = pt.dvector()
+    a = pt.dvector()
+    z = pt.dvector()
+    t = pt.dvector()
+    sv = pt.dvector()
+    size = cav_data_numpy.shape[0]
+    logp = logp_ddm(cav_data_numpy, v, a, z, t).sum()
+    grad = pt.grad(logp, wrt=[v, a, z, t])
+    grad_func = pytensor.function(
+        [v, a, z, t],
+        grad,
+        mode=NanGuardMode(nan_is_error=True, inf_is_error=True, big_is_error=False),
+    )
+    v_test = np.random.normal(size=size)
+    a_test = np.random.uniform(0.0001, 2, size=size)
+    z_test = np.random.uniform(0.1, 1.0, size=size)
+    t_test = np.random.uniform(0, 2, size=size)
+    sv_test = np.random.uniform(0.001, 1.0, size=size)
+    grad = np.array(grad_func(v_test, a_test, z_test, t_test))
+
+    assert np.all(np.isfinite(grad), axis=None), "Gradient contains non-finite values."
+
+    grad_func_sdv = pytensor.function(
+        [v, a, z, t, sv],
+        pt.grad(logp_ddm_sdv(cav_data_numpy, v, a, z, t, sv).sum(), wrt=[v, a, z, t]),
+        mode=NanGuardMode(nan_is_error=True, inf_is_error=True, big_is_error=False),
+    )
+
+    grad_sdv = np.array(grad_func_sdv(v_test, a_test, z_test, t_test, sv_test))
+
+    assert np.all(
+        np.isfinite(grad_sdv), axis=None
+    ), "Gradient contains non-finite values."
+
+
 @pytest.mark.parametrize("p_outlier, loglik_kind", param_matrix)
 def test_lapse_distribution_cav(p_outlier, loglik_kind):
     true_values = (0.5, 1.5, 0.5, 0.5)