658 lines
18 KiB
Python
658 lines
18 KiB
Python
r"""
|
|
The following constraints are implemented:
|
|
|
|
- ``constraints.boolean``
|
|
- ``constraints.cat``
|
|
- ``constraints.corr_cholesky``
|
|
- ``constraints.dependent``
|
|
- ``constraints.greater_than(lower_bound)``
|
|
- ``constraints.greater_than_eq(lower_bound)``
|
|
- ``constraints.independent(constraint, reinterpreted_batch_ndims)``
|
|
- ``constraints.integer_interval(lower_bound, upper_bound)``
|
|
- ``constraints.interval(lower_bound, upper_bound)``
|
|
- ``constraints.less_than(upper_bound)``
|
|
- ``constraints.lower_cholesky``
|
|
- ``constraints.lower_triangular``
|
|
- ``constraints.multinomial``
|
|
- ``constraints.nonnegative``
|
|
- ``constraints.nonnegative_integer``
|
|
- ``constraints.one_hot``
|
|
- ``constraints.positive_integer``
|
|
- ``constraints.positive``
|
|
- ``constraints.positive_semidefinite``
|
|
- ``constraints.positive_definite``
|
|
- ``constraints.real_vector``
|
|
- ``constraints.real``
|
|
- ``constraints.simplex``
|
|
- ``constraints.symmetric``
|
|
- ``constraints.stack``
|
|
- ``constraints.square``
|
|
- ``constraints.symmetric``
|
|
- ``constraints.unit_interval``
|
|
"""
|
|
|
|
import torch
|
|
|
|
__all__ = [
|
|
"Constraint",
|
|
"boolean",
|
|
"cat",
|
|
"corr_cholesky",
|
|
"dependent",
|
|
"dependent_property",
|
|
"greater_than",
|
|
"greater_than_eq",
|
|
"independent",
|
|
"integer_interval",
|
|
"interval",
|
|
"half_open_interval",
|
|
"is_dependent",
|
|
"less_than",
|
|
"lower_cholesky",
|
|
"lower_triangular",
|
|
"multinomial",
|
|
"nonnegative",
|
|
"nonnegative_integer",
|
|
"one_hot",
|
|
"positive",
|
|
"positive_semidefinite",
|
|
"positive_definite",
|
|
"positive_integer",
|
|
"real",
|
|
"real_vector",
|
|
"simplex",
|
|
"square",
|
|
"stack",
|
|
"symmetric",
|
|
"unit_interval",
|
|
]
|
|
|
|
|
|
class Constraint:
|
|
"""
|
|
Abstract base class for constraints.
|
|
|
|
A constraint object represents a region over which a variable is valid,
|
|
e.g. within which a variable can be optimized.
|
|
|
|
Attributes:
|
|
is_discrete (bool): Whether constrained space is discrete.
|
|
Defaults to False.
|
|
event_dim (int): Number of rightmost dimensions that together define
|
|
an event. The :meth:`check` method will remove this many dimensions
|
|
when computing validity.
|
|
"""
|
|
|
|
is_discrete = False # Default to continuous.
|
|
event_dim = 0 # Default to univariate.
|
|
|
|
def check(self, value):
|
|
"""
|
|
Returns a byte tensor of ``sample_shape + batch_shape`` indicating
|
|
whether each event in value satisfies this constraint.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def __repr__(self):
|
|
return self.__class__.__name__[1:] + "()"
|
|
|
|
|
|
class _Dependent(Constraint):
|
|
"""
|
|
Placeholder for variables whose support depends on other variables.
|
|
These variables obey no simple coordinate-wise constraints.
|
|
|
|
Args:
|
|
is_discrete (bool): Optional value of ``.is_discrete`` in case this
|
|
can be computed statically. If not provided, access to the
|
|
``.is_discrete`` attribute will raise a NotImplementedError.
|
|
event_dim (int): Optional value of ``.event_dim`` in case this
|
|
can be computed statically. If not provided, access to the
|
|
``.event_dim`` attribute will raise a NotImplementedError.
|
|
"""
|
|
|
|
def __init__(self, *, is_discrete=NotImplemented, event_dim=NotImplemented):
|
|
self._is_discrete = is_discrete
|
|
self._event_dim = event_dim
|
|
super().__init__()
|
|
|
|
@property
|
|
def is_discrete(self):
|
|
if self._is_discrete is NotImplemented:
|
|
raise NotImplementedError(".is_discrete cannot be determined statically")
|
|
return self._is_discrete
|
|
|
|
@property
|
|
def event_dim(self):
|
|
if self._event_dim is NotImplemented:
|
|
raise NotImplementedError(".event_dim cannot be determined statically")
|
|
return self._event_dim
|
|
|
|
def __call__(self, *, is_discrete=NotImplemented, event_dim=NotImplemented):
|
|
"""
|
|
Support for syntax to customize static attributes::
|
|
|
|
constraints.dependent(is_discrete=True, event_dim=1)
|
|
"""
|
|
if is_discrete is NotImplemented:
|
|
is_discrete = self._is_discrete
|
|
if event_dim is NotImplemented:
|
|
event_dim = self._event_dim
|
|
return _Dependent(is_discrete=is_discrete, event_dim=event_dim)
|
|
|
|
def check(self, x):
|
|
raise ValueError("Cannot determine validity of dependent constraint")
|
|
|
|
|
|
def is_dependent(constraint):
|
|
return isinstance(constraint, _Dependent)
|
|
|
|
|
|
class _DependentProperty(property, _Dependent):
|
|
"""
|
|
Decorator that extends @property to act like a `Dependent` constraint when
|
|
called on a class and act like a property when called on an object.
|
|
|
|
Example::
|
|
|
|
class Uniform(Distribution):
|
|
def __init__(self, low, high):
|
|
self.low = low
|
|
self.high = high
|
|
@constraints.dependent_property(is_discrete=False, event_dim=0)
|
|
def support(self):
|
|
return constraints.interval(self.low, self.high)
|
|
|
|
Args:
|
|
fn (Callable): The function to be decorated.
|
|
is_discrete (bool): Optional value of ``.is_discrete`` in case this
|
|
can be computed statically. If not provided, access to the
|
|
``.is_discrete`` attribute will raise a NotImplementedError.
|
|
event_dim (int): Optional value of ``.event_dim`` in case this
|
|
can be computed statically. If not provided, access to the
|
|
``.event_dim`` attribute will raise a NotImplementedError.
|
|
"""
|
|
|
|
def __init__(
|
|
self, fn=None, *, is_discrete=NotImplemented, event_dim=NotImplemented
|
|
):
|
|
super().__init__(fn)
|
|
self._is_discrete = is_discrete
|
|
self._event_dim = event_dim
|
|
|
|
def __call__(self, fn):
|
|
"""
|
|
Support for syntax to customize static attributes::
|
|
|
|
@constraints.dependent_property(is_discrete=True, event_dim=1)
|
|
def support(self):
|
|
...
|
|
"""
|
|
return _DependentProperty(
|
|
fn, is_discrete=self._is_discrete, event_dim=self._event_dim
|
|
)
|
|
|
|
|
|
class _IndependentConstraint(Constraint):
|
|
"""
|
|
Wraps a constraint by aggregating over ``reinterpreted_batch_ndims``-many
|
|
dims in :meth:`check`, so that an event is valid only if all its
|
|
independent entries are valid.
|
|
"""
|
|
|
|
def __init__(self, base_constraint, reinterpreted_batch_ndims):
|
|
assert isinstance(base_constraint, Constraint)
|
|
assert isinstance(reinterpreted_batch_ndims, int)
|
|
assert reinterpreted_batch_ndims >= 0
|
|
self.base_constraint = base_constraint
|
|
self.reinterpreted_batch_ndims = reinterpreted_batch_ndims
|
|
super().__init__()
|
|
|
|
@property
|
|
def is_discrete(self):
|
|
return self.base_constraint.is_discrete
|
|
|
|
@property
|
|
def event_dim(self):
|
|
return self.base_constraint.event_dim + self.reinterpreted_batch_ndims
|
|
|
|
def check(self, value):
|
|
result = self.base_constraint.check(value)
|
|
if result.dim() < self.reinterpreted_batch_ndims:
|
|
expected = self.base_constraint.event_dim + self.reinterpreted_batch_ndims
|
|
raise ValueError(
|
|
f"Expected value.dim() >= {expected} but got {value.dim()}"
|
|
)
|
|
result = result.reshape(
|
|
result.shape[: result.dim() - self.reinterpreted_batch_ndims] + (-1,)
|
|
)
|
|
result = result.all(-1)
|
|
return result
|
|
|
|
def __repr__(self):
|
|
return f"{self.__class__.__name__[1:]}({repr(self.base_constraint)}, {self.reinterpreted_batch_ndims})"
|
|
|
|
|
|
class _Boolean(Constraint):
|
|
"""
|
|
Constrain to the two values `{0, 1}`.
|
|
"""
|
|
|
|
is_discrete = True
|
|
|
|
def check(self, value):
|
|
return (value == 0) | (value == 1)
|
|
|
|
|
|
class _OneHot(Constraint):
|
|
"""
|
|
Constrain to one-hot vectors.
|
|
"""
|
|
|
|
is_discrete = True
|
|
event_dim = 1
|
|
|
|
def check(self, value):
|
|
is_boolean = (value == 0) | (value == 1)
|
|
is_normalized = value.sum(-1).eq(1)
|
|
return is_boolean.all(-1) & is_normalized
|
|
|
|
|
|
class _IntegerInterval(Constraint):
|
|
"""
|
|
Constrain to an integer interval `[lower_bound, upper_bound]`.
|
|
"""
|
|
|
|
is_discrete = True
|
|
|
|
def __init__(self, lower_bound, upper_bound):
|
|
self.lower_bound = lower_bound
|
|
self.upper_bound = upper_bound
|
|
super().__init__()
|
|
|
|
def check(self, value):
|
|
return (
|
|
(value % 1 == 0) & (self.lower_bound <= value) & (value <= self.upper_bound)
|
|
)
|
|
|
|
def __repr__(self):
|
|
fmt_string = self.__class__.__name__[1:]
|
|
fmt_string += (
|
|
f"(lower_bound={self.lower_bound}, upper_bound={self.upper_bound})"
|
|
)
|
|
return fmt_string
|
|
|
|
|
|
class _IntegerLessThan(Constraint):
|
|
"""
|
|
Constrain to an integer interval `(-inf, upper_bound]`.
|
|
"""
|
|
|
|
is_discrete = True
|
|
|
|
def __init__(self, upper_bound):
|
|
self.upper_bound = upper_bound
|
|
super().__init__()
|
|
|
|
def check(self, value):
|
|
return (value % 1 == 0) & (value <= self.upper_bound)
|
|
|
|
def __repr__(self):
|
|
fmt_string = self.__class__.__name__[1:]
|
|
fmt_string += f"(upper_bound={self.upper_bound})"
|
|
return fmt_string
|
|
|
|
|
|
class _IntegerGreaterThan(Constraint):
|
|
"""
|
|
Constrain to an integer interval `[lower_bound, inf)`.
|
|
"""
|
|
|
|
is_discrete = True
|
|
|
|
def __init__(self, lower_bound):
|
|
self.lower_bound = lower_bound
|
|
super().__init__()
|
|
|
|
def check(self, value):
|
|
return (value % 1 == 0) & (value >= self.lower_bound)
|
|
|
|
def __repr__(self):
|
|
fmt_string = self.__class__.__name__[1:]
|
|
fmt_string += f"(lower_bound={self.lower_bound})"
|
|
return fmt_string
|
|
|
|
|
|
class _Real(Constraint):
|
|
"""
|
|
Trivially constrain to the extended real line `[-inf, inf]`.
|
|
"""
|
|
|
|
def check(self, value):
|
|
return value == value # False for NANs.
|
|
|
|
|
|
class _GreaterThan(Constraint):
|
|
"""
|
|
Constrain to a real half line `(lower_bound, inf]`.
|
|
"""
|
|
|
|
def __init__(self, lower_bound):
|
|
self.lower_bound = lower_bound
|
|
super().__init__()
|
|
|
|
def check(self, value):
|
|
return self.lower_bound < value
|
|
|
|
def __repr__(self):
|
|
fmt_string = self.__class__.__name__[1:]
|
|
fmt_string += f"(lower_bound={self.lower_bound})"
|
|
return fmt_string
|
|
|
|
|
|
class _GreaterThanEq(Constraint):
|
|
"""
|
|
Constrain to a real half line `[lower_bound, inf)`.
|
|
"""
|
|
|
|
def __init__(self, lower_bound):
|
|
self.lower_bound = lower_bound
|
|
super().__init__()
|
|
|
|
def check(self, value):
|
|
return self.lower_bound <= value
|
|
|
|
def __repr__(self):
|
|
fmt_string = self.__class__.__name__[1:]
|
|
fmt_string += f"(lower_bound={self.lower_bound})"
|
|
return fmt_string
|
|
|
|
|
|
class _LessThan(Constraint):
|
|
"""
|
|
Constrain to a real half line `[-inf, upper_bound)`.
|
|
"""
|
|
|
|
def __init__(self, upper_bound):
|
|
self.upper_bound = upper_bound
|
|
super().__init__()
|
|
|
|
def check(self, value):
|
|
return value < self.upper_bound
|
|
|
|
def __repr__(self):
|
|
fmt_string = self.__class__.__name__[1:]
|
|
fmt_string += f"(upper_bound={self.upper_bound})"
|
|
return fmt_string
|
|
|
|
|
|
class _Interval(Constraint):
|
|
"""
|
|
Constrain to a real interval `[lower_bound, upper_bound]`.
|
|
"""
|
|
|
|
def __init__(self, lower_bound, upper_bound):
|
|
self.lower_bound = lower_bound
|
|
self.upper_bound = upper_bound
|
|
super().__init__()
|
|
|
|
def check(self, value):
|
|
return (self.lower_bound <= value) & (value <= self.upper_bound)
|
|
|
|
def __repr__(self):
|
|
fmt_string = self.__class__.__name__[1:]
|
|
fmt_string += (
|
|
f"(lower_bound={self.lower_bound}, upper_bound={self.upper_bound})"
|
|
)
|
|
return fmt_string
|
|
|
|
|
|
class _HalfOpenInterval(Constraint):
|
|
"""
|
|
Constrain to a real interval `[lower_bound, upper_bound)`.
|
|
"""
|
|
|
|
def __init__(self, lower_bound, upper_bound):
|
|
self.lower_bound = lower_bound
|
|
self.upper_bound = upper_bound
|
|
super().__init__()
|
|
|
|
def check(self, value):
|
|
return (self.lower_bound <= value) & (value < self.upper_bound)
|
|
|
|
def __repr__(self):
|
|
fmt_string = self.__class__.__name__[1:]
|
|
fmt_string += (
|
|
f"(lower_bound={self.lower_bound}, upper_bound={self.upper_bound})"
|
|
)
|
|
return fmt_string
|
|
|
|
|
|
class _Simplex(Constraint):
|
|
"""
|
|
Constrain to the unit simplex in the innermost (rightmost) dimension.
|
|
Specifically: `x >= 0` and `x.sum(-1) == 1`.
|
|
"""
|
|
|
|
event_dim = 1
|
|
|
|
def check(self, value):
|
|
return torch.all(value >= 0, dim=-1) & ((value.sum(-1) - 1).abs() < 1e-6)
|
|
|
|
|
|
class _Multinomial(Constraint):
|
|
"""
|
|
Constrain to nonnegative integer values summing to at most an upper bound.
|
|
|
|
Note due to limitations of the Multinomial distribution, this currently
|
|
checks the weaker condition ``value.sum(-1) <= upper_bound``. In the future
|
|
this may be strengthened to ``value.sum(-1) == upper_bound``.
|
|
"""
|
|
|
|
is_discrete = True
|
|
event_dim = 1
|
|
|
|
def __init__(self, upper_bound):
|
|
self.upper_bound = upper_bound
|
|
|
|
def check(self, x):
|
|
return (x >= 0).all(dim=-1) & (x.sum(dim=-1) <= self.upper_bound)
|
|
|
|
|
|
class _LowerTriangular(Constraint):
|
|
"""
|
|
Constrain to lower-triangular square matrices.
|
|
"""
|
|
|
|
event_dim = 2
|
|
|
|
def check(self, value):
|
|
value_tril = value.tril()
|
|
return (value_tril == value).view(value.shape[:-2] + (-1,)).min(-1)[0]
|
|
|
|
|
|
class _LowerCholesky(Constraint):
|
|
"""
|
|
Constrain to lower-triangular square matrices with positive diagonals.
|
|
"""
|
|
|
|
event_dim = 2
|
|
|
|
def check(self, value):
|
|
value_tril = value.tril()
|
|
lower_triangular = (
|
|
(value_tril == value).view(value.shape[:-2] + (-1,)).min(-1)[0]
|
|
)
|
|
|
|
positive_diagonal = (value.diagonal(dim1=-2, dim2=-1) > 0).min(-1)[0]
|
|
return lower_triangular & positive_diagonal
|
|
|
|
|
|
class _CorrCholesky(Constraint):
|
|
"""
|
|
Constrain to lower-triangular square matrices with positive diagonals and each
|
|
row vector being of unit length.
|
|
"""
|
|
|
|
event_dim = 2
|
|
|
|
def check(self, value):
|
|
tol = (
|
|
torch.finfo(value.dtype).eps * value.size(-1) * 10
|
|
) # 10 is an adjustable fudge factor
|
|
row_norm = torch.linalg.norm(value.detach(), dim=-1)
|
|
unit_row_norm = (row_norm - 1.0).abs().le(tol).all(dim=-1)
|
|
return _LowerCholesky().check(value) & unit_row_norm
|
|
|
|
|
|
class _Square(Constraint):
|
|
"""
|
|
Constrain to square matrices.
|
|
"""
|
|
|
|
event_dim = 2
|
|
|
|
def check(self, value):
|
|
return torch.full(
|
|
size=value.shape[:-2],
|
|
fill_value=(value.shape[-2] == value.shape[-1]),
|
|
dtype=torch.bool,
|
|
device=value.device,
|
|
)
|
|
|
|
|
|
class _Symmetric(_Square):
|
|
"""
|
|
Constrain to Symmetric square matrices.
|
|
"""
|
|
|
|
def check(self, value):
|
|
square_check = super().check(value)
|
|
if not square_check.all():
|
|
return square_check
|
|
return torch.isclose(value, value.mT, atol=1e-6).all(-2).all(-1)
|
|
|
|
|
|
class _PositiveSemidefinite(_Symmetric):
|
|
"""
|
|
Constrain to positive-semidefinite matrices.
|
|
"""
|
|
|
|
def check(self, value):
|
|
sym_check = super().check(value)
|
|
if not sym_check.all():
|
|
return sym_check
|
|
return torch.linalg.eigvalsh(value).ge(0).all(-1)
|
|
|
|
|
|
class _PositiveDefinite(_Symmetric):
|
|
"""
|
|
Constrain to positive-definite matrices.
|
|
"""
|
|
|
|
def check(self, value):
|
|
sym_check = super().check(value)
|
|
if not sym_check.all():
|
|
return sym_check
|
|
return torch.linalg.cholesky_ex(value).info.eq(0)
|
|
|
|
|
|
class _Cat(Constraint):
|
|
"""
|
|
Constraint functor that applies a sequence of constraints
|
|
`cseq` at the submatrices at dimension `dim`,
|
|
each of size `lengths[dim]`, in a way compatible with :func:`torch.cat`.
|
|
"""
|
|
|
|
def __init__(self, cseq, dim=0, lengths=None):
|
|
assert all(isinstance(c, Constraint) for c in cseq)
|
|
self.cseq = list(cseq)
|
|
if lengths is None:
|
|
lengths = [1] * len(self.cseq)
|
|
self.lengths = list(lengths)
|
|
assert len(self.lengths) == len(self.cseq)
|
|
self.dim = dim
|
|
super().__init__()
|
|
|
|
@property
|
|
def is_discrete(self):
|
|
return any(c.is_discrete for c in self.cseq)
|
|
|
|
@property
|
|
def event_dim(self):
|
|
return max(c.event_dim for c in self.cseq)
|
|
|
|
def check(self, value):
|
|
assert -value.dim() <= self.dim < value.dim()
|
|
checks = []
|
|
start = 0
|
|
for constr, length in zip(self.cseq, self.lengths):
|
|
v = value.narrow(self.dim, start, length)
|
|
checks.append(constr.check(v))
|
|
start = start + length # avoid += for jit compat
|
|
return torch.cat(checks, self.dim)
|
|
|
|
|
|
class _Stack(Constraint):
|
|
"""
|
|
Constraint functor that applies a sequence of constraints
|
|
`cseq` at the submatrices at dimension `dim`,
|
|
in a way compatible with :func:`torch.stack`.
|
|
"""
|
|
|
|
def __init__(self, cseq, dim=0):
|
|
assert all(isinstance(c, Constraint) for c in cseq)
|
|
self.cseq = list(cseq)
|
|
self.dim = dim
|
|
super().__init__()
|
|
|
|
@property
|
|
def is_discrete(self):
|
|
return any(c.is_discrete for c in self.cseq)
|
|
|
|
@property
|
|
def event_dim(self):
|
|
dim = max(c.event_dim for c in self.cseq)
|
|
if self.dim + dim < 0:
|
|
dim += 1
|
|
return dim
|
|
|
|
def check(self, value):
|
|
assert -value.dim() <= self.dim < value.dim()
|
|
vs = [value.select(self.dim, i) for i in range(value.size(self.dim))]
|
|
return torch.stack(
|
|
[constr.check(v) for v, constr in zip(vs, self.cseq)], self.dim
|
|
)
|
|
|
|
|
|
# Public interface.
|
|
dependent = _Dependent()
|
|
dependent_property = _DependentProperty
|
|
independent = _IndependentConstraint
|
|
boolean = _Boolean()
|
|
one_hot = _OneHot()
|
|
nonnegative_integer = _IntegerGreaterThan(0)
|
|
positive_integer = _IntegerGreaterThan(1)
|
|
integer_interval = _IntegerInterval
|
|
real = _Real()
|
|
real_vector = independent(real, 1)
|
|
positive = _GreaterThan(0.0)
|
|
nonnegative = _GreaterThanEq(0.0)
|
|
greater_than = _GreaterThan
|
|
greater_than_eq = _GreaterThanEq
|
|
less_than = _LessThan
|
|
multinomial = _Multinomial
|
|
unit_interval = _Interval(0.0, 1.0)
|
|
interval = _Interval
|
|
half_open_interval = _HalfOpenInterval
|
|
simplex = _Simplex()
|
|
lower_triangular = _LowerTriangular()
|
|
lower_cholesky = _LowerCholesky()
|
|
corr_cholesky = _CorrCholesky()
|
|
square = _Square()
|
|
symmetric = _Symmetric()
|
|
positive_semidefinite = _PositiveSemidefinite()
|
|
positive_definite = _PositiveDefinite()
|
|
cat = _Cat
|
|
stack = _Stack
|