pymc-extras 0.2.0__py3-none-any.whl

pymc_extras/preprocessing/standard_scaler.py

@@ -0,0 +1,17 @@
+import pandas as pd
+
+from sklearn.base import BaseEstimator, TransformerMixin
+from sklearn.preprocessing import StandardScaler
+
+
+class StandardScalerDF(StandardScaler, TransformerMixin, BaseEstimator):
+    def __init__(self, with_mean=True, with_std=True):
+        super().__init__(with_mean=with_mean, with_std=with_std)
+
+    def transform(self, X, y=None):
+        z = super().transform(X)
+        return pd.DataFrame(z, index=X.index, columns=X.columns)
+
+    def fit_transform(self, X, y=None):
+        z = super().fit_transform(X)
+        return pd.DataFrame(z, index=X.index, columns=X.columns)

pymc_extras/printing.py

@@ -0,0 +1,182 @@
+import numpy as np
+
+from pymc import Model
+from pymc.printing import str_for_dist, str_for_potential_or_deterministic
+from pytensor import Mode
+from pytensor.compile.sharedvalue import SharedVariable
+from pytensor.graph.type import Constant, Variable
+from rich.box import SIMPLE_HEAD
+from rich.table import Table
+
+
+def variable_expression(
+    model: Model,
+    var: Variable,
+    truncate_deterministic: int | None,
+) -> str:
+    """Get the expression of a variable in a human-readable format."""
+    if var in model.data_vars:
+        var_expr = "Data"
+    elif var in model.deterministics:
+        str_repr = str_for_potential_or_deterministic(var, dist_name="")
+        _, var_expr = str_repr.split(" ~ ")
+        var_expr = var_expr[1:-1]  # Remove outer parentheses (f(...))
+        if truncate_deterministic is not None and len(var_expr) > truncate_deterministic:
+            contents = var_expr[2:-1].split(", ")
+            str_len = 0
+            for show_n, content in enumerate(contents):
+                str_len += len(content) + 2
+                if str_len > truncate_deterministic:
+                    break
+            var_expr = f"f({', '.join(contents[:show_n])}, ...)"
+    elif var in model.potentials:
+        var_expr = str_for_potential_or_deterministic(var, dist_name="Potential").split(" ~ ")[1]
+    else:  # basic_RVs
+        var_expr = str_for_dist(var).split(" ~ ")[1]
+    return var_expr
+
+
+def _extract_dim_value(var: SharedVariable | Constant) -> np.ndarray:
+    if isinstance(var, SharedVariable):
+        return var.get_value(borrow=True)
+    else:
+        return var.data
+
+
+def dims_expression(model: Model, var: Variable) -> str:
+    """Get the dimensions of a variable in a human-readable format."""
+    if (dims := model.named_vars_to_dims.get(var.name)) is not None:
+        dim_sizes = {dim: _extract_dim_value(model.dim_lengths[dim]) for dim in dims}
+        return " × ".join(f"{dim}[{dim_size}]" for dim, dim_size in dim_sizes.items())
+    else:
+        dim_sizes = list(var.shape.eval(mode=Mode(linker="py", optimizer="fast_compile")))
+        return f"[{', '.join(map(str, dim_sizes))}]" if dim_sizes else ""
+
+
+def model_parameter_count(model: Model) -> int:
+    """Count the number of parameters in the model."""
+    rv_shapes = model.eval_rv_shapes()  # Includes transformed variables
+    return np.sum([np.prod(rv_shapes[free_rv.name]).astype(int) for free_rv in model.free_RVs])
+
+
+def model_table(
+    model: Model,
+    *,
+    split_groups: bool = True,
+    truncate_deterministic: int | None = None,
+    parameter_count: bool = True,
+) -> Table:
+    """Create a rich table with a summary of the model's variables and their expressions.
+
+    Parameters
+    ----------
+    model : Model
+        The PyMC model to summarize.
+    split_groups : bool
+        If True, each group of variables (data, free_RVs, deterministics, potentials, observed_RVs)
+        will be separated by a section.
+    truncate_deterministic : int | None
+        If not None, truncate the expression of deterministic variables that go beyond this length.
+    empty_dims : bool
+        If True, show the dimensions of scalar variables as an empty list.
+    parameter_count : bool
+        If True, add a row with the total number of parameters in the model.
+
+    Returns
+    -------
+    Table
+        A rich table with the model's variables, their expressions and dims.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        import numpy as np
+        import pymc as pm
+
+        from pymc_extras.printing import model_table
+
+        coords = {"subject": range(20), "param": ["a", "b"]}
+        with pm.Model(coords=coords) as m:
+            x = pm.Data("x", np.random.normal(size=(20, 2)), dims=("subject", "param"))
+            y = pm.Data("y", np.random.normal(size=(20,)), dims="subject")
+
+            beta = pm.Normal("beta", mu=0, sigma=1, dims="param")
+            mu = pm.Deterministic("mu", pm.math.dot(x, beta), dims="subject")
+            sigma = pm.HalfNormal("sigma", sigma=1)
+
+            y_obs = pm.Normal("y_obs", mu=mu, sigma=sigma, observed=y, dims="subject")
+
+        table = model_table(m)
+        table  # Displays the following table in an interactive environment
+        '''
+         Variable  Expression         Dimensions
+        ─────────────────────────────────────────────────────
+              x =  Data               subject[20] × param[2]
+              y =  Data               subject[20]
+
+           beta ~  Normal(0, 1)       param[2]
+          sigma ~  HalfNormal(0, 1)
+                                      Parameter count = 3
+
+             mu =  f(beta)            subject[20]
+
+          y_obs ~  Normal(mu, sigma)  subject[20]
+        '''
+
+    Output can be explicitly rendered in a rich console or exported to text, html or svg.
+
+    .. code-block:: python
+
+        from rich.console import Console
+
+        console = Console(record=True)
+        console.print(table)
+        text_export = console.export_text()
+        html_export = console.export_html()
+        svg_export = console.export_svg()
+
+    """
+    table = Table(
+        show_header=True,
+        show_edge=False,
+        box=SIMPLE_HEAD,
+        highlight=False,
+        collapse_padding=True,
+    )
+    table.add_column("Variable", justify="right")
+    table.add_column("Expression", justify="left")
+    table.add_column("Dimensions")
+
+    if split_groups:
+        groups = (
+            model.data_vars,
+            model.free_RVs,
+            model.deterministics,
+            model.potentials,
+            model.observed_RVs,
+        )
+    else:
+        # Show variables in the order they were defined
+        groups = (model.named_vars.values(),)
+
+    for group in groups:
+        if not group:
+            continue
+
+        for var in group:
+            var_name = var.name
+            sep = f'[b]{" ~" if (var in model.basic_RVs) else " ="}[/b]'
+            var_expr = variable_expression(model, var, truncate_deterministic)
+            dims_expr = dims_expression(model, var)
+            if dims_expr == "[]":
+                dims_expr = ""
+            table.add_row(var_name + sep, var_expr, dims_expr)
+
+        if parameter_count and (not split_groups or group == model.free_RVs):
+            n_parameters = model_parameter_count(model)
+            table.add_row("", "", f"[i]Parameter count = {n_parameters}[/i]")
+
+        table.add_section()
+
+    return table

pymc_extras/statespace/__init__.py

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

pymc_extras/statespace/core/__init__.py

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

pymc_extras/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_extras.statespace.core import PyMCStateSpace
+from pymc_extras.statespace.filters.distributions import LinearGaussianStateSpace
+from pymc_extras.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