Add ETS model

pymc_experimental/statespace/__init__.py

@@ -1,5 +1,7 @@
+from pymc_experimental.statespace.core.compile import compile_statespace
 from pymc_experimental.statespace.models import structural
+from pymc_experimental.statespace.models.ETS import BayesianETS
 from pymc_experimental.statespace.models.SARIMAX import BayesianSARIMA
 from pymc_experimental.statespace.models.VARMAX import BayesianVARMAX
 
-__all__ = ["structural", "BayesianSARIMA", "BayesianVARMAX"]
+__all__ = ["structural", "BayesianSARIMA", "BayesianVARMAX", "BayesianETS", "compile_statespace"]

pymc_experimental/statespace/core/__init__.py

@@ -1,4 +1,7 @@
+# ruff: noqa: I001
+
 from pymc_experimental.statespace.core.representation import PytensorRepresentation
 from pymc_experimental.statespace.core.statespace import PyMCStateSpace
+from pymc_experimental.statespace.core.compile import compile_statespace
 
-__all__ = ["PytensorRepresentation", "PyMCStateSpace"]
+__all__ = ["PytensorRepresentation", "PyMCStateSpace", "compile_statespace"]

pymc_experimental/statespace/core/compile.py

@@ -0,0 +1,48 @@
+import numpy as np
+import pymc as pm
+import pytensor
+import pytensor.tensor as pt
+
+from pymc_experimental.statespace.core import PyMCStateSpace
+from pymc_experimental.statespace.filters.distributions import LinearGaussianStateSpace
+from pymc_experimental.statespace.utils.constants import SHORT_NAME_TO_LONG
+
+
+def compile_statespace(
+    statespace_model: PyMCStateSpace, steps: int | None = None, **compile_kwargs
+):
+    if steps is None:
+        steps = pt.iscalar("steps")
+
+    x0, _, c, d, T, Z, R, H, Q = statespace_model._unpack_statespace_with_placeholders()
+
+    sequence_names = [x.name for x in [c, d] if x.ndim == 2]
+    sequence_names += [x.name for x in [T, Z, R, H, Q] if x.ndim == 3]
+
+    rename_dict = {v: k for k, v in SHORT_NAME_TO_LONG.items()}
+    sequence_names = list(map(rename_dict.get, sequence_names))
+
+    P0 = pt.zeros((x0.shape[0], x0.shape[0]))
+
+    outputs = LinearGaussianStateSpace.dist(
+        x0, P0, c, d, T, Z, R, H, Q, steps=steps, sequence_names=sequence_names
+    )
+
+    inputs = list(pytensor.graph.basic.explicit_graph_inputs(outputs))
+
+    _f = pm.compile_pymc(inputs, outputs, on_unused_input="ignore", **compile_kwargs)
+
+    def f(*, draws=1, **params):
+        if isinstance(steps, pt.Variable):
+            inner_steps = params.get("steps", 100)
+        else:
+            inner_steps = steps
+
+        output = [np.empty((draws, inner_steps + 1, x.type.shape[-1])) for x in outputs]
+        for i in range(draws):
+            draw = _f(**params)
+            for j, x in enumerate(draw):
+                output[j][i] = x
+        return [x.squeeze() for x in output]
+
+    return f

pymc_experimental/statespace/core/statespace.py

@@ -304,7 +304,19 @@ def _print_data_requirements(self) -> None:
             f"{out}"
         )
 
-    def _unpack_statespace_with_placeholders(self) -> tuple[pt.TensorVariable]:
+    def _unpack_statespace_with_placeholders(
+        self,
+    ) -> tuple[
+        pt.TensorVariable,
+        pt.TensorVariable,
+        pt.TensorVariable,
+        pt.TensorVariable,
+        pt.TensorVariable,
+        pt.TensorVariable,
+        pt.TensorVariable,
+        pt.TensorVariable,
+        pt.TensorVariable,
+    ]:
         """
         Helper function to quickly obtain all statespace matrices in the standard order. Matrices returned by this
         method will include pytensor placeholders.
@@ -448,8 +460,8 @@ def add_default_priors(self) -> None:
         raise NotImplementedError("The add_default_priors property has not been implemented!")
 
     def make_and_register_variable(
-        self, name, shape: int | tuple[int] | None = None, dtype=floatX
-    ) -> Variable:
+        self, name, shape: int | tuple[int, ...] | None = None, dtype=floatX
+    ) -> pt.TensorVariable:
         """
         Helper function to create a pytensor symbolic variable and register it in the _name_to_variable dictionary