openscvx 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

openscvx/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.1'
-__version_tuple__ = version_tuple = (0, 1, 1)
+__version__ = version = '0.1.3'
+__version_tuple__ = version_tuple = (0, 1, 3)

openscvx/augmentation/dynamics_augmentation.py

@@ -0,0 +1,122 @@
+from typing import Callable, List, Tuple
+
+import jax
+import jax.numpy as jnp
+
+from openscvx.constraints.violation import CTCSViolation
+from openscvx.dynamics import Dynamics
+
+def build_augmented_dynamics(
+    dynamics_non_augmented: Dynamics,
+    violations: List[CTCSViolation],
+    idx_x_true: slice,
+    idx_u_true: slice,
+) -> Dynamics:
+    dynamics_augmented = Dynamics(
+        f=get_augmented_dynamics(
+            dynamics_non_augmented.f, violations, idx_x_true, idx_u_true
+        ),
+    )
+    A, B = get_jacobians(
+        dynamics_augmented.f, dynamics_non_augmented, violations, idx_x_true, idx_u_true
+    )
+    dynamics_augmented.A = A
+    dynamics_augmented.B = B
+    return dynamics_augmented
+
+
+def get_augmented_dynamics(
+    dynamics: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray],
+    violations: List[CTCSViolation],
+    idx_x_true: slice,
+    idx_u_true: slice,
+) -> Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]:
+    def dynamics_augmented(x: jnp.array, u: jnp.array, node: int) -> jnp.array:
+        x_dot = dynamics(x[idx_x_true], u[idx_u_true])
+
+        # Iterate through the g_func dictionary and stack the output each function
+        # to x_dot
+        for v in violations:
+            x_dot = jnp.hstack([x_dot, v.g(x[idx_x_true], u[idx_u_true], node)])
+
+        return x_dot
+
+    return dynamics_augmented
+
+
+def get_jacobians(
+    dyn_augmented: Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray],
+    dynamics_non_augmented: Dynamics,
+    violations: List[CTCSViolation],
+    idx_x_true: slice,
+    idx_u_true: slice,
+) -> Tuple[
+    Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray],
+    Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray],
+]:
+    # 1) Early return if absolutely no custom grads anywhere
+    no_dyn_grads = dynamics_non_augmented.A is None and dynamics_non_augmented.B is None
+    no_vio_grads = all(v.g_grad_x is None and v.g_grad_u is None for v in violations)
+
+    if no_dyn_grads and no_vio_grads:
+        return (
+            jax.jacfwd(dyn_augmented, argnums=0),
+            jax.jacfwd(dyn_augmented, argnums=1),
+        )
+
+    # 2) Build the *true‐state* Jacobians of f(x_true,u_true)
+    f_fn = dynamics_non_augmented.f
+    if dynamics_non_augmented.A is None:
+        A_f = lambda x_true, u_true: jax.jacfwd(f_fn, argnums=0)(x_true, u_true)
+    else:
+        A_f = dynamics_non_augmented.A
+
+    if dynamics_non_augmented.B is None:
+        B_f = lambda x_true, u_true: jax.jacfwd(f_fn, argnums=1)(x_true, u_true)
+    else:
+        B_f = dynamics_non_augmented.B
+
+    # 3) Build per-violation gradients
+    def make_violation_grad_x(i: int) -> Callable:
+        viol = violations[i]
+        # use user‐provided if present, otherwise autodiff viol.g in argnum=0
+        return viol.g_grad_x or jax.jacfwd(viol.g, argnums=0)
+
+    def make_violation_grad_u(i: int) -> Callable:
+        viol = violations[i]
+        # use user‐provided if present, otherwise autodiff viol.g in argnum=0
+        return viol.g_grad_u or jax.jacfwd(viol.g, argnums=1)
+
+    # 4) Assemble A_aug, B_aug
+    def A(x_aug, u_aug, node):
+        # dynamics block + zero‐pad
+        Af = A_f(x_aug[idx_x_true], u_aug[idx_u_true])  # (n_f, n_x_true)
+        nv = sum(
+            v.g(x_aug[idx_x_true], u_aug[idx_u_true], node).shape[0] for v in violations
+        )  # total # rows of violations
+        zero_pad = jnp.zeros((Af.shape[0], nv))  # (n_f, n_v)
+        top = jnp.hstack([Af, zero_pad])  # (n_f, n_x_true + n_v)
+
+        # violation blocks
+        rows = [top]
+        for i in range(len(violations)):
+            dx_i = make_violation_grad_x(i)(
+                x_aug[idx_x_true], u_aug[idx_u_true], node
+            )  # (n_gi, n_x_true)
+            pad_i = jnp.zeros((dx_i.shape[0], nv))  # (n_gi, n_v)
+            rows.append(jnp.hstack([dx_i, pad_i]))  # (n_gi, n_x_true + n_v)
+
+        return jnp.vstack(rows)
+
+    def B(x_aug, u_aug, node):
+        Bf = B_f(x_aug[idx_x_true], u_aug[idx_u_true])  # (n_f, n_u_true)
+        rows = [Bf]
+        for i in range(len(violations)):
+            du_i = make_violation_grad_u(i)(
+                x_aug[idx_x_true], u_aug[idx_u_true], node
+            )  # (n_gi, n_u_true)
+            rows.append(du_i)
+
+        return jnp.vstack(rows)
+
+    return A, B

openscvx/config.py

@@ -2,6 +2,8 @@ import numpy as np
 from dataclasses import dataclass, field
 from typing import Dict, List
 
+from openscvx.constraints.boundary import BoundaryConstraint
+
 
 def get_affine_scaling_matrices(n, minimum, maximum):
     S = np.diag(np.maximum(np.ones(n), abs(minimum - maximum) / 2))
@@ -120,7 +122,8 @@ class PropagationConfig:
 class SimConfig:
     x_bar: np.ndarray
     u_bar: np.ndarray
-    initial_state: np.ndarray
+    initial_state: BoundaryConstraint
+    initial_state_prop: BoundaryConstraint
     final_state: np.ndarray
     max_state: np.ndarray
     min_state: np.ndarray
@@ -128,9 +131,11 @@ class SimConfig:
     min_control: np.ndarray
     total_time: float
     idx_x_true: slice
+    idx_x_true_prop: slice
     idx_u_true: slice
     idx_t: slice
     idx_y: slice
+    idx_y_prop: slice
     idx_s: slice
     ctcs_node_intervals: list = None
     constraints_ctcs: List[callable] = field(
@@ -138,6 +143,7 @@ class SimConfig:
     )  # TODO (norrisg): clean this up, consider moving to dedicated `constraints` dataclass
     constraints_nodal: List[callable] = field(default_factory=list)
     n_states: int = None
+    n_states_prop: int = None
     n_controls: int = None
     S_x: np.ndarray = None
     inv_S_x: np.ndarray = None

openscvx/constraints/ctcs.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 from typing import Callable, Sequence, Tuple, Optional
 import functools
+import types
 
 from jax.lax import cond
 import jax.numpy as jnp
@@ -10,8 +11,18 @@ import jax.numpy as jnp
 class CTCSConstraint:
     func: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]
     penalty: Callable[[jnp.ndarray], jnp.ndarray]
-    nodes: Optional[Sequence[Tuple[int, int]]] = None
+    nodes: Optional[Tuple[int, int]] = None
     idx: Optional[int] = None
+    grad_f_x: Optional[Callable] = None
+    grad_f_u: Optional[Callable] = None
+
+    def __post_init__(self):
+        if self.grad_f_x is not None:
+            _grad_f_x = self.grad_f_x
+            self.grad_f_x = lambda x, u, nodes: _grad_f_x(x, u)
+        if self.grad_f_u is not None:
+            _grad_f_u = self.grad_f_u
+            self.grad_f_u = lambda x, u, nodes: _grad_f_u(x, u)
 
     def __call__(self, x, u, node):
         return cond(
@@ -28,6 +39,8 @@ def ctcs(
     penalty: str = "squared_relu",
     nodes: Optional[Sequence[Tuple[int, int]]] = None,
     idx: Optional[int] = None,
+    grad_f_x: Optional[Callable] = None,
+    grad_f_u: Optional[Callable] = None,
 ):
     """Decorator to mark a function as a 'ctcs' constraint.
 
@@ -44,14 +57,25 @@ def ctcs(
         def pen(x):
             r = jnp.maximum(0, x)
             return jnp.where(r < delta, 0.5 * r**2, r - 0.5 * delta)
-
+    elif penalty == "smooth_relu":
+        c = 1e-8
+        pen = lambda x: (jnp.maximum(0, x) ** 2 + c**2) ** 0.5 - c
+    elif isinstance(penalty, types.LambdaType):
+        pen = penalty
     else:
         raise ValueError(f"Unknown penalty {penalty}")
 
     def decorator(f: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]):
         # wrap so name, doc, signature stay on f
         wrapped = functools.wraps(f)(f)
-        return CTCSConstraint(func=wrapped, penalty=pen, nodes=nodes, idx=idx)
+        return CTCSConstraint(
+            func=wrapped,
+            penalty=pen,
+            nodes=nodes,
+            idx=idx,
+            grad_f_x=grad_f_x,
+            grad_f_u=grad_f_u,
+        )
 
     # if called as @ctcs or @ctcs(...), _func will be None and we return decorator
     if _func is None:

openscvx/constraints/nodal.py

@@ -11,20 +11,22 @@ class NodalConstraint:
     nodes: Optional[List[int]] = None
     convex: bool = False
     vectorized: bool = False
+    grad_g_x: Optional[Callable] = None
+    grad_g_u: Optional[Callable] = None
 
     def __post_init__(self):
         if not self.convex:
-            # TODO: (haynec) switch to AOT instead of JIT
-            if self.vectorized:
-                # single-node but still using JAX
-                self.g = jit(self.func)
-                self.grad_g_x = jit(jacfwd(self.func, argnums=0))
-                self.grad_g_u = jit(jacfwd(self.func, argnums=1))
-            else:
-                self.g = vmap(jit(self.func), in_axes=(0, 0))
-                self.grad_g_x = jit(vmap(jacfwd(self.func, argnums=0), in_axes=(0, 0)))
-                self.grad_g_u = jit(vmap(jacfwd(self.func, argnums=1), in_axes=(0, 0)))
-        # if convex=True and inter_nodal=False, assume an external solver (e.g. CVX) will handle it
+            # single-node but still using JAX
+            self.g = self.func
+            if self.grad_g_x is None:
+                self.grad_g_x = jacfwd(self.func, argnums=0)
+            if self.grad_g_u is None:
+                self.grad_g_u = jacfwd(self.func, argnums=1)
+            if not self.vectorized:
+                self.g = vmap(self.g, in_axes=(0, 0))
+                self.grad_g_x = vmap(self.grad_g_x, in_axes=(0, 0))
+                self.grad_g_u = vmap(self.grad_g_u, in_axes=(0, 0))
+        # if convex=True assume an external solver (e.g. CVX) will handle it
 
     def __call__(self, x: jnp.ndarray, u: jnp.ndarray):
         return self.func(x, u)
@@ -36,6 +38,8 @@ def nodal(
     nodes: Optional[List[int]] = None,
     convex: bool = False,
     vectorized: bool = False,
+    grad_g_x: Optional[Callable] = None,
+    grad_g_u: Optional[Callable] = None,
 ):
     def decorator(f: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]):
         return NodalConstraint(
@@ -43,6 +47,8 @@ def nodal(
             nodes=nodes,
             convex=convex,
             vectorized=vectorized,
+            grad_g_x=grad_g_x,
+            grad_g_u=grad_g_u,
         )
 
     return decorator if _func is None else decorator(_func)

openscvx/constraints/violation.py

@@ -1,26 +1,67 @@
 from collections import defaultdict
+from typing import List, Optional, Tuple, Callable
+from dataclasses import dataclass
 
 import jax.numpy as jnp
 
-def get_g_func(constraints_ctcs: list[callable, callable]):
+from openscvx.constraints.ctcs import CTCSConstraint
+
+
+@dataclass
+class CTCSViolation:
+    g: Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]
+    g_grad_x: Optional[Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]] = None
+    g_grad_u: Optional[Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]] = None
+
+
+def get_g_grad_x(constraints_ctcs: List[CTCSConstraint]) -> Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]:
+    def g_grad_x(x: jnp.ndarray, u: jnp.ndarray, node: int) -> jnp.ndarray:
+        grads = [
+            c.grad_f_x(x, u, node) for c in constraints_ctcs if c.grad_f_x is not None
+        ]
+        return sum(grads) if grads else None
+
+    return g_grad_x
+
+
+def get_g_grad_u(constraints_ctcs: List[CTCSConstraint]) -> Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]:
+    def g_grad_u(x: jnp.ndarray, u: jnp.ndarray, node: int) -> jnp.ndarray:
+        grads = [
+            c.grad_f_u(x, u, node) for c in constraints_ctcs if c.grad_f_u is not None
+        ]
+        return sum(grads) if grads else None
+
+    return g_grad_u
+
+
+def get_g_func(constraints_ctcs: List[CTCSConstraint]) -> Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]:
     def g_func(x: jnp.array, u: jnp.array, node: int) -> jnp.array:
-        g_sum = 0
-        for g in constraints_ctcs:
-            g_sum += g(x,u, node)
-        return g_sum
+        return sum(c(x, u, node) for c in constraints_ctcs)
+
     return g_func
 
 
-def get_g_funcs(constraints_ctcs: list[callable]) -> list[callable]:
+def get_g_funcs(constraints_ctcs: List[CTCSConstraint]) -> List[CTCSViolation]:
     # Bucket by idx
-    groups: dict[int, list[callable]] = defaultdict(list)
+    groups: dict[int, List[CTCSConstraint]] = defaultdict(list)
     for c in constraints_ctcs:
         if c.idx is None:
             raise ValueError(f"CTCS constraint {c} has no .idx assigned")
         groups[c.idx].append(c)
 
-    # Build and return a list of get_g_func(funcs) in idx order
-    return [
-        get_g_func(funcs)
-        for idx, funcs in sorted(groups.items(), key=lambda kv: kv[0])
-    ]
+    # For each bucket, build one CTCSViolation
+    violations: List[CTCSViolation] = []
+    for idx, bucket in sorted(groups.items(), key=lambda kv: kv[0]):
+        g = get_g_func(bucket)
+        g_grad_x = get_g_grad_u(bucket) if all(c.grad_f_x for c in bucket) else None
+        g_grad_u = get_g_grad_x(bucket) if all(c.grad_f_u for c in bucket) else None
+
+        violations.append(
+            CTCSViolation(
+                g=g,
+                g_grad_x=g_grad_x,
+                g_grad_u=g_grad_u,
+            )
+        )
+
+    return violations

openscvx/discretization.py

@@ -1,6 +1,7 @@
 import jax.numpy as jnp
 import numpy as np
 
+from openscvx.dynamics import Dynamics
 from openscvx.integrators import solve_ivp_rk45, solve_ivp_diffrax
 
 
@@ -150,13 +151,13 @@ def calculate_discretization(
     return A_bar, B_bar, C_bar, z_bar, Vmulti
 
 
-def get_discretization_solver(state_dot, A, B, params):
+def get_discretization_solver(dyn: Dynamics, params):
     return lambda x, u: calculate_discretization(
         x=x,
         u=u,
-        state_dot=state_dot,
-        A=A,
-        B=B,
+        state_dot=dyn.f,
+        A=dyn.A,
+        B=dyn.B,
         n_x=params.sim.n_states,
         n_u=params.sim.n_controls,
         N=params.scp.n,

openscvx/dynamics.py

@@ -1,24 +1,41 @@
-import jax
+from dataclasses import dataclass
+from typing import Callable, Optional
+import functools
+
 import jax.numpy as jnp
 
 
-def get_augmented_dynamics(
-    dynamics: callable, g_funcs: list[callable], idx_x_true: slice, idx_u_true: slice
-) -> callable:
-    def dynamics_augmented(x: jnp.array, u: jnp.array, node: int) -> jnp.array:
-        x_dot = dynamics(x[idx_x_true], u[idx_u_true])
+@dataclass
+class Dynamics:
+    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
 
-        # Iterate through the g_func dictionary and stack the output each function
-        # to x_dot
-        for g in g_funcs:
-            x_dot = jnp.hstack([x_dot, g(x[idx_x_true], u[idx_u_true], node)])
+def dynamics(
+    _func=None,
+    *,
+    A: Optional[Callable] = None,
+    B: Optional[Callable] = None,):
+    """Decorator to mark a function as defining the system dynamics.
 
-        return x_dot
+    Use as:
+    @dynamics(A=my_grad_f_x, B=my_grad_f_u)')
+    def my_dynamics(x,u): ...
+    """
 
-    return dynamics_augmented
+    def decorator(f: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]):
+        # wrap so name, doc, signature stay on f
+        wrapped = functools.wraps(f)(f)
+        return Dynamics(
+            f=wrapped,
+            A=A,
+            B=B,
+        )
 
+    # if called as @dynamics or @dynamics(...), _func will be None and we return decorator
+    if _func is None:
+        return decorator
+    # if called as dynamics(func), we immediately decorate
+    else:
+        return decorator(_func)
 
-def get_jacobians(dyn: callable):
-    A = jax.jacfwd(dyn, argnums=0)
-    B = jax.jacfwd(dyn, argnums=1)
-    return A, B

openscvx/post_processing.py

@@ -1,4 +1,5 @@
 import numpy as np
+import jax.numpy as jnp
 
 from openscvx.propagation import s_to_t, t_to_tau, simulate_nonlinear_time
 from openscvx.config import Config
@@ -14,9 +15,13 @@ def propagate_trajectory_results(params: Config, result: dict, propagation_solve
 
     tau_vals, u_full = t_to_tau(u, t_full, u, t, params)
 
-    x_full = simulate_nonlinear_time(x[0], u, tau_vals, t, params, propagation_solver)
+    # Match free values from initial state to the initial value from the result
+    mask = jnp.array([t == "Free" for t in params.sim.initial_state_prop.types], dtype=bool)
+    params.sim.initial_state_prop.value = jnp.where(mask, x[0], params.sim.initial_state_prop.value)
 
-    print("Total CTCS Constraint Violation:", x_full[-1, params.sim.idx_y])
+    x_full = simulate_nonlinear_time(params.sim.initial_state_prop.value, u, tau_vals, t, params, propagation_solver)
+
+    print("Total CTCS Constraint Violation:", x_full[-1, params.sim.idx_y_prop])
     i = 0
     cost = np.zeros_like(x[-1, i])
     for type in params.sim.initial_state.type:
@@ -28,6 +33,16 @@ def propagate_trajectory_results(params: Config, result: dict, propagation_solve
         if type == "Minimize":
             cost += x[-1, i]
         i += 1
+    i=0
+    for type in params.sim.initial_state.type:
+        if type == "Maximize":
+            cost -= x[0, i]
+        i += 1
+    i = 0
+    for type in params.sim.final_state.type:
+        if type == "Maximize":
+            cost -= x[-1, i]
+        i += 1
     print("Cost: ", cost)
 
     more_result = dict(t_full=t_full, x_full=x_full, u_full=u_full)

openscvx/trajoptproblem.py

@@ -1,5 +1,5 @@
 import jax.numpy as jnp
-from typing import List
+from typing import List, Union, Optional
 import queue
 import threading
 import time
@@ -17,12 +17,13 @@ from openscvx.config import (
     DevConfig,
     Config,
 )
-from openscvx.dynamics import get_augmented_dynamics, get_jacobians
-from openscvx.constraints.violation import get_g_funcs
-from openscvx.augmentation import sort_ctcs_constraints
+from openscvx.dynamics import Dynamics
+from openscvx.augmentation.dynamics_augmentation import build_augmented_dynamics
+from openscvx.augmentation.ctcs import sort_ctcs_constraints
+from openscvx.constraints.violation import get_g_funcs, CTCSViolation
 from openscvx.discretization import get_discretization_solver
 from openscvx.propagation import get_propagation_solver
-from openscvx.constraints.boundary import BoundaryConstraint
+from openscvx.constraints.boundary import BoundaryConstraint, boundary
 from openscvx.constraints.ctcs import CTCSConstraint
 from openscvx.constraints.nodal import NodalConstraint
 from openscvx.ptr import PTR_init, PTR_main
@@ -35,8 +36,8 @@ from openscvx import io
 class TrajOptProblem:
     def __init__(
         self,
-        dynamics: callable,
-        constraints: List[callable],
+        dynamics: Dynamics,
+        constraints: List[Union[CTCSConstraint, NodalConstraint]],
         idx_time: int,
         N: int,
         time_init: float,
@@ -48,28 +49,34 @@ class TrajOptProblem:
         x_min: jnp.ndarray,
         u_max: jnp.ndarray,
         u_min: jnp.ndarray,
-        scp: ScpConfig = None,
-        dis: DiscretizationConfig = None,
-        prp: PropagationConfig = None,
-        sim: SimConfig = None,
-        dev: DevConfig = None,
-        cvx: ConvexSolverConfig = None,
+        dynamics_prop: callable = None,
+        initial_state_prop: BoundaryConstraint = None,
+        scp: Optional[ScpConfig] = None,
+        dis: Optional[DiscretizationConfig] = None,
+        prp: Optional[PropagationConfig] = None,
+        sim: Optional[SimConfig] = None,
+        dev: Optional[DevConfig] = None,
+        cvx: Optional[ConvexSolverConfig] = None,
         licq_min=0.0,
         licq_max=1e-4,
         time_dilation_factor_min=0.3,
         time_dilation_factor_max=3.0,
     ):
+        if dynamics_prop is None:
+            dynamics_prop = dynamics
+        
+        if initial_state_prop is None:
+            initial_state_prop = initial_state
 
         # TODO (norrisg) move this into some augmentation function, if we want to make this be executed after the init (i.e. within problem.initialize) need to rethink how problem is defined
         constraints_ctcs = []
         constraints_nodal = []
-        # TODO: (norrisg) change back to using isinstance once on PyPi
         for constraint in constraints:
-            if type(constraint).__name__ == CTCSConstraint.__name__:
+            if isinstance(constraint, CTCSConstraint):
                 constraints_ctcs.append(
                     constraint
                 )
-            elif type(constraint).__name__ == NodalConstraint.__name__:
+            elif isinstance(constraint, NodalConstraint):
                 constraints_nodal.append(
                     constraint
                 )
@@ -81,11 +88,15 @@ class TrajOptProblem:
         constraints_ctcs, node_intervals, num_augmented_states = sort_ctcs_constraints(constraints_ctcs, N)
 
         # Index tracking
-        idx_x_true = slice(0, len(x_max))
+        idx_x_true = slice(0, len(initial_state.value))
+        idx_x_true_prop = slice(0, len(initial_state_prop.value))
         idx_u_true = slice(0, len(u_max))
         idx_constraint_violation = slice(
             idx_x_true.stop, idx_x_true.stop + num_augmented_states
         )
+        idx_constraint_violation_prop = slice(
+            idx_x_true_prop.stop, idx_x_true_prop.stop + num_augmented_states
+        )
 
         idx_time_dilation = slice(idx_u_true.stop, idx_u_true.stop + 1)
 
@@ -106,6 +117,11 @@ class TrajOptProblem:
             [u_guess, np.full((u_guess.shape[0], 1), time_init)]
         )
 
+        initial_state_prop_values = np.hstack([initial_state_prop.value, np.repeat(licq_min, num_augmented_states)])
+        initial_state_prop_types = np.hstack([initial_state_prop.type, ["Fix"] * num_augmented_states])
+        initial_state_prop = boundary(initial_state_prop_values)
+        initial_state_prop.types = initial_state_prop_types
+
         if dis is None:
             dis = DiscretizationConfig()
 
@@ -114,17 +130,21 @@ class TrajOptProblem:
                 x_bar=x_bar_augmented,
                 u_bar=u_bar_augmented,
                 initial_state=initial_state,
+                initial_state_prop=initial_state_prop,
                 final_state=final_state,
                 max_state=x_max_augmented,
                 min_state=x_min_augmented,
                 max_control=u_max_augmented,
                 min_control=u_min_augmented,
                 total_time=time_init,
-                n_states=len(x_max),
+                n_states=len(initial_state.value),
+                n_states_prop=len(initial_state_prop.value),
                 idx_x_true=idx_x_true,
+                idx_x_true_prop=idx_x_true_prop,
                 idx_u_true=idx_u_true,
                 idx_t=idx_time,
                 idx_y=idx_constraint_violation,
+                idx_y_prop=idx_constraint_violation_prop,
                 idx_s=idx_time_dilation,
                 ctcs_node_intervals=node_intervals,
             )
@@ -160,9 +180,9 @@ class TrajOptProblem:
         sim.constraints_ctcs = constraints_ctcs
         sim.constraints_nodal = constraints_nodal
 
-        g_funcs = get_g_funcs(constraints_ctcs)
-        self.dynamics_augmented = get_augmented_dynamics(dynamics, g_funcs, idx_x_true, idx_u_true)
-        self.A_uncompiled, self.B_uncompiled = get_jacobians(self.dynamics_augmented)
+        ctcs_violation_funcs = get_g_funcs(constraints_ctcs)
+        self.dynamics_augmented = build_augmented_dynamics(dynamics, ctcs_violation_funcs, idx_x_true, idx_u_true)
+        self.dynamics_augmented_prop = build_augmented_dynamics(dynamics_prop, ctcs_violation_funcs, idx_x_true_prop, idx_u_true)
 
         self.params = Config(
             sim=sim,
@@ -212,18 +232,23 @@ class TrajOptProblem:
         self.params.sim.__post_init__()
 
         # Compile dynamics and jacobians
-        self.state_dot = jax.vmap(self.dynamics_augmented)
-        self.A = jax.jit(jax.vmap(self.A_uncompiled, in_axes=(0, 0, 0)))
-        self.B = jax.jit(jax.vmap(self.B_uncompiled, in_axes=(0, 0, 0)))
-        # TODO: (norrisg) Could consider using dataclass just to hold dynamics and jacobians
-        # TODO: (norrisg) Consider writing the compiled versions into the same variables?
-        # Otherwise if have a dataclass could have 2 instances, one for compied and one for uncompiled
+        self.dynamics_augmented.f = jax.vmap(self.dynamics_augmented.f)
+        self.dynamics_augmented.A = jax.jit(jax.vmap(self.dynamics_augmented.A, in_axes=(0, 0, 0)))
+        self.dynamics_augmented.B = jax.jit(jax.vmap(self.dynamics_augmented.B, in_axes=(0, 0, 0)))
+
+
+        self.dynamics_augmented_prop.f = jax.vmap(self.dynamics_augmented_prop.f)
+
+        for constraint in self.params.sim.constraints_nodal:
+            if not constraint.convex:
+                # TODO: (haynec) switch to AOT instead of JIT
+                constraint.g = jax.jit(constraint.g)
+                constraint.grad_g_x = jax.jit(constraint.grad_g_x)
+                constraint.grad_g_u = jax.jit(constraint.grad_g_u)
 
         # Generate solvers and optimal control problem
-        self.discretization_solver = get_discretization_solver(
-            self.state_dot, self.A, self.B, self.params
-        )
-        self.propagation_solver = get_propagation_solver(self.state_dot, self.params)
+        self.discretization_solver = get_discretization_solver(self.dynamics_augmented, self.params)
+        self.propagation_solver = get_propagation_solver(self.dynamics_augmented_prop.f, self.params)
         self.optimal_control_problem = OptimalControlProblem(self.params)
 
         # Initialize the PTR loop
@@ -247,7 +272,7 @@ class TrajOptProblem:
         self.propagation_solver = (
             jax.jit(self.propagation_solver)
             .lower(
-                np.ones((self.params.sim.n_states)),
+                np.ones((self.params.sim.n_states_prop)),
                 (0.0, 0.0),
                 np.ones((1, self.params.sim.n_controls)),
                 np.ones((1, self.params.sim.n_controls)),

{openscvx-0.1.1.dist-info → openscvx-0.1.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.1.1
+Version: 0.1.3
 Summary: A general Python-based successive convexification implementation which uses a JAX backend.
 Home-page: https://haynec.github.io/openscvx/
 Author: Chris Hayner and Griffin Norris
@@ -26,10 +26,15 @@ Dynamic: license-file
 
 <img src="figures/openscvx_logo.svg" width="1200"/>
 <p align="center">
-    <a href="https://github.com//haynec/OpenSCvx/actions/workflows/main.yml/badge.svg"><img src="https://github.com//haynec/OpenSCvx/actions/workflows/main.yml/badge.svg"/></a>
+    <a href="https://github.com//haynec/OpenSCvx/actions/workflows/website.yml/badge.svg"><img src="https://github.com//haynec/OpenSCvx/actions/workflows/website.yml/badge.svg"/></a>
     <a href="https://arxiv.org/abs/2410.22596"><img src="http://img.shields.io/badge/arXiv-2410.22596-B31B1B.svg"/></a>
     <a href="https://www.apache.org/licenses/LICENSE-2.0"><img src="https://img.shields.io/badge/License-Apache_2.0-blue.svg" alt="License: Apache 2.0"/></a>
 </p>
+<p align="center">
+    <a href="https://github.com//haynec/OpenSCvx/actions/workflows/ci.yml/badge.svg"><img src="https://github.com//haynec/OpenSCvx/actions/workflows/ci.yml/badge.svg"/></a>
+    <a href="https://github.com//haynec/OpenSCvx/actions/workflows/nightly.yml/badge.svg"><img src="https://github.com//haynec/OpenSCvx/actions/workflows/nightly.yml/badge.svg"/></a>
+    <a href="https://github.com//haynec/OpenSCvx/actions/workflows/release.yml/badge.svg"><img src="https://github.com//haynec/OpenSCvx/actions/workflows/release.yml/badge.svg"/></a>
+</p>
 
 <!-- PROJECT LOGO -->
 <br />
@@ -37,20 +42,50 @@ Dynamic: license-file
 <!-- GETTING STARTED -->
 ## Getting Started
 
-
 ### Installation
+
+To grab the latest stable release simply run
+
+```sh
+pip install openscvx
+```
+
+to install OpenSCVx in your python environment.
+
+<details>
+<summary>Install Development / Nightly Version</summary>
+
+If you want the pre-release version, you can install the latest `nightly` build with:
+
+```sh
+python3 -m pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ --pre --upgrade openscvx
+```
+
+This command will also upgrade an existing `nightly` install to the latest version
+Or if you want a specific pre-release version this can be installed with
+
+```sh
+python3 -m pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ openscvx==1.2.3.dev45
+```
+
+where `1.2.3.dev45 => <major>.<minor>.<patch>.dev<XY>` corresponds to your exact version
+</details>
+
+
+#### Dependencies
+
 The main packages are:
-- ```cvxpy``` - is used to formulate and solve the convex subproblems
-- ```jax``` - is used for determining the Jacobians using automatic differentiation, vectorization, and ahead-of-time (AOT) compilation of the dynamics and their Jacobians 
-- ```numpy``` - is used for numerical operations
-- ```diffrax``` - is used for the numerical integration of the dynamics
-- ```termcolor``` - is used for pretty command line output
-- ```plotly``` - is used for all visualizations
 
+- `cvxpy` - is used to formulate and solve the convex subproblems
+- `jax` - is used for determining the Jacobians using automatic differentiation, vectorization, and ahead-of-time (AOT) compilation of the dynamics and their Jacobians 
+- `numpy` - is used for numerical operations
+- `diffrax` - is used for the numerical integration of the dynamics
+- `termcolor` - is used for pretty command line output
+- `plotly` - is used for all visualizations
 
-These can be installed via conda or pip.
+These will be installed automatically, but can be installed via conda or pip if you are building from source.
 <details>
-<summary>Via Conda (Recommended) </summary>
+<summary>Via Conda</summary>
 
 1. Clone the repo
    ```sh
@@ -87,11 +122,11 @@ These can be installed via conda or pip.
 See `examples/` folder for several example trajectory optimization problems.
 To run a problem simply run `examples/main.py` with:
 
-```bash
-python3 -m examples.main
+```sh
+python3 examples/main.py
 ```
 
-To change which example is run by `main` simply replace the `params` import line:
+To change which example is run by `main` simply replace the `problem` import line:
 
 ```python
 # other imports
@@ -99,10 +134,11 @@ from examples.params.dr_vp import problem
 # rest of code
 ```
 
+and adjust the plotting as needed.
 Check out the problem definitions inside `examples/params` to see how to define your own problems.
 
-
 ## ToDos
+
 - [X] Standardized Vehicle and Constraint classes
 - [X] Implement QOCOGen with CVPYGEN
 - [X] Non-Dilated Time Propagation 
@@ -110,8 +146,12 @@ Check out the problem definitions inside `examples/params` to see how to define
 - [ ] Compiled at the subproblem level with JAX
 - [ ] Save and reload the compiled JAX code
 - [ ] Single Shot propagation
+- [ ] Unified Mathematical Interface
+
 ## What is implemented
+
 This repo has the following features:
+
 1. Free Final Time
 2. Fully adaptive time dilation (```s``` is appended to the control vector)
 3. Continuous-Time Constraint Satisfaction
@@ -122,11 +162,14 @@ This repo has the following features:
 <p align="right">(<a href="#readme-top">back to top</a>)</p>
 
 ## Acknowledgements
+
 This work was supported by a NASA Space Technology Graduate Research Opportunity and the Office of Naval Research under grant N00014-17-1-2433. The authors would like to acknowledge Natalia Pavlasek, Samuel Buckner, Abhi Kamath, Govind Chari, and Purnanand Elango as well as the other Autonomous Controls Laboratory members, for their many helpful discussions and support throughout this work.
 
 ## Citation
+
 Please cite the following works if you use the repository,
-```
+
+```tex
 @ARTICLE{hayner2025los,
         author={Hayner, Christopher R. and Carson III, John M. and Açıkmeşe, Behçet and Leung, Karen},
         journal={IEEE Robotics and Automation Letters}, 
@@ -139,7 +182,7 @@ Please cite the following works if you use the repository,
         doi={10.1109/LRA.2025.3545299}}
 ```
 
-```
+```tex
 @misc{elango2024ctscvx,
       title={Successive Convexification for Trajectory Optimization with Continuous-Time Constraint Satisfaction}, 
       author={Purnanand Elango and Dayou Luo and Abhinav G. Kamath and Samet Uzun and Taewan Kim and Behçet Açıkmeşe},

openscvx-0.1.3.dist-info/RECORD

@@ -0,0 +1,27 @@
+openscvx/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+openscvx/_version.py,sha256=NIzzV8ZM0W-CSLuEs1weG4zPrn_-8yr1AwwI1iuS6yo,511
+openscvx/config.py,sha256=lTdvJtU-Xiwq55zGtW3QvmldZj9Ne0xyKsSDSIKwqiI,10488
+openscvx/discretization.py,sha256=YF3mEeyYHgyTWQVNQsqpi1Mv72zDLyNfaMJSWqxj34c,4745
+openscvx/dynamics.py,sha256=X9sPpxUGGbdsnvQzgyrb_939N9ctBSsWVyI1eXtOKpc,1118
+openscvx/integrators.py,sha256=msIS-1Ehj-9TJLHfoCMs3vdyZ8NXz-TM0RII6aqRf4E,3821
+openscvx/io.py,sha256=fOvNWQWAegcN1gejeToaNbXenP5H5bAifNU8edJvdk4,4127
+openscvx/ocp.py,sha256=L_509EQiMsI6s5gBYlYyxKaHEzzRdpo-XAMjliCU3Rc,7544
+openscvx/plotting.py,sha256=fCvWJV4qWMhVyJlh18s12S_5xhj6EviF-_FuP0tWjx4,31207
+openscvx/post_processing.py,sha256=t3fUeDfA2PZq8S-WVnxDyN0xIgFuvpDP-5wbGQuIOVY,1618
+openscvx/propagation.py,sha256=XNezQnAM-NXb9L7aHUgKQOBn0CNUPeGGDL3_BbGoODU,3758
+openscvx/ptr.py,sha256=itDTR6RQUphnU226jaeRaAKuia-6v8U3MqAdw5-BYOk,5268
+openscvx/trajoptproblem.py,sha256=hcju3UB3iZ5peqhxh6NvDQYnsArnabR0mGUAKZ3sIaw,13360
+openscvx/utils.py,sha256=zmkKyto8Jowe_RAdOe8K0w6gzOu4JfxmX1RUL-3OFlY,2408
+openscvx/augmentation/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+openscvx/augmentation/ctcs.py,sha256=m1jdALXSqHq3WD6lCBAUI7FR0Sfs8aCYr66h0EwE4z4,1707
+openscvx/augmentation/dynamics_augmentation.py,sha256=7PL-mMfSmIIfMiXjdbXTolxOgUFSolzKpLu8WAmD384,4271
+openscvx/constraints/__init__.py,sha256=OOUcYEVoDWOSY50s2TbjpDjl3dRR3U04gRxmOyjbddY,258
+openscvx/constraints/boundary.py,sha256=yEhEnkKJ5f8NUeTksigEJjgBeE_YyuG_PJb_DWxg0L4,1541
+openscvx/constraints/ctcs.py,sha256=V763033aV82nAu7y4653KsAs11A7RpUysR_oUcnLfko,2572
+openscvx/constraints/nodal.py,sha256=YCS0cwUurA2OTQcHBb1EQqLxNt_w3MX8Nj8FH3GYClo,1726
+openscvx/constraints/violation.py,sha256=aIdDhHd-UndT0XB2QeuwLBKSNSAUWVkha_GeHOw9cQg,2362
+openscvx-0.1.3.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+openscvx-0.1.3.dist-info/METADATA,sha256=CVr7tDDV75qDlnuEuxf-IL0Arxkb1RXwv0Xj9oVxf_s,6911
+openscvx-0.1.3.dist-info/WHEEL,sha256=QZxptf4Y1BKFRCEDxD4h2V0mBFQOVFLFEpvxHmIs52A,91
+openscvx-0.1.3.dist-info/top_level.txt,sha256=nUT4Ybefzh40H8tVXqc1RzKESy_MAowElb-CIvAbd4Q,9
+openscvx-0.1.3.dist-info/RECORD,,

{openscvx-0.1.1.dist-info → openscvx-0.1.3.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.4.0)
+Generator: setuptools (80.6.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

openscvx-0.1.1.dist-info/RECORD

@@ -1,25 +0,0 @@
-openscvx/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-openscvx/_version.py,sha256=Mmxse1R0ki5tjz9qzU8AQyqUsLt8nTyCAbYQp8R87PU,511
-openscvx/augmentation.py,sha256=m1jdALXSqHq3WD6lCBAUI7FR0Sfs8aCYr66h0EwE4z4,1707
-openscvx/config.py,sha256=8Cl5O0ekf9MGNDTEeMHsp1C4XvY9NfJQkxd80lvnafM,10296
-openscvx/discretization.py,sha256=xqR6QPfp--9Nh0sjMfOkTc4OxOxeyHdZq7ip7ATzMNo,4704
-openscvx/dynamics.py,sha256=1GrjV75qxHa4bx2UK_Idul6qZa-IFjMUrqsne_gdV3E,684
-openscvx/integrators.py,sha256=msIS-1Ehj-9TJLHfoCMs3vdyZ8NXz-TM0RII6aqRf4E,3821
-openscvx/io.py,sha256=fOvNWQWAegcN1gejeToaNbXenP5H5bAifNU8edJvdk4,4127
-openscvx/ocp.py,sha256=L_509EQiMsI6s5gBYlYyxKaHEzzRdpo-XAMjliCU3Rc,7544
-openscvx/plotting.py,sha256=fCvWJV4qWMhVyJlh18s12S_5xhj6EviF-_FuP0tWjx4,31207
-openscvx/post_processing.py,sha256=TP1gi4TVlDS2HHpdqaIPCqfM5o4w7a7RCMU3Pu3czHw,1024
-openscvx/propagation.py,sha256=XNezQnAM-NXb9L7aHUgKQOBn0CNUPeGGDL3_BbGoODU,3758
-openscvx/ptr.py,sha256=itDTR6RQUphnU226jaeRaAKuia-6v8U3MqAdw5-BYOk,5268
-openscvx/trajoptproblem.py,sha256=fq68viMiS1UOcwquyKmWxcYMBQT3NvG3xdH-LHollHQ,11932
-openscvx/utils.py,sha256=zmkKyto8Jowe_RAdOe8K0w6gzOu4JfxmX1RUL-3OFlY,2408
-openscvx/constraints/__init__.py,sha256=OOUcYEVoDWOSY50s2TbjpDjl3dRR3U04gRxmOyjbddY,258
-openscvx/constraints/boundary.py,sha256=yEhEnkKJ5f8NUeTksigEJjgBeE_YyuG_PJb_DWxg0L4,1541
-openscvx/constraints/ctcs.py,sha256=05epAuo_mNm1AieNB6FWatkv0wOT1ebD4FdngIROljY,1788
-openscvx/constraints/nodal.py,sha256=a3CRI7sYBNoOk2wZz9n7nyUuQUjzAGIjRmuHlgBSidk,1592
-openscvx/constraints/violation.py,sha256=wKLNhInHoXXannf2J_nLtvm3dWOMZrhJy3mJLG4CTX0,809
-openscvx-0.1.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-openscvx-0.1.1.dist-info/METADATA,sha256=DiNKvRx7k9Z7mjnKgXvm5ClBgdIdBd66ub4kDPVy5q4,5384
-openscvx-0.1.1.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
-openscvx-0.1.1.dist-info/top_level.txt,sha256=nUT4Ybefzh40H8tVXqc1RzKESy_MAowElb-CIvAbd4Q,9
-openscvx-0.1.1.dist-info/RECORD,,