REF: Rename ar variables

Author: juanmpga

examples/time_series/Time_Series_Generative_Graph.ipynb

@@ -5,15 +5,15 @@ jupytext:
     format_name: myst
     format_version: 0.13
 kernelspec:
-  display_name: pymc_env
+  display_name: Python 3 (ipykernel)
   language: python
   name: python3
 ---
 
 (time_series_generative_graph)=
 # Time Series Models Derived From a Generative Graph
 
-:::{post} July, 2024
+:::{post} January, 2025
 :tags: time-series, 
 :category: intermediate, reference
 :author: Jesse Grabowski, Juan Orduz and Ricardo Vieira
@@ -70,7 +70,7 @@ rng = np.random.default_rng(42)
 
 ## Define AR(2) Process
 
-We start by encoding the generative graph of the AR(2) model as a function `ar_dist`. The strategy is to pass this function as a custom distribution via {class}`~pymc.CustomDist` inside a PyMC model. 
+We start by encoding the generative graph of the AR(2) model as a function `get_ar_steps`. The strategy is to pass this function as a custom distribution via {class}`~pymc.CustomDist` inside a PyMC model. 
 
 We need to specify the initial state (`ar_init`), the autoregressive coefficients (`rho`), and the standard deviation of the noise (`sigma`). Given such parameters, we can define the generative graph of the AR(2) model using the  {meth}`scan <pytensor.scan.basic.scan>` operation.
 
@@ -101,23 +101,23 @@ timeseries_length = 100  # Time series length
 # This is the transition function for the AR(2) model.
 # We take as inputs previous steps and then specify the autoregressive relationship.
 # Finally, we add Gaussian noise to the model.
-def ar_step(x_tm2, x_tm1, rho, sigma):
+def get_next_step(x_tm2, x_tm1, rho, sigma):
     mu = x_tm1 * rho[0] + x_tm2 * rho[1]
     x = mu + pm.Normal.dist(sigma=sigma)
     return x, collect_default_updates([x])
 
 
 # Here we actually "loop" over the time series.
-def ar_dist(ar_init, rho, sigma, size):
-    ar_innov, _ = pytensor.scan(
-        fn=ar_step,
+def get_ar_steps(ar_init, rho, sigma, size):
+    ar_steps, _ = pytensor.scan(
+        fn=get_next_step,
         outputs_info=[{"initial": ar_init, "taps": range(-lags, 0)}],
         non_sequences=[rho, sigma],
         n_steps=timeseries_length - lags,
         strict=True,
     )
 
-    return ar_innov
+    return ar_steps
 ```
 
 ## Generate AR(2) Graph
@@ -136,17 +136,17 @@ with pm.Model(coords=coords, check_bounds=False) as model:
 
     ar_init = pm.Normal(name="ar_init", sigma=0.5, dims=("lags",))
 
-    ar_innov = pm.CustomDist(
-        "ar_dist",
+    ar_steps = pm.CustomDist(
+        "ar_steps",
         ar_init,
         rho,
         sigma,
-        dist=ar_dist,
+        dist=get_ar_steps,
         dims=("steps",),
     )
 
     ar = pm.Deterministic(
-        name="ar", var=pt.concatenate([ar_init, ar_innov], axis=-1), dims=("timeseries_length",)
+        name="ar", var=pt.concatenate([ar_init, ar_steps], axis=-1), dims=("timeseries_length",)
     )
 
 
@@ -208,9 +208,9 @@ Next, we want to condition the AR(2) model on some observed data so that we can
 ```{code-cell} ipython3
 # select a random draw from the prior
 prior_draw = prior.prior.isel(chain=0, draw=chosen_draw)
-test_data = prior_draw["ar_dist"].to_numpy()
+test_data = prior_draw["ar_steps"].to_numpy()
 
-with pm.observe(model, {"ar_dist": test_data}) as observed_model:
+with pm.observe(model, {"ar_steps": test_data}) as observed_model:
     trace = pm.sample(chains=4, random_seed=rng)
 ```
 
@@ -319,17 +319,17 @@ $$
 Let's see how to do this in PyMC! The key observation is that we need to pass the observed data explicitly into out "for loop" in the generative graph. That is, we need to pass it into the {meth}`scan <pytensor.scan.basic.scan>` function.
 
 ```{code-cell} ipython3
-def conditional_ar_dist(y_data, rho, sigma, size):
+def get_conditional_ar_steps(y_data, rho, sigma, size):
     # Here we condition on the observed data by passing it through the `sequences` argument.
-    ar_innov, _ = pytensor.scan(
-        fn=ar_step,
+    ar_steps, _ = pytensor.scan(
+        fn=get_next_step,
         sequences=[{"input": y_data, "taps": list(range(-lags, 0))}],
         non_sequences=[rho, sigma],
         n_steps=timeseries_length - lags,
         strict=True,
     )
 
-    return ar_innov
+    return ar_steps
 ```
 
 Then we can simply generate samples from the posterior predictive distribution. Observe we need to "rewrite" the generative graph to include the conditioned transition step. When you call {meth}`~pm.sample_posterior_predictive`,PyMC will attempt to match the names of random variables in the active model context to names in the provided `idata.posterior`. If a match is found, the specified model prior is ignored, and replaced with draws from the posterior. This means we can put any prior we want on these parameters, because it will be ignored. We choose {class}`~pymc.distributions.continuous.Flat` because you cannot sample from it. This way, if PyMC does not find a match for one of our priors, we will get an error to let us know something isn't right. For a detailed explanation on these type of cross model predictions, see the great blog post [Out of model predictions with PyMC](https://www.pymc-labs.com/blog-posts/out-of-model-predictions-with-pymc/).
@@ -351,17 +351,17 @@ with pm.Model(coords=coords, check_bounds=False) as conditional_model:
     rho = pm.Flat(name="rho", dims=("lags",))
     sigma = pm.Flat(name="sigma")
 
-    ar_innov = pm.CustomDist(
-        "ar_dist",
+    ar_steps = pm.CustomDist(
+        "ar_steps",
         y_data,
         rho,
         sigma,
-        dist=conditional_ar_dist,
+        dist=get_conditional_ar_steps,
         dims=("steps",),
     )
 
     ar = pm.Deterministic(
-        name="ar", var=pt.concatenate([ar_init, ar_innov], axis=-1), dims=("timeseries_length",)
+        name="ar", var=pt.concatenate([ar_init, ar_steps], axis=-1), dims=("timeseries_length",)
     )
 
     post_pred_conditional = pm.sample_posterior_predictive(trace, var_names=["ar"], random_seed=rng)
@@ -446,34 +446,34 @@ with pm.Model(coords=coords, check_bounds=False) as forecasting_model:
     rho = pm.Flat(name="rho", dims=("lags",))
     sigma = pm.Flat(name="sigma")
 
-    def ar_dist_forecasting(forecast_initial_state, rho, sigma, size):
-        ar_innov, _ = pytensor.scan(
-            fn=ar_step,
+    def get_ar_steps_forecast(forecast_initial_state, rho, sigma, size):
+        ar_steps, _ = pytensor.scan(
+            fn=get_next_step,
             outputs_info=[{"initial": forecast_initial_state, "taps": range(-lags, 0)}],
             non_sequences=[rho, sigma],
             n_steps=forecast_steps,
             strict=True,
         )
-        return ar_innov
+        return ar_steps
 
-    ar_innov = pm.CustomDist(
-        "ar_dist",
+    ar_steps = pm.CustomDist(
+        "ar_steps",
         forecast_initial_state,
         rho,
         sigma,
-        dist=ar_dist_forecasting,
+        dist=get_ar_steps_forecast,
         dims=("steps",),
     )
 
     post_pred_forecast = pm.sample_posterior_predictive(
-        trace, var_names=["ar_dist"], random_seed=rng
+        trace, var_names=["ar_steps"], random_seed=rng
     )
 ```
 
 We can visualize the out-of-sample predictions and compare thee results wth the one from  `statsmodels`.
 
 ```{code-cell} ipython3
-forecast_post_pred_ar = post_pred_forecast.posterior_predictive["ar_dist"]
+forecast_post_pred_ar = post_pred_forecast.posterior_predictive["ar_steps"]
 
 _, ax = plt.subplots()
 for i, hdi_prob in enumerate((0.94, 0.64), 1):
@@ -497,7 +497,7 @@ ax.plot(
 )
 
 for i, hdi_prob in enumerate((0.94, 0.64), 1):
-    hdi = az.hdi(forecast_post_pred_ar, hdi_prob=hdi_prob)["ar_dist"]
+    hdi = az.hdi(forecast_post_pred_ar, hdi_prob=hdi_prob)["ar_steps"]
     lower = hdi.sel(hdi="lower")
     upper = hdi.sel(hdi="higher")
     ax.fill_between(