openscvx 0.1.2__py3-none-any.whl → 0.2.1.dev0__py3-none-any.whl

openscvx/constraints/ctcs.py

@@ -1,22 +1,98 @@
 from dataclasses import dataclass
-from typing import Callable, Sequence, Tuple, Optional
-import functools
-import types
+from typing import Callable, Optional, Tuple, Union
 
-from jax.lax import cond
 import jax.numpy as jnp
+from jax.lax import cond
+import functools
+import inspect
 
+from openscvx.backend.state import State, Variable
+from openscvx.backend.control import Control
+from openscvx.backend.parameter import Parameter
+
+# TODO: (norrisg) Unclear if should specify behavior for `idx`, `jacfwd` behavior for Jacobians, etc. since that logic is handled elsewhere and could change
 
 @dataclass
 class CTCSConstraint:
-    func: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]
+    """
+    Dataclass for continuous-time constraint satisfaction (CTCS) constraints over a trajectory interval.
+    
+    A `CTCSConstraint` wraps a residual function `func(x, u)`, applies a
+    pointwise `penalty` to its outputs, and accumulates the penalized sum
+    only within a specified node interval [nodes[0], nodes[1]).
+    
+    CTCS constraints are used for continuous-time constraints that need to be satisfied
+    over trajectory intervals rather than at specific nodes. The constraint function
+    should return residuals where positive values indicate constraint violations.
+
+    Usage examples:
+
+    ```python
+    @ctcs
+    def g(x_, u_):
+        return 2.0 - jnp.linalg.norm(x_[:3])  # ||x[:3]|| <= 2 constraint
+    ```
+    ```python
+    @ctcs(penalty="huber", nodes=(0, 10), idx=2)
+    def g(x_, u_):
+        return jnp.sin(x_) + u_  # sin(x) + u <= 0 constraint
+    ```
+    ```python
+    @ctcs(penalty="smooth_relu", scaling=0.5)
+    def g(x_, u_):
+        return x_[0]**2 + x_[1]**2 - 1.0  # ||x||^2 <= 1 constraint
+    ```
+
+    Or can directly wrap a function if a more lambda-function interface is desired:
+
+    ```python
+    constraint = ctcs(lambda x_, u_: jnp.maximum(0, x[0] - 1.0))
+    ```
+
+    Args:
+        func (Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]): 
+            Function computing constraint residuals g(x, u). 
+            - x: 1D array (state at a single node), shape (n_x,)
+            - u: 1D array (control at a single node), shape (n_u,)
+            - Additional parameters: passed as keyword arguments with names matching the parameter name plus an underscore (e.g., g_ for Parameter('g')).
+            Should return positive values for constraint violations (g(x,u) > 0 indicates violation).
+            If you want to use parameters, include them as extra arguments with the underscore naming convention.
+        penalty (Callable[[jnp.ndarray], jnp.ndarray]): 
+            Penalty function applied elementwise to g's output. Used to calculate and penalize
+            constraint violation during state augmentation. Common penalties include:
+            - "squared_relu": max(0, x)² (default)
+            - "huber": smooth approximation of absolute value
+            - "smooth_relu": differentiable version of ReLU
+        nodes (Optional[Tuple[int, int]]):
+            Half-open interval (start, end) of node indices where this constraint is active.
+            If None, the penalty applies at every node.
+        idx (Optional[int]): 
+            Optional index used to group CTCS constraints. Used during automatic state augmentation.
+            All CTCS constraints with the same index must be active over the same `nodes` interval
+        grad_f_x (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]):
+            User-supplied gradient of `func` w.r.t. state `x`, signature (x, u) -> jacobian.
+            If None, computed via `jax.jacfwd(func, argnums=0)` during state augmentation.
+        grad_f_u (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]):
+            User-supplied gradient of `func` w.r.t. input `u`, signature (x, u) -> jacobian.
+            If None, computed via `jax.jacfwd(func, argnums=1)` during state augmentation.
+        scaling (float):
+            Scaling factor to apply to the penalized sum.
+    """
+    func: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]  # takes (x_expr, u_expr, *param_exprs)
     penalty: Callable[[jnp.ndarray], jnp.ndarray]
     nodes: Optional[Tuple[int, int]] = None
     idx: Optional[int] = None
-    grad_f_x: Optional[Callable] = None
-    grad_f_u: Optional[Callable] = None
+    grad_f_x: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None
+    grad_f_u: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None
+    scaling: float = 1.0
 
     def __post_init__(self):
+        """
+        Adapt user-provided gradients to the three-argument signature (x, u, node).
+
+        If `grad_f_x` or `grad_f_u` are given as functions of (x, u), wrap them
+        so they accept the extra `node` argument to match `__call__`.
+        """
         if self.grad_f_x is not None:
             _grad_f_x = self.grad_f_x
             self.grad_f_x = lambda x, u, nodes: _grad_f_x(x, u)
@@ -24,62 +100,141 @@ class CTCSConstraint:
             _grad_f_u = self.grad_f_u
             self.grad_f_u = lambda x, u, nodes: _grad_f_u(x, u)
 
-    def __call__(self, x, u, node):
+    def __call__(self, x, u, node: int, *params):
+        """
+        Evaluate the penalized constraint at a given node index.
+        The penalty is summed only if `node` lies within the active interval.
+
+        Args:
+            x (jnp.ndarray): State vector at this node.
+            u (jnp.ndarray): Input vector at this node.
+            node (int): Trajectory time-step index.
+            *params (tuple): Sequence of (name, value) pairs for parameters.
+
+        Returns:
+            jnp.ndarray or float:
+                The total penalty (sum over selected residuals) if inside interval,
+                otherwise zero.
+        """
+        x_expr = x.expr if isinstance(x, (State, Variable)) else x
+        u_expr = u.expr if isinstance(u, (Control, Variable)) else u
+
+        # Inspect function signature for expected parameter names
+        func_signature = inspect.signature(self.func)
+        expected_args = set(func_signature.parameters.keys())
+
+        # Only include params whose name (with underscore) is in the function's signature
+        filtered_params = {
+            f"{name}_": value for name, value in params if f"{name}_" in expected_args
+        }
+
         return cond(
-            jnp.all((self.nodes[0] <= node) & (node < self.nodes[1])),
-            lambda _: jnp.sum(self.penalty(self.func(x, u))),
+            jnp.all((self.nodes[0] <= node) & (node < self.nodes[1]))
+            if self.nodes is not None else True,
+            lambda _: self.scaling * jnp.sum(self.penalty(self.func(x_expr, u_expr, **filtered_params))),
             lambda _: 0.0,
             operand=None,
         )
 
 
 def ctcs(
-    _func=None,
+    _func: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None,
     *,
-    penalty: str = "squared_relu",
-    nodes: Optional[Sequence[Tuple[int, int]]] = None,
+    penalty: Union[str, Callable[[jnp.ndarray], jnp.ndarray]] = "squared_relu",
+    nodes: Optional[Tuple[int, int]] = None,
     idx: Optional[int] = None,
-    grad_f_x: Optional[Callable] = None,
-    grad_f_u: Optional[Callable] = None,
-):
-    """Decorator to mark a function as a 'ctcs' constraint.
-
-    Use as:
-    @ctcs(nodes=[(0,10)], idx=2, penalty='huber')
-    def my_constraint(x,u): ...
+    grad_f_x: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None,
+    grad_f_u: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None,
+    scaling: float = 1.0,
+) -> Union[Callable, CTCSConstraint]:
+    """
+    Decorator to build a CTCSConstraint from a raw constraint function.
+
+    Supports built-in penalties by name or a custom penalty function.
+
+    Usage examples:
+
+    ```python
+    @ctcs
+    def g(x, u):
+        return jnp.maximum(0, x - 1)
+    ```
+    ```python
+    @ctcs("huber", nodes=[(0, 10)], idx=2)
+    def g2(x, u):
+        return jnp.sin(x) + u
+    ```
+
+    Or can directly wrap a function if a more lambda-function interface is desired:
+
+    ```python
+    constraint = [ctcs(lambda x, u: ...)]
+    ```
+
+    Args:
+        _func (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]):
+            The function to wrap; provided automatically when using bare @ctcs.
+        penalty (Union[str, Callable[[jnp.ndarray], jnp.ndarray]]):
+            Name of a built-in penalty ('squared_relu', 'huber', 'smooth_relu')
+            or a custom elementwise penalty function.
+        nodes (Optional[Tuple[int, int]]):
+            Half-open interval (start, end) of node indices where this constraint is active.
+            If None, the penalty applies at every node.
+        idx (Optional[int]):
+            Optional index used to group CTCS constraints. Used during automatic state augmentation.
+            All CTCS constraints with the same index must be active over the same `nodes` interval
+        grad_f_x (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]):
+            User-supplied gradient of `func` w.r.t state `x`.
+            If None, computed via `jax.jacfwd(func, argnums=0)` during state augmentation.
+        grad_f_u (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]):
+            User-supplied gradient of `func` w.r.t input `u`.
+            If None, computed via `jax.jacfwd(func, argnums=1)` during state augmentation.
+        scaling (float):
+            Scaling factor to apply to the penalized sum.
+
+    Returns:
+        Union[Callable, CTCSConstraint]
+            A decorator if called without a function, or a CTCSConstraint instance
+            when applied to a function.
+
+    Raises:
+        ValueError: If `penalty` string is not one of the supported names.
     """
     # prepare penalty function once
     if penalty == "squared_relu":
         pen = lambda x: jnp.maximum(0, x) ** 2
     elif penalty == "huber":
         delta = 0.25
-
-        def pen(x):
-            r = jnp.maximum(0, x)
-            return jnp.where(r < delta, 0.5 * r**2, r - 0.5 * delta)
+        def pen(x): return jnp.where(x < delta, 0.5 * x**2, x - 0.5 * delta)
     elif penalty == "smooth_relu":
         c = 1e-8
         pen = lambda x: (jnp.maximum(0, x) ** 2 + c**2) ** 0.5 - c
-    elif isinstance(penalty, types.LambdaType):
+    elif callable(penalty):
         pen = penalty
     else:
         raise ValueError(f"Unknown penalty {penalty}")
 
-    def decorator(f: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]):
-        # wrap so name, doc, signature stay on f
-        wrapped = functools.wraps(f)(f)
-        return CTCSConstraint(
-            func=wrapped,
+    def decorator(f: Callable):
+        @functools.wraps(f)  # preserves name, docstring, and signature on the wrapper
+        def wrapper(*args, **kwargs):
+            return f(*args, **kwargs)
+
+        # Now attach original function as attribute if needed
+        wrapper._original_func = f
+
+        # Return your CTCSConstraint with the original function, but keep the wrapper around it
+        constraint = CTCSConstraint(
+            func=wrapper,
             penalty=pen,
             nodes=nodes,
             idx=idx,
             grad_f_x=grad_f_x,
             grad_f_u=grad_f_u,
+            scaling=scaling,
         )
+        return constraint
 
-    # if called as @ctcs or @ctcs(...), _func will be None and we return decorator
     if _func is None:
         return decorator
-    # if called as ctcs(func), we immediately decorate
     else:
         return decorator(_func)

openscvx/constraints/nodal.py

@@ -1,20 +1,109 @@
 from dataclasses import dataclass
-from typing import Callable, Optional, List
+from typing import Callable, Optional, List, Union
 
 import jax.numpy as jnp
-from jax import jit, vmap, jacfwd
+from jax import vmap, jacfwd
 
 
 @dataclass
 class NodalConstraint:
+    """
+    Encapsulates a constraint function applied at specific trajectory nodes.
+
+    A `NodalConstraint` wraps a function `g(x, u)` that computes constraint residuals
+    for given state `x` and input `u`. It can optionally apply only at
+    a subset of trajectory nodes, support vectorized evaluation across nodes,
+    and integrate with convex solvers when `convex=True`.
+
+    **Expected input types:**
+
+    | Case                        | x, u type/shape                                 |
+    |-----------------------------|-------------------------------------------------|
+    | convex=False, vectorized=False | 1D arrays, shape (n_x,), (n_u,) (single node)   |
+    | convex=False, vectorized=True  | 2D arrays, shape (N, n_x), (N, n_u) (all nodes) |
+    | convex=True, vectorized=False  | list of cvxpy variables, one per node           |
+    | convex=True, vectorized=True   | list of cvxpy variables, one per node           |
+
+    **Expected output:**
+
+    | Case                        | Output type                                      |
+    |-----------------------------|--------------------------------------------------|
+    | convex=False, vectorized=False | float (single node)                              |
+    | convex=False, vectorized=True  | float array (per node)                           |
+    | convex=True, vectorized=False  | cvxpy expression (single node)                   |
+    | convex=True, vectorized=True   | list of cvxpy expressions (one per node)         |
+
+    Nonconvex examples:
+    
+    ```python
+    @nodal
+    def g(x_, u_):
+        return 1 - x_[0] <= 0
+    ```
+    ```python
+    @nodal(nodes=[0, 3])
+    def g(x_, u_):
+        return jnp.linalg.norm(x_) - 1.0 
+    ```
+
+    Or can directly wrap a function if a more lambda-function interface is desired:
+
+    ```python
+    constraint = nodal(lambda x_, u_: 1 - x_[0])
+    ```
+
+    Convex Examples:
+
+    ```python
+    @nodal(convex=True)
+    def g(x_, u_):
+        return cp.norm(x_) <= 1.0  # cvxpy expression following DPP rules
+    ```
+
+    Args:
+        func (Callable):
+            The user-supplied constraint function. The expected input and output types depend on the values of `convex` and `vectorized`:
+
+            | Case                          | x, u type/shape                                   | Output type                                      |
+            |-------------------------------|---------------------------------------------------|--------------------------------------------------|
+            | convex=False, vectorized=False | 1D arrays, shape (n_x,), (n_u,) (single node)     | float (single node)                              |
+            | convex=False, vectorized=True  | 2D arrays, shape (N, n_x), (N, n_u) (all nodes)   | float array (per node)                           |
+            | convex=True, vectorized=False  | list of cvxpy variables, one per node             | cvxpy expression (single node)                   |
+            | convex=True, vectorized=True   | list of cvxpy variables, one per node             | list of cvxpy expressions (one per node)         |
+
+            Additional parameters: always passed as keyword arguments with names matching the parameter name plus an underscore (e.g., `g_` for `Parameter('g')`).
+            For nonconvex constraints, the function should return constraint residuals (g(x,u) <= 0). For convex constraints, the function should return a cvxpy expression.
+        nodes (Optional[List[int]]):
+            Specific node indices where this constraint applies. If None, applies at all nodes.
+        convex (bool):
+            If True, the provided cvxpy.expression is directly passed to the cvxpy.problem.
+        vectorized (bool):
+            If False, automatically vectorizes `func` and its jacobians over
+            the node dimension using `jax.vmap`. If True, assumes `func` already
+            handles vectorization.
+        grad_g_x (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]):
+            User-supplied gradient of `func` wrt `x`. If None, computed via
+            `jax.jacfwd(func, argnums=0)`.
+        grad_g_u (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]): 
+            User-supplied gradient of `func` wrt `u`. If None, computed via
+            `jax.jacfwd(func, argnums=1)`.
+    """
+
     func: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]
     nodes: Optional[List[int]] = None
     convex: bool = False
     vectorized: bool = False
-    grad_g_x: Optional[Callable] = None
-    grad_g_u: Optional[Callable] = None
+    grad_g_x: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None
+    grad_g_u: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None
 
     def __post_init__(self):
+        """Initialize gradients and vectorization after instance creation.
+        
+        If the constraint is not convex, this method:
+        1. Sets up the constraint function
+        2. Computes gradients if not provided
+        3. Vectorizes the functions if needed
+        """
         if not self.convex:
             # single-node but still using JAX
             self.g = self.func
@@ -29,21 +118,66 @@ class NodalConstraint:
         # if convex=True assume an external solver (e.g. CVX) will handle it
 
     def __call__(self, x: jnp.ndarray, u: jnp.ndarray):
+        """Evaluate the constraint function at the given state and control.
+        
+        Args:
+            x (jnp.ndarray): The state vector.
+            u (jnp.ndarray): The control vector.
+            
+        Returns:
+            jnp.ndarray: The constraint violation values.
+        """
         return self.func(x, u)
 
 
 def nodal(
-    _func=None,
+    _func: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None,
     *,
     nodes: Optional[List[int]] = None,
     convex: bool = False,
     vectorized: bool = False,
-    grad_g_x: Optional[Callable] = None,
-    grad_g_u: Optional[Callable] = None,
-):
-    def decorator(f: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]):
+    grad_g_x: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None,
+    grad_g_u: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None,
+) -> Union[Callable, NodalConstraint]:
+    """
+    Decorator to build a `NodalConstraint` from a constraint function.
+
+    Usage examples:
+
+    ```python
+    @nodal
+    def g(x, u):
+        ...
+    ```
+    ```python
+    @nodal(nodes=[0, -1], convex=True, vectorized=False)
+    def g(x, u):
+        ...
+    ```
+
+    Or can directly wrap a function if a more lambda-function interface is desired:
+
+    ```python
+    constraint = nodal(lambda x, u: ...)
+    ```
+
+    Args:
+        _func (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]):
+            The function to wrap; populated automatically when using bare @nodal.
+            When `convex=False`, this is a standard function g(x, u) that should
+            return constraint residuals (g(x,u) <= 0). When `convex=True`, this
+            must be a cvxpy expression following the DPP ruleset.
+        nodes (Optional[List[int]]):
+            Node indices where the constraint applies; default None applies to all.
+        convex (bool):
+            If True, the provided cvxpy.expression is directly passed to the cvxpy.problem.
+        vectorized (bool):
+            If False, auto-vectorize over nodes using `jax.vmap`. If True, assumes
+            the function already handles vectorization.
+    """
+    def decorator(f: Callable):
         return NodalConstraint(
-            func=f,  # no wraps, just keep the original
+            func=f,
             nodes=nodes,
             convex=convex,
             vectorized=vectorized,
@@ -51,4 +185,9 @@ def nodal(
             grad_g_u=grad_g_u,
         )
 
-    return decorator if _func is None else decorator(_func)
+    if _func is None:
+        # Called with arguments, e.g., @nodal(nodes=[0, 1])
+        return decorator
+    else:
+        # Called as a bare decorator, e.g., @nodal
+        return decorator(_func)

openscvx/constraints/violation.py

@@ -9,6 +9,16 @@ from openscvx.constraints.ctcs import CTCSConstraint
 
 @dataclass
 class CTCSViolation:
+    """Class representing a continuous-time constraint satisfaction (CTCS) violation.
+    
+    This class holds the constraint function and its gradients for computing
+    constraint violations in continuous-time optimization problems.
+    
+    Attributes:
+        g (Callable): The constraint function that computes violations.
+        g_grad_x (Optional[Callable]): Gradient of g with respect to state x.
+        g_grad_u (Optional[Callable]): Gradient of g with respect to control u.
+    """
     g: Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]
     g_grad_x: Optional[Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]] = None
     g_grad_u: Optional[Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]] = None
@@ -35,8 +45,8 @@ def get_g_grad_u(constraints_ctcs: List[CTCSConstraint]) -> Callable[[jnp.ndarra
 
 
 def get_g_func(constraints_ctcs: List[CTCSConstraint]) -> Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]:
-    def g_func(x: jnp.array, u: jnp.array, node: int) -> jnp.array:
-        return sum(c(x, u, node) for c in constraints_ctcs)
+    def g_func(x: jnp.array, u: jnp.array, node: int, *params) -> jnp.array:
+        return sum(c(x, u, node, *params) for c in constraints_ctcs)
 
     return g_func