pymc-extras 0.2.6__py3-none-any.whl → 0.3.1__py3-none-any.whl

pymc_extras/statespace/core/statespace.py

@@ -2047,6 +2047,69 @@ class PyMCStateSpace:
 
         return scenario
 
+    def _build_forecast_model(
+        self, time_index, t0, forecast_index, scenario, filter_output, mvn_method
+    ):
+        filter_time_dim = TIME_DIM
+        temp_coords = self._fit_coords.copy()
+
+        dims = None
+        if all([dim in temp_coords for dim in [filter_time_dim, ALL_STATE_DIM, OBS_STATE_DIM]]):
+            dims = [TIME_DIM, ALL_STATE_DIM, OBS_STATE_DIM]
+
+        t0_idx = np.flatnonzero(time_index == t0)[0]
+
+        temp_coords["data_time"] = time_index
+        temp_coords[TIME_DIM] = forecast_index
+
+        mu_dims, cov_dims = None, None
+        if all([dim in self._fit_coords for dim in [TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM]]):
+            mu_dims = ["data_time", ALL_STATE_DIM]
+            cov_dims = ["data_time", ALL_STATE_DIM, ALL_STATE_AUX_DIM]
+
+        with pm.Model(coords=temp_coords) as forecast_model:
+            (_, _, *matrices), grouped_outputs = self._kalman_filter_outputs_from_dummy_graph(
+                data_dims=["data_time", OBS_STATE_DIM],
+            )
+
+            group_idx = FILTER_OUTPUT_TYPES.index(filter_output)
+            mu, cov = grouped_outputs[group_idx]
+
+            sub_dict = {
+                data_var: pt.as_tensor_variable(data_var.get_value(), name="data")
+                for data_var in forecast_model.data_vars
+            }
+
+            missing_data_vars = np.setdiff1d(
+                ar1=[*self.data_names, "data"], ar2=[k.name for k, _ in sub_dict.items()]
+            )
+            if missing_data_vars.size > 0:
+                raise ValueError(f"{missing_data_vars} data used for fitting not found!")
+
+            mu_frozen, cov_frozen = graph_replace([mu, cov], replace=sub_dict, strict=True)
+
+            x0 = pm.Deterministic(
+                "x0_slice", mu_frozen[t0_idx], dims=mu_dims[1:] if mu_dims is not None else None
+            )
+            P0 = pm.Deterministic(
+                "P0_slice", cov_frozen[t0_idx], dims=cov_dims[1:] if cov_dims is not None else None
+            )
+
+            _ = LinearGaussianStateSpace(
+                "forecast",
+                x0,
+                P0,
+                *matrices,
+                steps=len(forecast_index),
+                dims=dims,
+                sequence_names=self.kalman_filter.seq_names,
+                k_endog=self.k_endog,
+                append_x0=False,
+                method=mvn_method,
+            )
+
+        return forecast_model
+
     def forecast(
         self,
         idata: InferenceData,
@@ -2139,8 +2202,6 @@ class PyMCStateSpace:
                   the latent state trajectories: `y[t] = Z @ x[t] + nu[t]`, where `nu ~ N(0, H)`.
 
         """
-        filter_time_dim = TIME_DIM
-
         _validate_filter_arg(filter_output)
 
         compile_kwargs = kwargs.pop("compile_kwargs", {})
@@ -2185,58 +2246,23 @@ class PyMCStateSpace:
             use_scenario_index=use_scenario_index,
         )
         scenario = self._finalize_scenario_initialization(scenario, forecast_index)
-        temp_coords = self._fit_coords.copy()
-
-        dims = None
-        if all([dim in temp_coords for dim in [filter_time_dim, ALL_STATE_DIM, OBS_STATE_DIM]]):
-            dims = [TIME_DIM, ALL_STATE_DIM, OBS_STATE_DIM]
-
-        t0_idx = np.flatnonzero(time_index == t0)[0]
-
-        temp_coords["data_time"] = time_index
-        temp_coords[TIME_DIM] = forecast_index
-
-        mu_dims, cov_dims = None, None
-        if all([dim in self._fit_coords for dim in [TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM]]):
-            mu_dims = ["data_time", ALL_STATE_DIM]
-            cov_dims = ["data_time", ALL_STATE_DIM, ALL_STATE_AUX_DIM]
-
-        with pm.Model(coords=temp_coords) as forecast_model:
-            (_, _, *matrices), grouped_outputs = self._kalman_filter_outputs_from_dummy_graph(
-                scenario=scenario,
-                data_dims=["data_time", OBS_STATE_DIM],
-            )
-
-            for name in self.data_names:
-                if name in scenario.keys():
-                    pm.set_data(
-                        {"data": np.zeros((len(forecast_index), self.k_endog))},
-                        coords={"data_time": np.arange(len(forecast_index))},
-                    )
-                    break
 
-            group_idx = FILTER_OUTPUT_TYPES.index(filter_output)
-            mu, cov = grouped_outputs[group_idx]
-
-            x0 = pm.Deterministic(
-                "x0_slice", mu[t0_idx], dims=mu_dims[1:] if mu_dims is not None else None
-            )
-            P0 = pm.Deterministic(
-                "P0_slice", cov[t0_idx], dims=cov_dims[1:] if cov_dims is not None else None
-            )
+        forecast_model = self._build_forecast_model(
+            time_index=time_index,
+            t0=t0,
+            forecast_index=forecast_index,
+            scenario=scenario,
+            filter_output=filter_output,
+            mvn_method=mvn_method,
+        )
 
-            _ = LinearGaussianStateSpace(
-                "forecast",
-                x0,
-                P0,
-                *matrices,
-                steps=len(forecast_index),
-                dims=dims,
-                sequence_names=self.kalman_filter.seq_names,
-                k_endog=self.k_endog,
-                append_x0=False,
-                method=mvn_method,
-            )
+        with forecast_model:
+            if scenario is not None:
+                dummy_obs_data = np.zeros((len(forecast_index), self.k_endog))
+                pm.set_data(
+                    scenario | {"data": dummy_obs_data},
+                    coords={"data_time": np.arange(len(forecast_index))},
+                )
 
         forecast_model.rvs_to_initial_values = {
             k: None for k in forecast_model.rvs_to_initial_values.keys()

pymc_extras/statespace/filters/kalman_smoother.py

@@ -105,7 +105,7 @@ class KalmanSmoother:
         a_hat, P_hat = self.predict(a, P, T, R, Q)
 
         # Use pinv, otherwise P_hat is singular when there is missing data
-        smoother_gain = matrix_dot(pt.linalg.pinv(P_hat), T, P).T
+        smoother_gain = matrix_dot(pt.linalg.pinv(P_hat, hermitian=True), T, P).T
         a_smooth_next = a + smoother_gain @ (a_smooth - a_hat)
 
         P_smooth_next = P + quad_form_sym(smoother_gain, P_smooth - P_hat)

{pymc_extras-0.2.6.dist-info → pymc_extras-0.3.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pymc-extras
-Version: 0.2.6
+Version: 0.3.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
@@ -226,14 +226,14 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Scientific/Engineering :: Mathematics
-Requires-Python: >=3.10
-Requires-Dist: better-optimize>=0.1.2
+Requires-Python: >=3.11
+Requires-Dist: better-optimize>=0.1.4
+Requires-Dist: pydantic>=2.0.0
 Requires-Dist: pymc>=5.21.1
 Requires-Dist: scikit-learn
 Provides-Extra: complete
@@ -245,6 +245,8 @@ Requires-Dist: xhistogram; extra == 'dask-histogram'
 Provides-Extra: dev
 Requires-Dist: blackjax; extra == 'dev'
 Requires-Dist: dask[all]<2025.1.1; extra == 'dev'
+Requires-Dist: preliz>=0.5.0; extra == 'dev'
+Requires-Dist: pytest-mock; extra == 'dev'
 Requires-Dist: pytest>=6.0; extra == 'dev'
 Requires-Dist: statsmodels; extra == 'dev'
 Provides-Extra: docs

{pymc_extras-0.2.6.dist-info → pymc_extras-0.3.1.dist-info}/RECORD

@@ -1,7 +1,9 @@
 pymc_extras/__init__.py,sha256=YsR6OG72aW73y6dGS7w3nGGMV-V-ImHkmUOXKMPfMRA,1230
-pymc_extras/linearmodel.py,sha256=6eitl15Ec15mSZu7zoHZ7Wwy4U1DPwqfAgwEt6ILeIc,3920
+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/prior.py,sha256=0XbyRRVuS7aKY5gmvJr_iq4fGyHrRDeI_OjWu_O7CTA,39449
 pymc_extras/distributions/__init__.py,sha256=fDbrBt9mxEVp2CDPwnyCW3oiutzZ0PduB8EUH3fUrjI,1377
 pymc_extras/distributions/continuous.py,sha256=530wvcO-QcYVdiVN-iQRveImWfyJzzmxiZLMVShP7w4,11251
 pymc_extras/distributions/discrete.py,sha256=HNi-K0_hnNWTcfyBkWGh26sc71FwBgukQ_EjGAaAOjY,13036
@@ -13,21 +15,24 @@ 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=UH6S0bGfQKKyTSuqf7yezdy9PeE2bDU8U1v4eIRv4ZI,887
-pymc_extras/inference/find_map.py,sha256=g_qXZbMz6w-De9wCMbBx8yLNkQANdPVWxLN7nJ0O17I,18523
-pymc_extras/inference/fit.py,sha256=oe20RAajImZ-VD9Ucbzri8Bof4Y2KHNhNRG19v9O3lI,1336
-pymc_extras/inference/laplace.py,sha256=Rq_D6veUYmW93GEyU8UZXiQquvJw-lK1np7NPxKCFqU,22064
+pymc_extras/inference/__init__.py,sha256=YJIBqHoJnjglof7SVESH3u67li_ETmMy24zajld0DNE,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=DihZO3XgiIs7JFKgriluSQDdinMQZG1xjkrQ4OUREno,13859
+pymc_extras/inference/laplace_approx/idata.py,sha256=Ldkzny9qgYg9My2ydIaQfzKJ73y-xcgC7VJiGvbZhWE,13221
+pymc_extras/inference/laplace_approx/laplace.py,sha256=M8s8GfHE5SgNLOiyiHsDuMDvYGEaBMjjqthVnrxY248,18724
+pymc_extras/inference/laplace_approx/scipy_interface.py,sha256=qMxYodmmxaUGsOp1jc7HxBJc6L8NnmFT2Fd4UNNXu2c,8835
 pymc_extras/inference/pathfinder/__init__.py,sha256=FhAYrCWNx_dCrynEdjg2CZ9tIinvcVLBm67pNx_Y3kA,101
 pymc_extras/inference/pathfinder/importance_sampling.py,sha256=NwxepXOFit3cA5zEebniKdlnJ1rZWg56aMlH4MEOcG4,6264
 pymc_extras/inference/pathfinder/lbfgs.py,sha256=GOoJBil5Kft_iFwGNUGKSeqzI5x_shA4KQWDwgGuQtQ,7110
-pymc_extras/inference/pathfinder/pathfinder.py,sha256=GW04HQurj_3Nlo1C6_K2tEIeigo8x0buV3FqDLA88PQ,64439
+pymc_extras/inference/pathfinder/pathfinder.py,sha256=yme_wBHnREaT5gSOD6CZ0nb87oScmXjplERiQb0mcAg,64454
 pymc_extras/inference/smc/__init__.py,sha256=wyaT4NJl1YsSQRLiDy-i0Jq3CbJZ2BQd4nnCk-dIngY,603
 pymc_extras/inference/smc/sampling.py,sha256=AYwmKqGoV6pBtKnh9SUbBKbN7VcoFgb3MmNWV7SivMA,15365
 pymc_extras/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pymc_extras/model/model_api.py,sha256=UHMfQXxWBujeSiUySU0fDUC5Sd_BjT8FoVz3iBxQH_4,2400
 pymc_extras/model/marginal/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pymc_extras/model/marginal/distributions.py,sha256=iM1yT7_BmivgUSloQPKE2QXGPgjvLqDMY_OTBGsdAWg,15563
-pymc_extras/model/marginal/graph_analysis.py,sha256=0hWUH_PjfpgneQ3NaT__pWHS1fh50zNbI86kH4Nub0E,15693
+pymc_extras/model/marginal/graph_analysis.py,sha256=l_WSZHivm82297zMIm8i3G_h2F-4Tq397pQlcuEP-0I,15874
 pymc_extras/model/marginal/marginal_model.py,sha256=oIdikaSnefCkyMxmzAe222qGXNucxZpHYk7548fK6iA,23631
 pymc_extras/model/transforms/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pymc_extras/model/transforms/autoreparam.py,sha256=_NltGWmNqi_X9sHCqAvWcBveLTPxVy11-wENFTcN6kk,12377
@@ -37,11 +42,11 @@ pymc_extras/statespace/__init__.py,sha256=0MtZj7yT6jcyERvITnn-nkhyY8fO6Za4_vV53C
 pymc_extras/statespace/core/__init__.py,sha256=huHEiXAm8zV2MZyZ8GBHp6q7_fnWqveM7lC6ilpb3iE,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=9jCQ4odmLK3S33tQzKMqck2gsgVoo-C3hOCBX5dc9lA,104674
+pymc_extras/statespace/core/statespace.py,sha256=mr0jDeCEfclCAGfEfEzMsogiuonGO9k2BV67wH_YKio,105627
 pymc_extras/statespace/filters/__init__.py,sha256=N9Q4D0gAq_ZtT-GtrqiX1HkSg6Orv7o1TbrWUtnbTJE,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_smoother.py,sha256=SAjnqtiDdvV79Pp4jp6UzrdIMmH1lqXhCj5WLeHusr8,4167
+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
@@ -56,10 +61,9 @@ pymc_extras/statespace/utils/data_tools.py,sha256=01sz6XDtLYK9I5xghxYpD-PuDzGXv9
 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/pivoted_cholesky.py,sha256=QtnjP0pAl9b77fLAu-semwT4_9dcoiqx3dz1xKGBjMk,1871
 pymc_extras/utils/prior.py,sha256=QlWVr7uKIK9VncBw7Fz3YgaASKGDfqpORZHc-vz_9gQ,6841
 pymc_extras/utils/spline.py,sha256=qGq0gcoMG5dpdazKFzG0RXkkCWP8ADPPXN-653-oFn4,4820
-pymc_extras-0.2.6.dist-info/METADATA,sha256=zzdhVkdzXhL7MQH3R0uiCsrcl5i5uh1JLVdRBG6jJyY,18813
-pymc_extras-0.2.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-pymc_extras-0.2.6.dist-info/licenses/LICENSE,sha256=WjiLhUKEysJvy5e9jk6WwFv9tmAPtnov1uJ6gcH1kIs,11720
-pymc_extras-0.2.6.dist-info/RECORD,,
+pymc_extras-0.3.1.dist-info/METADATA,sha256=QxE1LJTwDeUfZzYZhDjhLA4sy0SX6UcXGDyeVC1dkmM,18881
+pymc_extras-0.3.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pymc_extras-0.3.1.dist-info/licenses/LICENSE,sha256=WjiLhUKEysJvy5e9jk6WwFv9tmAPtnov1uJ6gcH1kIs,11720
+pymc_extras-0.3.1.dist-info/RECORD,,