Implement utility to convert PyMC modelt to and from FunctionGraph

ricardoV94 · ricardoV94 · commit ba7968b0d4e4 · 2023-02-20T20:27:23.000+01:00
diff --git a/pymc_experimental/tests/utils/__init__.py b/pymc_experimental/tests/utils/__init__.py
diff --git a/pymc_experimental/tests/utils/test_model_fgraph.py b/pymc_experimental/tests/utils/test_model_fgraph.py
@@ -0,0 +1,144 @@
+import numpy as np
+import pymc as pm
+import pytensor.tensor as pt
+import pytest
+from pytensor.graph import FunctionGraph, node_rewriter
+from pytensor.graph.rewriting.basic import in2out
+from pytensor.tensor.exceptions import NotScalarConstantError
+
+from pymc_experimental.utils.model_fgraph import (
+    FreeRV,
+    deterministic,
+    fgraph_from_model,
+    free_rv,
+    model_from_fgraph,
+)
+
+
+def test_model_fgraph_conversion():
+    with pm.Model(coords={"test_dim": range(3)}) as m_old:
+        x = pm.Normal("x")
+        y = pm.Deterministic("y", x + 1)
+        w = pm.Normal("w", y)
+        z = pm.Normal("z", y, observed=[0, 1, 2], dims=("test_dim",))
+        pm.Potential("pot", x * 2)
+
+    m_fgraph = fgraph_from_model(m_old)
+    assert isinstance(m_fgraph, FunctionGraph)
+
+    m_new = model_from_fgraph(m_fgraph)
+    assert isinstance(m_new, pm.Model)
+
+    assert m_new.coords == {"test_dim": tuple(range(3))}
+    assert m_new.named_vars_to_dims == {"z": ["test_dim"]}
+
+    named_vars = {"x", "y", "w", "z", "pot"}
+    assert set(m_new.named_vars) == named_vars
+    for named_var in named_vars:
+        assert m_new[named_var] is not m_old[named_var]
+    assert m_new["x"] in m_new.free_RVs
+    assert m_new["w"] in m_new.free_RVs
+    assert m_new["y"] in m_new.deterministics
+    assert m_new["z"] in m_new.observed_RVs
+    assert m_new["pot"] in m_new.potentials
+
+    new_y_draw, new_z_draw = pm.draw([m_new["y"], m_new["z"]], draws=5, random_seed=1)
+    old_y_draw, old_z_draw = pm.draw([m_old["y"], m_old["z"]], draws=5, random_seed=1)
+    np.testing.assert_array_equal(new_y_draw, old_y_draw)
+    np.testing.assert_array_equal(new_z_draw, old_z_draw)
+
+    ip = m_new.initial_point()
+    np.testing.assert_equal(
+        m_new.compile_logp()(ip),
+        m_old.compile_logp()(ip),
+    )
+
+
+@pytest.fixture()
+def non_centered_rewrite():
+    @node_rewriter(tracks=[FreeRV])
+    def non_centered_param(fgraph: FunctionGraph, node):
+        """Rewrite that replaces centered normal by non-centered parametrization."""
+
+        rv, _, *dims = node.inputs
+        if not isinstance(rv.owner.op, pm.Normal):
+            return
+        rng, size, dtype, loc, scale = rv.owner.inputs
+
+        # Only apply rewrite if size information is explicit
+        if size.ndim == 0:
+            return None
+
+        try:
+            is_unit = (
+                pt.get_scalar_constant_value(loc) == 0 and pt.get_scalar_constant_value(scale) == 1
+            )
+        except NotScalarConstantError:
+            is_unit = False
+
+        # Nothing to do here
+        if is_unit:
+            return
+
+        raw_norm = pm.Normal.dist(0, 1, size=size, rng=rng)
+        raw_norm.name = f"{rv.name}_raw_"
+        raw_norm_value = raw_norm.clone()
+        fgraph.add_input(raw_norm_value)
+        raw_norm = free_rv(raw_norm, raw_norm_value, dims=dims)
+
+        new_norm = loc + raw_norm * scale
+        new_norm.name = rv.name
+        new_norm = deterministic(new_norm, dims=dims)
+
+        return [new_norm]
+
+    return in2out(non_centered_param)
+
+
+def test_fgraph_rewrite(non_centered_rewrite):
+
+    with pm.Model(coords={"subject": range(10)}) as m_old:
+        group_mean = pm.Normal("group_mean")
+        # FIXME: Transforms are not yet maintained across conversion
+        group_std = pm.HalfNormal("group_std", transform=None)
+        subject_mean = pm.Normal("subject_mean", group_mean, group_std, dims=("subject",))
+        obs = pm.Normal("obs", subject_mean, 1, observed=np.zeros(10), dims=("subject",))
+
+    fg = fgraph_from_model(m_old)
+    non_centered_rewrite.apply(fg)
+
+    m_new = model_from_fgraph(fg)
+    assert m_new.named_vars_to_dims == {
+        "subject_mean": ["subject"],
+        "subject_mean_raw_": ["subject"],
+        "obs": ["subject"],
+    }
+    assert set(m_new.named_vars) == {
+        "group_mean",
+        "group_std",
+        "subject_mean_raw_",
+        "subject_mean",
+        "obs",
+    }
+    assert {rv.name for rv in m_new.free_RVs} == {"group_mean", "group_std", "subject_mean_raw_"}
+    assert {rv.name for rv in m_new.observed_RVs} == {"obs"}
+    assert {rv.name for rv in m_new.deterministics} == {"subject_mean"}
+
+    with pm.Model() as m_ref:
+        group_mean = pm.Normal("group_mean")
+        # FIXME: Transforms are not yet maintained across conversion
+        group_std = pm.HalfNormal("group_std", transform=None)
+        subject_mean_raw = pm.Normal("subject_mean_raw_", 0, 1, shape=(10,))
+        subject_mean = pm.Deterministic("subject_mean", group_mean + subject_mean_raw * group_std)
+        obs = pm.Normal("obs", subject_mean, 1, observed=np.zeros(10))
+
+    np.testing.assert_array_equal(
+        pm.draw(m_new["subject_mean_raw_"], draws=7, random_seed=1),
+        pm.draw(m_ref["subject_mean_raw_"], draws=7, random_seed=1),
+    )
+
+    ip = m_new.initial_point()
+    np.testing.assert_equal(
+        m_new.compile_logp()(ip),
+        m_ref.compile_logp()(ip),
+    )
diff --git a/pymc_experimental/utils/model_fgraph.py b/pymc_experimental/utils/model_fgraph.py
@@ -0,0 +1,168 @@
+from typing import Optional, Sequence
+
+import pytensor
+from pymc.model import Model
+from pytensor.graph import Apply, FunctionGraph, Op
+from pytensor.tensor import TensorVariable
+
+from pymc_experimental.utils.pytensorf import StringType
+
+
+class ModelVar(Op):
+    """An Op that binds together a RV and its value"""
+
+    def make_node(self, rv, value=None, dims: Optional[Sequence[str]] = None):
+        assert isinstance(rv, TensorVariable)
+
+        if dims is not None:
+            dims = [pytensor.as_symbolic(dim) for dim in dims]
+            assert all(isinstance(dim.type, StringType) for dim in dims)
+            assert len(dims) == rv.type.ndim
+        else:
+            dims = ()
+
+        if value is not None:
+            assert isinstance(value, TensorVariable)
+            assert rv.type.in_same_class(value.type)
+            return Apply(self, [rv, value, *dims], [rv.type()])
+        else:
+            return Apply(self, [rv, *dims], [rv.type()])
+
+    def infer_shape(self, fgraph, node, inputs_shape):
+        return inputs_shape[0]
+
+    def do_constant_folding(self, fgraph, node):
+        return False
+
+    def perform(self, *args, **kwargs):
+        raise RuntimeError("ValuedRVs should never be evaluated!")
+
+
+class FreeRV(ModelVar):
+    pass
+
+
+class ObservedRV(ModelVar):
+    pass
+
+
+class Potential(ModelVar):
+    pass
+
+
+class Deterministic(ModelVar):
+    pass
+
+
+free_rv = FreeRV()
+observed_rv = ObservedRV()
+potential = Potential()
+deterministic = Deterministic()
+
+
+def toposort_replace(fgraph: FunctionGraph, replacements) -> None:
+    toposort = fgraph.toposort()
+    sorted_replacements = sorted(replacements, key=lambda pair: toposort.index(pair[0].owner))
+    fgraph.replace_all(tuple(sorted_replacements), import_missing=True)
+
+
+def fgraph_from_model(model: Model) -> FunctionGraph:
+
+    # Collect PyTensor variables
+    rvs_to_values = model.rvs_to_values
+    rvs = list(rvs_to_values.keys())
+    values = list(rvs_to_values.values())
+    free_rvs = model.free_RVs
+    deterministics = model.deterministics
+    potentials = model.potentials
+
+    # Collect PyMC meta-info
+    vars_to_dims = model.named_vars_to_dims
+    coords = model.coords
+
+    # TODO: Do something with these
+    dim_lengths = model.dim_lengths
+    rvs_to_transforms = model.rvs_to_transforms
+
+    # Not supported yet
+    if any(v is not None for v in model.rvs_to_total_sizes.values()):
+        raise NotImplementedError("Cannot convert models with total_sizes")
+    if any(v is not None for v in model.rvs_to_initial_values.values()):
+        raise NotImplementedError("Cannot convert models with non-default initial_values")
+
+    # We start the `dict` with mappings from the value variables to themselves,
+    # to prevent them from being cloned.
+    memo = {v: v for v in values}
+
+    fgraph = FunctionGraph(
+        outputs=rvs + potentials + deterministics,
+        clone=True,
+        memo=memo,
+        copy_orphans=False,
+        copy_inputs=False,
+    )
+    fgraph.coords = coords
+
+    # Introduce dummy Ops to label different types of ModelVariables
+    free_rvs_to_values = {memo[k]: v for k, v in rvs_to_values.items() if k in free_rvs}
+    observed_rvs_to_values = {memo[k]: v for k, v in rvs_to_values.items() if k not in free_rvs}
+    potentials = [memo[k] for k in potentials]
+    deterministics = [memo[k] for k in deterministics]
+
+    vars = fgraph.outputs
+    new_vars = []
+    for var in vars:
+        dims = vars_to_dims.get(var.name, None)
+        if var in free_rvs_to_values:
+            new_var = free_rv(var, free_rvs_to_values[var], dims)
+        elif var in observed_rvs_to_values:
+            new_var = observed_rv(var, observed_rvs_to_values[var], dims)
+        elif var in potentials:
+            new_var = potential(var, dims)
+        elif var in deterministics:
+            new_var = deterministic(var, dims)
+        else:
+            raise RuntimeError(f"Variable is not RV, Potential nor Deterministic: {new_var}")
+        new_vars.append(new_var)
+
+    toposort_replace(fgraph, tuple(zip(vars, new_vars)))
+    return fgraph
+
+
+def model_from_fgraph(fgraph: FunctionGraph) -> Model:
+    model = Model(coords=getattr(fgraph, "coords", None))
+
+    fgraph = fgraph.clone()
+    model_vars_to_vars = {
+        model_node.outputs[0]: model_node.inputs[0]
+        for model_node in fgraph.apply_nodes
+        if isinstance(model_node.op, ModelVar)
+    }
+    toposort_replace(fgraph, tuple(model_vars_to_vars.items()))
+
+    for model_var in model_vars_to_vars.keys():
+        if isinstance(model_var.owner.op, FreeRV):
+            var, value, *dims = model_var.owner.inputs
+            model.free_RVs.append(var)
+            model.create_value_var(var, transform=None, value_var=value)
+            model.set_initval(var, initval=None)
+        elif isinstance(model_var.owner.op, ObservedRV):
+            var, value, *dims = model_var.owner.inputs
+            model.observed_RVs.append(var)
+            model.create_value_var(var, transform=None, value_var=value)
+        elif isinstance(model_var.owner.op, Potential):
+            var, *dims = model_var.owner.inputs
+            model.potentials.append(var)
+        elif isinstance(model_var.owner.op, Deterministic):
+            var, *dims = model_var.owner.inputs
+            model.deterministics.append(var)
+        else:
+            continue  # Raise?
+
+        if not dims:
+            dims = None
+        else:
+            dims = [dim.data for dim in dims]
+        model.add_named_variable(var, dims=dims)
+
+    return model
diff --git a/pymc_experimental/utils/pytensorf.py b/pymc_experimental/utils/pytensorf.py
@@ -0,0 +1,33 @@
+import pytensor
+from pytensor.graph import Constant, Type
+
+
+class StringType(Type[str]):
+    def clone(self, **kwargs):
+        return type(self)()
+
+    def filter(self, x, strict=False, allow_downcast=None):
+        if isinstance(x, str):
+            return x
+        else:
+            raise TypeError("Expected a string!")
+
+    def __str__(self):
+        return "string"
+
+    @staticmethod
+    def may_share_memory(a, b):
+        return isinstance(a, str) and a is b
+
+
+stringtype = StringType()
+
+
+class StringConstant(Constant):
+    pass
+
+
+@pytensor._as_symbolic.register(str)
+def as_symbolic_string(x, **kwargs):
+
+    return StringConstant(stringtype, x)
