diff --git a/cognax/decisions/wfpt.py b/cognax/decisions/wfpt.py index 7be3c0f..f6562d6 100644 --- a/cognax/decisions/wfpt.py +++ b/cognax/decisions/wfpt.py @@ -237,20 +237,22 @@ def log_prob(self, value): RTs, w = jnp.broadcast_arrays(RTs, w) p = vmap_n(fnorm, n_times=RTs.ndim, x=(RTs - self.t0) / self.a**2, w=w) - return jnp.log( - jnp.exp( - jnp.log(p) - + ( - (self.a * w * self.v_scale) ** 2 - - 2 * self.a * v_loc * w - - (v_loc**2) * RTs - ) - / (2 * (self.v_scale**2) * (RTs - self.t0) + 2) + # decision time; clamped so the terms below stay finite before t0 + t = jnp.maximum(RTs - self.t0, 0.0) + log_p = ( + jnp.log(p) + + ( + (self.a * w * self.v_scale) ** 2 + - 2 * self.a * v_loc * w + - (v_loc**2) * t ) - / jnp.sqrt((self.v_scale**2) * (RTs - self.t0) + 1) - / (self.a**2) + / (2 * (self.v_scale**2) * t + 2) + - 0.5 * jnp.log1p((self.v_scale**2) * t) + - 2 * jnp.log(self.a) ) + return jnp.where(RTs > self.t0, log_p, -jnp.inf) + def __init__( self, v_loc, v_scale, a, w, t0, dt=0.01, deadline=5.0, *, validate_args=None ): diff --git a/test/test_dists.py b/test/test_dists.py index 6a72d27..d85c447 100644 --- a/test/test_dists.py +++ b/test/test_dists.py @@ -25,8 +25,8 @@ def get_choice_RTs(n_choice, dt, max_RT): - """choice_RTs with RTs from 0 -> max_RTs for each choice""" - t_range = np.arange(0, max_RT + dt, dt) + """choice_RTs with RTs from dt -> max_RT for each choice (RT = 0 is out of support)""" + t_range = np.arange(dt, max_RT + dt, dt) choices = np.repeat(np.arange(n_choice), repeats=len(t_range)) RTs = np.concatenate([t_range] * n_choice) choice_RTs = np.vstack([choices, RTs]).T @@ -54,6 +54,35 @@ def test_wfpt_integrate_to_one(): assert np.isclose(integral, 1, atol=0.01) +@pytest.mark.parametrize( + "v_loc, v_scale, t0", + [(1.0, 0.5, 0.0), (1.0, 0.5, 0.5), (2.0, 1.0, 0.8), (-1.0, 1.5, 0.8)], +) +def test_wfpt_normal_drift_integrate_to_one(v_loc, v_scale, t0): + dt = 0.0001 + choice_RTs = get_choice_RTs(n_choice=2, dt=dt, max_RT=30) + + wfpt = WFPTNormalDrift(v_loc=v_loc, v_scale=v_scale, a=1.5, w=0.5, t0=t0) + + probs = np.exp(wfpt.log_prob(value=choice_RTs)) + integral = np.sum(probs * dt) + + assert np.isclose(integral, 1, atol=0.01) + + +@pytest.mark.parametrize("t0", [0.0, 0.5]) +def test_wfpt_normal_drift_matches_wfpt_as_v_scale_vanishes(t0): + choice_RTs = get_choice_RTs(n_choice=2, dt=0.05, max_RT=5) + choice_RTs = choice_RTs[choice_RTs[:, 1] > t0] + + log_p = WFPTNormalDrift(v_loc=1.0, v_scale=1e-6, a=1.5, w=0.4, t0=t0).log_prob( + choice_RTs + ) + expected = WFPT(v=1.0, a=1.5, w=0.4, t0=t0).log_prob(choice_RTs) + + np.testing.assert_allclose(log_p, expected, rtol=1e-6) + + @pytest.mark.parametrize( "timer_args", [