adding a util function to calculate decaying reward-rate

src/aind_dynamic_foraging_basic_analysis/reward_rate.py

@@ -0,0 +1,91 @@
+import numpy as np
+from scipy.signal import fftconvolve
+
+def reward_rate(event_times,
+                t_start=None,
+                t_end=None,
+                dt=1/20, # matching photometry frame rate
+                tau=25.0, # to be optimised
+                kernel="exp", # exp/hyperbolic
+                hyper_p=1.0, # only for hyperbolic 
+                normalize_kernel=True):
+    """
+    Compute a reward-rate timeseries by convolving discrete reward events with a decay kernel.
+
+    Parameters
+    ----------
+    event_times : (N,) array-like
+        1-D array of reward timestamps [s].
+    t_start, t_end : float, optional
+        Output time range [s].  Defaults: t_start=0, t_end=max(event_times).
+    dt : float, default 0.001
+        Time step of the output signal [s].
+    tau : float, default 1.0
+        Time constant of the kernel [s].
+    kernel : {"exp", "hyperbolic"}, default "exp"
+        * "exp": k(t) = exp(-t / tau)                              (classic EWMA)  
+        * "hyperbolic": k(t) = (1 + t / tau)^{-hyper_p}            (hyperbolic discounting)
+    hyper_p : float, default 1.0
+        Power *p* of the hyperbolic kernel. Ignored if kernel="exp".
+    normalize_kernel : bool, default True
+        If True, scale the kernel so its integral equals 1.
+
+    Returns
+    -------
+    t : (T,) ndarray
+        Time axis [s].
+    r_rate : (T,) ndarray
+        Reward rate (events · s⁻¹).
+
+    Notes
+    -----
+    * Kernel is truncated at 5 × tau to keep computation fast.
+    * fftconvolve gives O(N log N) performance.
+    """
+
+    event_times = np.asarray(event_times, dtype=float)
+    if event_times.ndim != 1:
+        raise ValueError("event_times must be a 1-D array of timestamps.")
+
+    # ------------------------------------------------------------------
+    # Time axis: start at 0 so the series is zero until the first reward
+    # ------------------------------------------------------------------
+    if t_start is None:
+        t_start = 0.0
+    if t_end is None:
+        t_end = event_times.max()
+    t = np.arange(t_start, t_end + dt, dt)
+
+    # ------------------------------------------------------------------
+    # Binary event series aligned to the timeline
+    # ------------------------------------------------------------------
+    event_series = np.zeros_like(t)
+    idx = np.searchsorted(t, event_times - 1e-12)
+    idx = idx[(idx >= 0) & (idx < len(t))]
+    event_series[idx] = 1.0
+
+    # ------------------------------------------------------------------
+    # Build the decay kernel
+    # ------------------------------------------------------------------
+    t_kernel = np.arange(0, 5 * tau, dt)
+
+    if kernel == "exp":
+        k = np.exp(-t_kernel / tau)
+
+    elif kernel == "hyperbolic":
+        if hyper_p <= 0:
+            raise ValueError("hyper_p must be > 0 for a valid hyperbolic kernel.")
+        k = (1.0 + t_kernel / tau) ** (-hyper_p)
+
+    else:
+        raise ValueError("kernel must be 'exp' or 'hyperbolic'.")
+
+    if normalize_kernel:
+        k /= k.sum() * dt  # area = 1
+
+    # ------------------------------------------------------------------
+    # Convolution (truncate to match timeline length)
+    # ------------------------------------------------------------------
+    r_rate = fftconvolve(event_series, k, mode="full")[:len(t)]
+
+    return t, r_rate