Refactor convolve scanner factory functions

model/src/pyrenew/convolve.py

@@ -17,62 +17,110 @@
 
 import jax.numpy as jnp
 from jax.typing import ArrayLike
+from pyrenew.transformation import IdentityTransform
 
 
-def new_convolve_scanner(discrete_dist_flipped: ArrayLike) -> Callable:
-    """
-    Factory function to create a scanner function for convolving a discrete distribution
-    over a time series data subset.
+def new_convolve_scanner(
+    discrete_dist_flipped: ArrayLike, transform: Callable = None
+) -> Callable:
+    r"""
+    Factory function to create a "scanner" function
+    that can be used with :py:func:`jax.lax.scan` to
+    construct an array via backward-looking iterative
+    convolution.
 
     Parameters
     ----------
     discrete_dist_flipped : ArrayLike
-        A 1D jax array representing the discrete distribution flipped for convolution.
+        A 1D jax array representing the discrete
+        distribution flipped for convolution.
+
+    transform : Callable
+        A transformation to apply to the result
+        of the dot product and multiplication.
+        If None, use the identity transformation.
+        Default None.
 
     Returns
     -------
     Callable
-        A scanner function that can be used with jax.lax.scan for convolution.
-        This function takes a history subset and a multiplier, computes the dot product,
-        then updates and returns the new history subset and the convolution result.
+        A scanner function that can be used with
+        jax.lax.scan for convolution.
+        This function takes a history subset and
+        a multiplier, computes the dot product,
+        then updates and returns the new history
+        subset and the convolution result.
+
+    Notes
+    -----
+    The following iterative operation is found often
+    in renewal processes:
+
+    .. math::
+        X(t) = f\left(m(t) * \left[X(t - n),
+        X(t - n + 1), ... X(t - 1)\right] \dot \vec{d} \right)
+
+    Where `math`:\vec{d} is a vector of length `math`:n,
+    `math`:m(t) is a scalar for each value of time `math`:t,
+    and `math`:f is a scalar-valued function.
+
+    Given `math`:\vec{d}, and optionally `math`:f,
+    this factory function returns a new function that
+    peforms one step of this process while scanning along
+    an array of  multipliers (i.e. an array
+    giving the values of `math`:m(t)) using :py:func:jax.lax.scan.
     """
+    if transform is None:
+        transform = IdentityTransform()
 
     def _new_scanner(
         history_subset: ArrayLike, multiplier: float
     ) -> tuple[ArrayLike, float]:  # numpydoc ignore=GL08
-        new_val = multiplier * jnp.dot(discrete_dist_flipped, history_subset)
+        new_val = transform(
+            multiplier * jnp.dot(discrete_dist_flipped, history_subset)
+        )
         latest = jnp.hstack([history_subset[1:], new_val])
         return latest, new_val
 
     return _new_scanner
 
 
-def new_double_scanner(
+def new_double_convolve_scanner(
     dists: tuple[ArrayLike, ArrayLike],
-    transforms: tuple[Callable, Callable],
+    transforms: tuple[Callable, Callable] = (None, None),
 ) -> Callable:
     """
-    Factory function to create a scanner function that applies two sequential transformations
-    and convolutions using two discrete distributions.
+    Factory function to create a scanner function
+    that iteratively constructs arrays by applying
+    the dot-product/multiply/transform operation
+    twice per history subset, with the first yielding
+    operation yielding an additional scalar multiplier
+    for the second.
 
     Parameters
     ----------
     dists : tuple[ArrayLike, ArrayLike]
-        A tuple of two 1D jax arrays, each representing a discrete distribution for the
+        A tuple of two 1D jax arrays, each representing a
+        discrete distribution for the
         two stages of convolution.
     transforms : tuple[Callable, Callable]
-        A tuple of two functions, each transforming the output of the dot product at each
-        convolution stage.
+        A tuple of two functions, each transforming the
+        output of the dot product at each
+        convolution stage. If either is None,
+        the identity transformation will be used
+        at that step. Default (None, None)
 
     Returns
     -------
     Callable
-        A scanner function that applies two sequential convolutions and transformations.
-        It takes a history subset and a tuple of multipliers, computes the transformations
-        and dot products, and returns the updated history subset and a tuple of new values.
+        A scanner function that applies two sets of
+        convolution, multiply, and transform operations
+        in sequence to construct a new array by scanning
+        along a pair of input arrays that are equal in
+        length to each other.
     """
     d1, d2 = dists
-    t1, t2 = transforms
+    t1, t2 = [x if x is not None else IdentityTransform() for x in transforms]
 
     def _new_scanner(
         history_subset: ArrayLike,

model/src/pyrenew/latent/infection_functions.py

@@ -6,7 +6,8 @@
 import jax
 import jax.numpy as jnp
 from jax.typing import ArrayLike
-from pyrenew.convolve import new_convolve_scanner, new_double_scanner
+from pyrenew.convolve import new_convolve_scanner, new_double_convolve_scanner
+from pyrenew.transformation import ExpTransform, IdentityTransform
 
 
 def compute_infections_from_rt(
@@ -169,12 +170,12 @@ def compute_infections_from_rt_with_feedback(
     reductions in contact rate due to awareness of high incidence,
     et cetera.
     """
-    feedback_scanner = new_double_scanner(
+    feedback_scanner = new_double_convolve_scanner(
         dists=(
             reversed_infection_feedback_pmf,
             reversed_generation_interval_pmf,
         ),
-        transforms=(jnp.exp, lambda x: x),
+        transforms=(ExpTransform(), IdentityTransform()),
     )
     latest, infs_and_R_adj = jax.lax.scan(
         f=feedback_scanner,

model/src/test/test_convolve_scanners.py

@@ -0,0 +1,54 @@
+"""
+Unit tests for the iterative convolution
+scanner function factories found in pyrenew.convolve
+"""
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pyrenew.convolve as pc
+from numpy.testing import assert_array_equal
+
+
+def test_double_scanner_reduces_to_single():
+    """
+    Test that new_double_scanner() yields a function
+    that is equivalent to a single scanner if the first
+    scan is chosen appropriately
+    """
+    inits = jnp.array([0.352, 5.2, -3])
+    to_scan_a = jnp.array([0.5, 0.3, 0.2])
+
+    multipliers = jnp.array(np.random.normal(0, 0.5, size=500))
+
+    def transform_a(x: any):
+        """
+        transformation associated with
+        array to_scan_a
+
+        Parameters
+        ----------
+        x: any
+            input value
+
+        Returns
+        -------
+        The result of 4 * x + 0.025, where x is the input
+        value
+        """
+        return 4 * x + 0.025
+
+    scanner_a = pc.new_convolve_scanner(to_scan_a, transform_a)
+
+    double_scanner_a = pc.new_double_convolve_scanner(
+        (jnp.array([523, 2, -0.5233]), to_scan_a), (lambda x: 1, transform_a)
+    )
+
+    _, result_a = jax.lax.scan(f=scanner_a, init=inits, xs=multipliers)
+
+    _, result_a_double = jax.lax.scan(
+        f=double_scanner_a, init=inits, xs=(multipliers * 0.2352, multipliers)
+    )
+
+    assert_array_equal(result_a_double[1], jnp.ones_like(multipliers))
+    assert_array_equal(result_a_double[0], result_a)