Remove SingleTimeSeriesFilter

Author: jessegrabowski

pymc_experimental/statespace/core/statespace.py

@@ -19,7 +19,6 @@
 from pymc_experimental.statespace.core.representation import PytensorRepresentation
 from pymc_experimental.statespace.filters import (
     KalmanSmoother,
-    SingleTimeseriesFilter,
     SquareRootFilter,
     StandardFilter,
     UnivariateFilter,
@@ -52,7 +51,6 @@
 FILTER_FACTORY = {
     "standard": StandardFilter,
     "univariate": UnivariateFilter,
-    "single": SingleTimeseriesFilter,
     "cholesky": SquareRootFilter,
 }
 

pymc_experimental/statespace/filters/__init__.py

@@ -1,6 +1,5 @@
 from pymc_experimental.statespace.filters.distributions import LinearGaussianStateSpace
 from pymc_experimental.statespace.filters.kalman_filter import (
-    SingleTimeseriesFilter,
     SquareRootFilter,
     StandardFilter,
     UnivariateFilter,
@@ -11,7 +10,6 @@
     "StandardFilter",
     "UnivariateFilter",
     "KalmanSmoother",
-    "SingleTimeseriesFilter",
     "SquareRootFilter",
     "LinearGaussianStateSpace",
 ]

pymc_experimental/statespace/filters/kalman_filter.py

@@ -21,7 +21,6 @@
 MVN_CONST = pt.log(2 * pt.constant(np.pi, dtype="float64"))
 PARAM_NAMES = ["c", "d", "T", "Z", "R", "H", "Q"]
 
-assert_data_is_1d = Assert("UnivariateTimeSeries filter requires data be at most 1-dimensional")
 assert_time_varying_dim_correct = Assert(
     "The first dimension of a time varying matrix (the time dimension) must be "
     "equal to the first dimension of the data (the time dimension)."
@@ -751,50 +750,12 @@ def square_sequnece(L):
         ]
 
 
-class SingleTimeseriesFilter(BaseFilter):
-    """
-    Kalman filter optimized for univariate timeseries
-
-    If there is only a single observed timeseries, regardless of the number of hidden states, there is no need to
-    perform a matrix inversion anywhere in the filter.
-    """
-
-    # TODO: This class should eventually be made irrelevant by pytensor re-writes.
-    def check_params(self, data, a0, P0, c, d, T, Z, R, H, Q):
-        """
-        Wrap the data in an `Assert` `Op` to ensure there is only one observed state.
-        """
-        data = assert_data_is_1d(data, pt.eq(data.shape[1], 1))
-
-        return data, a0, P0, c, d, T, Z, R, H, Q
-
-    def update(self, a, P, y, d, Z, H, all_nan_flag):
-        y_hat = d + Z.dot(a)
-        v = y - y_hat.ravel()
-
-        PZT = P.dot(Z.T)
-
-        # F is scalar, K is a column vector
-        F = stabilize(Z.dot(PZT) + H, self.cov_jitter).ravel()
-
-        K = PZT / F
-        I_KZ = pt.eye(self.n_states) - K.dot(Z)
-
-        a_filtered = a + (K * v).ravel()
-
-        P_filtered = quad_form_sym(I_KZ, P) + quad_form_sym(K, H)
-
-        ll = pt.switch(all_nan_flag, 0.0, -0.5 * (MVN_CONST + pt.log(F) + v**2 / F)).ravel()[0]
-
-        return a_filtered, P_filtered, pt.atleast_1d(y_hat), pt.atleast_2d(F), ll
-
-
 class UnivariateFilter(BaseFilter):
     """
     The univariate kalman filter, described in [1], section 6.4.2, avoids inversion of the F matrix, as well as two
     matrix multiplications, at the cost of an additional loop. Note that the name doesn't mean there's only one
-    observed time series, that's the SingleTimeSeries filter. This is called univariate because it updates the state
-    mean and covariance matrices one variable at a time, using an inner-inner loop.
+    observed time series. This is called univariate because it updates the state mean and covariance matrices one
+    variable at a time, using an inner-inner loop.
 
     This is useful when states are perfectly observed, because the F matrix can easily become degenerate in these cases.
 

tests/statespace/test_kalman_filter.py

@@ -7,7 +7,6 @@
 
 from pymc_experimental.statespace.filters import (
     KalmanSmoother,
-    SingleTimeseriesFilter,
     SquareRootFilter,
     StandardFilter,
     UnivariateFilter,
@@ -34,20 +33,17 @@
 standard_inout = initialize_filter(StandardFilter())
 cholesky_inout = initialize_filter(SquareRootFilter())
 univariate_inout = initialize_filter(UnivariateFilter())
-single_inout = initialize_filter(SingleTimeseriesFilter())
 
 f_standard = pytensor.function(*standard_inout, on_unused_input="ignore")
 f_cholesky = pytensor.function(*cholesky_inout, on_unused_input="ignore")
 f_univariate = pytensor.function(*univariate_inout, on_unused_input="ignore")
-f_single_ts = pytensor.function(*single_inout, on_unused_input="ignore")
 
-filter_funcs = [f_standard, f_cholesky, f_univariate, f_single_ts]
+filter_funcs = [f_standard, f_cholesky, f_univariate]
 
 filter_names = [
     "StandardFilter",
     "CholeskyFilter",
     "UnivariateFilter",
-    "SingleTimeSeriesFilter",
 ]
 
 output_names = [
@@ -191,20 +187,12 @@ def test_output_with_multiple_observed(filter_func, filter_name, rng):
     p, m, r, n = 5, 5, 1, 10
     inputs = make_test_inputs(p, m, r, n, rng)
 
-    if filter_name == "SingleTimeSeriesFilter":
-        with pytest.raises(
-            AssertionError,
-            match="UnivariateTimeSeries filter requires data be at most 1-dimensional",
-        ):
-            filter_func(*inputs)
-
-    else:
-        outputs = filter_func(*inputs)
-        for output_idx, name in enumerate(output_names):
-            expected_output = get_expected_shape(name, p, m, r, n)
-            assert (
-                outputs[output_idx].shape == expected_output
-            ), f"Shape of {name} does not match expected"
+    outputs = filter_func(*inputs)
+    for output_idx, name in enumerate(output_names):
+        expected_output = get_expected_shape(name, p, m, r, n)
+        assert (
+            outputs[output_idx].shape == expected_output
+        ), f"Shape of {name} does not match expected"
 
 
 @pytest.mark.parametrize(
@@ -215,20 +203,12 @@ def test_missing_data(filter_func, filter_name, p, rng):
     m, r, n = 5, 1, 10
     inputs = make_test_inputs(p, m, r, n, rng, missing_data=1)
 
-    if p > 1 and filter_name == "SingleTimeSeriesFilter":
-        with pytest.raises(
-            AssertionError,
-            match="UnivariateTimeSeries filter requires data be at most 1-dimensional",
-        ):
-            filter_func(*inputs)
-
-    else:
-        outputs = filter_func(*inputs)
-        for output_idx, name in enumerate(output_names):
-            expected_output = get_expected_shape(name, p, m, r, n)
-            assert (
-                outputs[output_idx].shape == expected_output
-            ), f"Shape of {name} does not match expected"
+    outputs = filter_func(*inputs)
+    for output_idx, name in enumerate(output_names):
+        expected_output = get_expected_shape(name, p, m, r, n)
+        assert (
+            outputs[output_idx].shape == expected_output
+        ), f"Shape of {name} does not match expected"
 
 
 @pytest.mark.parametrize("filter_func", filter_funcs, ids=filter_names)
@@ -323,8 +303,8 @@ def test_all_covariance_matrices_are_PSD(filter_func, filter_name, n_missing, ob
 
 @pytest.mark.parametrize(
     "filter",
-    [StandardFilter, SingleTimeseriesFilter, SquareRootFilter],
-    ids=["standard", "single_ts", "cholesky"],
+    [StandardFilter, SquareRootFilter],
+    ids=["standard", "cholesky"],
 )
 def test_kalman_filter_jax(filter):
     pytest.importorskip("jax")