pymc-extras 0.2.0__py3-none-any.whl

pymc_extras/statespace/filters/kalman_smoother.py

@@ -0,0 +1,126 @@
+import pytensor
+import pytensor.tensor as pt
+
+from pytensor.compile import get_mode
+from pytensor.tensor.nlinalg import matrix_dot
+
+from pymc_extras.statespace.filters.utilities import (
+    quad_form_sym,
+    split_vars_into_seq_and_nonseq,
+    stabilize,
+)
+from pymc_extras.statespace.utils.constants import JITTER_DEFAULT
+
+
+class KalmanSmoother:
+    """
+    Kalman Smoother
+
+    """
+
+    def __init__(self, mode: str | None = None):
+        self.mode = mode
+        self.cov_jitter = JITTER_DEFAULT
+        self.seq_names = []
+        self.non_seq_names = []
+
+    def unpack_args(self, args):
+        """
+        The order of inputs to the inner scan function is not known, since some, all, or none of the input matrices
+        can be time varying. The order arguments are fed to the inner function is sequences, outputs_info,
+        non-sequences. This function works out which matrices are where, and returns a standardized order expected
+        by the kalman_step function.
+
+        The standard order is: a, P, a_smooth, P_smooth, T, R, Q
+        """
+        # If there are no sequence parameters (all params are static),
+        # no changes are needed, params will be in order.
+        args = list(args)
+        n_seq = len(self.seq_names)
+        if n_seq == 0:
+            return args
+
+        # The first two args are always a and P
+        a = args.pop(0)
+        P = args.pop(0)
+
+        # There are always two outputs_info wedged between the seqs and non_seqs
+        seqs, (a_smooth, P_smooth), non_seqs = (
+            args[:n_seq],
+            args[n_seq : n_seq + 2],
+            args[n_seq + 2 :],
+        )
+        return_ordered = []
+        for name in ["T", "R", "Q"]:
+            if name in self.seq_names:
+                idx = self.seq_names.index(name)
+                return_ordered.append(seqs[idx])
+            else:
+                idx = self.non_seq_names.index(name)
+                return_ordered.append(non_seqs[idx])
+
+        T, R, Q = return_ordered
+
+        return a, P, a_smooth, P_smooth, T, R, Q
+
+    def build_graph(
+        self, T, R, Q, filtered_states, filtered_covariances, mode=None, cov_jitter=JITTER_DEFAULT
+    ):
+        self.mode = mode
+        self.cov_jitter = cov_jitter
+
+        n, k = filtered_states.type.shape
+
+        a_last = pt.specify_shape(filtered_states[-1], (k,))
+        P_last = pt.specify_shape(filtered_covariances[-1], (k, k))
+
+        sequences, non_sequences, seq_names, non_seq_names = split_vars_into_seq_and_nonseq(
+            [T, R, Q], ["T", "R", "Q"]
+        )
+
+        self.seq_names = seq_names
+        self.non_seq_names = non_seq_names
+
+        smoother_result, updates = pytensor.scan(
+            self.smoother_step,
+            sequences=[filtered_states[:-1], filtered_covariances[:-1], *sequences],
+            outputs_info=[a_last, P_last],
+            non_sequences=non_sequences,
+            go_backwards=True,
+            name="kalman_smoother",
+            mode=get_mode(self.mode),
+        )
+
+        smoothed_states, smoothed_covariances = smoother_result
+        smoothed_states = pt.concatenate(
+            [smoothed_states[::-1], pt.expand_dims(a_last, axis=(0,))], axis=0
+        )
+        smoothed_covariances = pt.concatenate(
+            [smoothed_covariances[::-1], pt.expand_dims(P_last, axis=(0,))], axis=0
+        )
+
+        smoothed_states.name = "smoothed_states"
+        smoothed_covariances.name = "smoothed_covariances"
+
+        return smoothed_states, smoothed_covariances
+
+    def smoother_step(self, *args):
+        a, P, a_smooth, P_smooth, T, R, Q = self.unpack_args(args)
+        a_hat, P_hat = self.predict(a, P, T, R, Q)
+
+        # Use pinv, otherwise P_hat is singular when there is missing data
+        smoother_gain = matrix_dot(pt.linalg.pinv(P_hat), T, P).T
+        a_smooth_next = a + smoother_gain @ (a_smooth - a_hat)
+
+        P_smooth_next = P + quad_form_sym(smoother_gain, P_smooth - P_hat)
+        P_smooth_next = stabilize(P_smooth_next, self.cov_jitter)
+        P_smooth_next = pt.specify_shape(stabilize(P_smooth_next), P_smooth.type.shape)
+
+        return a_smooth_next, P_smooth_next
+
+    def predict(self, a, P, T, R, Q):
+        a_hat = T.dot(a)
+        P_hat = quad_form_sym(T, P) + quad_form_sym(R, Q)
+        P_hat = stabilize(P_hat, self.cov_jitter)
+
+        return a_hat, P_hat

pymc_extras/statespace/filters/utilities.py

@@ -0,0 +1,59 @@
+import pytensor.tensor as pt
+
+from pytensor.tensor.nlinalg import matrix_dot
+
+from pymc_extras.statespace.utils.constants import JITTER_DEFAULT, NEVER_TIME_VARYING, VECTOR_VALUED
+
+
+def decide_if_x_time_varies(x, name):
+    if name in NEVER_TIME_VARYING:
+        return False
+
+    ndim = x.ndim
+
+    if name in VECTOR_VALUED:
+        if ndim not in [1, 2]:
+            raise ValueError(
+                f"Vector {name} has {ndim} dimensions; it should have either 1 (static),"
+                f" or 2 (time varying )"
+            )
+
+        return ndim == 2
+
+    if ndim not in [2, 3]:
+        raise ValueError(
+            f"Matrix {name} has {ndim} dimensions; it should have either"
+            f" 2 (static), or 3 (time varying)."
+        )
+
+    return ndim == 3
+
+
+def split_vars_into_seq_and_nonseq(params, param_names):
+    """
+    Split inputs into those that are time varying and those that are not. This division is required by scan.
+    """
+    sequences, non_sequences = [], []
+    seq_names, non_seq_names = [], []
+
+    for param, name in zip(params, param_names):
+        if decide_if_x_time_varies(param, name):
+            sequences.append(param)
+            seq_names.append(name)
+        else:
+            non_sequences.append(param)
+            non_seq_names.append(name)
+
+    return sequences, non_sequences, seq_names, non_seq_names
+
+
+def stabilize(cov, jitter=JITTER_DEFAULT):
+    # Ensure diagonal is non-zero
+    cov = cov + pt.identity_like(cov) * jitter
+
+    return cov
+
+
+def quad_form_sym(A, B):
+    out = matrix_dot(A, B, A.T)
+    return 0.5 * (out + out.T)