openscvx 0.3.2.dev170__py3-none-any.whl

openscvx/discretization/discretization.py

@@ -0,0 +1,236 @@
+import jax.numpy as jnp
+import numpy as np
+
+from openscvx.config import Config
+from openscvx.integrators import solve_ivp_diffrax, solve_ivp_rk45
+from openscvx.lowered import Dynamics
+
+
+def dVdt(
+    tau: float,
+    V: jnp.ndarray,
+    u_cur: np.ndarray,
+    u_next: np.ndarray,
+    state_dot: callable,
+    A: callable,
+    B: callable,
+    n_x: int,
+    n_u: int,
+    N: int,
+    dis_type: str,
+    params: dict,
+) -> jnp.ndarray:
+    """Compute the time derivative of the augmented state vector.
+
+    This function computes the time derivative of the augmented state vector V,
+    which includes the state, state transition matrix, and control influence matrix.
+
+    Args:
+        tau (float): Current normalized time in [0,1].
+        V (jnp.ndarray): Augmented state vector.
+        u_cur (np.ndarray): Control input at current node.
+        u_next (np.ndarray): Control input at next node.
+        state_dot (callable): Function computing state derivatives.
+        A (callable): Function computing state Jacobian.
+        B (callable): Function computing control Jacobian.
+        n_x (int): Number of states.
+        n_u (int): Number of controls.
+        N (int): Number of nodes in trajectory.
+        dis_type (str): Discretization type ("ZOH" or "FOH").
+        **params: Additional parameters passed to state_dot, A, and B.
+
+    Returns:
+        jnp.ndarray: Time derivative of augmented state vector.
+    """
+    # Define the nodes
+    nodes = jnp.arange(0, N - 1)
+
+    # Define indices for slicing the augmented state vector
+    i0 = 0
+    i1 = n_x
+    i2 = i1 + n_x * n_x
+    i3 = i2 + n_x * n_u
+    i4 = i3 + n_x * n_u
+
+    # Unflatten V
+    V = V.reshape(-1, i4)
+
+    # Compute the interpolation factor based on the discretization type
+    if dis_type == "ZOH":
+        beta = 0.0
+    elif dis_type == "FOH":
+        beta = (tau) * N
+    alpha = 1 - beta
+
+    # Interpolate the control input
+    u = u_cur + beta * (u_next - u_cur)
+    s = u[:, -1]
+
+    # Initialize the augmented Jacobians
+    dfdx = jnp.zeros((V.shape[0], n_x, n_x))
+    dfdu = jnp.zeros((V.shape[0], n_x, n_u))
+
+    # Ensure x_seq and u have the same batch size
+    x = V[:, :n_x]
+    u = u[: x.shape[0]]
+
+    # Compute the nonlinear propagation term
+    f = state_dot(x, u[:, :-1], nodes, params)
+    F = s[:, None] * f
+
+    # Evaluate the State Jacobian
+    dfdx = A(x, u[:, :-1], nodes, params)
+    sdfdx = s[:, None, None] * dfdx
+
+    # Evaluate the Control Jacobian
+    dfdu_veh = B(x, u[:, :-1], nodes, params)
+    dfdu = dfdu.at[:, :, :-1].set(s[:, None, None] * dfdu_veh)
+    dfdu = dfdu.at[:, :, -1].set(f)
+
+    # Stack up the results into the augmented state vector
+    # fmt: off
+    dVdt = jnp.zeros_like(V)
+    dVdt = dVdt.at[:, i0:i1].set(F)
+    dVdt = dVdt.at[:, i1:i2].set(
+        jnp.matmul(sdfdx, V[:, i1:i2].reshape(-1, n_x, n_x)).reshape(-1, n_x * n_x)
+    )
+    dVdt = dVdt.at[:, i2:i3].set(
+        (jnp.matmul(sdfdx, V[:, i2:i3].reshape(-1, n_x, n_u)) + dfdu * alpha).reshape(-1, n_x * n_u)
+    )
+    dVdt = dVdt.at[:, i3:i4].set(
+        (jnp.matmul(sdfdx, V[:, i3:i4].reshape(-1, n_x, n_u)) + dfdu * beta).reshape(-1, n_x * n_u)
+    )
+    # fmt: on
+
+    return dVdt.reshape(-1)
+
+
+def calculate_discretization(
+    x,
+    u,
+    state_dot: callable,
+    A: callable,
+    B: callable,
+    settings: Config,
+    params: dict,
+):
+    """Calculate the discretized system matrices.
+
+    This function computes the discretized system matrices (A_bar, B_bar, C_bar)
+    and defect vector (z_bar) using numerical integration.
+
+    Args:
+        x: State trajectory.
+        u: Control trajectory.
+        state_dot (callable): Function computing state derivatives.
+        A (callable): Function computing state Jacobian.
+        B (callable): Function computing control Jacobian.
+        settings: Configuration settings for OpenSCvx.
+        custom_integrator (bool): Whether to use custom RK45 integrator.
+        debug (bool): Whether to use debug mode.
+        solver (str): Name of the solver to use.
+        rtol (float): Relative tolerance for integration.
+        atol (float): Absolute tolerance for integration.
+        dis_type (str): Discretization type ("ZOH" or "FOH").
+        **kwargs: Additional parameters passed to state_dot, A, and B.
+
+    Returns:
+        tuple: (A_bar, B_bar, C_bar, z_bar, Vmulti) where:
+            - A_bar: Discretized state transition matrix
+            - B_bar: Discretized control influence matrix
+            - C_bar: Discretized control influence matrix for next node
+            - z_bar: Defect vector
+            - Vmulti: Full augmented state trajectory
+    """
+    # Unpack settings
+    n_x = settings.sim.n_states
+    n_u = settings.sim.n_controls
+
+    N = settings.scp.n
+
+    # Define indices for slicing the augmented state vector
+    i0 = 0
+    i1 = n_x
+    i2 = i1 + n_x * n_x
+    i3 = i2 + n_x * n_u
+    i4 = i3 + n_x * n_u
+
+    # Initial augmented state
+    V0 = jnp.zeros((N - 1, i4))
+    V0 = V0.at[:, :n_x].set(x[:-1].astype(float))
+    V0 = V0.at[:, n_x : n_x + n_x * n_x].set(jnp.eye(n_x).reshape(1, -1).repeat(N - 1, axis=0))
+
+    # Choose integrator
+    integrator_args = dict(
+        u_cur=u[:-1].astype(float),
+        u_next=u[1:].astype(float),
+        state_dot=state_dot,
+        A=A,
+        B=B,
+        n_x=n_x,
+        n_u=n_u,
+        N=N,
+        dis_type=settings.dis.dis_type,
+        params=params,  # Pass params as single dict
+    )
+
+    # Define dVdt wrapper using named arguments
+    def dVdt_wrapped(t, y):
+        return dVdt(t, y, **integrator_args)
+
+    # Choose integrator
+    if settings.dis.custom_integrator:
+        sol = solve_ivp_rk45(
+            dVdt_wrapped,
+            1.0 / (N - 1),
+            V0.reshape(-1),
+            args=(),
+            is_not_compiled=settings.dev.debug,
+        )
+    else:
+        sol = solve_ivp_diffrax(
+            dVdt_wrapped,
+            1.0 / (N - 1),
+            V0.reshape(-1),
+            solver_name=settings.dis.solver,
+            rtol=settings.dis.rtol,
+            atol=settings.dis.atol,
+            args=(),
+            extra_kwargs=settings.dis.args,
+        )
+
+    Vend = sol[-1].T.reshape(-1, i4)
+    Vmulti = sol.T
+
+    x_prop = Vend[:, i0:i1]
+
+    # Return as 3D arrays: (N-1, n_x, n_x) for A_bar, (N-1, n_x, n_u) for B_bar/C_bar
+    A_bar = Vend[:, i1:i2].reshape(N - 1, n_x, n_x)
+    B_bar = Vend[:, i2:i3].reshape(N - 1, n_x, n_u)
+    C_bar = Vend[:, i3:i4].reshape(N - 1, n_x, n_u)
+
+    return A_bar, B_bar, C_bar, x_prop, Vmulti
+
+
+def get_discretization_solver(dyn: Dynamics, settings: Config):
+    """Create a discretization solver function.
+
+    This function creates a solver that computes the discretized system matrices
+    using the specified dynamics and settings.
+
+    Args:
+        dyn (Dynamics): System dynamics object.
+        settings: Configuration settings for discretization.
+
+    Returns:
+        callable: A function that computes the discretized system matrices.
+    """
+    return lambda x, u, params: calculate_discretization(
+        x=x,
+        u=u,
+        state_dot=dyn.f,
+        A=dyn.A,
+        B=dyn.B,
+        settings=settings,
+        params=params,  # Pass as single dict
+    )

openscvx/expert/__init__.py

@@ -0,0 +1,23 @@
+"""Expert-mode features for advanced users.
+
+This module contains features for expert users who need fine-grained control
+and are willing to bypass higher-level abstractions.
+"""
+
+from openscvx.expert.byof import (
+    ByofSpec,
+    CtcsConstraintSpec,
+    NodalConstraintSpec,
+    PenaltyFunction,
+)
+from openscvx.expert.lowering import apply_byof
+from openscvx.expert.validation import validate_byof
+
+__all__ = [
+    "ByofSpec",
+    "CtcsConstraintSpec",
+    "NodalConstraintSpec",
+    "PenaltyFunction",
+    "apply_byof",
+    "validate_byof",
+]

openscvx/expert/byof.py

@@ -0,0 +1,326 @@
+"""Bring-Your-Own-Functions (BYOF) - Expert User Mode.
+
+This module provides type definitions and documentation for expert users who want
+to bypass the symbolic layer and directly provide raw JAX functions.
+
+Important:
+    The unified state/control vectors include ALL states/controls in the order
+    they were provided, plus any augmented states from CTCS constraints. You are
+    responsible for correct indexing. Consider inspecting the symbolic problem
+    to understand the layout.
+
+Warning:
+    **Constraint Sign Convention**: All constraints follow g(x,u) <= 0 convention.
+    Return **negative when satisfied**, **positive when violated**.
+    Example: for x <= 10 return ``x - 10``, for x >= 5 return ``5 - x``.
+
+Function Signatures:
+    All byof functions must be JAX-compatible (use jax.numpy, avoid side effects).
+
+    - dynamics: ``(x, u, node, params) -> xdot_component``
+        - x: Full unified state vector (1D array)
+        - u: Full unified control vector (1D array)
+        - node: Integer node index
+        - params: Dict of parameters
+        - Returns: State derivative component (array matching state shape)
+
+    - nodal_constraints: ``(x, u, node, params) -> residual``
+        - Same arguments as dynamics
+        - Returns: Constraint residual (g <= 0: negative=satisfied, positive=violated)
+
+    - cross_nodal_constraints: ``(X, U, params) -> residual``
+        - X: State trajectory (N, n_x) where N is number of trajectory nodes,
+            n_x is unified state dimension
+        - U: Control trajectory (N, n_u) where N is number of trajectory nodes,
+            n_u is unified control dimension
+        - params: Dict of parameters
+        - Returns: Constraint residual (g <= 0: negative=satisfied, positive=violated)
+
+    - ctcs constraint_fn: ``(x, u, node, params) -> scalar``
+        - Same as nodal_constraints but MUST return scalar
+        - Returns: Constraint residual (g <= 0: negative=satisfied, positive=violated)
+
+    - ctcs penalty: ``(residual) -> penalty_value``
+        - residual: Scalar constraint residual
+        - Returns: Non-negative penalty value
+
+Example:
+    Basic usage mixing symbolic and byof::
+
+        import jax.numpy as jnp
+        import openscvx as ox
+        from openscvx import ByofSpec
+
+        # Define states
+        position = ox.State("position", shape=(2,))
+        velocity = ox.State("velocity", shape=(1,))
+        theta = ox.Control("theta", shape=(1,))
+
+        # Unified state: [position[0], position[1], velocity[0], time, augmented...]
+        # Unified control: [theta[0], time_dilation]
+
+        # Tip: Use the .slice property on State/Control objects for cleaner,
+        # more maintainable indexing instead of hardcoded indices.
+        byof: ByofSpec = {
+            "nodal_constraints": [
+                # Velocity bounds (applied to all nodes)
+                {
+                    "constraint_fn": lambda x, u, node, params: x[velocity.slice][0] - 10.0,
+                },
+                {
+                    "constraint_fn": lambda x, u, node, params: -x[velocity.slice][0],
+                },
+                # Velocity must be exactly 0 at start (selective enforcement)
+                {
+                    "constraint_fn": lambda x, u, node, params: x[velocity.slice][0],
+                    "nodes": [0],  # Only at first node
+                },
+            ],
+            "ctcs_constraints": [
+                {
+                    "constraint_fn": lambda x, u, node, params: x[position.slice][0] - 10.0,
+                    "penalty": "square",
+                    "bounds": (0.0, 1e-4),
+                }
+            ],
+        }
+
+        problem = ox.Problem(..., byof=byof)
+"""
+
+from typing import TYPE_CHECKING, Any, Callable, List, Literal, Tuple, TypedDict, Union
+
+if TYPE_CHECKING:
+    from jax import Array as JaxArray
+else:
+    JaxArray = Any
+
+__all__ = ["ByofSpec", "CtcsConstraintSpec", "NodalConstraintSpec", "PenaltyFunction"]
+
+
+# Type aliases for clarity
+DynamicsFunction = Callable[[JaxArray, JaxArray, int, dict], JaxArray]
+NodalConstraintFunction = Callable[[JaxArray, JaxArray, int, dict], JaxArray]
+CrossNodalConstraintFunction = Callable[[JaxArray, JaxArray, dict], JaxArray]
+CtcsConstraintFunction = Callable[[JaxArray, JaxArray, int, dict], float]
+PenaltyFunction = Union[Literal["square", "l1", "huber"], Callable[[float], float]]
+
+
+class NodalConstraintSpec(TypedDict, total=False):
+    """Specification for nodal constraint with optional node selection.
+
+    Nodal constraints are point-wise constraints evaluated at specific trajectory nodes.
+    By default, constraints apply to all nodes, but you can restrict enforcement to
+    specific nodes for boundary conditions, waypoints, or computational efficiency.
+
+    Attributes:
+        constraint_fn: Constraint function with signature ``(x, u, node, params) -> residual``.
+            Follows g(x,u) <= 0 convention (negative = satisfied). Required field.
+        nodes: List of integer node indices where constraint is enforced.
+            If omitted, applies to all nodes. Negative indices supported (e.g., -1 for last).
+            Optional field.
+
+    Example:
+        Boundary constraint only at first and last nodes::
+
+            nodal_spec: NodalConstraintSpec = {
+                "constraint_fn": lambda x, u, node, params: x[velocity.slice][0],
+                "nodes": [0, -1],  # Only at start and end
+            }
+
+        Waypoint constraint at middle of trajectory::
+
+            nodal_spec: NodalConstraintSpec = {
+                "constraint_fn": lambda x, u, node, params: jnp.linalg.norm(
+                    x[position.slice] - jnp.array([5.0, 7.5])
+                ) - 0.1,
+                "nodes": [N // 2],
+            }
+    """
+
+    constraint_fn: NodalConstraintFunction  # Required
+    nodes: List[int]
+
+
+class CtcsConstraintSpec(TypedDict, total=False):
+    """Specification for CTCS (Continuous-Time Constraint Satisfaction) constraint.
+
+    CTCS constraints are enforced by augmenting the dynamics with a penalty term that
+    accumulates violations over time. Useful for path constraints that must be satisfied
+    continuously, not just at discrete nodes.
+
+    Attributes:
+        constraint_fn: Function computing constraint residual with signature
+            ``(x, u, node, params) -> scalar``. Must return scalar.
+            Follows g(x,u) <= 0 convention (negative = satisfied). Required field.
+        penalty: Penalty function for positive residuals (violations).
+            Built-in options: "square" (max(r,0)^2, default), "l1" (max(r,0)),
+            "huber" (Huber loss). Custom: Callable ``(r) -> penalty`` (non-negative,
+            differentiable).
+        bounds: (min, max) bounds for augmented state accumulating penalties.
+            Default: (0.0, 1e-4). Max acts as soft constraint on total violation.
+        initial: Initial value for augmented state. Default: bounds[0] (usually 0.0).
+        over: Node interval (start, end) where constraint is active. The constraint
+            is enforced for nodes in [start, end). If omitted, constraint is active
+            over all nodes. Matches symbolic `.over()` method behavior.
+        idx: Constraint group index for sharing augmented states (default: 0).
+            All CTCS constraints (symbolic and byof) with the same idx share a single
+            augmented state. Their penalties are summed together. Use different idx values
+            to track different types of violations separately.
+
+    Warning:
+        If symbolic CTCS constraints exist with idx values [0, 1, 2], then byof idx **must** either:
+
+        - Match an existing idx (e.g., 0, 1, or 2) to add to that augmented state
+        - Be sequential after them (e.g., 3, 4, 5) to create new augmented states
+
+        You cannot use idx values that create gaps (e.g., if symbolic has [0, 1],
+        you cannot use byof idx=3 without also using idx=2).
+
+    Example:
+        Enforce position[0] <= 10.0 continuously::
+
+            # Assuming position = ox.State("position", shape=(2,))
+            ctcs_spec: CtcsConstraintSpec = {
+                "constraint_fn": lambda x, u, node, params: x[position.slice][0] - 10.0,
+                "penalty": "square",
+                "bounds": (0.0, 1e-4),
+                "initial": 0.0,
+                "idx": 0,  # Groups with other constraints having idx=0
+            }
+
+        Enforce constraint only over specific node range::
+
+            ctcs_spec: CtcsConstraintSpec = {
+                "constraint_fn": lambda x, u, node, params: x[position.slice][0] - 10.0,
+                "over": (10, 50),  # Active only for nodes 10-49
+                "penalty": "square",
+            }
+
+        Multiple constraints sharing an augmented state::
+
+            # If symbolic CTCS already has idx=[0, 1], then:
+
+            byof = {
+                "ctcs_constraints": [
+                    # Add to existing symbolic idx=0 augmented state
+                    {
+                        "constraint_fn": lambda x, u, node, params: x[pos.slice][0] - 10.0,
+                        "idx": 0,  # Shares with symbolic idx=0
+                    },
+                    # Add to existing symbolic idx=1 augmented state
+                    {
+                        "constraint_fn": lambda x, u, node, params: x[vel.slice][0] - 5.0,
+                        "idx": 1,  # Shares with symbolic idx=1
+                    },
+                    # Create NEW augmented state (sequential after symbolic)
+                    {
+                        "constraint_fn": lambda x, u, node, params: x[pos.slice][1] - 8.0,
+                        "idx": 2,  # New state (symbolic has 0,1, so next is 2)
+                    },
+                ]
+            }
+    """
+
+    constraint_fn: CtcsConstraintFunction  # Required
+    penalty: PenaltyFunction
+    bounds: Tuple[float, float]
+    initial: float
+    over: Tuple[int, int]
+    idx: int
+
+
+class ByofSpec(TypedDict, total=False):
+    """Bring-Your-Own-Functions specification for expert users.
+
+    Allows bypassing the symbolic layer and directly providing raw JAX functions.
+    All fields are optional - you can mix symbolic and byof as needed.
+
+    Warning:
+        You are responsible for:
+
+        - Correct indexing into unified state/control vectors
+        - Ensuring functions are JAX-compatible (use jax.numpy, no side effects)
+        - Ensuring functions are differentiable
+        - Following g(x,u) <= 0 convention for constraints
+
+    Tip:
+        Use the ``.slice`` property on State/Control objects for cleaner, more
+        maintainable indexing instead of hardcoded indices. For example, use
+        ``x[velocity.slice]`` instead of ``x[2:3]``. The slice property is set
+        after preprocessing and provides the correct indices into the unified
+        state/control vectors.
+
+    Attributes:
+        dynamics: Raw JAX functions for state derivatives. Maps state names to functions
+            with signature ``(x, u, node, params) -> xdot_component``. States here should
+            NOT appear in symbolic dynamics dict. You can mix: some states symbolic,
+            some in byof.
+        nodal_constraints: Point-wise constraints applied at specific nodes.
+            Each item is a :class:`NodalConstraintSpec` dict with:
+
+            - ``func``: Constraint function ``(x, u, node, params) -> residual`` (required)
+            - ``nodes``: List of node indices (optional, defaults to all nodes)
+
+            Follows g(x,u) <= 0 convention.
+        cross_nodal_constraints: Constraints coupling multiple nodes (smoothness, rate limits).
+            Signature: ``(X, U, params) -> residual`` where X is (N, n_x) and U is (N, n_u).
+            N is the number of trajectory nodes, n_x is state dimension, n_u is control dimension.
+            Follows g(X,U) <= 0 convention.
+        ctcs_constraints: Continuous-time constraint satisfaction via dynamics augmentation.
+            Each adds an augmented state accumulating violation penalties.
+            See :class:`CtcsConstraintSpec` for details.
+
+    Example:
+        Custom dynamics and constraints::
+
+            import jax.numpy as jnp
+            import openscvx as ox
+            from openscvx import ByofSpec
+
+            # Define states and controls
+            position = ox.State("position", shape=(2,))
+            velocity = ox.State("velocity", shape=(1,))
+            theta = ox.Control("theta", shape=(1,))
+
+            # Custom dynamics for one state using .slice property
+            def custom_velocity_dynamics(x, u, node, params):
+                # Use .slice property for clean indexing
+                return params["g"] * jnp.cos(u[theta.slice][0])
+
+            byof: ByofSpec = {
+                "dynamics": {
+                    "velocity": custom_velocity_dynamics,
+                },
+                "nodal_constraints": [
+                    # Applied to all nodes (no "nodes" field)
+                    {
+                        "constraint_fn": lambda x, u, node, params: x[velocity.slice][0] - 10.0,
+                    },
+                    {
+                        "constraint_fn": lambda x, u, node, params: -x[velocity.slice][0],
+                    },
+                    # Specify nodes for selective enforcement
+                    {
+                        "constraint_fn": lambda x, u, node, params: x[velocity.slice][0],
+                        "nodes": [0],  # Velocity must be exactly 0 at start
+                    },
+                ],
+                "cross_nodal_constraints": [
+                    # Constrain total velocity across trajectory: sum(velocities) >= 5
+                    # X.shape = (N, n_x), extract velocity column using slice
+                    lambda X, U, params: 5.0 - jnp.sum(X[:, velocity.slice]),
+                ],
+                "ctcs_constraints": [
+                    {
+                        "constraint_fn": lambda x, u, node, params: x[position.slice][0] - 5.0,
+                        "penalty": "square",
+                    }
+                ],
+            }
+    """
+
+    dynamics: dict[str, DynamicsFunction]
+    nodal_constraints: List[NodalConstraintSpec]
+    cross_nodal_constraints: List[CrossNodalConstraintFunction]
+    ctcs_constraints: List[CtcsConstraintSpec]