Added a PartialOrder transform

Author: velochy

pymc_extras/distributions/transforms/__init__.py

@@ -0,0 +1,3 @@
+from pymc_extras.distributions.transforms.partial_order import PartialOrder
+
+__all__ = ["PartialOrder"]

pymc_extras/distributions/transforms/partial_order.py

@@ -0,0 +1,144 @@
+#   Copyright 2025 The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+import numpy as np
+import pytensor.tensor as pt
+
+from pymc.logprob.transforms import Transform
+
+__all__ = ["PartialOrder"]
+
+
+# Find the minimum value for a given dtype
+def dtype_minval(dtype):
+    return np.iinfo(dtype).min if np.issubdtype(dtype, np.integer) else np.finfo(dtype).min
+
+
+# A padded version of np.where
+def padded_where(x, to_len, padval=-1):
+    w = np.where(x)
+    return np.concatenate([w[0], np.full(to_len - len(w[0]), padval)])
+
+
+# Partial order transform
+class PartialOrder(Transform):
+    name = "partial_order"
+
+    def __init__(self, adj_mat):
+        """Create a PartialOrder transform
+
+        This is a more flexible version of the pymc ordered transform that
+        allows specifying a (strict) partial order on the elements.
+
+        It works in O(N*D) in runtime, but takes O(N^3) in initialization,
+        where N is the number of nodes in the dag and
+        D is the maximum in-degree of a node in the transitive reduction.
+
+        Parameters
+        ----------
+        adj_mat: adjacency matrix for the DAG that generates the partial order,
+                 where adj_mat[i][j] = 1 denotes i<j.
+                 Note this also accepts multiple DAGs if RV is multidimensional
+        """
+
+        # Basic input checks
+        if adj_mat.ndim < 2:
+            raise ValueError("Adjacency matrix must have at least 2 dimensions")
+        if adj_mat.shape[-2] != adj_mat.shape[-1]:
+            raise ValueError("Adjacency matrix is not square")
+        if adj_mat.min() != 0 or adj_mat.max() != 1:
+            raise ValueError("Adjacency matrix must contain only 0s and 1s")
+
+        # Create index over the first ellipsis dimensions
+        idx = np.ix_(*[np.arange(s) for s in adj_mat.shape[:-2]])
+
+        # Transitive closure using Floyd-Warshall
+        tc = adj_mat.astype(bool)
+        for k in range(tc.shape[-1]):
+            tc |= np.logical_and(tc[..., :, k, None], tc[..., None, k, :])
+
+        # Check if the dag is acyclic
+        if np.any(tc.diagonal(axis1=-2, axis2=-1)):
+            raise ValueError("Partial order contains equalities")
+
+        # Transitive reduction using the closure
+        # This gives the minimum description of the partial order
+        # This is to minmax the input degree
+        adj_mat = tc * (1 - np.matmul(tc, tc))
+
+        # Find the maximum in-degree of the reduced dag
+        dag_idim = adj_mat.sum(axis=-2).max()
+
+        # Topological sort
+        ts_inds = np.zeros(adj_mat.shape[:-1], dtype=int)
+        dm = adj_mat.copy()
+        for i in range(adj_mat.shape[1]):
+            assert dm.sum(axis=-2).min() == 0  # DAG is acyclic
+            nind = np.argmin(dm.sum(axis=-2), axis=-1)
+            dm[(*idx, slice(None), nind)] = 1  # Make nind not show up again
+            dm[(*idx, nind, slice(None))] = 0  # Allow it's children to show
+            ts_inds[(*idx, i)] = nind
+        self.ts_inds = ts_inds
+
+        # Change the dag to adjacency lists (with -1 for NA)
+        dag_T = np.apply_along_axis(padded_where, axis=-2, arr=adj_mat, padval=-1, to_len=dag_idim)
+        self.dag = np.swapaxes(dag_T, -2, -1)
+        self.is_start = np.all(self.dag[..., :, :] == -1, axis=-1)
+
+    def initvals(self, lower=-1, upper=1):
+        vals = np.linspace(lower, upper, self.dag.shape[-2])
+        inds = np.argsort(self.ts_inds, axis=-1)
+        return vals[inds]
+
+    def backward(self, value, *inputs):
+        minv = dtype_minval(value.dtype)
+        x = pt.concatenate(
+            [pt.zeros_like(value), pt.full(value.shape[:-1], minv)[..., None]], axis=-1
+        )
+
+        # Indices to allow broadcasting the max over the last dimension
+        idx = np.ix_(*[np.arange(s) for s in self.dag.shape[:-2]])
+        idx2 = tuple(np.tile(i[:, None], self.dag.shape[-1]) for i in idx)
+
+        # Has to be done stepwise as next steps depend on previous values
+        # Also has to be done in topological order, hence the ts_inds
+        for i in range(self.dag.shape[-2]):
+            tsi = self.ts_inds[..., i]
+            if len(tsi.shape) == 0:
+                tsi = int(tsi)  # if shape 0, it's a scalar
+            ni = (*idx, tsi)  # i-th node in topological order
+            eni = (Ellipsis, *ni)
+            ist = self.is_start[ni]
+
+            mval = pt.max(x[(Ellipsis, *idx2, self.dag[ni])], axis=-1)
+            x = pt.set_subtensor(x[eni], ist * value[eni] + (1 - ist) * (mval + pt.exp(value[eni])))
+        return x[..., :-1]
+
+    def forward(self, value, *inputs):
+        y = pt.zeros_like(value)
+
+        minv = dtype_minval(value.dtype)
+        vx = pt.concatenate([value, pt.full(value.shape[:-1], minv)[..., None]], axis=-1)
+
+        # Indices to allow broadcasting the max over the last dimension
+        idx = np.ix_(*[np.arange(s) for s in self.dag.shape[:-2]])
+        idx = tuple(np.tile(i[:, None, None], self.dag.shape[-2:]) for i in idx)
+
+        y = self.is_start * value + (1 - self.is_start) * (
+            pt.log(value - pt.max(vx[(Ellipsis, *idx, self.dag[..., :])], axis=-1))
+        )
+
+        return y
+
+    def log_jac_det(self, value, *inputs):
+        return pt.sum(value * (1 - self.is_start), axis=-1)

tests/distributions/test_transform.py

@@ -0,0 +1,71 @@
+#   Copyright 2025 The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+import numpy as np
+import pymc as pm
+
+from pymc_extras.distributions.transforms import PartialOrder
+
+
+class TestPartialOrder:
+    adj_mats = np.array(
+        [
+            # 0 < {1, 2} < 3
+            [[0, 1, 1, 0], [0, 0, 0, 1], [0, 0, 0, 1], [0, 0, 0, 0]],
+            # 1 < 0 < 3 < 2
+            [[0, 0, 0, 1], [1, 0, 0, 0], [0, 0, 0, 0], [0, 0, 1, 0]],
+        ]
+    )
+
+    valid_values = np.array([[0, 2, 1, 3], [1, 0, 3, 2]], dtype=float)
+
+    # Test that forward and backward are inverses of eachother
+    # And that it works when extra dimensions are added in data
+    def test_forward_backward_dimensionality(self):
+        po = PartialOrder(self.adj_mats)
+        po0 = PartialOrder(self.adj_mats[0])
+        vv = self.valid_values
+        vv0 = self.valid_values[0]
+
+        testsets = [
+            (vv, po),
+            (po.initvals(), po),
+            (vv0, po0),
+            (po0.initvals(), po0),
+            (np.tile(vv0, (2, 1)), po0),
+            (np.tile(vv0, (2, 3, 2, 1)), po0),
+            (np.tile(vv, (2, 3, 2, 1, 1)), po),
+        ]
+
+        for vv, po in testsets:
+            fw = po.forward(vv)
+            bw = po.backward(fw)
+            np.testing.assert_allclose(bw.eval(), vv)
+
+    def test_sample_model(self):
+        po = PartialOrder(self.adj_mats)
+        with pm.Model() as model:
+            x = pm.Normal("x", size=(2, 4), transform=po, initval=po.initvals(-1, 1))
+            idata = pm.sample()
+
+        # Check that the order constraints are satisfied
+        xvs = idata.posterior.x.values.transpose(2, 3, 0, 1)
+        x0 = xvs[0]  # 0 < {1, 2} < 3
+        assert (
+            (x0[0] < x0[1]).all()
+            and (x0[0] < x0[2]).all()
+            and (x0[1] < x0[3]).all()
+            and (x0[2] < x0[3]).all()
+        )
+        x1 = xvs[1]  # 1 < 0 < 3 < 2
+        assert (x1[1] < x1[0]).all() and (x1[0] < x1[3]).all() and (x1[3] < x1[2]).all()