openscvx 0.4.1.dev88__tar.gz → 0.4.1.dev90__tar.gz

{openscvx-0.4.1.dev88/openscvx.egg-info → openscvx-0.4.1.dev90}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.4.1.dev88
+Version: 0.4.1.dev90
 Summary: A general Python-based successive convexification implementation which uses a JAX backend.
 Author-email: Chris Hayner and Griffin Norris <haynec@uw.edu>
 License: Apache Software License

{openscvx-0.4.1.dev88 → openscvx-0.4.1.dev90}/openscvx/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.4.1.dev88'
-__version_tuple__ = version_tuple = (0, 4, 1, 'dev88')
+__version__ = version = '0.4.1.dev90'
+__version_tuple__ = version_tuple = (0, 4, 1, 'dev90')
 
-__commit_id__ = commit_id = 'g6d7b113ab'
+__commit_id__ = commit_id = 'g7eea8c25b'

{openscvx-0.4.1.dev88 → openscvx-0.4.1.dev90}/openscvx/discretization/linearize_discretize.py

@@ -110,25 +110,26 @@ def _dVdt(
 
     Their derivatives follow from the variational equations:
 
-    - ``dx/dτ     = s · f(x, u)``
-    - ``dΦ/dτ     = s · A(x, u) · Φ``
-    - ``dB_d/dτ   = s · A(x, u) · B_d  +  α · s · B(x, u)``
-    - ``dC_d/dτ   = s · A(x, u) · C_d  +  β · s · B(x, u)``
+    - ``dx/dτ     = F(x, u)``          where F = s · f(x, u) is the time-dilated dynamics
+    - ``dΦ/dτ     = A(x, u) · Φ``
+    - ``dB_d/dτ   = A(x, u) · B_d  +  α · B(x, u)``
+    - ``dC_d/dτ   = A(x, u) · C_d  +  β · B(x, u)``
 
-    where ``s`` is the time-dilation factor (last column of ``u``),
-    ``A = ∂f/∂x``, ``B = ∂f/∂u``, and ``α, β`` are interpolation weights
-    determined by the hold type (ZOH: α=1, β=0; FOH: linear blend).
+    where ``A = ∂F/∂x`` and ``B = ∂F/∂u`` are Jacobians of the
+    time-dilated dynamics (which include the time-dilation factor ``s``
+    symbolically), and ``α, β`` are interpolation weights determined by the
+    hold type (ZOH: α=1, β=0; FOH: linear blend).
 
     Args:
         tau: Normalized time in [0, 1] within the current segment.
         V: Flattened augmented state vector, shape ``((N-1) * aug_dim,)``.
-        u_cur: Control at current node, shape ``(N-1, n_u+1)``.
-        u_next: Control at next node, shape ``(N-1, n_u+1)``.
-        state_dot: Vmapped dynamics ``f(x, u, node, params) -> x_dot``.
-        A: Vmapped state Jacobian ``∂f/∂x(x, u, node, params)``.
-        B: Vmapped control Jacobian ``∂f/∂u(x, u, node, params)``.
+        u_cur: Control at current node, shape ``(N-1, n_u)``.
+        u_next: Control at next node, shape ``(N-1, n_u)``.
+        state_dot: Vmapped time-dilated dynamics ``F(x, u, node, params) -> x_dot``.
+        A: Vmapped state Jacobian ``∂F/∂x(x, u, node, params)``.
+        B: Vmapped control Jacobian ``∂F/∂u(x, u, node, params)``.
         n_x: Number of states.
-        n_u: Number of controls (excluding time-dilation slack).
+        n_u: Number of controls (including time-dilation).
         N: Number of trajectory nodes.
         dis_type: ``"ZOH"`` (zero-order hold) or ``"FOH"`` (first-order hold).
         S_x: State scaling matrix (unused, reserved for future scaling).
@@ -165,46 +166,32 @@ def _dVdt(
         beta = (tau) * N
     alpha = 1 - beta
 
-    # TODO: (norrisg) integrate the multiplication with `s` into the symbolic layer
-    # This currently requires a hack to get the sparsity pattern to include the relation with s
-
     # Interpolate the control input
     u = u_cur + beta * (u_next - u_cur)
-    s = u[:, -1]
-
-    # Initialize the augmented Jacobians
-    dfdx = jnp.zeros((V.shape[0], n_x, n_x))
-    dfdu = jnp.zeros((V.shape[0], n_x, n_u))
 
     # Ensure x_seq and u have the same batch size
     x = V[:, :n_x]
     u = u[: x.shape[0]]
 
-    # Compute the nonlinear propagation term
-    f = state_dot(x, u[:, :-1], nodes, params)
-    F = s[:, None] * f
-
-    # Evaluate the State Jacobian
-    dfdx = A(x, u[:, :-1], nodes, params)
-    sdfdx = s[:, None, None] * dfdx
+    # Compute the time-dilated dynamics (s * f already included symbolically)
+    F = state_dot(x, u, nodes, params)
 
-    # Evaluate the Control Jacobian
-    dfdu_veh = B(x, u[:, :-1], nodes, params)
-    dfdu = dfdu.at[:, :, :-1].set(s[:, None, None] * dfdu_veh)
-    dfdu = dfdu.at[:, :, -1].set(f)
+    # Evaluate the Jacobians (already include time-dilation derivatives via autodiff)
+    dfdx = A(x, u, nodes, params)
+    dfdu = B(x, u, nodes, params)
 
     # Stack up the results into the augmented state vector
     # fmt: off
     dVdt = jnp.zeros_like(V)
     dVdt = dVdt.at[:, i0:i1].set(F)
     dVdt = dVdt.at[:, i1:i2].set(
-        jnp.matmul(sdfdx, V[:, i1:i2].reshape(-1, n_x, n_x)).reshape(-1, n_x * n_x)
+        jnp.matmul(dfdx, V[:, i1:i2].reshape(-1, n_x, n_x)).reshape(-1, n_x * n_x)
     )
     dVdt = dVdt.at[:, i2:i3].set(
-        (jnp.matmul(sdfdx, V[:, i2:i3].reshape(-1, n_x, n_u)) + dfdu * alpha).reshape(-1, n_x * n_u)
+        (jnp.matmul(dfdx, V[:, i2:i3].reshape(-1, n_x, n_u)) + dfdu * alpha).reshape(-1, n_x * n_u)
     )
     dVdt = dVdt.at[:, i3:i4].set(
-        (jnp.matmul(sdfdx, V[:, i3:i4].reshape(-1, n_x, n_u)) + dfdu * beta).reshape(-1, n_x * n_u)
+        (jnp.matmul(dfdx, V[:, i3:i4].reshape(-1, n_x, n_u)) + dfdu * beta).reshape(-1, n_x * n_u)
     )
     # fmt: on
 
@@ -231,11 +218,11 @@ def _calculate_discretization(
 
     Args:
         x: Reference state trajectory, shape ``(N, n_x)``.
-        u: Reference control trajectory, shape ``(N, n_u+1)`` (includes
-            time-dilation slack as last column).
-        state_dot: Vmapped dynamics ``f(x, u, node, params) -> x_dot``.
-        A: Vmapped state Jacobian ``∂f/∂x``.
-        B: Vmapped control Jacobian ``∂f/∂u``.
+        u: Reference control trajectory, shape ``(N, n_u)`` (includes
+            time-dilation as part of the unified control vector).
+        state_dot: Vmapped time-dilated dynamics ``F(x, u, node, params) -> x_dot``.
+        A: Vmapped state Jacobian ``∂F/∂x``.
+        B: Vmapped control Jacobian ``∂F/∂u``.
         settings: Configuration (integrator choice, tolerances, hold type, etc.).
         params: Parameters forwarded to ``state_dot``, ``A``, and ``B``.
 

{openscvx-0.4.1.dev88 → openscvx-0.4.1.dev90}/openscvx/expert/lowering.py

@@ -78,17 +78,23 @@ def apply_byof(
         state_slices = {state.name: state._slice for state in states}
         state_slices_prop = {state.name: state._slice for state in states_prop}
 
-        def _make_composite_dynamics(orig_f, byof_fns, slices_map):
+        # Time-dilation slice for multiplying byof outputs by s
+        td_slice = u_unified.time_dilation_slice
+
+        def _make_composite_dynamics(orig_f, byof_fns, slices_map, td_sl):
             """Create composite dynamics combining symbolic and byof state derivatives.
 
             This factory splices user-provided byof dynamics into the unified dynamics
             function at the appropriate slice indices, replacing the symbolic dynamics
-            for specific states while preserving the rest.
+            for specific states while preserving the rest. The byof outputs are
+            multiplied by the time-dilation factor s to match the symbolic dynamics
+            which already include s * f(x, u) via the Mul node.
 
             Args:
                 orig_f: Original unified dynamics (x, u, node, params) -> xdot
                 byof_fns: Dict mapping state names to byof dynamics functions
                 slices_map: Dict mapping state names to slice objects for indexing
+                td_sl: Slice for the time-dilation control in the unified u vector
 
             Returns:
                 Composite dynamics function with byof derivatives spliced in
@@ -98,11 +104,14 @@ def apply_byof(
                 # Start with symbolic/default dynamics for all states
                 xdot = orig_f(x, u, node, params)
 
-                # Splice in byof dynamics for specific states
+                # Time-dilation factor (symbolic dynamics already include s *)
+                s = u[td_sl]
+
+                # Splice in byof dynamics for specific states, multiplied by s
                 for state_name, byof_fn in byof_fns.items():
                     sl = slices_map[state_name]
-                    # Replace the derivative for this state with the byof result
-                    xdot = xdot.at[sl].set(byof_fn(x, u, node, params))
+                    # Replace the derivative for this state with s * byof result
+                    xdot = xdot.at[sl].set(s * byof_fn(x, u, node, params))
 
                 return xdot
 
@@ -110,12 +119,12 @@ def apply_byof(
 
         # Create composite optimization dynamics
         # Jacobians are computed by the discretizer, not here.
-        composite_f = _make_composite_dynamics(dynamics.f, byof_dynamics, state_slices)
+        composite_f = _make_composite_dynamics(dynamics.f, byof_dynamics, state_slices, td_slice)
         dynamics = Dynamics(f=composite_f)
 
         # Create composite propagation dynamics
         composite_f_prop = _make_composite_dynamics(
-            dynamics_prop.f, byof_dynamics, state_slices_prop
+            dynamics_prop.f, byof_dynamics, state_slices_prop, td_slice
         )
         dynamics_prop = Dynamics(f=composite_f_prop)
 
@@ -273,17 +282,24 @@ def apply_byof(
 
             penalty_fns.append(_make_penalty_fn(constraint_fn, penalty_func, over_interval))
 
+        # Time-dilation slice for multiplying byof CTCS penalties by s
+        td_slice = u_unified.time_dilation_slice
+
         if idx in idx_to_aug_slice:
             # This idx already exists from symbolic CTCS - add penalties to existing state
             aug_slice = idx_to_aug_slice[idx]
 
-            def _make_ctcs_addition(orig_f, pen_fns, aug_sl):
+            def _make_ctcs_addition(orig_f, pen_fns, aug_sl, td_sl):
                 """Create dynamics that adds penalties to existing augmented state.
 
+                The penalty is multiplied by the time-dilation factor s to match
+                the symbolic dynamics which already include s * f(x, u).
+
                 Args:
                     orig_f: Original dynamics function
                     pen_fns: List of penalty functions to add
                     aug_sl: Slice of the augmented state to modify
+                    td_sl: Slice for the time-dilation control
 
                 Returns:
                     Modified dynamics function
@@ -292,8 +308,9 @@ def apply_byof(
                 def modified_f(x, u, node, params):
                     xdot = orig_f(x, u, node, params)
 
-                    # Sum all penalties for this idx
-                    total_penalty = sum(pen_fn(x, u, node, params) for pen_fn in pen_fns)
+                    # Sum all penalties for this idx, scaled by time-dilation
+                    s = u[td_sl]
+                    total_penalty = s * sum(pen_fn(x, u, node, params) for pen_fn in pen_fns)
 
                     # Add to existing augmented state derivative
                     current_deriv = xdot[aug_sl]
@@ -305,8 +322,8 @@ def apply_byof(
 
             # Modify both optimization and propagation dynamics
             # Jacobians are computed by the discretizer, not here.
-            dynamics.f = _make_ctcs_addition(dynamics.f, penalty_fns, aug_slice)
-            dynamics_prop.f = _make_ctcs_addition(dynamics_prop.f, penalty_fns, aug_slice)
+            dynamics.f = _make_ctcs_addition(dynamics.f, penalty_fns, aug_slice, td_slice)
+            dynamics_prop.f = _make_ctcs_addition(dynamics_prop.f, penalty_fns, aug_slice, td_slice)
 
         else:
             # New idx - create new augmented state
@@ -315,12 +332,16 @@ def apply_byof(
             bounds = first_spec.get("bounds", (0.0, 1e-4))
             initial = first_spec.get("initial", bounds[0])
 
-            def _make_ctcs_new_state(orig_f, pen_fns):
+            def _make_ctcs_new_state(orig_f, pen_fns, td_sl):
                 """Create dynamics augmented with new CTCS state.
 
+                The penalty is multiplied by the time-dilation factor s to match
+                the symbolic dynamics which already include s * f(x, u).
+
                 Args:
                     orig_f: Original dynamics function
                     pen_fns: List of penalty functions to sum
+                    td_sl: Slice for the time-dilation control
 
                 Returns:
                     Augmented dynamics function
@@ -329,8 +350,9 @@ def apply_byof(
                 def augmented_f(x, u, node, params):
                     xdot = orig_f(x, u, node, params)
 
-                    # Sum all penalties for this new idx
-                    total_penalty = sum(pen_fn(x, u, node, params) for pen_fn in pen_fns)
+                    # Sum all penalties for this new idx, scaled by time-dilation
+                    s = u[td_sl]
+                    total_penalty = s * sum(pen_fn(x, u, node, params) for pen_fn in pen_fns)
 
                     # Append as new augmented state derivative
                     return jnp.concatenate([xdot, jnp.atleast_1d(total_penalty)])
@@ -339,11 +361,11 @@ def apply_byof(
 
             # Augment optimization dynamics
             # Jacobians are computed by the discretizer, not here.
-            aug_f = _make_ctcs_new_state(dynamics.f, penalty_fns)
+            aug_f = _make_ctcs_new_state(dynamics.f, penalty_fns, td_slice)
             dynamics = Dynamics(f=aug_f)
 
             # Augment propagation dynamics
-            aug_f_prop = _make_ctcs_new_state(dynamics_prop.f, penalty_fns)
+            aug_f_prop = _make_ctcs_new_state(dynamics_prop.f, penalty_fns, td_slice)
             dynamics_prop = Dynamics(f=aug_f_prop)
 
             # Create State objects for the new augmented states

{openscvx-0.4.1.dev88 → openscvx-0.4.1.dev90}/openscvx/propagation/propagation.py

@@ -12,7 +12,6 @@ def prop_aug_dy(
     u_next: np.ndarray,
     tau_init: float,
     node: int,
-    idx_s: int,
     state_dot: callable,
     dis_type: str,
     N: int,
@@ -20,8 +19,10 @@ def prop_aug_dy(
 ) -> np.ndarray:
     """Compute the augmented dynamics for propagation.
 
-    This function computes the time-scaled dynamics for propagating the system state,
-    taking into account the discretization type (ZOH or FOH) and time dilation.
+    This function computes the time-dilated dynamics for propagating the system
+    state, taking into account the discretization type (ZOH or FOH). The
+    time-dilation multiplication is already included in ``state_dot``
+    symbolically.
 
     Args:
         tau (float): Current normalized time in [0,1].
@@ -30,14 +31,13 @@ def prop_aug_dy(
         u_next (np.ndarray): Control input at next node.
         tau_init (float): Initial normalized time.
         node (int): Current node index.
-        idx_s (int): Index of time dilation variable in control vector.
-        state_dot (callable): Function computing state derivatives.
+        state_dot (callable): Function computing time-dilated state derivatives.
         dis_type (str): Discretization type ("ZOH" or "FOH").
         N (int): Number of nodes in trajectory.
         params: Dictionary of additional parameters passed to state_dot.
 
     Returns:
-        np.ndarray: Time-scaled state derivatives.
+        np.ndarray: Time-dilated state derivatives.
     """
     x = x[None, :]
 
@@ -47,7 +47,7 @@ def prop_aug_dy(
         beta = (tau - tau_init) * N
     u = u_current + beta * (u_next - u_current)
 
-    return u[:, idx_s] * state_dot(x, u[:, :-1], node, params).squeeze()
+    return state_dot(x, u, node, params).squeeze()
 
 
 def get_propagation_solver(state_dot: Dynamics, settings: Config) -> callable:
@@ -64,9 +64,7 @@ def get_propagation_solver(state_dot: Dynamics, settings: Config) -> callable:
         callable: A function that solves the propagation problem.
     """
 
-    def propagation_solver(
-        V0, tau_grid, u_cur, u_next, tau_init, node, idx_s, save_time, mask, params
-    ):
+    def propagation_solver(V0, tau_grid, u_cur, u_next, tau_init, node, save_time, mask, params):
         param_map_update = params
         return solve_ivp_diffrax_prop(
             f=prop_aug_dy,
@@ -77,7 +75,6 @@ def get_propagation_solver(state_dot: Dynamics, settings: Config) -> callable:
                 u_next,  # shape (1, n_controls)
                 tau_init,  # shape (1, 1)
                 node,  # shape (1, 1)
-                idx_s,  # int
                 state_dot,  # function or array
                 settings.dis.dis_type,
                 settings.scp.n,
@@ -242,7 +239,6 @@ def simulate_nonlinear_time(
             controls_next,
             np.array([[tau[k]]]),
             np.array([[k]]),
-            settings.sim.time_dilation_slice.stop,
             tau_cur_padded,
             mask_padded,
             params,

{openscvx-0.4.1.dev88 → openscvx-0.4.1.dev90}/openscvx/symbolic/augmentation.py

@@ -457,7 +457,7 @@ def augment_dynamics_with_ctcs(
     The augmented dynamics become:
         x_dot = f(x, u)
         aug_dot = penalty(g(x, u))  # For each constraint group
-        time_dot = time_dilation
+        time_dot = 1.0
 
     Args:
         xdot: Original dynamics expression for states

{openscvx-0.4.1.dev88 → openscvx-0.4.1.dev90}/openscvx/symbolic/builder.py

@@ -341,7 +341,19 @@ def preprocess_symbolic_problem(
             parameters=parameters,
         )
 
-    # ==================== PHASE 6: Process Algebraic Outputs ====================
+    # ==================== PHASE 6: Apply Time-Dilation Multiplication ===========
+    #
+    # Multiply all dynamics by the time-dilation control symbolically.
+    # This transforms f(x,u) -> s * f(x,u) so that JAX autodiff automatically
+    # computes the correct Jacobians including df/ds = f(x,u).
+    #
+    # This must happen AFTER propagation dynamics are assembled so that extra
+    # propagation states (e.g. distance) also get the s * multiplication.
+    time_dilation = next(c for c in controls_aug if c.name == "_time_dilation")
+    dynamics_aug = time_dilation * dynamics_aug
+    dynamics_prop = time_dilation * dynamics_prop
+
+    # ==================== PHASE 7: Process Algebraic Outputs ====================
 
     # Validate and canonicalize algebraic_prop expressions
     algebraic_prop_processed = None

{openscvx-0.4.1.dev88 → openscvx-0.4.1.dev90}/openscvx/symbolic/expr/expr.py

@@ -572,7 +572,7 @@ def traverse(expr: Expr, visit: Callable[[Expr], None]):
         traverse(child, visit)
 
 
-class Constant(Expr):
+class Constant(Leaf):
     """Constant value expression.
 
     Represents a constant numeric value in the expression tree. Constants are
@@ -601,40 +601,7 @@ class Constant(Expr):
         if not isinstance(value, np.ndarray):
             value = np.array(value, dtype=float)
         self.value = np.squeeze(value)
-
-    def canonicalize(self) -> "Expr":
-        """Constants are already in canonical form.
-
-        Returns:
-            Expr: Returns self since constants are already canonical
-        """
-        return self
-
-    def check_shape(self) -> Tuple[int, ...]:
-        """Return the shape of this constant's value.
-
-        Returns:
-            tuple: The shape of the constant's numpy array value
-        """
-        # Verify the invariant: constants should already be squeezed during construction
-        original_shape = self.value.shape
-        squeezed_shape = np.squeeze(self.value).shape
-        if original_shape != squeezed_shape:
-            raise ValueError(
-                f"Constant not properly normalized: has shape {original_shape} "
-                "but should have shape {squeezed_shape}. "
-                "Constants should be squeezed during construction."
-            )
-        return self.value.shape
-
-    def sparsity(self, n_x: int, n_u: int) -> Tuple[np.ndarray, np.ndarray]:
-        """Constants have no decision-variable dependence."""
-        shape = self.value.shape
-        n_out = int(np.prod(shape)) if shape else 1
-        return (
-            np.zeros((n_out, n_x), dtype=bool),
-            np.zeros((n_out, n_u), dtype=bool),
-        )
+        super().__init__(name="__constant__", shape=self.value.shape)
 
     def _hash_into(self, hasher: "hashlib._Hash") -> None:
         """Hash constant by its value.

{openscvx-0.4.1.dev88 → openscvx-0.4.1.dev90}/openscvx/symbolic/lower.py

@@ -717,27 +717,12 @@ def lower_symbolic_problem(
     # patterns (the superset) since dis_type isn't known yet.
     dynamics_sparsity = None
     if byof is None and problem.dynamics is not None:
-        from openscvx.symbolic.sparsity import _element_liveness, discrete_sparsity
+        from openscvx.symbolic.sparsity import discrete_sparsity
 
         n_x = sum(s.shape[0] for s in problem.states)
         n_u = sum(c.shape[0] for c in problem.controls)
         A_c, B_c = problem.dynamics.sparsity(n_x, n_u)
 
-        # TODO: (norrisg) tidy up how time-dilation is multiplied onto the
-        # dynamics so that it is automatically included into the sparsity.
-        # Requires fixing stuff in the discretizer.
-
-        # The discretization multiplies all dynamics by the time-dilation
-        # factor s and adds a column df/d(sigma) = f(x, u) to the control
-        # Jacobian (see linearize_discretize.py).  The symbolic dynamics
-        # don't reference sigma, so we patch B_c here: every row whose
-        # dynamics can be nonzero depends on sigma.
-        for ctrl in problem.controls:
-            if ctrl.name == "_time_dilation" and ctrl._slice is not None:
-                live = _element_liveness(problem.dynamics)
-                B_c[live, ctrl._slice] = True
-                break
-
         dynamics_sparsity = discrete_sparsity(A_c, B_c, dis_type="FOH")
 
     # Compute per-constraint Jacobian sparsity from the symbolic AST.

{openscvx-0.4.1.dev88 → openscvx-0.4.1.dev90}/openscvx/utils/caching.py

@@ -192,7 +192,6 @@ def load_or_compile_propagation_solver(
         np.ones((1, n_controls)),  # controls_next
         np.ones((1, 1)),  # tau_0
         np.ones((1, 1)).astype("int"),  # segment index
-        0,  # idx_s_stop
         np.ones((max_tau_len,)),  # save_time (tau_cur_padded)
         np.ones((max_tau_len,), dtype=bool),  # mask_padded (boolean mask)
         params,  # additional parameters as dict
@@ -223,7 +222,6 @@ def prime_propagation_solver(
         controls_next = np.ones((1, settings.sim.u.shape[0]), dtype=settings.sim.u.guess.dtype)
         tau_init = np.array([[0.0]], dtype=np.float64)
         node = np.array([[0]], dtype=np.int64)
-        idx_s_stop = settings.sim.time_dilation_slice.stop
         save_time = np.ones((settings.prp.max_tau_len,), dtype=np.float64)
         mask_padded = np.ones((settings.prp.max_tau_len,), dtype=bool)
         # Create dummy params dict with same structure
@@ -238,7 +236,6 @@ def prime_propagation_solver(
             controls_next,
             tau_init,
             node,
-            idx_s_stop,
             save_time,
             mask_padded,
             dummy_params,

{openscvx-0.4.1.dev88 → openscvx-0.4.1.dev90/openscvx.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.4.1.dev88
+Version: 0.4.1.dev90
 Summary: A general Python-based successive convexification implementation which uses a JAX backend.
 Author-email: Chris Hayner and Griffin Norris <haynec@uw.edu>
 License: Apache Software License

{openscvx-0.4.1.dev88 → openscvx-0.4.1.dev90}/tests/symbolic/test_augmentation.py

@@ -387,6 +387,9 @@ def test_augment_single_penalty_no_add():
         N,
     )
 
+    # Augmented dynamics should be Concat
+    assert isinstance(xdot_aug, Concat)
+
     # Single penalty should not be wrapped in Add
     # But it should be wrapped in CTCS
     ctcs_expr = xdot_aug.exprs[1]
@@ -418,6 +421,9 @@ def test_augment_multiple_penalties_create_add():
         N,
     )
 
+    # Augmented dynamics should be Concat
+    assert isinstance(xdot_aug, Concat)
+
     # Multiple penalties should be wrapped in Add
     penalty_expr = xdot_aug.exprs[1]
     assert isinstance(penalty_expr, Add)

{openscvx-0.4.1.dev88 → openscvx-0.4.1.dev90}/tests/test_discretization.py

@@ -19,7 +19,7 @@ def settings():
     p = Dummy()
     p.sim = Dummy()
     p.sim.n_states = 2
-    p.sim.n_controls = 1
+    p.sim.n_controls = 2  # 1 vehicle control + 1 time-dilation (unified)
     p.sim.S_x = jnp.eye(p.sim.n_states)
     p.sim.c_x = jnp.zeros(p.sim.n_states)
     p.sim.S_u = jnp.eye(p.sim.n_controls)
@@ -41,20 +41,12 @@ def settings():
 
 
 def state_dot(x, u, node, params):
-    # simple linear: x' = A_true x + B_true u
-    return x + u
-
-
-def A(x, u, node, params):
-    batch = x.shape[0]
-    eye = jnp.eye(2)
-    return jnp.broadcast_to(eye, (batch, 2, 2))
-
-
-def B(x, u, node, params):
-    batch = x.shape[0]
-    ones = jnp.ones((2, 1))
-    return jnp.broadcast_to(ones, (batch, 2, 1))
+    # simple time-dilated dynamics: x' = s * (x + u_vehicle)
+    # u = [u_vehicle, s] includes both vehicle control and time-dilation
+    # This is the un-vmapped version (single sample, not batched)
+    s = u[1]
+    u_v = u[0]
+    return s * (x + u_v)
 
 
 @pytest.fixture
@@ -72,9 +64,9 @@ def test_discretization_shapes(settings, dynamics):
     discretizer = LinearizeDiscretize()
     solver = discretizer.get_solver(dynamics, settings)
 
-    # dummy x,u
+    # dummy x,u (n_controls already includes time-dilation)
     x = jnp.ones((settings.scp.n, settings.sim.n_states))
-    u = jnp.ones((settings.scp.n, settings.sim.n_controls + 1))  # +1 slack
+    u = jnp.ones((settings.scp.n, settings.sim.n_controls))
 
     A_bar, B_bar, C_bar, x_prop, Vmulti = solver(x, u, {})
 
@@ -85,19 +77,23 @@ def test_discretization_shapes(settings, dynamics):
     assert B_bar.shape == ((N - 1), n_x, n_u)
     assert C_bar.shape == ((N - 1), n_x, n_u)
     assert x_prop.shape == ((N - 1), n_x)
-    # assert Vmulti.shape == (N, (n_x + n_x*n_x + 2*n_x*n_u + n_x))
 
 
 def test_jit_dVdt_compiles(settings):
-    # prepare trivial inputs
+    # prepare trivial inputs (n_u already includes time-dilation)
     n_x, n_u = settings.sim.n_states, settings.sim.n_controls
     N = settings.scp.n
     aug_dim = n_x + n_x * n_x + 2 * n_x * n_u
 
     tau = jnp.array(0.3)
     V_flat = jnp.ones((N - 1) * aug_dim)
-    u_cur = jnp.ones((N - 1, n_u + 1))
-    u_next = jnp.ones((N - 1, n_u + 1))
+    u_cur = jnp.ones((N - 1, n_u))
+    u_next = jnp.ones((N - 1, n_u))
+
+    # Create vmapped versions of dynamics and Jacobians (as _dVdt expects)
+    f_vmapped = jax.vmap(state_dot, in_axes=(0, 0, 0, None))
+    A_vmapped = jax.vmap(jax.jacfwd(state_dot, argnums=0), in_axes=(0, 0, 0, None))
+    B_vmapped = jax.vmap(jax.jacfwd(state_dot, argnums=1), in_axes=(0, 0, 0, None))
 
     # bind out the Python callables & settings
     def wrapped(tau_, V_):
@@ -106,20 +102,20 @@ def test_jit_dVdt_compiles(settings):
             V_,
             u_cur,
             u_next,
-            state_dot,
-            A,
-            B,
+            f_vmapped,
+            A_vmapped,
+            B_vmapped,
             n_x,
             n_u,
             N,
             settings.dis.dis_type,
-            {},
             settings.sim.S_x,
             settings.sim.c_x,
             settings.sim.S_u,
             settings.sim.c_u,
             settings.sim.inv_S_x,
             settings.sim.inv_S_u,
+            {},
         )
 
     # now JIT only over (tau_, V_)
@@ -141,9 +137,9 @@ def test_jit_discretization_solver_compiles(settings, dynamics, integrator):
     discretizer = LinearizeDiscretize()
     solver = discretizer.get_solver(dynamics, settings)
 
-    # dummy x,u (including slack column)
+    # dummy x,u (n_controls already includes time-dilation)
     x = jnp.ones((settings.scp.n, settings.sim.n_states))
-    u = jnp.ones((settings.scp.n, settings.sim.n_controls + 1))
+    u = jnp.ones((settings.scp.n, settings.sim.n_controls))
 
     # jit & lower & compile
     jitted = jax.jit(solver)