Add prediction intervals for filter/smoother

isssm/_modidx.py

@@ -54,10 +54,12 @@
                               'isssm.kalman._smooth_step': ('kalman_filter_smoother.html#_smooth_step', 'isssm/kalman.py'),
                               'isssm.kalman.account_for_nans': ('kalman_filter_smoother.html#account_for_nans', 'isssm/kalman.py'),
                               'isssm.kalman.disturbance_smoother': ('kalman_filter_smoother.html#disturbance_smoother', 'isssm/kalman.py'),
+                              'isssm.kalman.filter_intervals': ('kalman_filter_smoother.html#filter_intervals', 'isssm/kalman.py'),
                               'isssm.kalman.kalman': ('kalman_filter_smoother.html#kalman', 'isssm/kalman.py'),
                               'isssm.kalman.simulation_smoother': ('kalman_filter_smoother.html#simulation_smoother', 'isssm/kalman.py'),
                               'isssm.kalman.smoothed_signals': ('kalman_filter_smoother.html#smoothed_signals', 'isssm/kalman.py'),
-                              'isssm.kalman.smoother': ('kalman_filter_smoother.html#smoother', 'isssm/kalman.py')},
+                              'isssm.kalman.smoother': ('kalman_filter_smoother.html#smoother', 'isssm/kalman.py'),
+                              'isssm.kalman.smoother_intervals': ('kalman_filter_smoother.html#smoother_intervals', 'isssm/kalman.py')},
             'isssm.laplace_approximation': { 'isssm.laplace_approximation._initial_guess': ( 'laplace_approximation.html#_initial_guess',
                                                                                              'isssm/laplace_approximation.py'),
                                              'isssm.laplace_approximation.laplace_approximation': ( 'laplace_approximation.html#laplace_approximation',

isssm/kalman.py

@@ -1,8 +1,8 @@
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/10_kalman_filter_smoother.ipynb.
 
 # %% auto 0
-__all__ = ['State', 'StateCov', 'StateTransition', 'kalman', 'account_for_nans', 'smoother', 'FFBS', 'disturbance_smoother',
-           'smoothed_signals', 'simulation_smoother']
+__all__ = ['State', 'StateCov', 'StateTransition', 'kalman', 'account_for_nans', 'smoother', 'filter_intervals',
+           'smoother_intervals', 'FFBS', 'disturbance_smoother', 'smoothed_signals', 'simulation_smoother']
 
 # %% ../nbs/10_kalman_filter_smoother.ipynb 1
 import jax.numpy as jnp
@@ -141,6 +141,26 @@ def step(carry, inputs):
     return SmootherResult(x_smooth, Xi_smooth)
 
 # %% ../nbs/10_kalman_filter_smoother.ipynb 22
+from tensorflow_probability.substrates.jax.distributions import Normal
+def filter_intervals(result: FilterResult, alpha: Float=.05) -> Float[Array, "2 n+1 m"]:
+    x_filt, Xi_filt, *_ = result
+    marginal_variances = vmap(jnp.diag)(Xi_filt)
+    dist = Normal(x_filt, marginal_variances)
+    lower = dist.quantile(alpha / 2)
+    upper = dist.quantile(1 - alpha / 2)
+
+    return jnp.concatenate((lower[None], upper[None]))
+
+def smoother_intervals(result: SmootherResult, alpha: Float = .05) -> Float[Array, "2 n+1 m"]:
+    x_smooth, Xi_smooth = result
+    marginal_variances = vmap(jnp.diag)(Xi_smooth)
+    dist = Normal(x_smooth, marginal_variances)
+    lower = dist.quantile(alpha / 2)
+    upper = dist.quantile(1 - alpha / 2)
+
+    return jnp.concatenate((lower[None], upper[None]))
+
+# %% ../nbs/10_kalman_filter_smoother.ipynb 25
 def _simulate_smoothed_FW1994(
     x_filt: Float[Array, "n+1 m"],
     Xi_filt: Float[Array, "n+1 m m"],
@@ -192,7 +212,7 @@ def FFBS(
     key, subkey = jrn.split(key)
     return _simulate_smoothed_FW1994(x_filt, Xi_filt, Xi_pred, model.A, N, subkey)
 
-# %% ../nbs/10_kalman_filter_smoother.ipynb 27
+# %% ../nbs/10_kalman_filter_smoother.ipynb 30
 def disturbance_smoother(
     filtered: FilterResult, # filter result
     y: Observations, # observations
@@ -236,7 +256,7 @@ def smoothed_signals(
     eta_smooth, _ = disturbance_smoother(filtered, y, model)
     return y - eta_smooth
 
-# %% ../nbs/10_kalman_filter_smoother.ipynb 32
+# %% ../nbs/10_kalman_filter_smoother.ipynb 35
 from tensorflow_probability.substrates.jax.distributions import Chi2
 from .util import degenerate_cholesky
 from .util import location_antithetic, scale_antithethic