diff --git a/pymc_experimental/statespace/models/SARIMAX.py b/pymc_experimental/statespace/models/SARIMAX.py index 54b16fe60..52caa48e4 100644 --- a/pymc_experimental/statespace/models/SARIMAX.py +++ b/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 diff --git a/pymc_experimental/statespace/models/structural.py b/pymc_experimental/statespace/models/structural.py index 0a45c6f0d..bc8f4c429 100644 --- a/pymc_experimental/statespace/models/structural.py +++ b/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 diff --git a/pymc_experimental/tests/statespace/test_SARIMAX.py b/pymc_experimental/tests/statespace/test_SARIMAX.py index e4a01f86d..4d206076c 100644 --- a/pymc_experimental/tests/statespace/test_SARIMAX.py +++ b/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) diff --git a/pymc_experimental/tests/statespace/test_structural.py b/pymc_experimental/tests/statespace/test_structural.py index d6fa46e9f..d65ac972a 100644 --- a/pymc_experimental/tests/statespace/test_structural.py +++ b/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,10 +398,10 @@ 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 @@ -409,7 +409,7 @@ def create_structural_model_and_equivalent_statsmodel( 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