openscvx 0.5.3.dev49__tar.gz → 0.5.3.dev51__tar.gz

{openscvx-0.5.3.dev49/openscvx.egg-info → openscvx-0.5.3.dev51}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.5.3.dev49
+Version: 0.5.3.dev51
 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.5.3.dev49 → openscvx-0.5.3.dev51}/docs/Foundations/discretization.md

@@ -91,4 +91,41 @@ def dVdt(self, tau: float, V: jnp.ndarray, u_cur: np.ndarray, u_next: np.ndarray
     )
     return dVdt.flatten()
 ```
-    
+
+## Supplying your own Jacobian
+
+The state and control Jacobians above (`dfdx`, `dfdu`) come from automatic
+differentiation of the dynamics you wrote. Occasionally the exact Jacobian is
+not the one you want: a drag law, a contact model, or a lookup table can be
+smooth enough to integrate while its true derivative makes the convex
+subproblem badly conditioned, so the trust region collapses and the SCP loop
+crawls. The standard remedy is to linearize with a deliberately simplified
+(inexact) Jacobian.
+
+`Expr.with_jacobian` applies that remedy to a single term:
+
+``` py
+drag = -0.5 * rho * ox.Norm(vel) * vel
+J_drag = -0.5 * rho * ox.Norm(vel) * np.eye(3)  # drops the d||v||/dv term
+
+dynamics = {"pos": vel, "vel": thrust / m + drag.with_jacobian({vel: J_drag})}
+```
+
+The keys are the `State` or `Control` objects the derivative is taken with
+respect to, and each value is an expression (or constant array) of shape
+`(*term.shape, *variable.shape)`. Directions you do not name are still
+differentiated automatically, so the example above hands over `∂/∂vel` while
+`∂/∂thrust` still comes from autodiff.
+
+The *value* of the term is untouched — only the search direction changes — so a
+converged trajectory still satisfies the original nonlinear dynamics; the
+defect and the propagated solution are computed with the term as written. An
+inexact Jacobian trades linearization accuracy (and hence convergence rate) for
+conditioning, so reach for it when a term is a known source of stiffness, not
+as a default.
+
+The override lives inside the lowered JAX function, so every downstream
+derivative picks it up: the discretizer's `A`/`B`, nonconvex constraint
+linearization, and the sparsity pattern used for coloring. It has no meaning
+inside a constraint that is handed to the convex solver as written, and lowering
+one there raises.

{openscvx-0.5.3.dev49 → openscvx-0.5.3.dev51}/docs/UsersGuide/01_hello_world_brachistochrone.md

@@ -164,6 +164,35 @@ Here, we do _not_ specify an initial or final value for the control but we _do_
 The initial guess must be of shape $(n_u \times n_{\mathrm{nodes}})$, providing an initial guess for each control at every node.
 Technically, we can also provide a `.guess` for states. However, these default to a linear interpolation between initial and final conditions which is sufficient in most cases.
 
+#### Symbolic initial guesses
+
+`.guess` accepts three forms, on both states and controls:
+
+- an array of shape $(n_{\mathrm{nodes}} \times n)$, as above;
+- a **symbolic expression** in other states and controls;
+- a **callable of `tau`**, the normalized node coordinate, returning a symbolic expression.
+
+A guess expression is evaluated at every node, with each referenced variable at its guess value and `tau` running $0 \to 1$ across the grid — the same per-node mental model as dynamics and constraints.
+
+```python
+# Cross-variable: the speed guess follows whatever velocity's guess is
+speed.guess = ox.linalg.Norm(velocity)
+
+# Along the grid: a straight line from p0 to pf
+position.guess = lambda tau: p0 + (pf - p0) * tau
+```
+
+References between guesses are the dependency graph: `speed` above is resolved after `velocity`, and a circular reference raises. Everything is evaluated once, when the `Problem` is built (and again on `reset()` if you assign a new expression between solves), so the solver only ever sees ordinary arrays.
+
+Because evaluation is per node, guess expressions cannot express trajectory-level constructions such as finite differences or cumulative sums. Build those eagerly with NumPy instead:
+
+```python
+velocity.guess = np.gradient(position.guess, axis=0)
+```
+
+!!! Note
+    `time.guess` accepts all three forms, but `time_dilation_guess` remains array-only.
+
 ### Dynamics
 
 Dynamics are defined as a dictionary mapping state names to their time derivatives using symbolic expressions:

{openscvx-0.5.3.dev49 → openscvx-0.5.3.dev51}/openscvx/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.5.3.dev49'
-__version_tuple__ = version_tuple = (0, 5, 3, 'dev49')
+__version__ = version = '0.5.3.dev51'
+__version_tuple__ = version_tuple = (0, 5, 3, 'dev51')
 
-__commit_id__ = commit_id = 'g02b83a786'
+__commit_id__ = commit_id = 'g21d5581a4'

{openscvx-0.5.3.dev49 → openscvx-0.5.3.dev51}/openscvx/problem.py

@@ -75,6 +75,7 @@ from openscvx.symbolic.expr.control import Control
 from openscvx.symbolic.expr.state import State
 from openscvx.symbolic.expr.time import Time
 from openscvx.symbolic.lower import lower_symbolic_problem
+from openscvx.symbolic.preprocessing import resolve_guess_exprs
 from openscvx.symbolic.problem import SymbolicProblem
 from openscvx.utils import printing, profiling
 from openscvx.utils.caching import (
@@ -958,7 +959,8 @@ class Problem:
         This method reads the current `.guess` values from the original State and
         Control objects and updates the unified representations. This enables warm-starting
         workflows where the initial trajectory guess is updated between solves (e.g., shifting
-        the previous solution).
+        the previous solution). Guesses re-assigned symbolically between solves are
+        evaluated on the node grid first.
 
         Note:
             This only updates the unified representation. The AlgorithmState is
@@ -968,6 +970,8 @@ class Problem:
         if self._lowered is None:
             return
 
+        resolve_guess_exprs(self.symbolic.states, self.symbolic.controls, self.symbolic.N)
+
         # Sync optimization state guesses
         for state in self.symbolic.states:
             if state.guess is not None:

{openscvx-0.5.3.dev49 → openscvx-0.5.3.dev51}/openscvx/symbolic/builder.py

@@ -48,6 +48,7 @@ from openscvx.symbolic.preprocessing import (
     collect_and_assign_slices,
     convert_dynamics_dict_to_expr,
     fill_default_guesses,
+    resolve_guess_exprs,
     validate_and_normalize_constraint_nodes,
     validate_boundary_conditions,
     validate_bounds,
@@ -220,8 +221,9 @@ def preprocess_symbolic_problem(
         states = list(states) + [time]
         time_state = time
 
-    # Fill in default guess if needed (for Time instances)
-    if isinstance(time_state, Time) and time_state.guess is None:
+    # Fill in default guess if needed (for Time instances); a symbolic time guess
+    # is resolved in Phase 2, so it must not be overwritten here.
+    if isinstance(time_state, Time) and time_state.guess is None and time_state._guess_expr is None:
         time_state.guess = time_state._generate_default_guess(N)
 
     # Add CTCS constraints for time bounds
@@ -287,6 +289,9 @@ def preprocess_symbolic_problem(
     # ==================== PHASE 2: Expression Validation ====================
     # Validate all expressions (use unsorted constraints)
     collect_and_assign_slices(states, controls)
+    # Evaluate symbolic guesses now that slices exist and time has been added, and
+    # before augmentation derives the time-dilation guess from time's guess numerically.
+    resolve_guess_exprs(states, controls, N)
     validate_constraints_at_root(constraints.unsorted)
     validate_and_normalize_constraint_nodes(constraints.unsorted, N)
     validate_dynamics_dimension(dynamics_concat, states)
@@ -521,6 +526,17 @@ def add_propagation_states(
 
     # ==================== PHASE 1: Validate Extra States ====================
 
+    # Propagation-only states are integrated from their initial condition, never
+    # seeded from a guess trajectory, so a symbolic guess has nothing to resolve on.
+    for state in states_extra:
+        if state._guess_expr is not None:
+            raise ValueError(
+                f"State '{state.name}': symbolic initial guesses are not supported for "
+                f"propagation-only states. These states are integrated forward from "
+                f"their initial condition after the solve, so they carry no guess "
+                f"trajectory."
+            )
+
     # Validate that extra states don't conflict with optimization state names
     opt_state_names = {s.name for s in states_opt}
     extra_state_names = {s.name for s in states_extra}

{openscvx-0.5.3.dev49 → openscvx-0.5.3.dev51}/openscvx/symbolic/expr/__init__.py

@@ -33,6 +33,11 @@ Module Organization:
         Array manipulation operations including `Index`, `Concat`, `Stack`, `Hstack`,
         and `Vstack` for indexing, slicing, and combining arrays.
 
+    User-Supplied Jacobians (autodiff.py):
+        `WithJacobian`, built by `Expr.with_jacobian`, replaces a subexpression's
+        derivative with respect to chosen variables — the inexact-Jacobian knob for
+        SCP conditioning.
+
     Constraints (constraint.py):
         Constraint types including `Constraint`, `Equality`, `Inequality`,
         `NodalConstraint`, and `CTCS` (Continuous-Time Constraint Satisfaction).
@@ -80,6 +85,9 @@ from .arithmetic import Add, Div, MatMul, Mul, Neg, Power, Sub
 # Array operations
 from .array import Block, Concat, Hstack, Index, Stack, Vstack
 
+# User-supplied Jacobians
+from .autodiff import WithJacobian
+
 # Specialized constraints
 from .constraint import (
     CTCS,
@@ -258,4 +266,6 @@ __all__ = [
     "ctcs",
     # Vmap
     "Vmap",
+    # User-supplied Jacobians
+    "WithJacobian",
 ]

openscvx-0.5.3.dev51/openscvx/symbolic/expr/autodiff.py

@@ -0,0 +1,159 @@
+"""User-supplied Jacobians for a subexpression.
+
+Successive convexification linearizes the dynamics and constraints at every
+iteration, so the *Jacobian* — not the value — is what shapes each convex
+subproblem. Occasionally the exact Jacobian is the wrong one to hand the
+solver: an aerodynamic drag term, a contact model, or a lookup table can be
+smooth enough to integrate but stiff enough that its true derivative wrecks the
+conditioning of the subproblem. The standard remedy is a deliberately
+simplified (inexact) Jacobian — the trajectory still converges to a solution of
+the *nonlinear* problem, because the value used for the defect is untouched;
+only the search direction changes.
+
+:class:`WithJacobian` is that remedy, written at sub-expression granularity::
+
+    drag = -0.5 * rho * ox.Norm(vel) * vel
+    accel = thrust / m + drag.with_jacobian({vel: J_drag})
+
+The node wraps a subexpression and carries an override Jacobian per variable.
+Lowering to JAX emits a :func:`jax.custom_jvp` rule around the wrapped
+subexpression, so every downstream differentiation — the discretizer's
+``jacfwd``, constraint linearization, sparse coloring — sees the user's
+Jacobian in the overridden directions and autodiff everywhere else, with no
+further special-casing. The primal value is always the wrapped expression
+itself.
+"""
+
+import hashlib
+from typing import Dict, List, Tuple
+
+import numpy as np
+
+from .expr import Expr, to_expr
+from .variable import Variable
+
+
+def _without_unit_axes(shape: Tuple[int, ...]) -> Tuple[int, ...]:
+    """Drop length-1 axes, the way :class:`Constant` squeezes a wrapped array."""
+    return tuple(d for d in shape if d != 1)
+
+
+class WithJacobian(Expr):
+    """A subexpression whose derivative w.r.t. chosen variables is user-supplied.
+
+    Construct this through :meth:`Expr.with_jacobian` rather than directly. The
+    node evaluates exactly like the expression it wraps; only differentiation
+    differs, and only in the overridden directions (see the module docstring for
+    the motivation).
+
+    Attributes:
+        expr: The wrapped expression, and the node's value.
+        overrides: ``(variable, jacobian)`` pairs in the order given, each
+            Jacobian shaped ``(*expr.shape, *variable.shape)``.
+    """
+
+    def __init__(self, expr: Expr, jacobians: Dict[Variable, Expr]):
+        """Wrap ``expr`` with an override Jacobian per variable.
+
+        Args:
+            expr: Expression whose derivative is being overridden.
+            jacobians: Map from a ``State`` or ``Control`` to the Jacobian of
+                ``expr`` with respect to it. Values are coerced with
+                :func:`to_expr`, so constant matrices are accepted.
+
+        Raises:
+            TypeError: If a key is not a decision variable.
+            ValueError: If no Jacobian is given.
+        """
+        if not jacobians:
+            raise ValueError(
+                "with_jacobian needs at least one {variable: jacobian} entry; an empty "
+                "mapping would wrap the expression without changing anything."
+            )
+        self.expr = to_expr(expr)
+        overrides = []
+        for var, jac in jacobians.items():
+            if not isinstance(var, Variable):
+                raise TypeError(
+                    f"with_jacobian keys must be the State or Control the derivative is "
+                    f"taken with respect to, got {type(var).__name__}. Write "
+                    f"`expr.with_jacobian({{vel: J_vel}})` with `vel` the variable itself."
+                )
+            overrides.append((var, to_expr(jac)))
+        self.overrides: List[Tuple[Variable, Expr]] = overrides
+
+    def children(self) -> List[Expr]:
+        """Return the wrapped expression followed by each override Jacobian.
+
+        The overridden variables are not children: they identify *directions*,
+        not operands, and appear as leaves inside the Jacobian expressions only
+        if the user put them there.
+        """
+        return [self.expr, *(jac for _, jac in self.overrides)]
+
+    def canonicalize(self) -> Expr:
+        """Canonicalize the wrapped expression and every override Jacobian."""
+        return WithJacobian(
+            self.expr.canonicalize(),
+            {var: jac.canonicalize() for var, jac in self.overrides},
+        )
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Return the wrapped expression's shape, validating each Jacobian.
+
+        A Jacobian must be shaped ``(*expr.shape, *var.shape)``. Unit axes are
+        ignored in the comparison: a constant is squeezed when it is wrapped, so
+        an honest ``(3, 1)`` Jacobian arrives as ``(3,)``. The JAX lowering
+        reshapes it back before contracting.
+
+        Returns:
+            The shape of the wrapped expression — the wrapper is value-transparent.
+
+        Raises:
+            ValueError: If a Jacobian's shape is not ``(*expr.shape, *var.shape)``.
+        """
+        shape = self.expr.check_shape()
+        for var, jac in self.overrides:
+            expected = (*shape, *var.shape)
+            got = jac.check_shape()
+            if _without_unit_axes(got) != _without_unit_axes(expected):
+                raise ValueError(
+                    f"Jacobian for {var.__class__.__name__} '{var.name}' has shape {got}, "
+                    f"expected {expected}: the derivative of an expression of shape "
+                    f"{shape} with respect to a variable of shape {var.shape}."
+                )
+        return shape
+
+    def sparsity(self, n_x: int, n_u: int) -> Tuple[np.ndarray, np.ndarray]:
+        """Conservative pattern: the children's, widened to the overridden columns.
+
+        The default union over children is not enough here. A user Jacobian may
+        be a constant (no dependence of its own) yet still couple outputs to
+        every entry of the overridden variable, so those columns are marked
+        dense on top of the children's pattern.
+        """
+        from .control import Control
+
+        S_x, S_u = super().sparsity(n_x, n_u)
+        for var, _ in self.overrides:
+            if var._slice is None:
+                continue
+            S = S_u if isinstance(var, Control) else S_x
+            S[:, var._slice] = True
+        return S_x, S_u
+
+    def _hash_into(self, hasher: "hashlib._Hash") -> None:
+        """Hash the wrapper, the wrapped expression, and each override in order.
+
+        Args:
+            hasher: A hashlib hash object to update
+        """
+        hasher.update(b"WithJacobian")
+        self.expr._hash_into(hasher)
+        for var, jac in self.overrides:
+            var._hash_into(hasher)
+            jac._hash_into(hasher)
+
+    def __repr__(self) -> str:
+        vars_ = ", ".join(var.name for var, _ in self.overrides)
+        return f"{self.expr!r}.with_jacobian({{{vars_}}})"

{openscvx-0.5.3.dev49 → openscvx-0.5.3.dev51}/openscvx/symbolic/expr/expr.py

@@ -308,6 +308,47 @@ class Expr:
         """
         return NodeReference(self, k)
 
+    def with_jacobian(self, jacobians: dict) -> "Expr":
+        """Override this expression's derivative w.r.t. chosen variables.
+
+        Successive convexification linearizes at every iteration, so the Jacobian
+        of a term — not its value — is what conditions the convex subproblem. A
+        deliberately simplified (inexact) Jacobian is a standard remedy for a term
+        that is smooth enough to integrate but whose true derivative makes the
+        subproblem badly conditioned. The value of the expression is untouched, so
+        the converged trajectory still satisfies the original nonlinear problem;
+        only the search direction changes.
+
+        Directions that are not overridden are still differentiated automatically,
+        so a single term can hand-specify its state Jacobian and let the control
+        Jacobian fall out of autodiff.
+
+        Args:
+            jacobians: Map from a ``State`` or ``Control`` to the Jacobian of this
+                expression with respect to it, shaped ``(*self.shape, *var.shape)``
+                (a variable of shape ``(n,)`` and a scalar expression give shape
+                ``(n,)``). Constant matrices are accepted alongside expressions.
+
+        Returns:
+            Expr: A ``WithJacobian`` node wrapping this expression.
+
+        Example:
+            Damp the drag Jacobian without changing the drag force:
+
+                drag = -0.5 * rho * ox.Norm(vel) * vel
+                J_drag = -0.5 * rho * ox.Norm(vel) * np.eye(3)  # drops the d|v|/dv term
+
+                dynamics = {"vel": thrust / m + drag.with_jacobian({vel: J_drag})}
+
+        Note:
+            The override applies on the JAX path — dynamics linearization and
+            constraint linearization. Inside a constraint lowered to CVXPy the
+            wrapper has no meaning and raises.
+        """
+        from .autodiff import WithJacobian
+
+        return WithJacobian(self, jacobians)
+
     def children(self) -> List["Expr"]:
         """Return the child expressions of this node.
 

{openscvx-0.5.3.dev49 → openscvx-0.5.3.dev51}/openscvx/symbolic/expr/time.py

@@ -1,11 +1,57 @@
-from typing import Any, List, Optional, Union
+from typing import Any, List, Optional, Tuple, Union
 
 import numpy as np
 from pydantic import ConfigDict, field_validator
 
+from openscvx.symbolic.expr.expr import Expr, Leaf
 from openscvx.symbolic.expr.state import State, StateSpec
 from openscvx.symbolic.expr.variable import Variable
 
+#: Key under which the normalized node grid is passed to a lowered guess expression.
+TAU_PARAM = "__tau__"
+
+
+class Tau(Leaf):
+    """Normalized node coordinate, running 0 at the first node to 1 at the last.
+
+    ``Tau`` is internal plumbing for initial guesses, not part of the user-facing
+    vocabulary. A user who needs the node coordinate writes their guess as a
+    callable of one argument::
+
+        pos.guess = lambda tau: p0 + (pf - p0) * tau
+
+    and the library applies that callable once, at assignment, with a ``Tau``
+    instance. The result is an ordinary expression that the guess-resolution pass
+    evaluates node by node (see
+    :func:`openscvx.symbolic.preprocessing.resolve_guess_exprs`).
+
+    Physical time is a ``Time`` state; ``Tau`` is the dimensionless grid
+    coordinate the trajectory is discretized on, and is meaningful only while
+    building a guess.
+    """
+
+    def __init__(self):
+        super().__init__("tau", shape=())
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Reject tau outside of an initial-guess expression.
+
+        Raises:
+            ValueError: Always. Dynamics, constraints, and costs are shape-checked
+                eagerly, so this fires exactly when a captured ``tau`` has leaked
+                out of the guess it was created for.
+        """
+        raise ValueError(
+            "tau is only meaningful inside initial-guess expressions. It was received "
+            "by a `guess = lambda tau: ...` callable and then used in dynamics, a "
+            "constraint, or a cost, where there is no node grid to evaluate it on. "
+            "Use the `time` state for physical time, or index nodes explicitly with "
+            "`.at(k)`."
+        )
+
+    def __repr__(self) -> str:
+        return "Tau()"
+
 
 class Time(State):
     """Time state variable for trajectory optimization.
@@ -92,15 +138,17 @@ class Time(State):
             min: Minimum time bound.
             max: Maximum time bound.
             guess: Initial guess for the time trajectory. 1D array of shape
-                (N,) or 2D of shape (N, 1). If not provided, a linear
+                (N,), 2D of shape (N, 1), or any symbolic form accepted by
+                :attr:`Variable.guess`. If not provided, a linear
                 interpolation from initial to final is generated.
             time_dilation_min: Absolute minimum bound for time dilation.
                 Defaults to `0.3 * time_final` if not set.
             time_dilation_max: Absolute maximum bound for time dilation.
                 Defaults to `3.0 * time_final` if not set.
             time_dilation_guess: Initial guess for time dilation.
-                1D array of shape (N,) or 2D of shape (N, 1). Overrides
-                the default finite-difference computation.
+                1D array of shape (N,) or 2D of shape (N, 1) — arrays only,
+                symbolic guesses are not accepted here. Overrides the default
+                finite-difference computation.
             uniform_time_grid: If True, constrain the time dilation to be
                 the same across all nodes (uniform time steps). Defaults
                 to False.
@@ -125,11 +173,7 @@ class Time(State):
 
         super().__init__("time", shape=(1,), min=min, max=max, initial=initial, final=final)
 
-        # guess: normalize 1D → 2D before assigning
         if guess is not None:
-            guess = np.asarray(guess, dtype=float)
-            if guess.ndim == 1:
-                guess = guess.reshape(-1, 1)
             self.guess = guess
         if time_dilation_min is not None:
             self.time_dilation_min = time_dilation_min
@@ -167,12 +211,25 @@ class Time(State):
         State.final.fset(self, val)
 
     @Variable.guess.setter
-    def guess(self, arr):
-        """Set the time guess. Accepts 1D (N,) or 2D (N, 1) arrays."""
-        arr = np.asarray(arr, dtype=float)
-        if arr.ndim == 1:
-            arr = arr.reshape(-1, 1)
-        Variable.guess.fset(self, arr)
+    def guess(self, value):
+        """Set the time guess.
+
+        Accepts everything :attr:`Variable.guess` accepts — an array, a symbolic
+        expression, or a callable of tau — and additionally reshapes a 1D ``(N,)``
+        array to ``(N, 1)`` since Time is always shape ``(1,)``. A symbolic guess
+        resolving to one value per node is reshaped the same way at build time::
+
+            time.guess = lambda tau: t0 + (tf - t0) * tau**2
+
+        Note:
+            ``time_dilation_guess`` is a separate, array-only setting; it is not
+            resolved from expressions.
+        """
+        if not isinstance(value, Expr) and not callable(value):
+            value = np.asarray(value, dtype=float)
+            if value.ndim == 1:
+                value = value.reshape(-1, 1)
+        Variable.guess.fset(self, value)
 
     @property
     def time_dilation_min(self) -> Optional[float]:
@@ -199,7 +256,7 @@ class Time(State):
 
     @property
     def time_dilation_guess(self) -> Optional[np.ndarray]:
-        """Initial guess for time dilation, shape (N, 1)."""
+        """Initial guess for time dilation, shape (N, 1). Arrays only."""
         return self._time_dilation_guess
 
     @time_dilation_guess.setter