Iterate on proposal

Author: Jesse Grabowski

pymc_extras/statespace/core/properties.py

@@ -0,0 +1,176 @@
+from collections.abc import Iterator
+from dataclasses import dataclass, fields
+from typing import Generic, Self, TypeVar
+
+from pymc_extras.statespace.core import PyMCStateSpace
+from pymc_extras.statespace.models.structural.core import Component
+from pymc_extras.statespace.utils.constants import (
+    ALL_STATE_AUX_DIM,
+    ALL_STATE_DIM,
+    OBS_STATE_AUX_DIM,
+    OBS_STATE_DIM,
+    SHOCK_AUX_DIM,
+    SHOCK_DIM,
+)
+
+
+@dataclass(frozen=True)
+class Property:
+    def __str__(self) -> str:
+        return "\n".join(f"{f.name}: {getattr(self, f.name)}" for f in fields(self))
+
+
+T = TypeVar("T", bound=Property)
+
+
+@dataclass(frozen=True)
+class Info(Generic[T]):
+    items: tuple[T, ...]
+    key_field: str = "name"
+    _index: dict[str, T] | None = None
+
+    def __post_init__(self):
+        index = {}
+        missing_attr = []
+        for item in self.items:
+            if not hasattr(item, self.key_field):
+                missing_attr.append(item)
+                continue
+            key = getattr(item, self.key_field)
+            if key in index:
+                raise ValueError(f"Duplicate {self.key_field} '{key}' detected.")
+            index[key] = item
+        if missing_attr:
+            raise AttributeError(f"Items missing attribute '{self.key_field}': {missing_attr}")
+        object.__setattr__(self, "_index", index)
+
+    def _key(self, item: T) -> str:
+        return getattr(item, self.key_field)
+
+    def get(self, key: str, default=None) -> T | None:
+        return self._index.get(key, default)
+
+    def __getitem__(self, key: str) -> T:
+        try:
+            return self._index[key]
+        except KeyError as e:
+            available = ", ".join(self._index.keys())
+            raise KeyError(f"No {self.key_field} '{key}'. Available: [{available}]") from e
+
+    def __contains__(self, key: object) -> bool:
+        return key in self._index
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._index)
+
+    def __len__(self) -> int:
+        return len(self._index)
+
+    def __str__(self) -> str:
+        return f"{self.key_field}s: {list(self._index.keys())}"
+
+    @property
+    def names(self) -> tuple[str, ...]:
+        return tuple(self._index.keys())
+
+
+@dataclass(frozen=True)
+class Parameter(Property):
+    name: str
+    shape: tuple[int, ...]
+    dims: tuple[str, ...]
+    constraints: str | None = None
+
+
+@dataclass(frozen=True)
+class ParameterInfo(Info[Parameter]):
+    def __init__(self, parameters: list[Parameter]):
+        super().__init__(items=tuple(parameters), key_field="name")
+
+
+@dataclass(frozen=True)
+class Data(Property):
+    name: str
+    shape: tuple[int, ...]
+    dims: tuple[str, ...]
+    is_exogenous: bool
+
+
+@dataclass(frozen=True)
+class DataInfo(Info[Data]):
+    def __init__(self, data: list[Data]):
+        super().__init__(items=tuple(data), key_field="name")
+
+    @property
+    def needs_exogenous_data(self) -> bool:
+        return any(d.is_exogenous for d in self.items)
+
+    def __str__(self) -> str:
+        return f"data: {[d.name for d in self.items]}\nneeds exogenous data: {self.needs_exogenous_data}"
+
+
+@dataclass(frozen=True)
+class Coord(Property):
+    dimension: str
+    labels: tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class CoordInfo(Info[Coord]):
+    def __init__(self, coords: list[Coord]):
+        super().__init__(items=tuple(coords), key_field="dimension")
+
+    def __str__(self) -> str:
+        base = "coordinates:"
+        for coord in self.items:
+            coord_str = str(coord)
+            indented = "\n".join("  " + line for line in coord_str.splitlines())
+            base += "\n" + indented + "\n"
+        return base
+
+    @classmethod
+    def default_coords_from_model(cls, model: Component | PyMCStateSpace) -> Self:
+        states = tuple(model.state_names)
+        obs_states = tuple(model.observed_state_names)
+        shocks = tuple(model.shock_names)
+
+        dim_to_labels = (
+            (ALL_STATE_DIM, states),
+            (ALL_STATE_AUX_DIM, states),
+            (OBS_STATE_DIM, obs_states),
+            (OBS_STATE_AUX_DIM, obs_states),
+            (SHOCK_DIM, shocks),
+            (SHOCK_AUX_DIM, shocks),
+        )
+
+        coords = [Coord(dimension=dim, labels=labels) for dim, labels in dim_to_labels]
+        return cls(coords)
+
+
+@dataclass(frozen=True)
+class State(Property):
+    name: str
+    observed: bool
+    shared: bool
+
+
+@dataclass(frozen=True)
+class StateInfo(Info[State]):
+    def __init__(self, states: list[State]):
+        super().__init__(items=tuple(states), key_field="name")
+
+    def __str__(self) -> str:
+        return (
+            f"states: {[s.name for s in self.items]}\nobserved: {[s.observed for s in self.items]}"
+        )
+
+
+@dataclass(frozen=True)
+class Shock(Property):
+    name: str
+
+
+@dataclass(frozen=True)
+class ShockInfo(Info[Shock]):
+    def __init__(self, shocks: list[Shock]):
+        super().__init__(items=tuple(shocks), key_field="name")

tests/statespace/core/test_properties.py

@@ -0,0 +1,119 @@
+import pytest
+
+from pymc_extras.statespace.core.properties import (
+    CoordInfo,
+    Data,
+    DataInfo,
+    Parameter,
+    ParameterInfo,
+    Shock,
+    ShockInfo,
+    State,
+    StateInfo,
+)
+from pymc_extras.statespace.utils.constants import (
+    ALL_STATE_AUX_DIM,
+    ALL_STATE_DIM,
+    OBS_STATE_AUX_DIM,
+    OBS_STATE_DIM,
+    SHOCK_AUX_DIM,
+    SHOCK_DIM,
+)
+
+
+def test_property_str_formats_fields():
+    p = Parameter(name="alpha", shape=(2,), dims=("param",))
+    s = str(p).splitlines()
+    assert s == [
+        "name: alpha",
+        "shape: (2,)",
+        "dims: ('param',)",
+        "constraints: None",
+    ]
+
+
+def test_info_lookup_contains_and_missing_key():
+    params = [
+        Parameter("a", (1,), ("d",)),
+        Parameter("b", (2,), ("d",)),
+        Parameter("c", (3,), ("d",)),
+    ]
+    info = ParameterInfo(params)
+
+    assert info.get("b").name == "b"
+    assert info["a"].shape == (1,)
+    assert "c" in info
+
+    with pytest.raises(KeyError) as e:
+        _ = info["missing"]
+    assert "No name 'missing'" in str(e.value)
+
+
+def test_data_info_needs_exogenous_and_str():
+    data = [
+        Data("price", (10,), ("time",), is_exogenous=False),
+        Data("x", (10,), ("time",), is_exogenous=True),
+    ]
+    info = DataInfo(data)
+
+    assert info.needs_exogenous_data is True
+    s = str(info)
+    assert "data: ['price', 'x']" in s
+    assert "needs exogenous data: True" in s
+
+    no_exog = DataInfo([Data("y", (10,), ("time",), is_exogenous=False)])
+    assert no_exog.needs_exogenous_data is False
+
+
+def test_coord_info_make_defaults_from_component_and_types():
+    class DummyComponent:
+        state_names = ["x1", "x2"]
+        observed_state_names = ["x2"]
+        shock_names = ["eps1"]
+
+    ci = CoordInfo.default_coords_from_model(DummyComponent())
+
+    expected = [
+        (ALL_STATE_DIM, ("x1", "x2")),
+        (ALL_STATE_AUX_DIM, ("x1", "x2")),
+        (OBS_STATE_DIM, ("x2",)),
+        (OBS_STATE_AUX_DIM, ("x2",)),
+        (SHOCK_DIM, ("eps1",)),
+        (SHOCK_AUX_DIM, ("eps1",)),
+    ]
+
+    assert len(ci.items) == 6
+    for dim, labels in expected:
+        assert dim in ci
+        assert ci[dim].labels == labels
+        assert isinstance(ci[dim].labels, tuple)
+
+
+def test_state_info_and_shockinfo_basic():
+    states = [
+        State("x1", observed=True, shared=False),
+        State("x2", observed=False, shared=True),
+    ]
+    state_info = StateInfo(states)
+    assert state_info["x1"].observed is True
+    s = str(state_info)
+
+    assert "states: ['x1', 'x2']" in s
+    assert "observed: [True, False]" in s
+
+    shocks = [Shock("s1"), Shock("s2")]
+    shock_info = ShockInfo(shocks)
+
+    assert "s1" in shock_info
+    assert shock_info["s2"].name == "s2"
+
+
+def test_info_is_iterable_and_unpackable():
+    items = [Parameter("p1", (1,), ("d",)), Parameter("p2", (2,), ("d",))]
+    info = ParameterInfo(items)
+
+    names = info.names
+    assert names == ("p1", "p2")
+
+    a, b = info.items
+    assert a.name == "p1" and b.name == "p2"