Skip to content

Commit c9458e7

Browse files
jessegrabowskijessegrabowski
authored
Merge pull request #4 from Dekermanjian/multivariate-structural
bug fix multivariate regression component univariate case
2 parents 4fc8db2 + 1c8b61a commit c9458e7

File tree

11 files changed

+1767
-788
lines changed

11 files changed

+1767
-788
lines changed

‎notebooks/multivariate_ssm.ipynb

Lines changed: 1443 additions & 703 deletions
Large diffs are not rendered by default.

‎pymc_extras/statespace/models/structural/components/level_trend.py

Lines changed: 19 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -164,44 +164,45 @@ def populate_component_properties(self):
164164
k_posdef = self.k_posdef // k_endog
165165

166166
name_slice = POSITION_DERIVATIVE_NAMES[:k_states]
167-
self.param_names = [f"{self.name}_initial"]
167+
self.param_names = [f"initial_{self.name}"]
168168
base_names = [name for name, mask in zip(name_slice, self._order_mask) if mask]
169169
self.state_names = [
170170
f"{name}[{obs_name}]" for obs_name in self.observed_state_names for name in base_names
171171
]
172-
self.param_dims = {f"{self.name}_initial": (f"{self.name}_state",)}
173-
self.coords = {f"{self.name}_state": base_names}
172+
self.param_dims = {f"initial_{self.name}": (f"state_{self.name}",)}
173+
self.coords = {f"state_{self.name}": base_names}
174174

175175
if k_endog > 1:
176-
self.param_dims[f"{self.name}_state"] = (
177-
f"{self.name}_endog",
178-
f"{self.name}_state",
176+
self.param_dims[f"state_{self.name}"] = (
177+
f"endog_{self.name}",
178+
f"state_{self.name}",
179179
)
180-
self.param_dims = {f"{self.name}_initial": (f"{self.name}_endog", f"{self.name}_state")}
181-
self.coords[f"{self.name}_endog"] = self.observed_state_names
180+
self.param_dims = {f"initial_{self.name}": (f"endog_{self.name}", f"state_{self.name}")}
181+
self.coords[f"endog_{self.name}"] = self.observed_state_names
182182

183183
shape = (k_endog, k_states) if k_endog > 1 else (k_states,)
184-
self.param_info = {f"{self.name}_initial": {"shape": shape, "constraints": None}}
184+
self.param_info = {f"initial_{self.name}": {"shape": shape, "constraints": None}}
185185

186186
if self.k_posdef > 0:
187-
self.param_names += [f"{self.name}_sigma"]
187+
self.param_names += [f"sigma_{self.name}"]
188188

189-
shock_base_names = [
189+
base_shock_names = [
190190
name for name, mask in zip(name_slice, self.innovations_order) if mask
191191
]
192+
192193
self.shock_names = [
193194
f"{name}[{obs_name}]"
194195
for obs_name in self.observed_state_names
195-
for name in shock_base_names
196+
for name in base_shock_names
196197
]
197198

198-
self.param_dims[f"{self.name}_sigma"] = (
199+
self.param_dims[f"sigma_{self.name}"] = (
199200
(f"{self.name}_shock",)
200201
if k_endog == 1
201-
else (f"{self.name}_endog", f"{self.name}_shock")
202+
else (f"endog_{self.name}", f"{self.name}_shock")
202203
)
203-
self.coords[f"{self.name}_shock"] = self.shock_names
204-
self.param_info[f"{self.name}_sigma"] = {
204+
self.coords[f"{self.name}_shock"] = base_shock_names
205+
self.param_info[f"sigma_{self.name}"] = {
205206
"shape": (k_posdef,) if k_endog == 1 else (k_endog, k_posdef),
206207
"constraints": "Positive",
207208
}
@@ -215,7 +216,7 @@ def make_symbolic_graph(self) -> None:
215216
k_posdef = self.k_posdef // k_endog
216217

217218
initial_trend = self.make_and_register_variable(
218-
f"{self.name}_initial",
219+
f"initial_{self.name}",
219220
shape=(k_states,) if k_endog == 1 else (k_endog, k_states),
220221
)
221222
self.ssm["initial_state", :] = initial_trend.ravel()
@@ -242,7 +243,7 @@ def make_symbolic_graph(self) -> None:
242243

243244
if k_posdef > 0:
244245
sigma_trend = self.make_and_register_variable(
245-
f"{self.name}_sigma",
246+
f"sigma_{self.name}",
246247
shape=(k_posdef,) if k_endog == 1 else (k_endog, k_posdef),
247248
)
248249
diag_idx = np.diag_indices(k_posdef * k_endog)

‎pymc_extras/statespace/models/structural/components/regression.py

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,9 @@ def populate_component_properties(self) -> None:
9292
self.param_names = [f"beta_{self.name}"]
9393
self.data_names = [f"data_{self.name}"]
9494
self.param_dims = {
95-
f"beta_{self.name}": (f"{self.name}_endog", f"{self.name}_state"),
95+
f"beta_{self.name}": (f"endog_{self.name}", f"state_{self.name}")
96+
if k_endog > 1
97+
else (f"state_{self.name}",)
9698
}
9799

98100
base_names = self.state_names
@@ -104,30 +106,30 @@ def populate_component_properties(self) -> None:
104106
f"beta_{self.name}": {
105107
"shape": (k_endog, k_states) if k_endog > 1 else (k_states,),
106108
"constraints": None,
107-
"dims": (f"{self.name}_endog", f"{self.name}_state")
109+
"dims": (f"endog_{self.name}", f"state_{self.name}")
108110
if k_endog > 1
109-
else (f"{self.name}_state",),
111+
else (f"state_{self.name}",),
110112
},
111113
}
112114

113115
self.data_info = {
114116
f"data_{self.name}": {
115117
"shape": (None, k_states),
116-
"dims": (TIME_DIM, "exog_state"),
118+
"dims": (TIME_DIM, f"state_{self.name}"),
117119
},
118120
}
119121
self.coords = {
120-
f"{self.name}_state": self.state_names,
121-
f"{self.name}_endog": self.observed_state_names,
122+
f"state_{self.name}": base_names,
123+
f"endog_{self.name}": self.observed_state_names,
122124
}
123125

124126
if self.innovations:
125127
self.param_names += [f"sigma_beta_{self.name}"]
126-
self.param_dims[f"sigma_beta_{self.name}"] = f"{self.name}_state"
128+
self.param_dims[f"sigma_beta_{self.name}"] = (f"state_{self.name}",)
127129
self.param_info[f"sigma_beta_{self.name}"] = {
128-
"shape": (),
130+
"shape": (k_states,),
129131
"constraints": "Positive",
130-
"dims": (f"{self.name}_state",)
132+
"dims": (f"state_{self.name}",)
131133
if k_endog == 1
132-
else (f"{self.name}_endog", f"{self.name}_state"),
134+
else (f"endog_{self.name}", f"state_{self.name}"),
133135
}

‎pymc_extras/statespace/models/structural/components/seasonality.py

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -191,8 +191,21 @@ def populate_component_properties(self):
191191
else (f"{self.name}_endog", f"{self.name}_state"),
192192
}
193193
}
194-
self.param_dims = {f"{self.name}_coefs": (f"{self.name}_state",)}
195-
self.coords = {f"{self.name}_state": self.state_names}
194+
195+
self.param_dims = {
196+
f"{self.name}_coefs": (f"{self.name}_state",)
197+
if k_endog == 1
198+
else (f"{self.name}_endog", f"{self.name}_state")
199+
}
200+
201+
self.coords = (
202+
{f"{self.name}_state": self.provided_state_names}
203+
if k_endog == 1
204+
else {
205+
f"{self.name}_endog": self.observed_state_names,
206+
f"{self.name}_state": self.provided_state_names,
207+
}
208+
)
196209

197210
if self.innovations:
198211
self.param_names += [f"sigma_{self.name}"]

0 commit comments

Comments
 (0)