Initialize a prior from a fitted posterior

pymc_experimental/tests/test_prior_from_trace.py

@@ -0,0 +1,102 @@
+import pymc_experimental as pmx
+from pymc.distributions import transforms
+import pytest
+import arviz as az
+import numpy as np
+
+
+@pytest.mark.parametrize(
+    "case",
+    [
+        (("a", dict(name="b")), dict(name="b", transform=None, dims=None)),
+        (("a", None), dict(name="a", transform=None, dims=None)),
+        (("a", transforms.log), dict(name="a", transform=transforms.log, dims=None)),
+        (
+            ("a", dict(transform=transforms.log)),
+            dict(name="a", transform=transforms.log, dims=None),
+        ),
+        (("a", dict(name="b")), dict(name="b", transform=None, dims=None)),
+        (("a", dict(name="b", dims="test")), dict(name="b", transform=None, dims="test")),
+        (("a", ("test",)), dict(name="a", transform=None, dims=("test",))),
+    ],
+)
+def test_parsing_arguments(case):
+    inp, out = case
+    test = pmx.utils.prior._arg_to_param_cfg(*inp)
+    assert test == out
+
+
+@pytest.fixture
+def coords():
+    return dict(test=range(3), simplex=range(4))
+
+
+@pytest.fixture
+def param_cfg():
+    return dict(
+        a=pmx.utils.prior._arg_to_param_cfg("d"),
+        b=pmx.utils.prior._arg_to_param_cfg("b", dict(transform=transforms.log, dims=("test",))),
+        c=pmx.utils.prior._arg_to_param_cfg(
+            "c", dict(transform=transforms.simplex, dims=("simplex",))
+        ),
+    )
+
+
+@pytest.fixture
+def transformed_data(param_cfg, coords):
+    vars = dict()
+    for k, cfg in param_cfg.items():
+        if cfg["dims"] is not None:
+            extra_dims = [len(coords[d]) for d in cfg["dims"]]
+        else:
+            extra_dims = []
+        orig = np.random.randn(4, 100, *extra_dims)
+        vars[k] = orig
+    return vars
+
+
+@pytest.fixture
+def idata(transformed_data, param_cfg):
+    vars = dict()
+    for k, orig in transformed_data.items():
+        cfg = param_cfg[k]
+        if cfg["transform"] is not None:
+            var = cfg["transform"].backward(orig).eval()
+        else:
+            var = orig
+        assert not np.isnan(var).any()
+        vars[k] = var
+    return az.convert_to_inference_data(vars)
+
+
+def test_idata_for_tests(idata, param_cfg):
+    assert set(idata.posterior.keys()) == set(param_cfg)
+    assert len(idata.posterior.coords["chain"]) == 4
+    assert len(idata.posterior.coords["draw"]) == 100
+
+
+def test_args_compose():
+    cfg = pmx.utils.prior._parse_args(
+        var_names=["a"],
+        b=("test",),
+        c=transforms.log,
+        d="e",
+        f=dict(dims="test"),
+        g=dict(name="h", dims="test", transform=transforms.log),
+    )
+    assert cfg == dict(
+        a=dict(name="a", dims=None, transform=None),
+        b=dict(name="b", dims=("test",), transform=None),
+        c=dict(name="c", dims=None, transform=transforms.log),
+        d=dict(name="e", dims=None, transform=None),
+        f=dict(name="f", dims="test", transform=None),
+        g=dict(name="h", dims="test", transform=transforms.log),
+    )
+
+
+def test_transform_idata(transformed_data, idata, param_cfg):
+    flat_info = pmx.utils.prior._flatten(idata, **param_cfg)
+    expected_shape = 0
+    for v in transformed_data.values():
+        expected_shape += int(np.prod(v.shape[2:]))
+    assert flat_info["data"].shape[1] == expected_shape

pymc_experimental/utils/__init__.py

@@ -1 +1,2 @@
 from pymc_experimental.utils import spline
+from pymc_experimental.utils import prior

pymc_experimental/utils/prior.py

@@ -0,0 +1,76 @@
+from typing import TypedDict, Optional, Union, Tuple, Sequence, Dict, List
+import aeppl.transforms
+import arviz
+import numpy as np
+
+
+class ParamCfg(TypedDict):
+    name: str
+    transform: Optional[aeppl.transforms.RVTransform]
+    dims: Optional[Union[str, Tuple[str]]]
+
+
+class ShapeInfo(TypedDict):
+    # shape might not match slice due to a transform
+    shape: Tuple[int]
+    slice: slice
+
+
+class VarInfo(TypedDict):
+    sinfo: ShapeInfo
+    vinfo: ParamCfg
+
+
+class FlatInfo(TypedDict):
+    data: np.ndarray
+    info: List[VarInfo]
+
+
+def _arg_to_param_cfg(
+    key, value: Optional[Union[ParamCfg, aeppl.transforms.RVTransform, str, Tuple]] = None
+):
+    if value is None:
+        cfg = ParamCfg(name=key, transform=None, dims=None)
+    elif isinstance(value, Tuple):
+        cfg = ParamCfg(name=key, transform=None, dims=value)
+    elif isinstance(value, str):
+        cfg = ParamCfg(name=value, transform=None, dims=None)
+    elif isinstance(value, aeppl.transforms.RVTransform):
+        cfg = ParamCfg(name=key, transform=value, dims=None)
+    else:
+        cfg = value.copy()
+        cfg.setdefault("name", key)
+        cfg.setdefault("transform", None)
+        cfg.setdefault("dims", None)
+    return cfg
+
+
+def _parse_args(
+    var_names: Sequence[str], **kwargs: Union[ParamCfg, aeppl.transforms.RVTransform, str, Tuple]
+) -> Dict[str, ParamCfg]:
+    results = dict()
+    for var in var_names:
+        results[var] = _arg_to_param_cfg(var)
+    for key, val in kwargs.items():
+        results[key] = _arg_to_param_cfg(key, val)
+    return results
+
+
+def _flatten(idata: arviz.InferenceData, **kwargs: ParamCfg) -> FlatInfo:
+    posterior = idata.posterior
+    vars = list()
+    info = list()
+    begin = 0
+    for key, cfg in kwargs.items():
+        data = posterior[key].values
+        # omitting chain, draw
+        shape = data.shape[2:]
+        if cfg["transform"] is not None:
+            data = cfg["transform"].forward(data).eval()
+        data = data.reshape(*data.shape[:2], -1)
+        data = data.reshape(-1, data.shape[2])
+        end = begin + data.shape[1]
+        vars.append(data)
+        info.append(dict(shape=shape, slice=slice(begin, end)))
+        begin = end
+    return dict(data=np.concatenate(vars, axis=-1), infp=info)