wehub-resource-sync
153 lines
4.4 KiB
Python
153 lines
4.4 KiB
Python
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# 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.
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
import paddle
|
|
|
|
if TYPE_CHECKING:
|
|
from paddle import Tensor
|
|
|
|
|
|
class Constraint:
|
|
"""Constraint condition for random variable."""
|
|
|
|
def __call__(self, value: Tensor) -> Tensor:
|
|
raise NotImplementedError
|
|
|
|
def check(self, value: Tensor) -> Tensor:
|
|
return self(value)
|
|
|
|
|
|
class Real(Constraint):
|
|
def __call__(self, value: Tensor) -> Tensor:
|
|
return value == value
|
|
|
|
|
|
class RealVector(Constraint):
|
|
event_dim = 1
|
|
|
|
def __call__(self, value: Tensor) -> Tensor:
|
|
if value.dim() < 1:
|
|
return paddle.zeros(value.shape[:-1], dtype='bool')
|
|
return (value == value).reshape((*value.shape[:-1], -1)).all(-1)
|
|
|
|
|
|
class Range(Constraint):
|
|
def __init__(self, lower: float | Tensor, upper: float | Tensor) -> None:
|
|
self._lower = lower
|
|
self._upper = upper
|
|
super().__init__()
|
|
|
|
def __call__(self, value: Tensor) -> Tensor:
|
|
return self._lower <= value <= self._upper
|
|
|
|
|
|
class IntegerInterval(Constraint):
|
|
event_dim = 0
|
|
is_discrete = True
|
|
|
|
def __init__(self, lower: int, upper: int) -> None:
|
|
self._lower = lower
|
|
self._upper = upper
|
|
super().__init__()
|
|
|
|
def __call__(self, value: Tensor) -> Tensor:
|
|
return (
|
|
(value >= self._lower) & (value <= self._upper) & (value % 1 == 0)
|
|
)
|
|
|
|
|
|
class Positive(Constraint):
|
|
def __call__(self, value: Tensor) -> Tensor:
|
|
return value >= 0.0
|
|
|
|
|
|
class LowerTriangular(Constraint):
|
|
event_dim = 2
|
|
|
|
def __call__(self, value: Tensor) -> Tensor:
|
|
if value.dim() < 2:
|
|
return paddle.zeros(value.shape[:-2], dtype='bool')
|
|
value_tril = paddle.tril(value)
|
|
return (value_tril == value).reshape((*value.shape[:-2], -1)).all(-1)
|
|
|
|
|
|
class LowerCholesky(Constraint):
|
|
event_dim = 2
|
|
|
|
def __call__(self, value: Tensor) -> Tensor:
|
|
if value.dim() < 2:
|
|
return paddle.zeros(value.shape[:-2], dtype='bool')
|
|
value_tril = paddle.tril(value)
|
|
lower_triangular = (
|
|
(value_tril == value).reshape((*value.shape[:-2], -1)).all(-1)
|
|
)
|
|
positive_diagonal = (value.diagonal(axis1=-2, axis2=-1) > 0).all(-1)
|
|
return lower_triangular & positive_diagonal
|
|
|
|
|
|
class Square(Constraint):
|
|
event_dim = 2
|
|
|
|
def __call__(self, value: Tensor) -> Tensor:
|
|
if value.dim() < 2:
|
|
return paddle.full_like(value.sum(), False, dtype='bool')
|
|
batch_value = value.reshape((*value.shape[:-2], -1)).sum(-1)
|
|
return paddle.full_like(
|
|
batch_value, value.shape[-2] == value.shape[-1], dtype='bool'
|
|
)
|
|
|
|
|
|
class Symmetric(Square):
|
|
def __call__(self, value: Tensor) -> Tensor:
|
|
square_check = super().__call__(value)
|
|
if value.dim() < 2:
|
|
return square_check
|
|
if value.shape[-2] != value.shape[-1]:
|
|
return square_check
|
|
return square_check & paddle.isclose(value, value.mT, atol=1e-6).all(
|
|
-2
|
|
).all(-1)
|
|
|
|
|
|
class PositiveDefinite(Symmetric):
|
|
def __call__(self, value: Tensor) -> Tensor:
|
|
if value.dim() < 2:
|
|
return paddle.zeros(value.shape[:-2], dtype='bool')
|
|
sym_check = super().__call__(value)
|
|
if value.shape[-2] != value.shape[-1]:
|
|
return sym_check
|
|
return sym_check & (paddle.linalg.eigvalsh(value) > 0).all(-1)
|
|
|
|
|
|
class Simplex(Constraint):
|
|
def __call__(self, value: Tensor) -> Tensor:
|
|
return paddle.all(value >= 0, axis=-1) and (
|
|
(value.sum(-1) - 1).abs() < 1e-6
|
|
)
|
|
|
|
|
|
real = Real()
|
|
real_vector = RealVector()
|
|
integer_interval = IntegerInterval
|
|
positive = Positive()
|
|
lower_triangular = LowerTriangular()
|
|
lower_cholesky = LowerCholesky()
|
|
square = Square()
|
|
symmetric = Symmetric()
|
|
positive_definite = PositiveDefinite()
|
|
simplex = Simplex()
|