Generalize ordered transform

Author: velochy

pymc/distributions/transforms.py

@@ -89,26 +89,48 @@ def log_jac_det(self, value, *inputs):
 
 
 class Ordered(Transform):
+    """
+    Transforms a vector of values into a vector of ordered values.
+
+    Parameters
+    ----------
+    positive: If True, all values are positive. This has better geometry than just chaining with a log transform.
+    ascending: If True, the values are in ascending order (default). If False, the values are in descending order.
+    """
+
     name = "ordered"
 
-    def __init__(self, ndim_supp=None):
+    def __init__(self, ndim_supp=None, positive=False, ascending=True):
         if ndim_supp is not None:
             warnings.warn("ndim_supp argument is deprecated and has no effect", FutureWarning)
+        self.positive = positive
+        self.ascending = ascending
 
     def backward(self, value, *inputs):
-        x = pt.zeros(value.shape)
-        x = pt.set_subtensor(x[..., 0], value[..., 0])
-        x = pt.set_subtensor(x[..., 1:], pt.exp(value[..., 1:]))
-        return pt.cumsum(x, axis=-1)
+        if self.positive:
+            x = pt.exp(value)
+        else:  # Everything except the first element is positive
+            x = pt.zeros(value.shape)
+            x = pt.set_subtensor(x[..., 0], value[..., 0])
+            x = pt.set_subtensor(x[..., 1:], pt.exp(value[..., 1:]))
+        x = pt.cumsum(x, axis=-1)
+        if not self.ascending:
+            x = x[..., ::-1]
+        return x
 
     def forward(self, value, *inputs):
+        if not self.ascending:
+            value = value[..., ::-1]
         y = pt.zeros(value.shape)
-        y = pt.set_subtensor(y[..., 0], value[..., 0])
+        y = pt.set_subtensor(y[..., 0], pt.log(value[..., 0]) if self.positive else value[..., 0])
         y = pt.set_subtensor(y[..., 1:], pt.log(value[..., 1:] - value[..., :-1]))
         return y
 
     def log_jac_det(self, value, *inputs):
-        return pt.sum(value[..., 1:], axis=-1)
+        if self.positive:
+            return pt.sum(value, axis=-1)
+        else:
+            return pt.sum(value[..., 1:], axis=-1)
 
 
 class SumTo1(Transform):

tests/distributions/test_transform.py

@@ -281,6 +281,18 @@ def test_ordered():
     vals = get_values(tr.ordered, Vector(R, 3), pt.vector, floatX(np.zeros(3)))
     assert_array_equal(np.diff(vals) >= 0, True)
 
+    # Check that positive=True creates positive and still ordered values
+    vals = get_values(tr.Ordered(positive=True), Vector(R, 3), pt.vector, floatX(np.zeros(3)))
+    assert_array_equal(vals > 0, True)
+    assert_array_equal(np.diff(vals) >= 0, True)
+
+    # Check that positive=True and ascending=False creates descending values
+    vals = get_values(
+        tr.Ordered(positive=True, ascending=False), Vector(R, 3), pt.vector, floatX(np.zeros(3))
+    )
+    assert_array_equal(vals > 0, True)
+    assert_array_equal(np.diff(vals) <= 0, True)
+
 
 def test_chain_values():
     chain_tranf = tr.Chain([tr.logodds, tr.ordered])