Source code for rlaopt.atoms.halfspace

"""Halfspace constraint atom for optimization."""

import torch

from rlaopt.atoms.polyhedron import Polyhedron
from rlaopt.expression import Variable
from rlaopt.ext_tensordict import TensorDict




class Halfspace(Polyhedron):
    """Halfspace constraint atom representing a linear inequality.

    A halfspace constraint restricts a variable to satisfy a linear inequality:
    c^T x <= upper, which defines a half-space in the variable space.

    This is a special case of Polyhedra with a single linear inequality
    constraint, but with an efficient closed-form proximal operator (projection
    onto the halfspace).

    Args:
        x: Variable to constrain.
        c: Normal vector defining the halfspace orientation.
        upper: Upper bound for the linear form c^T x.

    Examples:
        >>> # Constraint: c^T x <= 1
        >>> x = Variable((5,), name='x')
        >>> c = torch.randn(5)
        >>> halfspace = Halfspace(x, c=c, upper=1.0)

        >>> # Non-negativity for first coordinate: x[0] >= 0
        >>> # Rewritten as: -x[0] <= 0, so c = [-1, 0, 0, ...], upper = 0
        >>> c = torch.zeros(5)
        >>> c[0] = -1.0
        >>> nonneg = Halfspace(x, c=c, upper=torch.tensor(0.0))

        >>> # Use proximal operator for projection
        >>> violating_point = torch.randn(5)
        >>> projected = halfspace.prox(violating_point, prox_scaling=1.0)
    """



    def __init__(self, x: Variable, c: torch.Tensor, upper: float):
        """Initialize the halfspace constraint atom.

        Args:
            x: Variable to constrain.
            c: Normal vector defining the halfspace orientation.
            upper: Upper bound for the linear form c^T x.
        """
        super().__init__(x, A=None, b=None, C=c, lower=None, upper=upper)




    def is_proxable(self) -> bool:
        """Check if the halfspace constraint has a computable proximal operator.

        Returns:
            bool: Always True, as halfspace constraints have a closed-form
                proximal operator (projection onto the halfspace).
        """
        return True


    def _prox(
        self, relevant_variable_values: TensorDict, prox_scaling: float
    ) -> TensorDict:
        """Compute the proximal operator of the halfspace constraint.

        The proximal operator projects onto the halfspace by moving the point
        perpendicular to the boundary until it satisfies c^T x <= upper.
        """
        c = self.get_buffer("C")
        upper = self.get_buffer("upper")

        def project_onto_halfspace(x: torch.Tensor) -> torch.Tensor:
            r = torch.dot(c, x) - upper
            # x is feasible
            if r <= 0:
                return x
            # x is infeasible
            else:
                c_norm = torch.linalg.norm(c, 2)
                return x - torch.nn.functional.relu(r) * c / c_norm**2

        return relevant_variable_values.apply(project_onto_halfspace)