Skip to content

lowered

Lowered problem dataclasses.

This module contains dataclasses representing the outputs of the lowering phase, where symbolic expressions are converted to executable JAX and CVXPy code.

CVXPyVariables dataclass

CVXPy variables and parameters for the optimal control problem.

This dataclass holds all CVXPy Variable and Parameter objects needed to construct and solve the optimal control problem. It replaces the previous untyped dictionary approach with a typed, self-documenting structure.

The variables are organized into logical groups
  • SCP weights: Parameters controlling trust region and penalty weights
  • State: Variables and parameters for the state trajectory
  • Control: Variables and parameters for the control trajectory
  • Dynamics: Parameters for the discretized dynamics constraints
  • Nodal constraints: Parameters for linearized non-convex nodal constraints
  • Cross-node constraints: Parameters for linearized cross-node constraints
  • Scaling: Affine scaling matrices and offset vectors
  • Scaled expressions: CVXPy expressions for scaled state/control at each node

Attributes:

Name Type Description
w_tr Parameter

Trust region weight parameter (scalar, nonneg)
lam_cost Parameter

Cost function weight parameter (scalar, nonneg)
lam_vc Parameter

Virtual control penalty weights (N-1 x n_states, nonneg)
lam_vb Parameter

Virtual buffer penalty weight (scalar, nonneg)
x Variable

State variable (N x n_states)
dx Variable

State error variable (N x n_states)
x_bar Parameter

Previous SCP state parameter (N x n_states)
x_init Parameter

Initial state parameter (n_states,)
x_term Parameter

Terminal state parameter (n_states,)
u Variable

Control variable (N x n_controls)
du Variable

Control error variable (N x n_controls)
u_bar Parameter

Previous SCP control parameter (N x n_controls)
A_d Parameter

Discretized state Jacobian parameter (N-1 x n_states*n_states)
B_d Parameter

Discretized control Jacobian parameter (N-1 x n_states*n_controls)
C_d Parameter

Discretized control Jacobian (next node) parameter
x_prop Parameter

Propagated state parameter (N-1 x n_states)
nu Variable

Virtual control variable (N-1 x n_states)
g List[Parameter]

List of constraint value parameters (one per nodal constraint)
grad_g_x List[Parameter]

List of state gradient parameters (one per nodal constraint)
grad_g_u List[Parameter]

List of control gradient parameters (one per nodal constraint)
nu_vb List[Variable]

List of virtual buffer variables (one per nodal constraint)
g_cross List[Parameter]

List of cross-node constraint value parameters
grad_g_X_cross List[Parameter]

List of trajectory state gradient parameters
grad_g_U_cross List[Parameter]

List of trajectory control gradient parameters
nu_vb_cross List[Variable]

List of cross-node virtual buffer variables
S_x ndarray

State scaling matrix (n_states x n_states)
inv_S_x ndarray

Inverse state scaling matrix
c_x ndarray

State offset vector (n_states,)
S_u ndarray

Control scaling matrix (n_controls x n_controls)
inv_S_u ndarray

Inverse control scaling matrix
c_u ndarray

Control offset vector (n_controls,)
x_nonscaled List

List of scaled state expressions at each node
u_nonscaled List

List of scaled control expressions at each node
dx_nonscaled List

List of scaled state error expressions at each node
du_nonscaled List

List of scaled control error expressions at each node
Source code in openscvx/lowered/cvxpy_variables.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
@dataclass
class CVXPyVariables:
    """CVXPy variables and parameters for the optimal control problem.

    This dataclass holds all CVXPy Variable and Parameter objects needed to
    construct and solve the optimal control problem. It replaces the previous
    untyped dictionary approach with a typed, self-documenting structure.

    The variables are organized into logical groups:
        - SCP weights: Parameters controlling trust region and penalty weights
        - State: Variables and parameters for the state trajectory
        - Control: Variables and parameters for the control trajectory
        - Dynamics: Parameters for the discretized dynamics constraints
        - Nodal constraints: Parameters for linearized non-convex nodal constraints
        - Cross-node constraints: Parameters for linearized cross-node constraints
        - Scaling: Affine scaling matrices and offset vectors
        - Scaled expressions: CVXPy expressions for scaled state/control at each node

    Attributes:
        w_tr: Trust region weight parameter (scalar, nonneg)
        lam_cost: Cost function weight parameter (scalar, nonneg)
        lam_vc: Virtual control penalty weights (N-1 x n_states, nonneg)
        lam_vb: Virtual buffer penalty weight (scalar, nonneg)

        x: State variable (N x n_states)
        dx: State error variable (N x n_states)
        x_bar: Previous SCP state parameter (N x n_states)
        x_init: Initial state parameter (n_states,)
        x_term: Terminal state parameter (n_states,)

        u: Control variable (N x n_controls)
        du: Control error variable (N x n_controls)
        u_bar: Previous SCP control parameter (N x n_controls)

        A_d: Discretized state Jacobian parameter (N-1 x n_states*n_states)
        B_d: Discretized control Jacobian parameter (N-1 x n_states*n_controls)
        C_d: Discretized control Jacobian (next node) parameter
        x_prop: Propagated state parameter (N-1 x n_states)
        nu: Virtual control variable (N-1 x n_states)

        g: List of constraint value parameters (one per nodal constraint)
        grad_g_x: List of state gradient parameters (one per nodal constraint)
        grad_g_u: List of control gradient parameters (one per nodal constraint)
        nu_vb: List of virtual buffer variables (one per nodal constraint)

        g_cross: List of cross-node constraint value parameters
        grad_g_X_cross: List of trajectory state gradient parameters
        grad_g_U_cross: List of trajectory control gradient parameters
        nu_vb_cross: List of cross-node virtual buffer variables

        S_x: State scaling matrix (n_states x n_states)
        inv_S_x: Inverse state scaling matrix
        c_x: State offset vector (n_states,)
        S_u: Control scaling matrix (n_controls x n_controls)
        inv_S_u: Inverse control scaling matrix
        c_u: Control offset vector (n_controls,)

        x_nonscaled: List of scaled state expressions at each node
        u_nonscaled: List of scaled control expressions at each node
        dx_nonscaled: List of scaled state error expressions at each node
        du_nonscaled: List of scaled control error expressions at each node
    """

    # SCP weight parameters
    w_tr: "cp.Parameter"
    lam_cost: "cp.Parameter"
    lam_vc: "cp.Parameter"
    lam_vb: "cp.Parameter"

    # State variables and parameters
    x: "cp.Variable"
    dx: "cp.Variable"
    x_bar: "cp.Parameter"
    x_init: "cp.Parameter"
    x_term: "cp.Parameter"

    # Control variables and parameters
    u: "cp.Variable"
    du: "cp.Variable"
    u_bar: "cp.Parameter"

    # Dynamics discretization parameters
    A_d: "cp.Parameter"
    B_d: "cp.Parameter"
    C_d: "cp.Parameter"
    x_prop: "cp.Parameter"
    nu: "cp.Variable"

    # Nodal constraint linearization (lists, one per constraint)
    g: List["cp.Parameter"] = field(default_factory=list)
    grad_g_x: List["cp.Parameter"] = field(default_factory=list)
    grad_g_u: List["cp.Parameter"] = field(default_factory=list)
    nu_vb: List["cp.Variable"] = field(default_factory=list)

    # Cross-node constraint linearization (lists, one per constraint)
    g_cross: List["cp.Parameter"] = field(default_factory=list)
    grad_g_X_cross: List["cp.Parameter"] = field(default_factory=list)
    grad_g_U_cross: List["cp.Parameter"] = field(default_factory=list)
    nu_vb_cross: List["cp.Variable"] = field(default_factory=list)

    # Scaling matrices and offsets (numpy arrays)
    S_x: np.ndarray = field(default_factory=lambda: np.array([]))
    inv_S_x: np.ndarray = field(default_factory=lambda: np.array([]))
    c_x: np.ndarray = field(default_factory=lambda: np.array([]))
    S_u: np.ndarray = field(default_factory=lambda: np.array([]))
    inv_S_u: np.ndarray = field(default_factory=lambda: np.array([]))
    c_u: np.ndarray = field(default_factory=lambda: np.array([]))

    # Scaled CVXPy expressions at each node (lists of length N)
    x_nonscaled: List = field(default_factory=list)
    u_nonscaled: List = field(default_factory=list)
    dx_nonscaled: List = field(default_factory=list)
    du_nonscaled: List = field(default_factory=list)

Dynamics dataclass

Dataclass to hold a system dynamics function and its Jacobians.

This dataclass is used internally by openscvx to store the compiled dynamics function and its gradients after symbolic expressions are lowered to JAX. Users typically don't instantiate this class directly.

Attributes:

Name Type Description
f Callable[[ndarray, ndarray], ndarray]

Function defining the continuous time nonlinear system dynamics as x_dot = f(x, u, ...params). - 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')). If you use vectorized integration or batch evaluation, x and u may be 2D arrays (N, n_x) and (N, n_u).
A Optional[Callable[[ndarray, ndarray], ndarray]]

Jacobian of f w.r.t. x.
B Optional[Callable[[ndarray, ndarray], ndarray]]

Jacobian of f w.r.t. u.
Source code in openscvx/lowered/dynamics.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
@dataclass
class Dynamics:
    """Dataclass to hold a system dynamics function and its Jacobians.

    This dataclass is used internally by openscvx to store the compiled dynamics
    function and its gradients after symbolic expressions are lowered to JAX.
    Users typically don't instantiate this class directly.

    Attributes:
        f (Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]):
            Function defining the continuous time nonlinear system dynamics
            as x_dot = f(x, u, ...params).
            - 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')).
            If you use vectorized integration or batch evaluation, x and u
            may be 2D arrays (N, n_x) and (N, n_u).
        A (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]):
            Jacobian of ``f`` w.r.t. ``x``.
        B (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]):
            Jacobian of ``f`` w.r.t. ``u``.
    """

    f: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]
    A: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None
    B: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None

LoweredCrossNodeConstraint dataclass

Lowered cross-node constraint with trajectory-level evaluation.

Unlike regular LoweredNodalConstraint which operates on single-node vectors and is vmapped across the trajectory, LoweredCrossNodeConstraint operates on full trajectory arrays to relate multiple nodes simultaneously.

This is necessary for constraints like: - Rate limits: x[k] - x[k-1] <= max_rate - Multi-step dependencies: x[k] = 2*x[k-1] - x[k-2] - Periodic boundaries: x[0] = x[N-1]

The function signatures differ from LoweredNodalConstraint: - Regular: f(x, u, node, params) -> scalar (vmapped to handle (N, n_x)) - Cross-node: f(X, U, params) -> scalar (single constraint with fixed node indices)

Attributes:

Name Type Description
func Callable[[ndarray, ndarray, dict], ndarray]

Function (X, U, params) -> scalar residual where X: (N, n_x), U: (N, n_u) Returns constraint residual following g(X, U) <= 0 convention The constraint references fixed trajectory nodes (e.g., X[5] - X[4])
grad_g_X Callable[[ndarray, ndarray, dict], ndarray]

Function (X, U, params) -> (N, n_x) Jacobian wrt full state trajectory This is typically sparse - most constraints only couple nearby nodes
grad_g_U Callable[[ndarray, ndarray, dict], ndarray]

Function (X, U, params) -> (N, n_u) Jacobian wrt full control trajectory Often zero or very sparse for cross-node state constraints
Example

For rate constraint x[5] - x[4] <= r:

func(X, U, params) -> scalar residual
grad_g_X(X, U, params) -> (N, n_x) sparse Jacobian
    where grad_g_X[5, :] = ∂g/∂x[5] (derivative wrt node 5)
    and grad_g_X[4, :] = ∂g/∂x[4] (derivative wrt node 4)
    all other entries are zero
Performance Note - Dense Jacobian Storage

The Jacobian matrices grad_g_X and grad_g_U are stored as DENSE arrays with shape (N, n_x) and (N, n_u), but most cross-node constraints only couple a small number of nearby nodes, making these matrices extremely sparse.

For example, a rate limit constraint x[k] - x[k-1] <= r only has non-zero Jacobian entries at positions [k, :] and [k-1, :]. All other N-2 rows are zero but still stored in memory.

Memory impact for large problems: - A single constraint with N=100 nodes, n_x=10 states requires ~8KB for grad_g_X (compared to ~160 bytes if sparse with 2 non-zero rows) - Multiple cross-node constraints multiply this overhead - May cause issues for N > 1000 with many constraints

Performance impact: - Slower autodiff (computes many zero gradients) - Inefficient constraint linearization in the SCP solver - Potential GPU memory limitations for very large problems

The current implementation prioritizes simplicity and compatibility with JAX's autodiff over memory efficiency. Future versions may support sparse Jacobian formats (COO, CSR, or custom sparse representations) if this becomes a bottleneck in practice.

Source code in openscvx/lowered/jax_constraints.py 
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
@dataclass
class LoweredCrossNodeConstraint:
    """Lowered cross-node constraint with trajectory-level evaluation.

    Unlike regular LoweredNodalConstraint which operates on single-node vectors
    and is vmapped across the trajectory, LoweredCrossNodeConstraint operates
    on full trajectory arrays to relate multiple nodes simultaneously.

    This is necessary for constraints like:
    - Rate limits: x[k] - x[k-1] <= max_rate
    - Multi-step dependencies: x[k] = 2*x[k-1] - x[k-2]
    - Periodic boundaries: x[0] = x[N-1]

    The function signatures differ from LoweredNodalConstraint:
    - Regular: f(x, u, node, params) -> scalar (vmapped to handle (N, n_x))
    - Cross-node: f(X, U, params) -> scalar (single constraint with fixed node indices)

    Attributes:
        func: Function (X, U, params) -> scalar residual
            where X: (N, n_x), U: (N, n_u)
            Returns constraint residual following g(X, U) <= 0 convention
            The constraint references fixed trajectory nodes (e.g., X[5] - X[4])
        grad_g_X: Function (X, U, params) -> (N, n_x) Jacobian wrt full state trajectory
            This is typically sparse - most constraints only couple nearby nodes
        grad_g_U: Function (X, U, params) -> (N, n_u) Jacobian wrt full control trajectory
            Often zero or very sparse for cross-node state constraints

    Example:
        For rate constraint x[5] - x[4] <= r:

            func(X, U, params) -> scalar residual
            grad_g_X(X, U, params) -> (N, n_x) sparse Jacobian
                where grad_g_X[5, :] = ∂g/∂x[5] (derivative wrt node 5)
                and grad_g_X[4, :] = ∂g/∂x[4] (derivative wrt node 4)
                all other entries are zero

    Performance Note - Dense Jacobian Storage:
        The Jacobian matrices grad_g_X and grad_g_U are stored as DENSE arrays with
        shape (N, n_x) and (N, n_u), but most cross-node constraints only couple a
        small number of nearby nodes, making these matrices extremely sparse.

        For example, a rate limit constraint x[k] - x[k-1] <= r only has non-zero
        Jacobian entries at positions [k, :] and [k-1, :]. All other N-2 rows are
        zero but still stored in memory.

        Memory impact for large problems:
        - A single constraint with N=100 nodes, n_x=10 states requires ~8KB for
          grad_g_X (compared to ~160 bytes if sparse with 2 non-zero rows)
        - Multiple cross-node constraints multiply this overhead
        - May cause issues for N > 1000 with many constraints

        Performance impact:
        - Slower autodiff (computes many zero gradients)
        - Inefficient constraint linearization in the SCP solver
        - Potential GPU memory limitations for very large problems

        The current implementation prioritizes simplicity and compatibility with
        JAX's autodiff over memory efficiency. Future versions may support sparse
        Jacobian formats (COO, CSR, or custom sparse representations) if this
        becomes a bottleneck in practice.
    """

    func: Callable[[jnp.ndarray, jnp.ndarray, dict], jnp.ndarray]
    grad_g_X: Callable[[jnp.ndarray, jnp.ndarray, dict], jnp.ndarray]
    grad_g_U: Callable[[jnp.ndarray, jnp.ndarray, dict], jnp.ndarray]

LoweredCvxpyConstraints dataclass

CVXPy-lowered convex constraints.

Contains constraints that have been lowered to CVXPy constraint objects. These are added directly to the optimal control problem without linearization.

Attributes:

Name Type Description
constraints list[Constraint]

List of CVXPy constraint objects (cp.Constraint). Includes both nodal and cross-node convex constraints.
Source code in openscvx/lowered/cvxpy_constraints.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
@dataclass
class LoweredCvxpyConstraints:
    """CVXPy-lowered convex constraints.

    Contains constraints that have been lowered to CVXPy constraint objects.
    These are added directly to the optimal control problem without
    linearization.

    Attributes:
        constraints: List of CVXPy constraint objects (cp.Constraint).
            Includes both nodal and cross-node convex constraints.
    """

    constraints: list["cp.Constraint"] = field(default_factory=list)

LoweredJaxConstraints dataclass

JAX-lowered non-convex constraints with gradient functions.

Contains constraints that have been lowered to JAX callable functions with automatically computed gradients. These are used for linearization in the SCP (Sequential Convex Programming) loop.

Attributes:

Name Type Description
nodal list[LoweredNodalConstraint]

List of LoweredNodalConstraint objects. Each has func, grad_g_x, grad_g_u callables and nodes list.
cross_node list[LoweredCrossNodeConstraint]

List of LoweredCrossNodeConstraint objects. Each has func, grad_g_X, grad_g_U for trajectory-level constraints.
ctcs list[CTCS]

CTCS constraints (unchanged from input, not lowered here).
Source code in openscvx/lowered/jax_constraints.py 
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
@dataclass
class LoweredJaxConstraints:
    """JAX-lowered non-convex constraints with gradient functions.

    Contains constraints that have been lowered to JAX callable functions
    with automatically computed gradients. These are used for linearization
    in the SCP (Sequential Convex Programming) loop.

    Attributes:
        nodal: List of LoweredNodalConstraint objects. Each has `func`,
            `grad_g_x`, `grad_g_u` callables and `nodes` list.
        cross_node: List of LoweredCrossNodeConstraint objects. Each has
            `func`, `grad_g_X`, `grad_g_U` for trajectory-level constraints.
        ctcs: CTCS constraints (unchanged from input, not lowered here).
    """

    nodal: list[LoweredNodalConstraint] = field(default_factory=list)
    cross_node: list[LoweredCrossNodeConstraint] = field(default_factory=list)
    ctcs: list["CTCS"] = field(default_factory=list)

LoweredNodalConstraint dataclass

Dataclass to hold a lowered symbolic constraint function and its jacobians.

This is a simplified drop-in replacement for NodalConstraint that holds only the essential lowered JAX functions and their jacobians, without the complexity of convex/vectorized flags or post-initialization logic.

Designed for use with symbolic expressions that have been lowered to JAX and will be linearized for sequential convex programming.

Parameters:

Name Type Description Default
func Callable[[ndarray, ndarray], ndarray]

The lowered constraint function g(x, u, ...params) that returns constraint residuals. Should follow g(x, u) <= 0 convention. - 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

 required
grad_g_x Optional[Callable[[ndarray, ndarray], ndarray]]

Jacobian of g w.r.t. x. If None, should be computed using jax.jacfwd.

 None
grad_g_u Optional[Callable[[ndarray, ndarray], ndarray]]

Jacobian of g w.r.t. u. If None, should be computed using jax.jacfwd.

 None
nodes Optional[List[int]]

List of node indices where this constraint applies. Set after lowering from NodalConstraint.

 None
Source code in openscvx/lowered/jax_constraints.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
@dataclass
class LoweredNodalConstraint:
    """
    Dataclass to hold a lowered symbolic constraint function and its jacobians.

    This is a simplified drop-in replacement for NodalConstraint that holds
    only the essential lowered JAX functions and their jacobians, without
    the complexity of convex/vectorized flags or post-initialization logic.

    Designed for use with symbolic expressions that have been lowered to JAX
    and will be linearized for sequential convex programming.

    Args:
        func (Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]):
            The lowered constraint function g(x, u, ...params) that returns
            constraint residuals. Should follow g(x, u) <= 0 convention.
            - 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

        grad_g_x (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]):
            Jacobian of g w.r.t. x. If None, should be computed using jax.jacfwd.

        grad_g_u (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]):
            Jacobian of g w.r.t. u. If None, should be computed using jax.jacfwd.

        nodes (Optional[List[int]]): List of node indices where this constraint applies.
            Set after lowering from NodalConstraint.
    """

    func: 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
    nodes: Optional[List[int]] = None

LoweredProblem dataclass

Container for all outputs from symbolic problem lowering.

This dataclass holds all the results of lowering symbolic expressions to executable JAX and CVXPy code. It provides a clean, typed interface for accessing the various components needed for optimization.

Attributes:

Name Type Description
dynamics Dynamics

Optimization dynamics with fields f, A, B (JAX functions)
dynamics_prop Dynamics

Propagation dynamics with fields f, A, B
jax_constraints LoweredJaxConstraints

Non-convex constraints lowered to JAX with gradients
cvxpy_constraints LoweredCvxpyConstraints

Convex constraints lowered to CVXPy
x_unified UnifiedState

Aggregated optimization state interface
u_unified UnifiedControl

Aggregated optimization control interface
x_prop_unified UnifiedState

Aggregated propagation state interface
ocp_vars CVXPyVariables

Typed CVXPy variables and parameters for OCP construction
cvxpy_params Dict[str, Parameter]

Dict mapping user parameter names to CVXPy Parameter objects
Example

After lowering a symbolic problem::

lowered = lower_symbolic_problem(
    dynamics_aug=dynamics,
    states_aug=states,
    controls_aug=controls,
    constraints=constraint_set,
    parameters=params,
    N=50,
)

# Access components
dx_dt = lowered.dynamics.f(x, u, node, params)
jacobian_A = lowered.dynamics.A(x, u, node, params)

# Use CVXPy objects
ocp = OptimalControlProblem(settings, lowered)
Source code in openscvx/lowered/problem.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
@dataclass
class LoweredProblem:
    """Container for all outputs from symbolic problem lowering.

    This dataclass holds all the results of lowering symbolic expressions
    to executable JAX and CVXPy code. It provides a clean, typed interface
    for accessing the various components needed for optimization.

    Attributes:
        dynamics: Optimization dynamics with fields f, A, B (JAX functions)
        dynamics_prop: Propagation dynamics with fields f, A, B
        jax_constraints: Non-convex constraints lowered to JAX with gradients
        cvxpy_constraints: Convex constraints lowered to CVXPy
        x_unified: Aggregated optimization state interface
        u_unified: Aggregated optimization control interface
        x_prop_unified: Aggregated propagation state interface
        ocp_vars: Typed CVXPy variables and parameters for OCP construction
        cvxpy_params: Dict mapping user parameter names to CVXPy Parameter objects

    Example:
        After lowering a symbolic problem::

            lowered = lower_symbolic_problem(
                dynamics_aug=dynamics,
                states_aug=states,
                controls_aug=controls,
                constraints=constraint_set,
                parameters=params,
                N=50,
            )

            # Access components
            dx_dt = lowered.dynamics.f(x, u, node, params)
            jacobian_A = lowered.dynamics.A(x, u, node, params)

            # Use CVXPy objects
            ocp = OptimalControlProblem(settings, lowered)
    """

    # JAX dynamics
    dynamics: Dynamics
    dynamics_prop: Dynamics

    # Lowered constraints (separate types for JAX vs CVXPy)
    jax_constraints: LoweredJaxConstraints
    cvxpy_constraints: LoweredCvxpyConstraints

    # Unified interfaces
    x_unified: UnifiedState
    u_unified: UnifiedControl
    x_prop_unified: UnifiedState

    # CVXPy objects
    ocp_vars: CVXPyVariables
    cvxpy_params: Dict[str, "cp.Parameter"]

ParameterDict

Bases: dict

Dictionary that syncs to both internal _parameters dict and CVXPy parameters.

This allows users to naturally update parameters like

problem.parameters["obs_radius"] = 2.0

Changes automatically propagate to: 1. Internal _parameters dict (plain dict for JAX) 2. CVXPy parameters (for optimization)

Source code in openscvx/lowered/parameters.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
class ParameterDict(dict):
    """Dictionary that syncs to both internal _parameters dict and CVXPy parameters.

    This allows users to naturally update parameters like:
        problem.parameters["obs_radius"] = 2.0

    Changes automatically propagate to:
    1. Internal _parameters dict (plain dict for JAX)
    2. CVXPy parameters (for optimization)
    """

    def __init__(self, problem, internal_dict, *args, **kwargs):
        self._problem = problem
        self._internal_dict = internal_dict  # Reference to plain dict for JAX
        super().__init__()
        # Initialize with float enforcement by using __setitem__
        if args:
            other = args[0]
            if hasattr(other, "items"):
                for key, value in other.items():
                    self[key] = value
            else:
                for key, value in other:
                    self[key] = value
        for key, value in kwargs.items():
            self[key] = value

    def __setitem__(self, key, value):
        # Enforce float dtype to prevent int/float mismatch bugs
        value = np.asarray(value, dtype=float)
        super().__setitem__(key, value)
        # Sync to internal dict for JAX
        self._internal_dict[key] = value
        # Sync to CVXPy if it exists
        lowered = getattr(self._problem, "_lowered", None)
        if lowered is not None and key in lowered.cvxpy_params:
            lowered.cvxpy_params[key].value = value

    def update(self, other=None, **kwargs):
        """Update multiple parameters and sync to internal dict and CVXPy."""
        if other is not None:
            if hasattr(other, "items"):
                for key, value in other.items():
                    self[key] = value
            else:
                for key, value in other:
                    self[key] = value
        for key, value in kwargs.items():
            self[key] = value
update(other=None, **kwargs)

Update multiple parameters and sync to internal dict and CVXPy.

Source code in openscvx/lowered/parameters.py 
44
45
46
47
48
49
50
51
52
53
54
def update(self, other=None, **kwargs):
    """Update multiple parameters and sync to internal dict and CVXPy."""
    if other is not None:
        if hasattr(other, "items"):
            for key, value in other.items():
                self[key] = value
        else:
            for key, value in other:
                self[key] = value
    for key, value in kwargs.items():
        self[key] = value

UnifiedControl dataclass

Unified control vector aggregating multiple Control objects.

UnifiedControl is a drop-in replacement for individual Control objects that holds aggregated data from multiple Control instances. It maintains compatibility with optimization infrastructure while providing access to individual control components through slicing.

The unified control separates user-defined "true" controls from augmented controls added internally (e.g., for time dilation). This separation allows clean access to physical control inputs while supporting advanced features.

Attributes:

Name Type Description
name str

Name identifier for the unified control vector
shape tuple

Combined shape (total_dim,) of all aggregated controls
min ndarray

Lower bounds for all control variables, shape (total_dim,)
max ndarray

Upper bounds for all control variables, shape (total_dim,)
guess ndarray

Initial guess trajectory, shape (num_nodes, total_dim)
_true_dim int

Number of user-defined control dimensions (excludes augmented)
_true_slice slice

Slice for extracting true controls from unified vector
_augmented_slice slice

Slice for extracting augmented controls
time_dilation_slice Optional[slice]

Slice for time dilation control, if present
Properties

true: Returns UnifiedControl view containing only true (user-defined) controls augmented: Returns UnifiedControl view containing only augmented controls

Example

Creating a unified control from multiple Control objects::

from openscvx.symbolic.unified import unify_controls

thrust = ox.Control("thrust", shape=(3,), min=0, max=10)
torque = ox.Control("torque", shape=(3,), min=-1, max=1)

unified = unify_controls([thrust, torque], name="u")
print(unified.shape)        # (6,)
print(unified.min)          # [0, 0, 0, -1, -1, -1]
print(unified.true.shape)   # (6,) - all are true controls
print(unified.augmented.shape)  # (0,) - no augmented controls

Appending controls dynamically::

unified = UnifiedControl(name="u", shape=(0,), _true_dim=0)
unified.append(min=-1, max=1, guess=0.0)  # Add scalar control
print(unified.shape)  # (1,)
See Also
  • unify_controls(): Factory function for creating UnifiedControl from Control list
  • Control: Individual symbolic control variable
  • UnifiedState: Analogous unified state vector
Source code in openscvx/lowered/unified.py 
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
@dataclass
class UnifiedControl:
    """Unified control vector aggregating multiple Control objects.

    UnifiedControl is a drop-in replacement for individual Control objects that holds
    aggregated data from multiple Control instances. It maintains compatibility with
    optimization infrastructure while providing access to individual control components
    through slicing.

    The unified control separates user-defined "true" controls from augmented controls
    added internally (e.g., for time dilation). This separation allows clean access to
    physical control inputs while supporting advanced features.

    Attributes:
        name (str): Name identifier for the unified control vector
        shape (tuple): Combined shape (total_dim,) of all aggregated controls
        min (np.ndarray): Lower bounds for all control variables, shape (total_dim,)
        max (np.ndarray): Upper bounds for all control variables, shape (total_dim,)
        guess (np.ndarray): Initial guess trajectory, shape (num_nodes, total_dim)
        _true_dim (int): Number of user-defined control dimensions (excludes augmented)
        _true_slice (slice): Slice for extracting true controls from unified vector
        _augmented_slice (slice): Slice for extracting augmented controls
        time_dilation_slice (Optional[slice]): Slice for time dilation control, if present

    Properties:
        true: Returns UnifiedControl view containing only true (user-defined) controls
        augmented: Returns UnifiedControl view containing only augmented controls

    Example:
        Creating a unified control from multiple Control objects::

            from openscvx.symbolic.unified import unify_controls

            thrust = ox.Control("thrust", shape=(3,), min=0, max=10)
            torque = ox.Control("torque", shape=(3,), min=-1, max=1)

            unified = unify_controls([thrust, torque], name="u")
            print(unified.shape)        # (6,)
            print(unified.min)          # [0, 0, 0, -1, -1, -1]
            print(unified.true.shape)   # (6,) - all are true controls
            print(unified.augmented.shape)  # (0,) - no augmented controls

        Appending controls dynamically::

            unified = UnifiedControl(name="u", shape=(0,), _true_dim=0)
            unified.append(min=-1, max=1, guess=0.0)  # Add scalar control
            print(unified.shape)  # (1,)

    See Also:
        - unify_controls(): Factory function for creating UnifiedControl from Control list
        - Control: Individual symbolic control variable
        - UnifiedState: Analogous unified state vector
    """

    name: str
    shape: tuple
    min: Optional[np.ndarray] = None
    max: Optional[np.ndarray] = None
    guess: Optional[np.ndarray] = None
    _true_dim: int = 0
    _true_slice: Optional[slice] = None
    _augmented_slice: Optional[slice] = None
    time_dilation_slice: Optional[slice] = None  # Slice for time dilation control
    scaling_min: Optional[np.ndarray] = None  # Scaling minimum bounds for unified control
    scaling_max: Optional[np.ndarray] = None  # Scaling maximum bounds for unified control

    def __post_init__(self):
        """Initialize slices after dataclass creation."""
        if self._true_slice is None:
            self._true_slice = slice(0, self._true_dim)
        if self._augmented_slice is None:
            self._augmented_slice = slice(self._true_dim, self.shape[0])

    @property
    def true(self) -> "UnifiedControl":
        """Get the true (user-defined) control variables.

        Returns a view of the unified control containing only user-defined controls,
        excluding internal augmented controls added for time dilation, etc.

        Returns:
            UnifiedControl: Sliced view containing only true control variables

        Example:
            Get true user defined controls::

                unified = unify_controls([thrust, torque, time_dilation], name="u")
                true_controls = unified.true  # Only thrust and torque
        """
        return self[self._true_slice]

    @property
    def augmented(self) -> "UnifiedControl":
        """Get the augmented (internal) control variables.

        Returns a view of the unified control containing only augmented controls
        added internally by the optimization framework (e.g., time dilation control).

        Returns:
            UnifiedControl: Sliced view containing only augmented control variables

        Example:
            Get augmented controls::

                unified = unify_controls([thrust, time_dilation], name="u")
                aug_controls = unified.augmented  # Only time dilation
        """
        return self[self._augmented_slice]

    def append(
        self,
        other: "Optional[Control | UnifiedControl]" = None,
        *,
        min=-np.inf,
        max=np.inf,
        guess=0.0,
        augmented=False,
    ):
        """Append another control or create a new control variable.

        This method allows dynamic extension of the unified control, either by appending
        another Control/UnifiedControl object or by creating a new scalar control variable
        with specified properties. Modifies the unified control in-place.

        Args:
            other (Optional[Control | UnifiedControl]): Control object to append. If None,
                creates a new scalar control variable with properties from keyword args.
            min (float): Lower bound for new scalar control (default: -inf)
            max (float): Upper bound for new scalar control (default: inf)
            guess (float): Initial guess value for new scalar control (default: 0.0)
            augmented (bool): Whether the appended control is augmented (internal) rather
                than true (user-defined). Affects _true_dim tracking. Default: False

        Returns:
            None: Modifies the unified control in-place

        Example:
            Appending a Control object::

                unified = unify_controls([thrust], name="u")
                torque = ox.Control("torque", shape=(3,), min=-1, max=1)
                unified.append(torque)
                print(unified.shape)  # (6,) - thrust (3) + torque (3)

            Creating new scalar control variables::

                unified = UnifiedControl(name="u", shape=(0,), _true_dim=0)
                unified.append(min=-1, max=1, guess=0.0)  # Add scalar control
                print(unified.shape)  # (1,)
        """
        # Import here to avoid circular imports at module level
        from openscvx.symbolic.expr.control import Control

        if isinstance(other, (Control, UnifiedControl)):
            # Append another control object
            new_shape = (self.shape[0] + other.shape[0],)

            # Update bounds
            if self.min is not None and other.min is not None:
                new_min = np.concatenate([self.min, other.min])
            else:
                new_min = self.min

            if self.max is not None and other.max is not None:
                new_max = np.concatenate([self.max, other.max])
            else:
                new_max = self.max

            # Update guess
            if self.guess is not None and other.guess is not None:
                new_guess = np.concatenate([self.guess, other.guess], axis=1)
            else:
                new_guess = self.guess

            # Update true dimension
            if not augmented:
                new_true_dim = self._true_dim + getattr(other, "_true_dim", other.shape[0])
            else:
                new_true_dim = self._true_dim

            # Update all attributes in place
            self.shape = new_shape
            self.min = new_min
            self.max = new_max
            self.guess = new_guess
            self._true_dim = new_true_dim
            self._true_slice = slice(0, self._true_dim)
            self._augmented_slice = slice(self._true_dim, self.shape[0])

        else:
            # Create a single new variable
            new_shape = (self.shape[0] + 1,)

            # Extend arrays
            if self.min is not None:
                self.min = np.concatenate([self.min, np.array([min])])
            if self.max is not None:
                self.max = np.concatenate([self.max, np.array([max])])
            if self.guess is not None:
                guess_arr = np.full((self.guess.shape[0], 1), guess)
                self.guess = np.concatenate([self.guess, guess_arr], axis=1)

            # Update dimensions
            self.shape = new_shape
            if not augmented:
                self._true_dim += 1
            self._true_slice = slice(0, self._true_dim)
            self._augmented_slice = slice(self._true_dim, self.shape[0])

    def __getitem__(self, idx):
        """Get a subset of the unified control variables.

        Enables slicing of the unified control to extract subsets of control variables.
        Returns a new UnifiedControl containing only the sliced dimensions.

        Args:
            idx (slice): Slice object specifying which control dimensions to extract.
                Only simple slices with step=1 are supported.

        Returns:
            UnifiedControl: New unified control containing only the sliced dimensions

        Raises:
            NotImplementedError: If idx is not a slice, or if step != 1

        Example:
            Generate unified control object::

                unified = unify_controls([thrust, torque], name="u")

            thrust has shape (3,), torque has shape (3,)::

                first_three = unified[0:3]  # Extract thrust only
                print(first_three.shape)  # (3,)

        Note:
            The sliced control maintains all properties (bounds, guesses, etc.) for
            the selected dimensions. The _true_dim is recalculated based on which
            dimensions fall within the original true control range.
        """
        if isinstance(idx, slice):
            start, stop, step = idx.indices(self.shape[0])
            if step != 1:
                raise NotImplementedError("Step slicing not supported")

            new_shape = (stop - start,)
            new_name = f"{self.name}[{start}:{stop}]"

            # Slice all arrays
            new_min = self.min[idx] if self.min is not None else None
            new_max = self.max[idx] if self.max is not None else None
            new_guess = self.guess[:, idx] if self.guess is not None else None

            # Calculate new true dimension
            new_true_dim = max(0, min(stop, self._true_dim) - max(start, 0))

            return UnifiedControl(
                name=new_name,
                shape=new_shape,
                min=new_min,
                max=new_max,
                guess=new_guess,
                _true_dim=new_true_dim,
                _true_slice=slice(0, new_true_dim),
                _augmented_slice=slice(new_true_dim, new_shape[0]),
            )
        else:
            raise NotImplementedError("Only slice indexing is supported")

    def __repr__(self):
        """String representation of the UnifiedControl object."""
        return f"UnifiedControl('{self.name}', shape={self.shape})"
augmented: UnifiedControl property

Get the augmented (internal) control variables.

Returns a view of the unified control containing only augmented controls added internally by the optimization framework (e.g., time dilation control).

Returns:

Name Type Description
UnifiedControl UnifiedControl

Sliced view containing only augmented control variables
Example

Get augmented controls::

unified = unify_controls([thrust, time_dilation], name="u")
aug_controls = unified.augmented  # Only time dilation
true: UnifiedControl property

Get the true (user-defined) control variables.

Returns a view of the unified control containing only user-defined controls, excluding internal augmented controls added for time dilation, etc.

Returns:

Name Type Description
UnifiedControl UnifiedControl

Sliced view containing only true control variables
Example

Get true user defined controls::

unified = unify_controls([thrust, torque, time_dilation], name="u")
true_controls = unified.true  # Only thrust and torque
append(other: Optional[Control | UnifiedControl] = None, *, min=-np.inf, max=np.inf, guess=0.0, augmented=False)

Append another control or create a new control variable.

This method allows dynamic extension of the unified control, either by appending another Control/UnifiedControl object or by creating a new scalar control variable with specified properties. Modifies the unified control in-place.

Parameters:

Name Type Description Default
other Optional[Control | UnifiedControl]

Control object to append. If None, creates a new scalar control variable with properties from keyword args.

 None
min float

Lower bound for new scalar control (default: -inf)

 -inf
max float

Upper bound for new scalar control (default: inf)

 inf
guess float

Initial guess value for new scalar control (default: 0.0)

 0.0
augmented bool

Whether the appended control is augmented (internal) rather than true (user-defined). Affects _true_dim tracking. Default: False

 False

Returns:

Name Type Description
None

Modifies the unified control in-place
Example

Appending a Control object::

unified = unify_controls([thrust], name="u")
torque = ox.Control("torque", shape=(3,), min=-1, max=1)
unified.append(torque)
print(unified.shape)  # (6,) - thrust (3) + torque (3)

Creating new scalar control variables::

unified = UnifiedControl(name="u", shape=(0,), _true_dim=0)
unified.append(min=-1, max=1, guess=0.0)  # Add scalar control
print(unified.shape)  # (1,)
Source code in openscvx/lowered/unified.py 
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
def append(
    self,
    other: "Optional[Control | UnifiedControl]" = None,
    *,
    min=-np.inf,
    max=np.inf,
    guess=0.0,
    augmented=False,
):
    """Append another control or create a new control variable.

    This method allows dynamic extension of the unified control, either by appending
    another Control/UnifiedControl object or by creating a new scalar control variable
    with specified properties. Modifies the unified control in-place.

    Args:
        other (Optional[Control | UnifiedControl]): Control object to append. If None,
            creates a new scalar control variable with properties from keyword args.
        min (float): Lower bound for new scalar control (default: -inf)
        max (float): Upper bound for new scalar control (default: inf)
        guess (float): Initial guess value for new scalar control (default: 0.0)
        augmented (bool): Whether the appended control is augmented (internal) rather
            than true (user-defined). Affects _true_dim tracking. Default: False

    Returns:
        None: Modifies the unified control in-place

    Example:
        Appending a Control object::

            unified = unify_controls([thrust], name="u")
            torque = ox.Control("torque", shape=(3,), min=-1, max=1)
            unified.append(torque)
            print(unified.shape)  # (6,) - thrust (3) + torque (3)

        Creating new scalar control variables::

            unified = UnifiedControl(name="u", shape=(0,), _true_dim=0)
            unified.append(min=-1, max=1, guess=0.0)  # Add scalar control
            print(unified.shape)  # (1,)
    """
    # Import here to avoid circular imports at module level
    from openscvx.symbolic.expr.control import Control

    if isinstance(other, (Control, UnifiedControl)):
        # Append another control object
        new_shape = (self.shape[0] + other.shape[0],)

        # Update bounds
        if self.min is not None and other.min is not None:
            new_min = np.concatenate([self.min, other.min])
        else:
            new_min = self.min

        if self.max is not None and other.max is not None:
            new_max = np.concatenate([self.max, other.max])
        else:
            new_max = self.max

        # Update guess
        if self.guess is not None and other.guess is not None:
            new_guess = np.concatenate([self.guess, other.guess], axis=1)
        else:
            new_guess = self.guess

        # Update true dimension
        if not augmented:
            new_true_dim = self._true_dim + getattr(other, "_true_dim", other.shape[0])
        else:
            new_true_dim = self._true_dim

        # Update all attributes in place
        self.shape = new_shape
        self.min = new_min
        self.max = new_max
        self.guess = new_guess
        self._true_dim = new_true_dim
        self._true_slice = slice(0, self._true_dim)
        self._augmented_slice = slice(self._true_dim, self.shape[0])

    else:
        # Create a single new variable
        new_shape = (self.shape[0] + 1,)

        # Extend arrays
        if self.min is not None:
            self.min = np.concatenate([self.min, np.array([min])])
        if self.max is not None:
            self.max = np.concatenate([self.max, np.array([max])])
        if self.guess is not None:
            guess_arr = np.full((self.guess.shape[0], 1), guess)
            self.guess = np.concatenate([self.guess, guess_arr], axis=1)

        # Update dimensions
        self.shape = new_shape
        if not augmented:
            self._true_dim += 1
        self._true_slice = slice(0, self._true_dim)
        self._augmented_slice = slice(self._true_dim, self.shape[0])

UnifiedState dataclass

Unified state vector aggregating multiple State objects.

UnifiedState is a drop-in replacement for individual State objects that holds aggregated data from multiple State instances. It maintains compatibility with optimization infrastructure while providing access to individual state components through slicing.

The unified state separates user-defined "true" states from augmented states added internally (e.g., for CTCS constraints or time variables). This separation allows clean access to physical states while supporting advanced features.

Attributes:

Name Type Description
name str

Name identifier for the unified state vector
shape tuple

Combined shape (total_dim,) of all aggregated states
min ndarray

Lower bounds for all state variables, shape (total_dim,)
max ndarray

Upper bounds for all state variables, shape (total_dim,)
guess ndarray

Initial guess trajectory, shape (num_nodes, total_dim)
initial ndarray

Initial boundary conditions, shape (total_dim,)
final ndarray

Final boundary conditions, shape (total_dim,)
_initial ndarray

Internal initial values, shape (total_dim,)
_final ndarray

Internal final values, shape (total_dim,)
initial_type ndarray

Boundary condition types at t0 ("Fix" or "Free"), shape (total_dim,), dtype=object
final_type ndarray

Boundary condition types at tf ("Fix" or "Free"), shape (total_dim,), dtype=object
_true_dim int

Number of user-defined state dimensions (excludes augmented)
_true_slice slice

Slice for extracting true states from unified vector
_augmented_slice slice

Slice for extracting augmented states
time_slice Optional[slice]

Slice for time state variable, if present
ctcs_slice Optional[slice]

Slice for CTCS augmented states, if present
Properties

true: Returns UnifiedState view containing only true (user-defined) states augmented: Returns UnifiedState view containing only augmented states

Example

Creating a unified state from multiple State objects::

from openscvx.symbolic.unified import unify_states

position = ox.State("pos", shape=(3,), min=-10, max=10)
velocity = ox.State("vel", shape=(3,), min=-5, max=5)

unified = unify_states([position, velocity], name="x")
print(unified.shape)        # (6,)
print(unified.min)          # [-10, -10, -10, -5, -5, -5]
print(unified.true.shape)   # (6,) - all are true states
print(unified.augmented.shape)  # (0,) - no augmented states

Appending states dynamically::

unified = UnifiedState(name="x", shape=(0,), _true_dim=0)
unified.append(min=-1, max=1, guess=0.5)  # Add scalar state
print(unified.shape)  # (1,)
See Also
  • unify_states(): Factory function for creating UnifiedState from State list
  • State: Individual symbolic state variable
  • UnifiedControl: Analogous unified control vector
Source code in openscvx/lowered/unified.py 
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
@dataclass
class UnifiedState:
    """Unified state vector aggregating multiple State objects.

    UnifiedState is a drop-in replacement for individual State objects that holds
    aggregated data from multiple State instances. It maintains compatibility with
    optimization infrastructure while providing access to individual state components
    through slicing.

    The unified state separates user-defined "true" states from augmented states
    added internally (e.g., for CTCS constraints or time variables). This separation
    allows clean access to physical states while supporting advanced features.

    Attributes:
        name (str): Name identifier for the unified state vector
        shape (tuple): Combined shape (total_dim,) of all aggregated states
        min (np.ndarray): Lower bounds for all state variables, shape (total_dim,)
        max (np.ndarray): Upper bounds for all state variables, shape (total_dim,)
        guess (np.ndarray): Initial guess trajectory, shape (num_nodes, total_dim)
        initial (np.ndarray): Initial boundary conditions, shape (total_dim,)
        final (np.ndarray): Final boundary conditions, shape (total_dim,)
        _initial (np.ndarray): Internal initial values, shape (total_dim,)
        _final (np.ndarray): Internal final values, shape (total_dim,)
        initial_type (np.ndarray): Boundary condition types at t0 ("Fix" or "Free"),
            shape (total_dim,), dtype=object
        final_type (np.ndarray): Boundary condition types at tf ("Fix" or "Free"),
            shape (total_dim,), dtype=object
        _true_dim (int): Number of user-defined state dimensions (excludes augmented)
        _true_slice (slice): Slice for extracting true states from unified vector
        _augmented_slice (slice): Slice for extracting augmented states
        time_slice (Optional[slice]): Slice for time state variable, if present
        ctcs_slice (Optional[slice]): Slice for CTCS augmented states, if present

    Properties:
        true: Returns UnifiedState view containing only true (user-defined) states
        augmented: Returns UnifiedState view containing only augmented states

    Example:
        Creating a unified state from multiple State objects::

            from openscvx.symbolic.unified import unify_states

            position = ox.State("pos", shape=(3,), min=-10, max=10)
            velocity = ox.State("vel", shape=(3,), min=-5, max=5)

            unified = unify_states([position, velocity], name="x")
            print(unified.shape)        # (6,)
            print(unified.min)          # [-10, -10, -10, -5, -5, -5]
            print(unified.true.shape)   # (6,) - all are true states
            print(unified.augmented.shape)  # (0,) - no augmented states

        Appending states dynamically::

            unified = UnifiedState(name="x", shape=(0,), _true_dim=0)
            unified.append(min=-1, max=1, guess=0.5)  # Add scalar state
            print(unified.shape)  # (1,)

    See Also:
        - unify_states(): Factory function for creating UnifiedState from State list
        - State: Individual symbolic state variable
        - UnifiedControl: Analogous unified control vector
    """

    name: str
    shape: tuple
    min: Optional[np.ndarray] = None
    max: Optional[np.ndarray] = None
    guess: Optional[np.ndarray] = None
    initial: Optional[np.ndarray] = None
    final: Optional[np.ndarray] = None
    _initial: Optional[np.ndarray] = None
    _final: Optional[np.ndarray] = None
    initial_type: Optional[np.ndarray] = None
    final_type: Optional[np.ndarray] = None
    _true_dim: int = 0
    _true_slice: Optional[slice] = None
    _augmented_slice: Optional[slice] = None
    time_slice: Optional[slice] = None  # Slice for time state
    ctcs_slice: Optional[slice] = None  # Slice for CTCS augmented states
    scaling_min: Optional[np.ndarray] = None  # Scaling minimum bounds for unified state
    scaling_max: Optional[np.ndarray] = None  # Scaling maximum bounds for unified state

    def __post_init__(self):
        """Initialize slices after dataclass creation."""
        if self._true_slice is None:
            self._true_slice = slice(0, self._true_dim)
        if self._augmented_slice is None:
            self._augmented_slice = slice(self._true_dim, self.shape[0])

    @property
    def true(self) -> "UnifiedState":
        """Get the true (user-defined) state variables.

        Returns a view of the unified state containing only user-defined states,
        excluding internal augmented states added for CTCS, time, etc.

        Returns:
            UnifiedState: Sliced view containing only true state variables

        Example:
            Get true user-defined state::

                unified = unify_states([position, velocity, ctcs_aug], name="x")
                true_states = unified.true  # Only position and velocity
                true_states.shape  # (6,) if position and velocity are 3D each
        """
        return self[self._true_slice]

    @property
    def augmented(self) -> "UnifiedState":
        """Get the augmented (internal) state variables.

        Returns a view of the unified state containing only augmented states
        added internally by the optimization framework (e.g., CTCS penalty states,
        time variables).

        Returns:
            UnifiedState: Sliced view containing only augmented state variables

        Example:
            Get augmented state::

                unified = unify_states([position, ctcs_aug], name="x")
                aug_states = unified.augmented  # Only CTCS states
        """
        return self[self._augmented_slice]

    def append(
        self,
        other: "Optional[State | UnifiedState]" = None,
        *,
        min=-np.inf,
        max=np.inf,
        guess=0.0,
        initial=0.0,
        final=0.0,
        augmented=False,
    ):
        """Append another state or create a new state variable.

        This method allows dynamic extension of the unified state, either by appending
        another State/UnifiedState object or by creating a new scalar state variable
        with specified properties. Modifies the unified state in-place.

        Args:
            other (Optional[State | UnifiedState]): State object to append. If None,
                creates a new scalar state variable with properties from keyword args.
            min (float): Lower bound for new scalar state (default: -inf)
            max (float): Upper bound for new scalar state (default: inf)
            guess (float): Initial guess value for new scalar state (default: 0.0)
            initial (float): Initial boundary condition for new scalar state (default: 0.0)
            final (float): Final boundary condition for new scalar state (default: 0.0)
            augmented (bool): Whether the appended state is augmented (internal) rather
                than true (user-defined). Affects _true_dim tracking. Default: False

        Returns:
            None: Modifies the unified state in-place

        Example:
            Appending a State object::

                unified = unify_states([position], name="x")
                velocity = ox.State("vel", shape=(3,), min=-5, max=5)
                unified.append(velocity)
                print(unified.shape)  # (6,) - position (3) + velocity (3)

            Creating new scalar state variables::

                unified = UnifiedState(name="x", shape=(0,), _true_dim=0)
                unified.append(min=-1, max=1, guess=0.5)  # Add scalar state
                unified.append(min=-2, max=2, augmented=True)  # Add augmented state
                print(unified.shape)  # (2,)
                print(unified._true_dim)  # 1 (only first is true)

        Note:
            Maintains the invariant that true states appear before augmented states
            in the unified vector. When appending augmented states, they are added
            to the end but don't increment _true_dim.
        """
        # Import here to avoid circular imports at module level
        from openscvx.symbolic.expr.state import State

        if isinstance(other, (State, UnifiedState)):
            # Append another state object
            new_shape = (self.shape[0] + other.shape[0],)

            # Update bounds
            if self.min is not None and other.min is not None:
                new_min = np.concatenate([self.min, other.min])
            else:
                new_min = self.min

            if self.max is not None and other.max is not None:
                new_max = np.concatenate([self.max, other.max])
            else:
                new_max = self.max

            # Update guess
            if self.guess is not None and other.guess is not None:
                new_guess = np.concatenate([self.guess, other.guess], axis=1)
            else:
                new_guess = self.guess

            # Update initial/final conditions
            if self.initial is not None and other.initial is not None:
                new_initial = np.concatenate([self.initial, other.initial])
            else:
                new_initial = self.initial

            if self.final is not None and other.final is not None:
                new_final = np.concatenate([self.final, other.final])
            else:
                new_final = self.final

            # Update internal arrays
            if self._initial is not None and other._initial is not None:
                new__initial = np.concatenate([self._initial, other._initial])
            else:
                new__initial = self._initial

            if self._final is not None and other._final is not None:
                new__final = np.concatenate([self._final, other._final])
            else:
                new__final = self._final

            # Update types
            if self.initial_type is not None and other.initial_type is not None:
                new_initial_type = np.concatenate([self.initial_type, other.initial_type])
            else:
                new_initial_type = self.initial_type

            if self.final_type is not None and other.final_type is not None:
                new_final_type = np.concatenate([self.final_type, other.final_type])
            else:
                new_final_type = self.final_type

            # Update true dimension
            if not augmented:
                new_true_dim = self._true_dim + getattr(other, "_true_dim", other.shape[0])
            else:
                new_true_dim = self._true_dim

            # Update all attributes in place
            self.shape = new_shape
            self.min = new_min
            self.max = new_max
            self.guess = new_guess
            self.initial = new_initial
            self.final = new_final
            self._initial = new__initial
            self._final = new__final
            self.initial_type = new_initial_type
            self.final_type = new_final_type
            self._true_dim = new_true_dim
            self._true_slice = slice(0, self._true_dim)
            self._augmented_slice = slice(self._true_dim, self.shape[0])

        else:
            # Create a single new variable
            new_shape = (self.shape[0] + 1,)

            # Extend arrays
            if self.min is not None:
                self.min = np.concatenate([self.min, np.array([min])])
            if self.max is not None:
                self.max = np.concatenate([self.max, np.array([max])])
            if self.guess is not None:
                guess_arr = np.full((self.guess.shape[0], 1), guess)
                self.guess = np.concatenate([self.guess, guess_arr], axis=1)
            if self.initial is not None:
                self.initial = np.concatenate([self.initial, np.array([initial])])
            if self.final is not None:
                self.final = np.concatenate([self.final, np.array([final])])
            if self._initial is not None:
                self._initial = np.concatenate([self._initial, np.array([initial])])
            if self._final is not None:
                self._final = np.concatenate([self._final, np.array([final])])
            if self.initial_type is not None:
                self.initial_type = np.concatenate(
                    [self.initial_type, np.array(["Fix"], dtype=object)]
                )
            if self.final_type is not None:
                self.final_type = np.concatenate([self.final_type, np.array(["Fix"], dtype=object)])

            # Update dimensions
            self.shape = new_shape
            if not augmented:
                self._true_dim += 1
            self._true_slice = slice(0, self._true_dim)
            self._augmented_slice = slice(self._true_dim, self.shape[0])

    def __getitem__(self, idx):
        """Get a subset of the unified state variables.

        Enables slicing of the unified state to extract subsets of state variables.
        Returns a new UnifiedState containing only the sliced dimensions.

        Args:
            idx (slice): Slice object specifying which state dimensions to extract.
                Only simple slices with step=1 are supported.

        Returns:
            UnifiedState: New unified state containing only the sliced dimensions

        Raises:
            NotImplementedError: If idx is not a slice, or if step != 1

        Example:
            Generate unified state object::

                unified = unify_states([position, velocity], name="x")

            position has shape (3,), velocity has shape (3,)::

                first_three = unified[0:3]  # Extract position only
                print(first_three.shape)  # (3,)
                last_three = unified[3:6]  # Extract velocity only
                print(last_three.shape)  # (3,)

        Note:
            The sliced state maintains all properties (bounds, guesses, etc.) for
            the selected dimensions. The _true_dim is recalculated based on which
            dimensions fall within the original true state range.
        """
        if isinstance(idx, slice):
            start, stop, step = idx.indices(self.shape[0])
            if step != 1:
                raise NotImplementedError("Step slicing not supported")

            new_shape = (stop - start,)
            new_name = f"{self.name}[{start}:{stop}]"

            # Slice all arrays
            new_min = self.min[idx] if self.min is not None else None
            new_max = self.max[idx] if self.max is not None else None
            new_guess = self.guess[:, idx] if self.guess is not None else None
            new_initial = self.initial[idx] if self.initial is not None else None
            new_final = self.final[idx] if self.final is not None else None
            new__initial = self._initial[idx] if self._initial is not None else None
            new__final = self._final[idx] if self._final is not None else None
            new_initial_type = self.initial_type[idx] if self.initial_type is not None else None
            new_final_type = self.final_type[idx] if self.final_type is not None else None

            # Calculate new true dimension
            new_true_dim = max(0, min(stop, self._true_dim) - max(start, 0))

            return UnifiedState(
                name=new_name,
                shape=new_shape,
                min=new_min,
                max=new_max,
                guess=new_guess,
                initial=new_initial,
                final=new_final,
                _initial=new__initial,
                _final=new__final,
                initial_type=new_initial_type,
                final_type=new_final_type,
                _true_dim=new_true_dim,
                _true_slice=slice(0, new_true_dim),
                _augmented_slice=slice(new_true_dim, new_shape[0]),
            )
        else:
            raise NotImplementedError("Only slice indexing is supported")

    def __repr__(self):
        """String representation of the UnifiedState object."""
        return f"UnifiedState('{self.name}', shape={self.shape})"
augmented: UnifiedState property

Get the augmented (internal) state variables.

Returns a view of the unified state containing only augmented states added internally by the optimization framework (e.g., CTCS penalty states, time variables).

Returns:

Name Type Description
UnifiedState UnifiedState

Sliced view containing only augmented state variables
Example

Get augmented state::

unified = unify_states([position, ctcs_aug], name="x")
aug_states = unified.augmented  # Only CTCS states
true: UnifiedState property

Get the true (user-defined) state variables.

Returns a view of the unified state containing only user-defined states, excluding internal augmented states added for CTCS, time, etc.

Returns:

Name Type Description
UnifiedState UnifiedState

Sliced view containing only true state variables
Example

Get true user-defined state::

unified = unify_states([position, velocity, ctcs_aug], name="x")
true_states = unified.true  # Only position and velocity
true_states.shape  # (6,) if position and velocity are 3D each
append(other: Optional[State | UnifiedState] = None, *, min=-np.inf, max=np.inf, guess=0.0, initial=0.0, final=0.0, augmented=False)

Append another state or create a new state variable.

This method allows dynamic extension of the unified state, either by appending another State/UnifiedState object or by creating a new scalar state variable with specified properties. Modifies the unified state in-place.

Parameters:

Name Type Description Default
other Optional[State | UnifiedState]

State object to append. If None, creates a new scalar state variable with properties from keyword args.

 None
min float

Lower bound for new scalar state (default: -inf)

 -inf
max float

Upper bound for new scalar state (default: inf)

 inf
guess float

Initial guess value for new scalar state (default: 0.0)

 0.0
initial float

Initial boundary condition for new scalar state (default: 0.0)

 0.0
final float

Final boundary condition for new scalar state (default: 0.0)

 0.0
augmented bool

Whether the appended state is augmented (internal) rather than true (user-defined). Affects _true_dim tracking. Default: False

 False

Returns:

Name Type Description
None

Modifies the unified state in-place
Example

Appending a State object::

unified = unify_states([position], name="x")
velocity = ox.State("vel", shape=(3,), min=-5, max=5)
unified.append(velocity)
print(unified.shape)  # (6,) - position (3) + velocity (3)

Creating new scalar state variables::

unified = UnifiedState(name="x", shape=(0,), _true_dim=0)
unified.append(min=-1, max=1, guess=0.5)  # Add scalar state
unified.append(min=-2, max=2, augmented=True)  # Add augmented state
print(unified.shape)  # (2,)
print(unified._true_dim)  # 1 (only first is true)
Note

Maintains the invariant that true states appear before augmented states in the unified vector. When appending augmented states, they are added to the end but don't increment _true_dim.

Source code in openscvx/lowered/unified.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
def append(
    self,
    other: "Optional[State | UnifiedState]" = None,
    *,
    min=-np.inf,
    max=np.inf,
    guess=0.0,
    initial=0.0,
    final=0.0,
    augmented=False,
):
    """Append another state or create a new state variable.

    This method allows dynamic extension of the unified state, either by appending
    another State/UnifiedState object or by creating a new scalar state variable
    with specified properties. Modifies the unified state in-place.

    Args:
        other (Optional[State | UnifiedState]): State object to append. If None,
            creates a new scalar state variable with properties from keyword args.
        min (float): Lower bound for new scalar state (default: -inf)
        max (float): Upper bound for new scalar state (default: inf)
        guess (float): Initial guess value for new scalar state (default: 0.0)
        initial (float): Initial boundary condition for new scalar state (default: 0.0)
        final (float): Final boundary condition for new scalar state (default: 0.0)
        augmented (bool): Whether the appended state is augmented (internal) rather
            than true (user-defined). Affects _true_dim tracking. Default: False

    Returns:
        None: Modifies the unified state in-place

    Example:
        Appending a State object::

            unified = unify_states([position], name="x")
            velocity = ox.State("vel", shape=(3,), min=-5, max=5)
            unified.append(velocity)
            print(unified.shape)  # (6,) - position (3) + velocity (3)

        Creating new scalar state variables::

            unified = UnifiedState(name="x", shape=(0,), _true_dim=0)
            unified.append(min=-1, max=1, guess=0.5)  # Add scalar state
            unified.append(min=-2, max=2, augmented=True)  # Add augmented state
            print(unified.shape)  # (2,)
            print(unified._true_dim)  # 1 (only first is true)

    Note:
        Maintains the invariant that true states appear before augmented states
        in the unified vector. When appending augmented states, they are added
        to the end but don't increment _true_dim.
    """
    # Import here to avoid circular imports at module level
    from openscvx.symbolic.expr.state import State

    if isinstance(other, (State, UnifiedState)):
        # Append another state object
        new_shape = (self.shape[0] + other.shape[0],)

        # Update bounds
        if self.min is not None and other.min is not None:
            new_min = np.concatenate([self.min, other.min])
        else:
            new_min = self.min

        if self.max is not None and other.max is not None:
            new_max = np.concatenate([self.max, other.max])
        else:
            new_max = self.max

        # Update guess
        if self.guess is not None and other.guess is not None:
            new_guess = np.concatenate([self.guess, other.guess], axis=1)
        else:
            new_guess = self.guess

        # Update initial/final conditions
        if self.initial is not None and other.initial is not None:
            new_initial = np.concatenate([self.initial, other.initial])
        else:
            new_initial = self.initial

        if self.final is not None and other.final is not None:
            new_final = np.concatenate([self.final, other.final])
        else:
            new_final = self.final

        # Update internal arrays
        if self._initial is not None and other._initial is not None:
            new__initial = np.concatenate([self._initial, other._initial])
        else:
            new__initial = self._initial

        if self._final is not None and other._final is not None:
            new__final = np.concatenate([self._final, other._final])
        else:
            new__final = self._final

        # Update types
        if self.initial_type is not None and other.initial_type is not None:
            new_initial_type = np.concatenate([self.initial_type, other.initial_type])
        else:
            new_initial_type = self.initial_type

        if self.final_type is not None and other.final_type is not None:
            new_final_type = np.concatenate([self.final_type, other.final_type])
        else:
            new_final_type = self.final_type

        # Update true dimension
        if not augmented:
            new_true_dim = self._true_dim + getattr(other, "_true_dim", other.shape[0])
        else:
            new_true_dim = self._true_dim

        # Update all attributes in place
        self.shape = new_shape
        self.min = new_min
        self.max = new_max
        self.guess = new_guess
        self.initial = new_initial
        self.final = new_final
        self._initial = new__initial
        self._final = new__final
        self.initial_type = new_initial_type
        self.final_type = new_final_type
        self._true_dim = new_true_dim
        self._true_slice = slice(0, self._true_dim)
        self._augmented_slice = slice(self._true_dim, self.shape[0])

    else:
        # Create a single new variable
        new_shape = (self.shape[0] + 1,)

        # Extend arrays
        if self.min is not None:
            self.min = np.concatenate([self.min, np.array([min])])
        if self.max is not None:
            self.max = np.concatenate([self.max, np.array([max])])
        if self.guess is not None:
            guess_arr = np.full((self.guess.shape[0], 1), guess)
            self.guess = np.concatenate([self.guess, guess_arr], axis=1)
        if self.initial is not None:
            self.initial = np.concatenate([self.initial, np.array([initial])])
        if self.final is not None:
            self.final = np.concatenate([self.final, np.array([final])])
        if self._initial is not None:
            self._initial = np.concatenate([self._initial, np.array([initial])])
        if self._final is not None:
            self._final = np.concatenate([self._final, np.array([final])])
        if self.initial_type is not None:
            self.initial_type = np.concatenate(
                [self.initial_type, np.array(["Fix"], dtype=object)]
            )
        if self.final_type is not None:
            self.final_type = np.concatenate([self.final_type, np.array(["Fix"], dtype=object)])

        # Update dimensions
        self.shape = new_shape
        if not augmented:
            self._true_dim += 1
        self._true_slice = slice(0, self._true_dim)
        self._augmented_slice = slice(self._true_dim, self.shape[0])