Constraints for multiple intervals

numpyro/distributions/constraints.py

@@ -60,6 +60,8 @@
     "unit_interval",
     "zero_sum",
     "Constraint",
+    "union_of_closed_intervals",
+    "intersection_of_closed_intervals",
 ]
 
 import math
@@ -766,6 +768,140 @@ def tree_flatten(self):
         return (self.event_dim,), (("event_dim",), dict())
 
 
+class _UnionOfClosedIntervals(Constraint):
+    """A constraint representing the union of multiple intervals."""
+
+    event_dim = 1
+
+    def __init__(self, lower_bounds, upper_bounds) -> None:
+        assert isinstance(lower_bounds, (list, tuple, jnp.ndarray))
+        assert isinstance(upper_bounds, (list, tuple, jnp.ndarray))
+        assert len(lower_bounds) == len(upper_bounds), (
+            f"lower_bounds and upper_bounds must have the same length, "
+            f"but got {len(lower_bounds)} and {len(upper_bounds)}"
+        )
+        self.lower_bounds = jnp.asarray(lower_bounds)
+        self.upper_bounds = jnp.asarray(upper_bounds)
+
+    def __call__(self, x):
+        r"""Check if the input is within the union of intervals
+
+        .. math::
+            x \in \bigcup_{i=1}^{n} [a_i, b_i] \implies \bigvee_{i=1}^{n} (x \in [a_i, b_i])
+
+        :param x: The input to be checked.
+        """
+        return jnp.any((x >= self.lower_bounds) & (x <= self.upper_bounds), axis=-1)
+
+    def feasible_like(self, prototype):
+        return jnp.broadcast_to(
+            (self.lower_bounds + self.upper_bounds) / 2, jnp.shape(prototype)
+        )
+
+    def tree_flatten(self):
+        return (self.lower_bounds, self.upper_bounds), (
+            ("lower_bounds", "upper_bounds"),
+            dict(),
+        )
+
+    def __eq__(self, other):
+        if not isinstance(other, _UnionOfClosedIntervals):
+            return False
+        return jnp.array_equal(self.lower_bounds, other.lower_bounds) & jnp.array_equal(
+            self.upper_bounds, other.upper_bounds
+        )
+
+
+class _IntersectionOfClosedIntervals(Constraint):
+    """A constraint representing the intersection of multiple intervals."""
+
+    event_dim = 1
+
+    def __init__(self, lower_bounds, upper_bounds) -> None:
+        assert isinstance(lower_bounds, (list, tuple, jnp.ndarray))
+        assert isinstance(upper_bounds, (list, tuple, jnp.ndarray))
+        assert len(lower_bounds) == len(upper_bounds), (
+            f"lower_bounds and upper_bounds must have the same length, "
+            f"but got {len(lower_bounds)} and {len(upper_bounds)}"
+        )
+        self.lower_bounds = jnp.asarray(lower_bounds)
+        self.upper_bounds = jnp.asarray(upper_bounds)
+
+    def __call__(self, x):
+        r"""Check if the input is within the union of intervals
+
+        .. math::
+            x \in \bigcap_{i=1}^{n} [a_i, b_i] \implies \bigwedge_{i=1}^{n} (x \in [a_i, b_i])
+
+        :param x: The input to be checked.
+        """
+        return jnp.all((x >= self.lower_bounds) & (x <= self.upper_bounds), axis=-1)
+
+    def feasible_like(self, prototype):
+        return jnp.broadcast_to(
+            (self.lower_bounds + self.upper_bounds) / 2, jnp.shape(prototype)
+        )
+
+    def tree_flatten(self):
+        return (self.lower_bounds, self.upper_bounds), (
+            ("lower_bounds", "upper_bounds"),
+            dict(),
+        )
+
+    def __eq__(self, other):
+        if not isinstance(other, _IntersectionOfClosedIntervals):
+            return False
+        return jnp.array_equal(self.lower_bounds, other.lower_bounds) & jnp.array_equal(
+            self.upper_bounds, other.upper_bounds
+        )
+
+
+class _UniqueIntervals(Constraint):
+    """A constraint representing a set of unique intervals for a single dimension."""
+
+    event_dim = 1
+
+    def __init__(self, lower_bounds, upper_bounds) -> None:
+        assert isinstance(lower_bounds, (list, tuple, jnp.ndarray))
+        assert isinstance(upper_bounds, (list, tuple, jnp.ndarray))
+        assert len(lower_bounds) == len(upper_bounds), (
+            f"lower_bounds and upper_bounds must have the same length, "
+            f"but got {len(lower_bounds)} and {len(upper_bounds)}"
+        )
+        self.lower_bounds = jnp.asarray(lower_bounds)
+        self.upper_bounds = jnp.asarray(upper_bounds)
+
+    def __call__(self, x):
+        r"""Check if the input is within the specified intervals
+
+        .. math::
+            \bigwedge_{i=1}^{n} (x_i \in [a_i, b_i])
+
+        :param x: The input to be checked.
+        """
+        less_than = jnp.all(x <= self.upper_bounds, axis=-1)
+        greater_than = jnp.all(x >= self.lower_bounds, axis=-1)
+        return less_than & greater_than
+
+    def feasible_like(self, prototype):
+        return jnp.broadcast_to(
+            (self.lower_bounds + self.upper_bounds) / 2, jnp.shape(prototype)
+        )
+
+    def tree_flatten(self):
+        return (self.lower_bounds, self.upper_bounds), (
+            ("lower_bounds", "upper_bounds"),
+            dict(),
+        )
+
+    def __eq__(self, other):
+        if not isinstance(other, _UniqueIntervals):
+            return False
+        return jnp.array_equal(self.lower_bounds, other.lower_bounds) & jnp.array_equal(
+            self.upper_bounds, other.upper_bounds
+        )
+
+
 # TODO: Make types consistent
 # See https://github.com/pytorch/pytorch/issues/50616
 
@@ -805,3 +941,6 @@ def tree_flatten(self):
 unit_interval = _UnitInterval()
 open_interval = _OpenInterval
 zero_sum = _ZeroSum
+union_of_closed_intervals = _UnionOfClosedIntervals
+intersection_of_closed_intervals = _IntersectionOfClosedIntervals
+unique_intervals = _UniqueIntervals

numpyro/distributions/transforms.py

@@ -1636,3 +1636,29 @@ def _transform_to_simplex(constraint):
 @biject_to.register(constraints.zero_sum)
 def _transform_to_zero_sum(constraint):
     return ZeroSumTransform(constraint.event_dim)
+
+
+@biject_to.register(constraints.union_of_closed_intervals)
+@biject_to.register(constraints.intersection_of_closed_intervals)
+def _transform_to_union_of_intervals(constraint):
+    return ComposeTransform(
+        [
+            ReshapeTransform(
+                forward_shape=constraint.lower_bounds.shape,
+                inverse_shape=(constraint.lower_bounds.size,),
+            ),
+            ExpTransform(),
+        ]
+    )
+
+
+@biject_to.register(constraints.unique_intervals)
+def _transform_to_unique_intervals(constraint):
+    return ComposeTransform(
+        [
+            ReshapeTransform(
+                forward_shape=constraint.lower_bounds.shape,
+                inverse_shape=(constraint.lower_bounds.size,),
+            ),
+        ]
+    )

test/test_constraints.py

@@ -68,6 +68,19 @@ class T(namedtuple("TestCase", ["constraint_cls", "params", "kwargs"])):
     ),
     "open_interval": T(constraints.open_interval, (_a(-1.0), _a(1.0)), dict()),
     "zero_sum": T(constraints.zero_sum, (), dict(event_dim=1)),
+    "union_of_closed_intervals": T(
+        constraints.union_of_closed_intervals,
+        (_a((-100, -50, 0, 50)), _a((-50, 0, 50, 100))),
+        dict(),
+    ),
+    "intersection_of_closed_intervals": T(
+        constraints.intersection_of_closed_intervals,
+        (_a((-100, -130, -150, -250)), _a((200, 100, 250, 100))),
+        dict(),
+    ),
+    "unique_intervals": T(
+        constraints.unique_intervals, (_a((-10, -8, -6, -4)), _a((4, 6, 8, 10))), dict()
+    ),
 }
 
 # TODO: BijectorConstraint

test/test_transforms.py

@@ -383,6 +383,22 @@ def test_batched_recursive_linear_transform():
         (constraints.softplus_positive, (2,)),
         (constraints.unit_interval, (4,)),
         (constraints.nonnegative, (7,)),
+        (
+            constraints.union_of_closed_intervals(
+                (-100, -50, 0, 50), (-50, 0, 50, 100)
+            ),
+            (4,),
+        ),
+        (
+            constraints.intersection_of_closed_intervals(
+                (-100, -130, -150, -250), (200, 100, 250, 100)
+            ),
+            (4,),
+        ),
+        (
+            constraints.unique_intervals((-10, -8, -6, -4), (4, 6, 8, 10)),
+            (4,),
+        ),
     ],
     ids=str,
 )