SSM hurricane Notebooks update (#804)

* StateSpace Module Hurricane Case Study WIP

rebased main to resolve conflicts

* attempt to fix plot rendering on ReviewNB

* attempt to fix plot rendering on ReviewNB try jupyterlab renderer

* attempt to fix plot rendering on ReviewNB try notebook renderer

* 1. Made changes in accordance to reviewer comments
2. Improved 24-hour ahead forecasts of simplest model
3. Added a section and model with how to add exogenous variables to the model

* 1. added forecast to exogenous variables section
2. cleaned up typos and figure sizes/zoom

* 1. updated the exogenous covariates section
2. Added a non-linear B-Splines section

* 1. Added text to spline section
2. Added closing remarks section
3. Updated references bib file

* 1. fixed typos, latex errors, and improved readability of text
2. Restructured function docstrings to comply with numpy's guidestyle
3. Added ODE derivation of Newtonian equations of motion
4. Added legend to origins scatterplot
5. Fixed issue with passing in future covariates into exogenous model when forecasting
6. Added model evaluations and comparisons
7. Fixed issue with diffuse initialization of P0

* updated notebook removing exogenous forecasts hack

* made some figures static to reduce file size

* 1. updated docstring indentation
2. updated generate forecast utility
3. updated how mode=JAX is specified
4. renamed simple model to Newtonian model
5. added explanations to exogenous data duplicating and mirroring
6. added note on how the newtonian model differs from typical parameteric models

* attempt to fix latex rendering

* missed spot in cleaning up latex

* updated equation number to attempt to fix latex rendering

* removed equation numbering

* fixed broken reference in closing remarks

* updated headers. Only one H1 header in notebook

* updated authoring dates and fixed missed h3 header

* made change to utilize conditional posteriors for predictive state means and covariances. Updated plotly setting to automatically select renderer

* Revert "made change to utilize conditional posteriors for predictive state means and covariances. Updated plotly setting to automatically select renderer"

This reverts commit 07355d2.

Revert using conditional posterior to using save_kalman_filter_outputs_in_idata

* try mimetype plotly renderer

* try sphinx_gallery plotly renderer

* try modifying javascript to render correctly

* try disabling plotly mathjax

* try iframe renderer

* changed plotly to static rendering

* fixed inconsistency where transition matrix was sometimes referred to with A and sometimes with T to only being referred to using T

* try remove-cell tag to change thumbnail

* updated the way in which we sample filter_outputs using the sample_filter_outputs method

* updated to use new arviz

* added arviz-variat style to arviz plots

Author: Dekermanjian

examples/case_studies/ssm_hurricane_tracking.ipynb

@@ -5,7 +5,7 @@ jupytext:
     format_name: myst
     format_version: 0.13
 kernelspec:
-  display_name: pymc_examples_dev
+  display_name: pymc-examples
   language: python
   name: python3
 myst:
@@ -91,8 +91,8 @@ import warnings
 
 warnings.filterwarnings("ignore", message="The RandomType SharedVariables", category=UserWarning)
 
-import arviz as az
 import arviz.labels as azl
+import arviz.preview as az
 import numpy as np
 import pymc as pm
 import pytensor.tensor as pt
@@ -119,6 +119,8 @@ from pymc_extras.statespace.utils.constants import (
 
 # make all plotly figures static
 pio.renderers.default = "svg"
+# set arviz style
+az.style.use("arviz-variat")
 ```
 
 ## Helper Functions
@@ -313,9 +315,7 @@ def generate_period_forecasts(
     longitude_cppc = az.extract(forecasts["forecast_observed"].sel(observed_state="x"))
     latitude_cppc = az.extract(forecasts["forecast_observed"].sel(observed_state="y"))
     cppc_var = forecasts["forecast_observed"].var(("chain", "draw"))
-    cppc_covs = xr.cov(
-        latitude_cppc["forecast_observed"], longitude_cppc["forecast_observed"], dim="sample"
-    )
+    cppc_covs = xr.cov(latitude_cppc, longitude_cppc, dim="sample")
     covs_list = []
     for i in range(cppc_covs.shape[0]):
         covs_list.append(
@@ -964,23 +964,32 @@ with pm.Model(coords=n_ssm.coords) as newtonian:
 
     n_ssm.build_statespace_graph(
         data=fiona_df.select("longitude", "latitude").to_numpy(),
-        save_kalman_filter_outputs_in_idata=True,
     )
     newtonian_idata = pm.sample(
         nuts_sampler="nutpie", nuts_sampler_kwargs={"backend": "jax", "gradient_backend": "jax"}
     )
 ```
 
 ```{code-cell} ipython3
-az.summary(newtonian_idata, var_names="acceleration_innovations", kind="stats")
+az.summary(newtonian_idata, var_names="acceleration_innovations", kind="stats", round_to=4)
+```
+
+```{code-cell} ipython3
+n_ssm_filter_outputs = n_ssm.sample_filter_outputs(
+    newtonian_idata, filter_output_names=["predicted_covariances", "predicted_observed_states"]
+)
 ```
 
 ```{code-cell} ipython3
-predicted_covs = newtonian_idata.posterior["predicted_covariance"].mean(("chain", "draw"))
+predicted_covs = n_ssm_filter_outputs.posterior_predictive["predicted_covariances"].mean(
+    ("chain", "draw")
+)
 ```
 
 ```{code-cell} ipython3
-post_mean = newtonian_idata.posterior["predicted_observed_state"].mean(("chain", "draw"))
+post_mean = n_ssm_filter_outputs.posterior_predictive["predicted_observed_states"].mean(
+    ("chain", "draw")
+)
 ```
 
 Not bad for a model with only one parameter. We can see that the forecast gets wonky in the middle where the trajectory of the Hurricane changes directions over short time periods. Again, it is important to keep in mind that what we are plotting are the one-step/period ahead forecast. In our case, our periods are six hours apart. Unfortunately, a 6-hour ahead hurricane forecast is not very practical. Let's see what we get when we generate a 4-period (24-hour) ahead forecast.
@@ -1343,10 +1352,7 @@ with pm.Model(coords=exog_ssm.coords) as exogenous:
 
     acceleration_innovations = pm.Gamma("acceleration_innovations", 0.1, 5, shape=(1,))
 
-    exog_ssm.build_statespace_graph(
-        data=fiona_df.select("longitude", "latitude").to_numpy(),
-        save_kalman_filter_outputs_in_idata=True,
-    )
+    exog_ssm.build_statespace_graph(data=fiona_df.select("longitude", "latitude").to_numpy())
     exogenous_idata = pm.sample(
         nuts_sampler="nutpie", nuts_sampler_kwargs={"backend": "jax", "gradient_backend": "jax"}
     )
@@ -1355,7 +1361,7 @@ with pm.Model(coords=exog_ssm.coords) as exogenous:
 Typically, the surface wind speed and the central pressure of a hurricane carry little information on the path the hurricane will take. The path of a hurricane is, generally, influenced by surrounding atmospheric conditions like pressure gradients. Knowing this, it makes sense to see that many of our beta parameters are close to zero, indicating little to no influence on the hurricanes' path.
 
 ```{code-cell} ipython3
-az.plot_trace(exogenous_idata, var_names="acceleration_innovations");
+az.plot_trace_dist(exogenous_idata, var_names="acceleration_innovations");
 ```
 
 ```{code-cell} ipython3
@@ -1367,8 +1373,18 @@ az.plot_forest(
 ### Make in-sample forecasts with new exogenous model
 
 ```{code-cell} ipython3
-predicted_covs = exogenous_idata.posterior["predicted_covariance"].mean(("chain", "draw"))
-post_mean = exogenous_idata.posterior["predicted_observed_state"].mean(("chain", "draw"))
+exog_ssm_filter_outputs = exog_ssm.sample_filter_outputs(
+    exogenous_idata, filter_output_names=["predicted_covariances", "predicted_observed_states"]
+)
+```
+
+```{code-cell} ipython3
+predicted_covs = exog_ssm_filter_outputs.posterior_predictive["predicted_covariances"].mean(
+    ("chain", "draw")
+)
+post_mean = exog_ssm_filter_outputs.posterior_predictive["predicted_observed_states"].mean(
+    ("chain", "draw")
+)
 ```
 
 Our one-period ahead forecasts seem to be slightly worse than our Newtonian model. You will notice that at the end of the forecast we see that our trajectory is erroneously more north rather than north-east. Since the exogenous variables we added to the model don't carry additional information with respect to the hurricane's trajectory, this results are expected.
@@ -1805,10 +1821,7 @@ with pm.Model(coords=spline_ssm.coords) as spline_model:
 
     acceleration_innovations = pm.Gamma("acceleration_innovations", 0.1, 5, shape=(1,))
 
-    spline_ssm.build_statespace_graph(
-        data=fiona_df.select("longitude", "latitude").to_numpy(),
-        save_kalman_filter_outputs_in_idata=True,
-    )
+    spline_ssm.build_statespace_graph(data=fiona_df.select("longitude", "latitude").to_numpy())
     spline_idata = pm.sample(
         nuts_sampler="nutpie", nuts_sampler_kwargs={"backend": "jax", "gradient_backend": "jax"}
     )
@@ -1817,11 +1830,11 @@ with pm.Model(coords=spline_ssm.coords) as spline_model:
 Most of our spline parameters are around zero, with a handful of exceptions. Let's take a look at how these effect our forecasts.
 
 ```{code-cell} ipython3
-az.plot_trace(spline_idata, var_names="acceleration_innovations");
+az.plot_trace_dist(spline_idata, var_names="acceleration_innovations");
 ```
 
 ```{code-cell} ipython3
-az.plot_trace(spline_idata, var_names=["beta_exog"], compact=True, figsize=(20, 8));
+az.plot_trace_dist(spline_idata, var_names=["beta_exog"], compact=True);
 ```
 
 ### Make in-sample forecasts with new spline model
@@ -1831,8 +1844,18 @@ az.plot_trace(spline_idata, var_names=["beta_exog"], compact=True, figsize=(20,
 Our one-period ahead forecasts, look better than the ones we generated from the Exogenous covariates model, but worse than the original model that purely follows Newtonian kinematics.
 
 ```{code-cell} ipython3
-predicted_covs = spline_idata.posterior["predicted_covariance"].mean(("chain", "draw"))
-post_mean = spline_idata.posterior["predicted_observed_state"].mean(("chain", "draw"))
+spline_ssm_filter_outputs = spline_ssm.sample_filter_outputs(
+    spline_idata, filter_output_names=["predicted_covariances", "predicted_observed_states"]
+)
+```
+
+```{code-cell} ipython3
+predicted_covs = spline_ssm_filter_outputs.posterior_predictive["predicted_covariances"].mean(
+    ("chain", "draw")
+)
+post_mean = spline_ssm_filter_outputs.posterior_predictive["predicted_observed_states"].mean(
+    ("chain", "draw")
+)
 ```
 
 ```{code-cell} ipython3
@@ -1941,6 +1964,7 @@ fig.show(width=1000, renderer="png", config={"displayModeBar": False})
 ```
 
 ## Authors
+* Updated by Jonathan Dekermanjian in August, 2025 to use the `sample_filter_outputs` method
 * Authored by Jonathan Dekermanjian in June, 2025 
 
 +++