pymc-extras 0.4.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

pymc_extras/statespace/models/structural/core.py

@@ -120,7 +120,7 @@ class StructuralTimeSeries(PyMCStateSpace):
             initial_trend = pm.Normal('initial_trend', sigma=10, dims=ss_mod.param_dims['initial_trend'])
             sigma_trend = pm.HalfNormal('sigma_trend', sigma=1, dims=ss_mod.param_dims['sigma_trend'])
 
-            seasonal_coefs = pm.Normal('seasonal_coefs', sigma=1, dims=ss_mod.param_dims['seasonal_coefs'])
+            seasonal_coefs = pm.Normal('params_seasonal', sigma=1, dims=ss_mod.param_dims['params_seasonal'])
             sigma_seasonal = pm.HalfNormal('sigma_seasonal', sigma=1)
 
             sigma_obs = pm.Exponential('sigma_obs', 1, dims=ss_mod.param_dims['sigma_obs'])
@@ -310,16 +310,20 @@ class StructuralTimeSeries(PyMCStateSpace):
 
         for i, (name, s) in enumerate(zip(names, state_slices)):
             obs_idx = info[name]["obs_state_idx"]
+
             if obs_idx is None:
                 continue
 
             X = data[..., s]
+
             if info[name]["combine_hidden_states"]:
-                sum_idx = np.flatnonzero(obs_idx)
-                result.append(X[..., sum_idx].sum(axis=-1)[..., None])
+                sum_idx_joined = np.flatnonzero(obs_idx)
+                sum_idx_split = np.split(sum_idx_joined, info[name]["k_endog"])
+                for sum_idx in sum_idx_split:
+                    result.append(X[..., sum_idx].sum(axis=-1)[..., None])
             else:
-                comp_names = self.state_names[s]
-                for j, state_name in enumerate(comp_names):
+                n_components = len(self.state_names[s])
+                for j in range(n_components):
                     result.append(X[..., j, None])
 
         return np.concatenate(result, axis=-1)
@@ -332,7 +336,15 @@ class StructuralTimeSeries(PyMCStateSpace):
 
         for i, (name, s) in enumerate(zip(names, state_slices)):
             if info[name]["combine_hidden_states"]:
-                result.append(name)
+                if self.k_endog == 1:
+                    result.append(name)
+                else:
+                    # If there are multiple observed states, we will combine per hidden state, preserving the
+                    # observed state names. Note this happens even if this *component* has only 1 state for consistency,
+                    # as long as the statespace model has multiple observed states.
+                    result.extend(
+                        [f"{name}[{obs_name}]" for obs_name in info[name]["observed_state_names"]]
+                    )
             else:
                 comp_names = self.state_names[s]
                 result.extend([f"{name}[{comp_name}]" for comp_name in comp_names])
@@ -350,7 +362,7 @@ class StructuralTimeSeries(PyMCStateSpace):
         Returns
         -------
         idata: Dataset
-            An Dataset object with hidden states transformed to represent only the "interpretable" subcomponents
+            A Dataset object with hidden states transformed to represent only the "interpretable" subcomponents
             of the structural model.
 
         Notes
@@ -459,6 +471,10 @@ class Component:
     obs_state_idxs : np.ndarray | None, optional
         Indices indicating which states contribute to observed variables. If None,
         defaults to None.
+    share_states : bool, optional
+        Whether states are shared across multiple endogenous variables in multivariate
+        models. When True, the same latent states affect all observed variables.
+        Default is False.
 
     Examples
     --------
@@ -500,10 +516,12 @@ class Component:
         combine_hidden_states=True,
         component_from_sum=False,
         obs_state_idxs=None,
+        share_states: bool = False,
     ):
         self.name = name
         self.k_endog = k_endog
         self.k_states = k_states
+        self.share_states = share_states
         self.k_posdef = k_posdef
         self.measurement_error = measurement_error
 
@@ -540,11 +558,12 @@ class Component:
         self._component_info = {
             self.name: {
                 "k_states": self.k_states,
-                "k_enodg": self.k_endog,
+                "k_endog": self.k_endog,
                 "k_posdef": self.k_posdef,
                 "observed_state_names": self.observed_state_names,
                 "combine_hidden_states": combine_hidden_states,
                 "obs_state_idx": obs_state_idxs,
+                "share_states": self.share_states,
             }
         }
 

pymc_extras/statespace/utils/constants.py

@@ -12,6 +12,7 @@ MA_PARAM_DIM = "lag_ma"
 SEASONAL_AR_PARAM_DIM = "seasonal_lag_ar"
 SEASONAL_MA_PARAM_DIM = "seasonal_lag_ma"
 ETS_SEASONAL_DIM = "seasonal_lag"
+EXOGENOUS_DIM = "exogenous"
 
 NEVER_TIME_VARYING = ["initial_state", "initial_state_cov", "a0", "P0"]
 VECTOR_VALUED = ["initial_state", "state_intercept", "obs_intercept", "a0", "c", "d"]
@@ -38,14 +39,16 @@ SHORT_NAME_TO_LONG = dict(zip(MATRIX_NAMES, LONG_MATRIX_NAMES))
 LONG_NAME_TO_SHORT = dict(zip(LONG_MATRIX_NAMES, MATRIX_NAMES))
 
 FILTER_OUTPUT_NAMES = [
-    "filtered_state",
-    "predicted_state",
-    "filtered_covariance",
-    "predicted_covariance",
+    "filtered_states",
+    "predicted_states",
+    "filtered_covariances",
+    "predicted_covariances",
+    "predicted_observed_states",
+    "predicted_observed_covariances",
 ]
 
-SMOOTHER_OUTPUT_NAMES = ["smoothed_state", "smoothed_covariance"]
-OBSERVED_OUTPUT_NAMES = ["predicted_observed_state", "predicted_observed_covariance"]
+SMOOTHER_OUTPUT_NAMES = ["smoothed_states", "smoothed_covariances"]
+OBSERVED_OUTPUT_NAMES = ["predicted_observed_states", "predicted_observed_covariances"]
 
 MATRIX_DIMS = {
     "x0": (ALL_STATE_DIM,),
@@ -60,14 +63,14 @@ MATRIX_DIMS = {
 }
 
 FILTER_OUTPUT_DIMS = {
-    "filtered_state": (TIME_DIM, ALL_STATE_DIM),
-    "smoothed_state": (TIME_DIM, ALL_STATE_DIM),
-    "predicted_state": (TIME_DIM, ALL_STATE_DIM),
-    "filtered_covariance": (TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM),
-    "smoothed_covariance": (TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM),
-    "predicted_covariance": (TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM),
-    "predicted_observed_state": (TIME_DIM, OBS_STATE_DIM),
-    "predicted_observed_covariance": (TIME_DIM, OBS_STATE_DIM, OBS_STATE_AUX_DIM),
+    "filtered_states": (TIME_DIM, ALL_STATE_DIM),
+    "smoothed_states": (TIME_DIM, ALL_STATE_DIM),
+    "predicted_states": (TIME_DIM, ALL_STATE_DIM),
+    "filtered_covariances": (TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM),
+    "smoothed_covariances": (TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM),
+    "predicted_covariances": (TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM),
+    "predicted_observed_states": (TIME_DIM, OBS_STATE_DIM),
+    "predicted_observed_covariances": (TIME_DIM, OBS_STATE_DIM, OBS_STATE_AUX_DIM),
 }
 
 POSITION_DERIVATIVE_NAMES = ["level", "trend", "acceleration", "jerk", "snap", "crackle", "pop"]

pymc_extras/statespace/utils/data_tools.py

@@ -53,7 +53,7 @@ def _validate_data_shape(data_shape, n_obs, obs_coords=None, check_col_names=Fal
         if len(missing_cols) > 0:
             raise ValueError(
                 "Columns of DataFrame provided as data do not match state names. The following states were"
-                f'not found: {", ".join(missing_cols)}. This may result in unexpected results in complex'
+                f"not found: {', '.join(missing_cols)}. This may result in unexpected results in complex"
                 f"statespace models"
             )
 

{pymc_extras-0.4.0.dist-info → pymc_extras-0.4.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pymc-extras
-Version: 0.4.0
+Version: 0.4.1
 Summary: A home for new additions to PyMC, which may include unusual probability distribitions, advanced model fitting algorithms, or any code that may be inappropriate to include in the pymc repository, but may want to be made available to users.
 Project-URL: Documentation, https://pymc-extras.readthedocs.io/
 Project-URL: Repository, https://github.com/pymc-devs/pymc-extras.git

{pymc_extras-0.4.0.dist-info → pymc_extras-0.4.1.dist-info}/RECORD

@@ -2,12 +2,12 @@ pymc_extras/__init__.py,sha256=YsR6OG72aW73y6dGS7w3nGGMV-V-ImHkmUOXKMPfMRA,1230
 pymc_extras/deserialize.py,sha256=dktK5gsR96X3zAUoRF5udrTiconknH3uupiAWqkZi0M,5937
 pymc_extras/linearmodel.py,sha256=KkvZ_DBXOD6myPgVNzu742YV0OzDK449_pDqNC5yae4,3975
 pymc_extras/model_builder.py,sha256=sAw77fxdiy046BvDPjocuMlbJ0Efj-CDAGtmcwYmoG0,26361
-pymc_extras/printing.py,sha256=G8mj9dRd6i0PcsbcEWZm56ek6V8mmil78RI4MUhywBs,6506
+pymc_extras/printing.py,sha256=bFOANgsOWDk0vbRMvm2h_D5TsT7OiSojdG7tvyfCw28,6506
 pymc_extras/prior.py,sha256=0XbyRRVuS7aKY5gmvJr_iq4fGyHrRDeI_OjWu_O7CTA,39449
-pymc_extras/distributions/__init__.py,sha256=fDbrBt9mxEVp2CDPwnyCW3oiutzZ0PduB8EUH3fUrjI,1377
+pymc_extras/distributions/__init__.py,sha256=Cge3AP7gzD6qTJY7v2tYRtSgn-rlnIo7wQBgf3IfKQ8,1377
 pymc_extras/distributions/continuous.py,sha256=530wvcO-QcYVdiVN-iQRveImWfyJzzmxiZLMVShP7w4,11251
 pymc_extras/distributions/discrete.py,sha256=HNi-K0_hnNWTcfyBkWGh26sc71FwBgukQ_EjGAaAOjY,13036
-pymc_extras/distributions/histogram_utils.py,sha256=5RTvlGCUrp2qzshrchmPyWxjhs6RIYL62SMikjDM1TU,5814
+pymc_extras/distributions/histogram_utils.py,sha256=xvCc19nlOmeb9PLZDcsR5PRdmcr5sRefZlPlCvxmGfM,5814
 pymc_extras/distributions/timeseries.py,sha256=M5MZ-nik_tgkaoZ1hdUGEZ9g04DQyVLwszVJqSKwNcY,12719
 pymc_extras/distributions/multivariate/__init__.py,sha256=E8OeLW9tTotCbrUjEo4um76-_WQD56PehsPzkKmhfyA,93
 pymc_extras/distributions/multivariate/r2d2m2cp.py,sha256=bUj9bB-hQi6CpaJfvJjgNPi727uTbvAdxl9fm1zNBqY,16005
@@ -15,7 +15,7 @@ pymc_extras/distributions/transforms/__init__.py,sha256=FUp2vyRE6_2eUcQ_FVt5Dn0-
 pymc_extras/distributions/transforms/partial_order.py,sha256=oEZlc9WgnGR46uFEjLzKEUxlhzIo2vrUUbBE3vYrsfQ,8404
 pymc_extras/gp/__init__.py,sha256=sFHw2y3lEl5tG_FDQHZUonQ_k0DF1JRf0Rp8dpHmge0,745
 pymc_extras/gp/latent_approx.py,sha256=cDEMM6H1BL2qyKg7BZU-ISrKn2HJe7hDaM4Y8GgQDf4,6682
-pymc_extras/inference/__init__.py,sha256=YJIBqHoJnjglof7SVESH3u67li_ETmMy24zajld0DNE,917
+pymc_extras/inference/__init__.py,sha256=sy1JYQGNZNvPs-3jVFfbFQTW0iCIrbjH3aHBpx1HQi0,917
 pymc_extras/inference/fit.py,sha256=U_jfzuyjk5bV6AvOxtOKzBg-q4z-_BOR06Hn38T0W6E,1328
 pymc_extras/inference/laplace_approx/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pymc_extras/inference/laplace_approx/find_map.py,sha256=fP8DQ21OZbkUiBaq-TXGe7CtH0umupFacRC3qReoiKU,14022
@@ -38,41 +38,41 @@ pymc_extras/model/transforms/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NM
 pymc_extras/model/transforms/autoreparam.py,sha256=_NltGWmNqi_X9sHCqAvWcBveLTPxVy11-wENFTcN6kk,12377
 pymc_extras/preprocessing/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pymc_extras/preprocessing/standard_scaler.py,sha256=Vajp33ma6OkwlU54JYtSS8urHbMJ3CRiRFxZpvFNuus,600
-pymc_extras/statespace/__init__.py,sha256=0MtZj7yT6jcyERvITnn-nkhyY8fO6Za4_vV53CF6ND0,429
-pymc_extras/statespace/core/__init__.py,sha256=huHEiXAm8zV2MZyZ8GBHp6q7_fnWqveM7lC6ilpb3iE,309
+pymc_extras/statespace/__init__.py,sha256=PxV8i4aa2XJarRM6aKU14_bEY1AoLu4bNXIBy_E1rRw,431
+pymc_extras/statespace/core/__init__.py,sha256=LEhkqdMZzzcTyzYml45IM4ykWoCdbWWj2c29IpM_ey8,309
 pymc_extras/statespace/core/compile.py,sha256=9FZfE8Bi3VfElxujfOIKRVvmyL9M5R0WfNEqPc5kbVQ,1603
-pymc_extras/statespace/core/representation.py,sha256=DwNIun6wdeEA20oWBx5M4govyWTf5JI87aGQ_E6Mb4U,18956
-pymc_extras/statespace/core/statespace.py,sha256=mr0jDeCEfclCAGfEfEzMsogiuonGO9k2BV67wH_YKio,105627
-pymc_extras/statespace/filters/__init__.py,sha256=N9Q4D0gAq_ZtT-GtrqiX1HkSg6Orv7o1TbrWUtnbTJE,420
+pymc_extras/statespace/core/representation.py,sha256=boY-jjlkd3KuuO2XiSuV-GwEAyEqRJ9267H72AmE3BU,18956
+pymc_extras/statespace/core/statespace.py,sha256=yu7smA5w7l1LFNjTwuKLnGarGLx4HEPJKQ9ZMDbWhDY,108161
+pymc_extras/statespace/filters/__init__.py,sha256=F0EtZUhArp23lj3upy6zB0mDTjLIjwGh0pKmMny0QfY,420
 pymc_extras/statespace/filters/distributions.py,sha256=-s1c5s2zm6FMc0UqKSrWnJzIF4U5bvJT_3mMNTyV_ak,11927
-pymc_extras/statespace/filters/kalman_filter.py,sha256=Z6kxsbW8_VQ6ZcPjDMA5d_XPfdUY1-4GfRwKbBNfVZs,31438
+pymc_extras/statespace/filters/kalman_filter.py,sha256=rgpgF4KZXX5M8yRwblrt2SEINKgoXgiKNfKkbl7ZU9Y,31464
 pymc_extras/statespace/filters/kalman_smoother.py,sha256=5jlSZAPveJzD5Q8omnpn7Gb1jgElBMgixGR7H9zoH8U,4183
 pymc_extras/statespace/filters/utilities.py,sha256=iwdaYnO1cO06t_XUjLLRmqb8vwzzVH6Nx1iyZcbJL2k,1584
 pymc_extras/statespace/models/ETS.py,sha256=08sbiuNvKdxcgKzS7jWj-z4jf-su73WFkYc8sKkGdEs,28538
-pymc_extras/statespace/models/SARIMAX.py,sha256=aXR6KYuqtSBOk-jvm9NvnOX5vu4QesBgCIL-KR89SXs,22207
-pymc_extras/statespace/models/VARMAX.py,sha256=CX8Gs4GP3u6aD2b8nxonA5zCoQzIca668iLF0SkKF00,16456
-pymc_extras/statespace/models/__init__.py,sha256=U79b8rTHBNijVvvGOd43nLu4PCloPUH1rwlN87-n88c,317
+pymc_extras/statespace/models/SARIMAX.py,sha256=Yppz_k1ZyZuKPC62WIye6K7luw44cP-dog73VVkw0L4,25096
+pymc_extras/statespace/models/VARMAX.py,sha256=7obJFXES9t9NONlcUQoeJ9TCqyoDlVat9FkPviQhAq0,25947
+pymc_extras/statespace/models/__init__.py,sha256=DUwPrwfnz9AUbgZOFvZeUpWEw5FiPAK5X9x7vZrRWqY,319
 pymc_extras/statespace/models/utilities.py,sha256=jpUYByAy6rMFP7l56uST1SEYchRa-clsFQ-At_1NLSw,27123
-pymc_extras/statespace/models/structural/__init__.py,sha256=lplt7FYxeV4jt3zx0jAkaqeCyk-bVrYSk7z6Ubz86O8,811
-pymc_extras/statespace/models/structural/core.py,sha256=Vm61KhlgTHLvTE_kCSZu63G-UJRgfD6KtzVjwqYsrvk,34350
+pymc_extras/statespace/models/structural/__init__.py,sha256=jvbczE1IeNkhW7gMQ2vF2BhhKHeYyfD90mV-Awko-Vs,811
+pymc_extras/statespace/models/structural/core.py,sha256=n0cbP8_-NFLmflFF4x37AyOOIHcY5iylRrgTzjyOAhM,35374
 pymc_extras/statespace/models/structural/utils.py,sha256=Eze34Z0iXJzDC_gZEY2mHrp2VIYu8rHV915vM4U5Sn4,359
 pymc_extras/statespace/models/structural/components/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pymc_extras/statespace/models/structural/components/autoregressive.py,sha256=vaFmfJ0EljUQKiTTqqMthMztK4jVXuxBu5sxhu9esPU,6861
-pymc_extras/statespace/models/structural/components/cycle.py,sha256=CGuV0TJSj9KkO7eHHZUxAhaGKS1erODkv-dLXiVkNUI,12907
-pymc_extras/statespace/models/structural/components/level_trend.py,sha256=rfRHuWlI84QA_as1mWUnp9Z-KeHp9rGLyEL-GUSsfVQ,9895
-pymc_extras/statespace/models/structural/components/measurement_error.py,sha256=49oN77hUpZMS-1M5DxDvKREAi4Tp041JrrF33wuNeAs,4628
-pymc_extras/statespace/models/structural/components/regression.py,sha256=D1NJHlVf8bOYxmQMVDZ1dE1JiVwnBr73rvvmYym4i9c,8571
-pymc_extras/statespace/models/structural/components/seasonality.py,sha256=ZByEMpovVO1Y-124OLB894td0YVy6HaM3vjYhVYANC4,19093
+pymc_extras/statespace/models/structural/components/autoregressive.py,sha256=HkS5an5fuNOBGcjHFNMUVNJrF1BNnlpQxvmPq_5dD0s,8021
+pymc_extras/statespace/models/structural/components/cycle.py,sha256=qEiGFGMEXKS2Tl_zgzKIp77ijGXCVq6UIHEZp_ErHSQ,13931
+pymc_extras/statespace/models/structural/components/level_trend.py,sha256=7glYX_tKOJPq6uB1NBuPQFFZGkhcwK4GMZUBTcU0xIY,11357
+pymc_extras/statespace/models/structural/components/measurement_error.py,sha256=5LHDx3IplNrWSGcsY3xJLywKPosTqr42jlrvm80ZApM,5316
+pymc_extras/statespace/models/structural/components/regression.py,sha256=27PRV9I64_VXIyjUi7pRr_gbk7sSI5DfJ4FBAbq5WCM,9856
+pymc_extras/statespace/models/structural/components/seasonality.py,sha256=soXJIZ2xewUhSUb5s2MGnxvnQCcir7ZgbgkSr94xEvc,26987
 pymc_extras/statespace/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pymc_extras/statespace/utils/constants.py,sha256=S2J4FGqMKNQJQnDjE9y4ul0D_pwSKgiu5kii7Dem_cA,2403
+pymc_extras/statespace/utils/constants.py,sha256=-4vCXo7-X3IuzdcplWBrAV9m9tm8JngcgoE-8imGmj0,2518
 pymc_extras/statespace/utils/coord_tools.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pymc_extras/statespace/utils/data_tools.py,sha256=QDBcELbE5Dyng24tE3Pl5IiPOxl4EDEH73lRKMymIOE,6620
+pymc_extras/statespace/utils/data_tools.py,sha256=Tomur7d8WCKlMXUCrPqufqVTKUe_nLLCHdipsM9pmaI,6620
 pymc_extras/utils/__init__.py,sha256=yxI9cJ7fCtVQS0GFw0y6mDGZIQZiK53vm3UNKqIuGSk,758
 pymc_extras/utils/linear_cg.py,sha256=KkXhuimFsrKtNd_0By2ApxQQQNm5FdBtmDQJOVbLYkA,10056
 pymc_extras/utils/model_equivalence.py,sha256=8QIftID2HDxD659i0RXHazQ-l2Q5YegCRLcDqb2p9Pc,2187
 pymc_extras/utils/prior.py,sha256=QlWVr7uKIK9VncBw7Fz3YgaASKGDfqpORZHc-vz_9gQ,6841
 pymc_extras/utils/spline.py,sha256=qGq0gcoMG5dpdazKFzG0RXkkCWP8ADPPXN-653-oFn4,4820
-pymc_extras-0.4.0.dist-info/METADATA,sha256=NilZkqZEg7D4NfORAYoHEKpGnv1E0P5_GjmAB42zqzA,18898
-pymc_extras-0.4.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-pymc_extras-0.4.0.dist-info/licenses/LICENSE,sha256=WjiLhUKEysJvy5e9jk6WwFv9tmAPtnov1uJ6gcH1kIs,11720
-pymc_extras-0.4.0.dist-info/RECORD,,
+pymc_extras-0.4.1.dist-info/METADATA,sha256=TpuX_8nEFjQfPlC51u_2EvQV3XwHAvgYCQMKYzeVU_E,18898
+pymc_extras-0.4.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pymc_extras-0.4.1.dist-info/licenses/LICENSE,sha256=WjiLhUKEysJvy5e9jk6WwFv9tmAPtnov1uJ6gcH1kIs,11720
+pymc_extras-0.4.1.dist-info/RECORD,,