Use model variables named "sigma" as standard deviations rather than variances (#296)

* Model variables named "sigma" no longer expect variances

* Rename variables "sigma" to "sigma2" in `create_structural_model_and_equivalent_statsmodel` test function

Author: jessegrabowski

pymc_experimental/statespace/models/SARIMAX.py

@@ -514,14 +514,14 @@ def make_symbolic_graph(self) -> None:
         state_cov = self.make_and_register_variable(
             "sigma_state", shape=(self.k_posdef,), dtype=floatX
         )
-        self.ssm[state_cov_idx] = state_cov
+        self.ssm[state_cov_idx] = state_cov**2
 
         if self.measurement_error:
             obs_cov_idx = ("obs_cov",) + np.diag_indices(self.k_endog)
             obs_cov = self.make_and_register_variable(
                 "sigma_obs", shape=(self.k_endog,), dtype=floatX
             )
-            self.ssm[obs_cov_idx] = obs_cov
+            self.ssm[obs_cov_idx] = obs_cov**2
 
         # The initial conditions have to be done last in the case of stationary initialization, because it will depend
         # on c, T, R and Q

pymc_experimental/statespace/models/structural.py

@@ -818,7 +818,7 @@ def make_symbolic_graph(self) -> None:
             sigma_trend = self.make_and_register_variable("sigma_trend", shape=(self.k_posdef,))
             diag_idx = np.diag_indices(self.k_posdef)
             idx = np.s_["state_cov", diag_idx[0], diag_idx[1]]
-            self.ssm[idx] = sigma_trend
+            self.ssm[idx] = sigma_trend**2
 
 
 class MeasurementError(Component):
@@ -884,7 +884,7 @@ def make_symbolic_graph(self) -> None:
         error_sigma = self.make_and_register_variable(f"sigma_{self.name}", shape=(self.k_endog,))
         diag_idx = np.diag_indices(self.k_endog)
         idx = np.s_["obs_cov", diag_idx[0], diag_idx[1]]
-        self.ssm[idx] = error_sigma
+        self.ssm[idx] = error_sigma**2
 
 
 class AutoregressiveComponent(Component):
@@ -991,7 +991,7 @@ def make_symbolic_graph(self) -> None:
         self.ssm[ar_idx] = ar_params
 
         cov_idx = ("state_cov", *np.diag_indices(1))
-        self.ssm[cov_idx] = sigma_ar
+        self.ssm[cov_idx] = sigma_ar**2
 
 
 class TimeSeasonality(Component):
@@ -1175,7 +1175,7 @@ def make_symbolic_graph(self) -> None:
             self.ssm["selection", 0, 0] = 1
             season_sigma = self.make_and_register_variable(f"sigma_{self.name}", shape=(1,))
             cov_idx = ("state_cov", *np.diag_indices(1))
-            self.ssm[cov_idx] = season_sigma
+            self.ssm[cov_idx] = season_sigma**2
 
 
 class FrequencySeasonality(Component):
@@ -1273,7 +1273,7 @@ def make_symbolic_graph(self) -> None:
 
         if self.innovations:
             sigma_season = self.make_and_register_variable(f"sigma_{self.name}", shape=(1,))
-            self.ssm["state_cov", :, :] = pt.eye(self.k_posdef) * sigma_season
+            self.ssm["state_cov", :, :] = pt.eye(self.k_posdef) * sigma_season**2
             self.ssm["selection", :, :] = np.eye(self.k_states)
 
     def populate_component_properties(self):
@@ -1453,7 +1453,7 @@ def make_symbolic_graph(self) -> None:
 
         if self.innovations:
             sigma_cycle = self.make_and_register_variable(f"sigma_{self.name}", shape=(1,))
-            self.ssm["state_cov", :, :] = pt.eye(self.k_posdef) * sigma_cycle
+            self.ssm["state_cov", :, :] = pt.eye(self.k_posdef) * sigma_cycle**2
 
     def populate_component_properties(self):
         self.state_names = [f"{self.name}_{f}" for f in ["Sin", "Cos"]]
@@ -1556,7 +1556,7 @@ def make_symbolic_graph(self) -> None:
                 f"sigma_beta_{self.name}", (self.k_states,)
             )
             row_idx, col_idx = np.diag_indices(self.k_states)
-            self.ssm["state_cov", row_idx, col_idx] = sigma_beta
+            self.ssm["state_cov", row_idx, col_idx] = sigma_beta**2
 
     def populate_component_properties(self) -> None:
         self.shock_names = self.state_names

pymc_experimental/tests/statespace/test_SARIMAX.py

@@ -296,7 +296,9 @@ def test_SARIMAX_update_matches_statsmodels(p, d, q, P, D, Q, S, data, rng):
                 ),
             )
 
-        pm.Deterministic("sigma_state", pt.as_tensor_variable(np.array([param_d["sigma2"]])))
+        pm.Deterministic(
+            "sigma_state", pt.as_tensor_variable(np.sqrt(np.array([param_d["sigma2"]])))
+        )
 
         mod._insert_random_variables()
         matrices = pm.draw(mod.subbed_ssm)

pymc_experimental/tests/statespace/test_structural.py

@@ -196,9 +196,9 @@ def create_structural_model_and_equivalent_statsmodel(
     components = []
 
     if irregular:
-        sigma = np.abs(rng.normal(size=(1,))).astype(floatX)
-        params["sigma_irregular"] = sigma
-        sm_params["sigma2.irregular"] = sigma.item()
+        sigma2 = np.abs(rng.normal(size=(1,))).astype(floatX)
+        params["sigma_irregular"] = np.sqrt(sigma2)
+        sm_params["sigma2.irregular"] = sigma2.item()
         expected_param_dims["sigma_irregular"] += ("observed_state",)
 
         comp = st.MeasurementError("irregular")
@@ -255,7 +255,7 @@ def create_structural_model_and_equivalent_statsmodel(
             ).astype(floatX),
             np.zeros(2, dtype=floatX),
         )
-        sigma_level_value = np.abs(rng.normal(size=(2,)))[
+        sigma_level_value2 = np.abs(rng.normal(size=(2,)))[
             np.array(level_trend_innov_order, dtype="bool")
         ]
         max_order = np.flatnonzero(level_value)[-1].item() + 1
@@ -267,9 +267,9 @@ def create_structural_model_and_equivalent_statsmodel(
 
         if sum(level_trend_innov_order) > 0:
             expected_param_dims["sigma_trend"] += ("trend_shock",)
-            params["sigma_trend"] = sigma_level_value
+            params["sigma_trend"] = np.sqrt(sigma_level_value2)
 
-        sigma_level_value = sigma_level_value.tolist()
+        sigma_level_value = sigma_level_value2.tolist()
         if stochastic_level:
             sigma = sigma_level_value.pop(0)
             sm_params["sigma2.level"] = sigma
@@ -298,9 +298,9 @@ def create_structural_model_and_equivalent_statsmodel(
         sm_init.update(seasonal_dict)
 
         if stochastic_seasonal:
-            sigma = np.abs(rng.normal(size=(1,))).astype(floatX)
-            params["sigma_seasonal"] = sigma
-            sm_params["sigma2.seasonal"] = sigma
+            sigma2 = np.abs(rng.normal(size=(1,))).astype(floatX)
+            params["sigma_seasonal"] = np.sqrt(sigma2)
+            sm_params["sigma2.seasonal"] = sigma2
             expected_coords[SHOCK_DIM] += [
                 "seasonal",
             ]
@@ -343,9 +343,9 @@ def create_structural_model_and_equivalent_statsmodel(
                 state_count += 1
 
             if has_innov:
-                sigma = np.abs(rng.normal(size=(1,))).astype(floatX)
-                params[f"sigma_seasonal_{s}"] = sigma
-                sm_params[f"sigma2.freq_seasonal_{s}({n})"] = sigma
+                sigma2 = np.abs(rng.normal(size=(1,))).astype(floatX)
+                params[f"sigma_seasonal_{s}"] = np.sqrt(sigma2)
+                sm_params[f"sigma2.freq_seasonal_{s}({n})"] = sigma2
                 expected_coords[SHOCK_DIM] += state_names
                 expected_coords[SHOCK_AUX_DIM] += state_names
 
@@ -374,12 +374,12 @@ def create_structural_model_and_equivalent_statsmodel(
         sm_init["cycle.auxilliary"] = init_cycle[1]
 
         if stochastic_cycle:
-            sigma = np.abs(rng.normal(size=(1,))).astype(floatX)
-            params["sigma_cycle"] = sigma
+            sigma2 = np.abs(rng.normal(size=(1,))).astype(floatX)
+            params["sigma_cycle"] = np.sqrt(sigma2)
             expected_coords[SHOCK_DIM] += state_names
             expected_coords[SHOCK_AUX_DIM] += state_names
 
-            sm_params["sigma2.cycle"] = sigma
+            sm_params["sigma2.cycle"] = sigma2
 
         if damped_cycle:
             rho = rng.beta(1, 1, size=(1,)).astype(floatX)
@@ -398,18 +398,18 @@ def create_structural_model_and_equivalent_statsmodel(
     if autoregressive is not None:
         ar_names = [f"L{i+1}.data" for i in range(autoregressive)]
         ar_params = rng.normal(size=(autoregressive,)).astype(floatX)
-        sigma = np.abs(rng.normal(size=(1,))).astype(floatX)
+        sigma2 = np.abs(rng.normal(size=(1,))).astype(floatX)
 
         params["ar_params"] = ar_params
-        params["sigma_ar"] = sigma
+        params["sigma_ar"] = np.sqrt(sigma2)
         expected_param_dims["ar_params"] += (AR_PARAM_DIM,)
         expected_coords[AR_PARAM_DIM] += tuple(list(range(1, autoregressive + 1)))
         expected_coords[ALL_STATE_DIM] += ar_names
         expected_coords[ALL_STATE_AUX_DIM] += ar_names
         expected_coords[SHOCK_DIM] += ["ar_innovation"]
         expected_coords[SHOCK_AUX_DIM] += ["ar_innovation"]
 
-        sm_params["sigma2.ar"] = sigma
+        sm_params["sigma2.ar"] = sigma2
         for i, rho in enumerate(ar_params):
             sm_init[f"ar.L{i+1}"] = 0
             sm_params[f"ar.L{i+1}"] = rho