algorithms

Successive convexification algorithms for trajectory optimization.

This module provides implementations of SCvx (Successive Convexification) algorithms for solving non-convex trajectory optimization problems through iterative convex approximation.

All algorithms inherit from :class:Algorithm, enabling pluggable algorithm implementations and custom SCvx variants. An algorithm builds its fused, JAX-pure iteration body (build_iteration), stores it along with the diagnostics emitter (initialize), drives it one step at a time on the Python path (step), and owns its convergence policy (converged) — honored identically by solve(), solve_jax(), and solve_batched().

The iterate carry is split into two objects:

• :class:AlgorithmState — a frozen, JAX-registered pytree holding only the current iterate (x, u, weights, propagated states, diagnostic scalars). Every leaf is a jnp.ndarray so the state composes with jax.vmap / jax.jit / jax.grad.

• :class:AlgorithmHistory — a CPU-side mutable record of every iteration's trajectories, discretizations, weights, and diagnostics. Grown by the SCP loop via record_iteration; never crosses the JAX boundary.

Autotuners follow the same contract: :class:AutotuningBase declares update_weights as a pure functional update on the :class:AlgorithmState pytree (no mutation, no string returns, no list appends, no Python-level branching on iterate values). The next-iterate adaptive state is reported via the :class:AdaptiveStateCode IntEnum and converted to a human-readable string on the printing path with :func:adaptive_state_code_to_str.

Current Implementations:

• :class:PenalizedTrustRegion: Penalized Trust Region (PTR) algorithm

AcceptanceRatioAutotuner

Bases: AutotuningBase

Four-bucket acceptance-ratio update on lam_prox.

The acceptance ratio :math:\rho (actual over predicted nonlinear reduction) selects one of four buckets — REJECT / ACCEPT_HIGHER / ACCEPT_CONSTANT / ACCEPT_LOWER — via a jnp.where cascade so the whole update traces under jax.jit. Subclasses supply the multiplier rule (how lam_vc / lam_vb_* respond on acceptance) through :meth:_update_multipliers, and must assign an :class:AcceptanceRatioHyper (or subclass) instance to self.hyper in __init__ — the shared body reads its knobs off state.hyper.

update_weights is a pure functional update on the :class:AlgorithmState pytree; see the base-class contract.

Source code in openscvx/algorithms/autotuner/acceptance_ratio.py

class AcceptanceRatioAutotuner(AutotuningBase):
    """Four-bucket acceptance-ratio update on ``lam_prox``.

    The acceptance ratio :math:`\\rho` (actual over predicted nonlinear
    reduction) selects one of four buckets — ``REJECT`` / ``ACCEPT_HIGHER`` /
    ``ACCEPT_CONSTANT`` / ``ACCEPT_LOWER`` — via a ``jnp.where`` cascade so the
    whole update traces under ``jax.jit``. Subclasses supply the multiplier
    rule (how ``lam_vc`` / ``lam_vb_*`` respond on acceptance) through
    :meth:`_update_multipliers`, and must assign an
    :class:`AcceptanceRatioHyper` (or subclass) instance to ``self.hyper`` in
    ``__init__`` — the shared body reads its knobs off ``state.hyper``.

    ``update_weights`` is a pure functional update on the
    :class:`AlgorithmState` pytree; see the base-class contract.
    """

    COLUMNS: List[Column] = [
        Column("J_nonlin", "J_nonlin", 8, "{: .1e}", None, Verbosity.STANDARD),
        Column("J_lin", "J_lin", 8, "{: .1e}", None, Verbosity.STANDARD),
        Column("pred_reduction", "pred_red", 9, "{: .1e}", min_verbosity=Verbosity.FULL),
        Column("actual_reduction", "act_red", 9, "{: .1e}", min_verbosity=Verbosity.FULL),
        Column(
            "acceptance_ratio",
            "acc_ratio",
            9,
            "{: .2e}",
            color_acceptance_ratio,
            Verbosity.STANDARD,
        ),
        Column("lam_prox", "lam_prox", 8, "{: .1e}", min_verbosity=Verbosity.FULL),
        Column("adaptive_state", "Adaptive", 16, "{}", color_adaptive_state, Verbosity.FULL),
    ]

    def citation(self) -> List[str]:
        """Return BibTeX citations for the SCvx acceptance-ratio update rule."""
        return [
            r"""@inproceedings{mao2016scvx,
  title={Successive convexification of non-convex optimal control problems and its
    convergence properties},
  author={Mao, Yuanqi and Szmuk, Michael and A{\c{c}}{\i}kme{\c{s}}e, Beh{\c{c}}et},
  booktitle={2016 IEEE 55th Conference on Decision and Control (CDC)},
  pages={3636--3641},
  year={2016},
  publisher={IEEE}
}""",
            r"""@misc{mao2019scvx,
  title={Successive Convexification: A Superlinearly Convergent Algorithm for
    Non-convex Optimal Control Problems},
  author={Mao, Yuanqi and Szmuk, Michael and Xu, Xiangru and A{\c{c}}{\i}kme{\c{s}}e, Beh{\c{c}}et},
  year={2019},
  eprint={1804.06539},
  archivePrefix={arXiv},
  primaryClass={math.OC}
}""",
        ]

    def _update_multipliers(
        self,
        state: "AlgorithmState",
        candidate: "CandidateIterate",
        nodal_constraints: "LoweredJaxConstraints",
        params: dict,
        settings: Config,
        new_lam_prox: jnp.ndarray,
    ) -> dict:
        """Return the multiplier-weight updates to apply on an accepted iterate.

        Each returned entry maps an :class:`AlgorithmState` field name
        (``lam_vc`` / ``lam_vb_nodal`` / ``lam_vb_cross``) to its updated
        value, computed against the *candidate* trajectory; the shared body
        gates each by acceptance (a rejected iterate keeps the previous value).
        The default holds every multiplier constant — return ``{}``.
        """
        return {}

    def update_weights(
        self,
        state: "AlgorithmState",
        candidate: "CandidateIterate",
        nodal_constraints: "LoweredJaxConstraints",
        settings: Config,
        params: dict,
    ) -> "AlgorithmState":
        """Return the next-iterate state per the acceptance-ratio rules.

        Pure functional update — see class docstring.
        """
        candidate_x_prop = candidate.x_prop_plus[1:]
        nonlin_cost, nonlin_pen, nodal_pen = calculate_nonlinear_penalty(
            candidate_x_prop,
            candidate.x,
            candidate.u,
            state.lam_vc,
            state.lam_vb_nodal,
            state.lam_vb_cross,
            state.lam_cost,
            nodal_constraints,
            params,
            settings,
        )
        J_nonlin = nonlin_cost + nonlin_pen + nodal_pen

        lam_cost_next = self._relaxed_lam_cost(state)

        def first_iter(state):
            # Iter 1: accept unconditionally, leave weights at their init values,
            # only refresh trajectory + propagation fields.
            return state.replace(
                x=candidate.x,
                u=candidate.u,
                x_prop=candidate.x_prop,
                x_prop_plus=candidate.x_prop_plus,
                lam_cost=lam_cost_next,
                J_nonlin=J_nonlin,
                adaptive_state_code=jnp.asarray(int(AdaptiveStateCode.INITIAL), dtype=jnp.int32),
            )

        def later_iter(state):
            # Recompute the previous iterate's J_nonlin from the pytree fields
            # (state.x/state.u were the *previous* accepted iterate).
            prev_x_prop = state.x_prop_plus[1:]
            prev_cost, prev_pen, prev_nodal_pen = calculate_nonlinear_penalty(
                prev_x_prop,
                state.x,
                state.u,
                state.lam_vc,
                state.lam_vb_nodal,
                state.lam_vb_cross,
                state.lam_cost,
                nodal_constraints,
                params,
                settings,
            )
            prev_J_nonlin = prev_cost + prev_pen + prev_nodal_pen

            actual = prev_J_nonlin - J_nonlin
            predicted = prev_J_nonlin - candidate.J_lin
            # If predicted reduction is exactly zero, force the reject bucket
            # (rho = -inf) deterministically instead of raising.
            safe_pred = jnp.where(predicted == 0.0, 1.0, predicted)
            rho = jnp.where(predicted == 0.0, -jnp.inf, actual / safe_pred)

            is_reject = rho < state.hyper.eta_0
            is_accept_higher = (rho >= state.hyper.eta_0) & (rho < state.hyper.eta_1)
            is_accept_constant = (rho >= state.hyper.eta_1) & (rho < state.hyper.eta_2)
            # is_accept_lower implicit (else)
            accepted = ~is_reject

            # Compute both lam_prox candidates and gate.
            lp_higher = jnp.minimum(state.hyper.lam_prox_max, state.hyper.gamma_1 * state.lam_prox)
            lp_lower = jnp.maximum(state.hyper.lam_prox_min, state.hyper.gamma_2 * state.lam_prox)
            new_lam_prox = jnp.where(
                is_reject | is_accept_higher,
                lp_higher,
                jnp.where(is_accept_constant, state.lam_prox, lp_lower),
            )

            # Multiplier (virtual-control / virtual-buffer) updates: computed
            # against the *candidate* trajectory and gated by `accepted`. Reject
            # keeps the previous values (so the next subproblem doesn't see a
            # bumped penalty that wasn't earned). The default hook returns {} —
            # multipliers held constant.
            multiplier_updates = self._update_multipliers(
                state, candidate, nodal_constraints, params, settings, new_lam_prox
            )
            gated_multipliers = {
                name: jnp.where(accepted, value, getattr(state, name))
                for name, value in multiplier_updates.items()
            }

            code = jnp.where(
                is_reject,
                jnp.int32(AdaptiveStateCode.REJECT),
                jnp.where(
                    is_accept_higher,
                    jnp.int32(AdaptiveStateCode.ACCEPT_HIGHER),
                    jnp.where(
                        is_accept_constant,
                        jnp.int32(AdaptiveStateCode.ACCEPT_CONSTANT),
                        jnp.int32(AdaptiveStateCode.ACCEPT_LOWER),
                    ),
                ),
            )

            return state.replace(
                x=jnp.where(accepted, candidate.x, state.x),
                u=jnp.where(accepted, candidate.u, state.u),
                x_prop=jnp.where(accepted, candidate.x_prop, state.x_prop),
                x_prop_plus=jnp.where(accepted, candidate.x_prop_plus, state.x_prop_plus),
                lam_prox=new_lam_prox,
                lam_cost=lam_cost_next,
                J_nonlin=jnp.where(accepted, J_nonlin, state.J_nonlin),
                predicted_reduction=predicted,
                actual_reduction=actual,
                acceptance_ratio=rho,
                adaptive_state_code=code,
                **gated_multipliers,
            )

        return jax.lax.cond(state.k == 1, first_iter, later_iter, state)

citation() -> List[str]

Return BibTeX citations for the SCvx acceptance-ratio update rule.

Source code in openscvx/algorithms/autotuner/acceptance_ratio.py

    def citation(self) -> List[str]:
        """Return BibTeX citations for the SCvx acceptance-ratio update rule."""
        return [
            r"""@inproceedings{mao2016scvx,
  title={Successive convexification of non-convex optimal control problems and its
    convergence properties},
  author={Mao, Yuanqi and Szmuk, Michael and A{\c{c}}{\i}kme{\c{s}}e, Beh{\c{c}}et},
  booktitle={2016 IEEE 55th Conference on Decision and Control (CDC)},
  pages={3636--3641},
  year={2016},
  publisher={IEEE}
}""",
            r"""@misc{mao2019scvx,
  title={Successive Convexification: A Superlinearly Convergent Algorithm for
    Non-convex Optimal Control Problems},
  author={Mao, Yuanqi and Szmuk, Michael and Xu, Xiangru and A{\c{c}}{\i}kme{\c{s}}e, Beh{\c{c}}et},
  year={2019},
  eprint={1804.06539},
  archivePrefix={arXiv},
  primaryClass={math.OC}
}""",
        ]

update_weights(state: AlgorithmState, candidate: CandidateIterate, nodal_constraints: LoweredJaxConstraints, settings: Config, params: dict) -> AlgorithmState

Return the next-iterate state per the acceptance-ratio rules.

Pure functional update — see class docstring.

Source code in openscvx/algorithms/autotuner/acceptance_ratio.py

def update_weights(
    self,
    state: "AlgorithmState",
    candidate: "CandidateIterate",
    nodal_constraints: "LoweredJaxConstraints",
    settings: Config,
    params: dict,
) -> "AlgorithmState":
    """Return the next-iterate state per the acceptance-ratio rules.

    Pure functional update — see class docstring.
    """
    candidate_x_prop = candidate.x_prop_plus[1:]
    nonlin_cost, nonlin_pen, nodal_pen = calculate_nonlinear_penalty(
        candidate_x_prop,
        candidate.x,
        candidate.u,
        state.lam_vc,
        state.lam_vb_nodal,
        state.lam_vb_cross,
        state.lam_cost,
        nodal_constraints,
        params,
        settings,
    )
    J_nonlin = nonlin_cost + nonlin_pen + nodal_pen

    lam_cost_next = self._relaxed_lam_cost(state)

    def first_iter(state):
        # Iter 1: accept unconditionally, leave weights at their init values,
        # only refresh trajectory + propagation fields.
        return state.replace(
            x=candidate.x,
            u=candidate.u,
            x_prop=candidate.x_prop,
            x_prop_plus=candidate.x_prop_plus,
            lam_cost=lam_cost_next,
            J_nonlin=J_nonlin,
            adaptive_state_code=jnp.asarray(int(AdaptiveStateCode.INITIAL), dtype=jnp.int32),
        )

    def later_iter(state):
        # Recompute the previous iterate's J_nonlin from the pytree fields
        # (state.x/state.u were the *previous* accepted iterate).
        prev_x_prop = state.x_prop_plus[1:]
        prev_cost, prev_pen, prev_nodal_pen = calculate_nonlinear_penalty(
            prev_x_prop,
            state.x,
            state.u,
            state.lam_vc,
            state.lam_vb_nodal,
            state.lam_vb_cross,
            state.lam_cost,
            nodal_constraints,
            params,
            settings,
        )
        prev_J_nonlin = prev_cost + prev_pen + prev_nodal_pen

        actual = prev_J_nonlin - J_nonlin
        predicted = prev_J_nonlin - candidate.J_lin
        # If predicted reduction is exactly zero, force the reject bucket
        # (rho = -inf) deterministically instead of raising.
        safe_pred = jnp.where(predicted == 0.0, 1.0, predicted)
        rho = jnp.where(predicted == 0.0, -jnp.inf, actual / safe_pred)

        is_reject = rho < state.hyper.eta_0
        is_accept_higher = (rho >= state.hyper.eta_0) & (rho < state.hyper.eta_1)
        is_accept_constant = (rho >= state.hyper.eta_1) & (rho < state.hyper.eta_2)
        # is_accept_lower implicit (else)
        accepted = ~is_reject

        # Compute both lam_prox candidates and gate.
        lp_higher = jnp.minimum(state.hyper.lam_prox_max, state.hyper.gamma_1 * state.lam_prox)
        lp_lower = jnp.maximum(state.hyper.lam_prox_min, state.hyper.gamma_2 * state.lam_prox)
        new_lam_prox = jnp.where(
            is_reject | is_accept_higher,
            lp_higher,
            jnp.where(is_accept_constant, state.lam_prox, lp_lower),
        )

        # Multiplier (virtual-control / virtual-buffer) updates: computed
        # against the *candidate* trajectory and gated by `accepted`. Reject
        # keeps the previous values (so the next subproblem doesn't see a
        # bumped penalty that wasn't earned). The default hook returns {} —
        # multipliers held constant.
        multiplier_updates = self._update_multipliers(
            state, candidate, nodal_constraints, params, settings, new_lam_prox
        )
        gated_multipliers = {
            name: jnp.where(accepted, value, getattr(state, name))
            for name, value in multiplier_updates.items()
        }

        code = jnp.where(
            is_reject,
            jnp.int32(AdaptiveStateCode.REJECT),
            jnp.where(
                is_accept_higher,
                jnp.int32(AdaptiveStateCode.ACCEPT_HIGHER),
                jnp.where(
                    is_accept_constant,
                    jnp.int32(AdaptiveStateCode.ACCEPT_CONSTANT),
                    jnp.int32(AdaptiveStateCode.ACCEPT_LOWER),
                ),
            ),
        )

        return state.replace(
            x=jnp.where(accepted, candidate.x, state.x),
            u=jnp.where(accepted, candidate.u, state.u),
            x_prop=jnp.where(accepted, candidate.x_prop, state.x_prop),
            x_prop_plus=jnp.where(accepted, candidate.x_prop_plus, state.x_prop_plus),
            lam_prox=new_lam_prox,
            lam_cost=lam_cost_next,
            J_nonlin=jnp.where(accepted, J_nonlin, state.J_nonlin),
            predicted_reduction=predicted,
            actual_reduction=actual,
            acceptance_ratio=rho,
            adaptive_state_code=code,
            **gated_multipliers,
        )

    return jax.lax.cond(state.k == 1, first_iter, later_iter, state)

AcceptanceRatioHyper

Bases: LamCostRelaxHyper

Knobs shared by every acceptance-ratio autotuner.

Extends the shared :class:LamCostRelaxHyper lam_cost relaxation knobs with the acceptance-ratio decision thresholds and the lam_prox clips.

Source code in openscvx/algorithms/autotuner/acceptance_ratio.py

class AcceptanceRatioHyper(LamCostRelaxHyper):
    """Knobs shared by every acceptance-ratio autotuner.

    Extends the shared :class:`LamCostRelaxHyper` ``lam_cost`` relaxation knobs
    with the acceptance-ratio decision thresholds and the ``lam_prox`` clips.
    """

    gamma_1: float = 2.0
    gamma_2: float = 0.5
    eta_0: float = 1e-2
    eta_1: float = 1e-1
    eta_2: float = 0.8
    lam_prox_min: float = 1e-3
    lam_prox_max: float = 1e4

AdaptiveProximalWeight

Bases: AcceptanceRatioAutotuner

PTR-style proximal adaptation with fixed virtual-penalty weights.

Inherits the four-bucket acceptance-ratio update on lam_prox from :class:~openscvx.algorithms.autotuner.acceptance_ratio.AcceptanceRatioAutotuner and keeps the default multiplier hook, so the virtual-control and virtual-buffer weights are carried unchanged.

update_weights is a pure functional update on the :class:AlgorithmState pytree; see the base-class contract.

Source code in openscvx/algorithms/autotuner/adaptive_proximal_weight.py

class AdaptiveProximalWeight(AcceptanceRatioAutotuner):
    """PTR-style proximal adaptation with fixed virtual-penalty weights.

    Inherits the four-bucket acceptance-ratio update on ``lam_prox`` from
    :class:`~openscvx.algorithms.autotuner.acceptance_ratio.AcceptanceRatioAutotuner`
    and keeps the default multiplier hook, so the virtual-control and
    virtual-buffer weights are carried unchanged.

    ``update_weights`` is a pure functional update on the
    :class:`AlgorithmState` pytree; see the base-class contract.
    """

    def __init__(
        self,
        gamma_1: float = 2.0,
        gamma_2: float = 0.5,
        eta_0: float = 1e-2,
        eta_1: float = 1e-1,
        eta_2: float = 0.8,
        lam_prox_min: float = 1e-3,
        lam_prox_max: float = 1e4,
        lam_cost_drop: int = -1,
        lam_cost_relax: float = 1.0,
    ):
        self.hyper = AcceptanceRatioHyper(
            gamma_1=gamma_1,
            gamma_2=gamma_2,
            eta_0=eta_0,
            eta_1=eta_1,
            eta_2=eta_2,
            lam_prox_min=lam_prox_min,
            lam_prox_max=lam_prox_max,
            lam_cost_drop=lam_cost_drop,
            lam_cost_relax=lam_cost_relax,
        )

AdaptiveStateCode

Bases: IntEnum

Outcome of one autotuner update_weights call.

The autotuner stores its decision as an int32 on attr:AlgorithmState.adaptive_state_code so the pytree remains JAX-traceable. The SCP printing path maps the code back to a human-readable string via :func:adaptive_state_code_to_str.

Source code in openscvx/algorithms/state.py

class AdaptiveStateCode(IntEnum):
    """Outcome of one autotuner ``update_weights`` call.

    The autotuner stores its decision as an ``int32`` on
    :py:attr:`AlgorithmState.adaptive_state_code` so the pytree remains
    JAX-traceable. The SCP printing path maps the code back to a human-readable
    string via :func:`adaptive_state_code_to_str`.
    """

    REJECT = 0
    ACCEPT_HIGHER = 1
    ACCEPT_CONSTANT = 2
    ACCEPT_LOWER = 3
    INITIAL = 4

Algorithm

Bases: ABC

Abstract base class for successive convexification algorithms.

An Algorithm owns the SCP iteration end to end: it builds the JAX-pure iteration body (:meth:build_iteration), stores it (:meth:initialize), and drives it one step at a time (:meth:step). step takes the current :class:AlgorithmState and a CPU-side :class:AlgorithmHistory, returns the next state plus a convergence flag, and appends to history for diagnostics.

Statefullness

Avoid storing mutable iteration state on self. All iteration state lives on :class:AlgorithmState (JAX-traceable) and :class:AlgorithmHistory (Python-side), threaded explicitly through step().

The surface :class:~openscvx.problem.Problem relies on, beyond the abstract methods below:

• autotuner — the :class:AutotuningBase whose declared hyper fields :class:~openscvx.problem.Problem enumerates to build the per-solve / batched override channel.
• ep_tr / ep_vb / ep_vc — convergence thresholds :meth:Problem._sync_scp_constants snapshots onto :class:AlgorithmState before each solve.
• :meth:get_columns — the iteration-table columns to print.
• weights.build_vb_arrays — called once in :meth:Problem.initialize to size the virtual-buffer weight arrays.
• :meth:_hash_into — contributes the algorithm's identity to the solve_batched export cache key.

t_max is honored only on the Python solve() path; the JAX lax.while_loop behind solve_jax / solve_batched terminates on k_max and the convergence predicate alone (no wall-clock probe inside a trace).

Source code in openscvx/algorithms/base.py

class Algorithm(ABC):
    """Abstract base class for successive convexification algorithms.

    An ``Algorithm`` owns the SCP iteration end to end: it *builds* the
    JAX-pure iteration body (:meth:`build_iteration`), *stores* it
    (:meth:`initialize`), and *drives* it one step at a time (:meth:`step`).
    ``step`` takes the current :class:`AlgorithmState` and a CPU-side
    :class:`AlgorithmHistory`, returns the next state plus a convergence flag,
    and appends to ``history`` for diagnostics.

    !!! tip "Statefullness"
        Avoid storing mutable iteration state on ``self``. All iteration
        state lives on :class:`AlgorithmState` (JAX-traceable) and
        :class:`AlgorithmHistory` (Python-side), threaded explicitly through
        ``step()``.

    The surface :class:`~openscvx.problem.Problem` relies on, beyond the
    abstract methods below:

    - ``autotuner`` — the :class:`AutotuningBase` whose declared ``hyper``
      fields :class:`~openscvx.problem.Problem` enumerates to build the
      per-solve / batched override channel.
    - ``ep_tr`` / ``ep_vb`` / ``ep_vc`` — convergence thresholds
      :meth:`Problem._sync_scp_constants` snapshots onto
      :class:`AlgorithmState` before each solve.
    - :meth:`get_columns` — the iteration-table columns to print.
    - ``weights.build_vb_arrays`` — called once in
      :meth:`Problem.initialize` to size the virtual-buffer weight arrays.
    - :meth:`_hash_into` — contributes the algorithm's identity to the
      ``solve_batched`` export cache key.

    ``t_max`` is honored only on the Python ``solve()`` path; the JAX
    ``lax.while_loop`` behind ``solve_jax`` / ``solve_batched`` terminates on
    ``k_max`` and the convergence predicate alone (no wall-clock probe inside a
    trace).
    """

    def __init__(
        self,
        weights: "Weights",
        autotuner: "AutotuningBase",
        k_max: int,
        t_max: Optional[float],
        ep_tr: float,
        ep_vb: float,
        ep_vc: float,
    ):
        """Record the SCP weights, autotuner, and convergence parameters.

        Subclasses build the user-facing defaults (PTR owns them) and end their
        own ``__init__`` with ``super().__init__(...)``; every parameter here is
        required so the ABC is not a fourth place those defaults are declared.

        Args:
            weights: SCP weights used by the algorithm and autotuner.
            autotuner: The penalty-weight update rule (:class:`AutotuningBase`).
            k_max: Maximum number of SCP iterations.
            t_max: Optional wall-clock time limit in seconds (``solve()`` only;
                ``None`` means no limit).
            ep_tr: Trust-region convergence threshold.
            ep_vb: Virtual-buffer convergence threshold.
            ep_vc: Virtual-control convergence threshold.
        """
        self.weights = weights
        self.autotuner = autotuner
        self.k_max = k_max
        self.t_max = t_max
        self.ep_tr = ep_tr
        self.ep_vb = ep_vb
        self.ep_vc = ep_vc

    @abstractmethod
    def build_iteration(
        self,
        dis_continuous: Callable,
        dis_impulsive: Callable,
        jax_constraints: "LoweredJaxConstraints",
        solver_callback: Callable,
        settings: "Config",
    ) -> Callable:
        """Build the JAX-pure SCP iteration body for this algorithm.

        :class:`~openscvx.problem.Problem` assembles the discretization
        solvers, lowered constraints, and convex-solver callback, then asks
        the algorithm to fuse them into one step. The algorithm owns this
        because the fusion is algorithm-specific (which autotuner runs, which
        penalty terms are assembled); ``Problem`` stays algorithm-agnostic.

        Args:
            dis_continuous: Continuous-dynamics discretization solver,
                ``(x, u, params) -> (A_d, B_d, C_d, x_prop, V)``.
            dis_impulsive: Impulsive/discrete-dynamics discretization solver,
                ``(x_nodes, u, params) -> (x_prop_plus, D_d, E_d, W)``.
            jax_constraints: Lowered JAX constraints the body operates over.
            solver_callback: The convex backend's ``iteration_callback``,
                ``(state, SubproblemData) -> SubproblemSolution`` (see
                :class:`~openscvx.solvers.ptr_solver.PTRSolver`).
            settings: Problem configuration.

        Returns:
            The JAX-pure ``(state, params) -> (next_state, diagnostics)``
            iteration body. It must advance ``state.k`` by one per call —
            both ``Problem.solve()``'s Python loop and the ``lax.while_loop``
            behind ``solve_jax`` terminate on ``k`` reaching ``k_max``.
            :class:`~openscvx.problem.Problem` wraps the body in ``jax.jit``
            and hands it back via :meth:`initialize`.
        """
        raise NotImplementedError

    @abstractmethod
    def initialize(self, iteration_fn: Callable, emitter: Callable) -> None:
        """Store the fused SCP iteration body and per-iteration infrastructure.

        Args:
            iteration_fn: The JAX-pure ``(state, params) -> (next_state,
                diagnostics)`` body returned by :meth:`build_iteration`,
                already wrapped in ``jax.jit`` by
                :class:`~openscvx.problem.Problem`.
            emitter: Per-iteration diagnostics sink (printing queue / no-op).
        """
        raise NotImplementedError

    @abstractmethod
    def step(
        self,
        state: "AlgorithmState",
        history: "AlgorithmHistory",
        params: dict,
        settings: "Config",
    ) -> Tuple["AlgorithmState", bool]:
        """Execute one iteration.

        Args:
            state: Current-iterate pytree.
            history: CPU-side append-only iteration log.
            params: Problem parameters (may change between steps).
            settings: Configuration object (may change between steps).

        Returns:
            ``(next_state, converged)`` — the post-iteration state pytree
            and a flag indicating whether the SCP convergence criteria are
            satisfied.
        """
        raise NotImplementedError

    @abstractmethod
    def get_columns(self, verbosity: int) -> List[Column]:
        """Return the iteration-table columns to print at ``verbosity``.

        :class:`~openscvx.problem.Problem` calls this once per solve to lay
        out the progress table. Implementations typically concatenate the
        algorithm's own columns with the autotuner's ``COLUMNS`` and filter by
        each column's ``min_verbosity``.
        """
        raise NotImplementedError

    def converged(self, state: "AlgorithmState") -> jnp.ndarray:
        """Boolean SCP convergence test from the metrics and thresholds on ``state``.

        The default — every metric below its threshold — is algorithm-agnostic:
        :class:`AlgorithmState` carries ``J_tr`` / ``J_vb`` / ``J_vc`` and the
        ``ep_*`` tolerances generically, so it serves any SCP algorithm. Override
        to change the convergence policy; the override is honored on all three
        solve paths (``solve`` / ``solve_jax`` / ``solve_batched``), since
        :meth:`step` and the ``lax.while_loop`` harness both route through it.

        Must be JAX-traceable: it runs inside the ``lax.while_loop`` cond of
        ``solve_jax`` / ``solve_batched`` (where it is ``jax.vmap``'d per batch
        element) as well as on the Python ``solve()`` path.

        Under ``save_compiled``, the predicate is baked into the exported
        ``solve_batched`` artifact, but the cache key sees only the algorithm's
        class *name* (:meth:`_hash_into`): introducing an override — a new
        subclass — is a clean cache miss, while *editing the body* of an
        existing override is invisible to the key and silently loads the stale
        artifact. Clear the solver cache (or rename the class) when iterating
        on a predicate with ``save_compiled`` enabled.
        """
        return (state.J_tr < state.ep_tr) & (state.J_vb < state.ep_vb) & (state.J_vc < state.ep_vc)

    def _hash_into(self, hasher: "hashlib._Hash") -> None:
        """Contribute this algorithm's identity to the ``solve_batched`` cache key.

        The exported batched loop bakes in the initial penalty weights and the
        autotuner's update rule — none of which the symbolic problem hash
        covers. The default folds in the concrete class name, then every
        :class:`~openscvx.algorithms.weights.Weights` field (derived from the
        dataclass, so a new field is hashed automatically), then the autotuner
        via its own ``_hash_into`` — mirroring the symbolic ``_hash_into``
        protocol. The :class:`~openscvx.algorithms.state.AlgorithmState` field
        schema is folded in too: the state pytree is the artifact's input, its
        leaves bind positionally to the exported avals, and many are
        identically-shaped scalars — so a field reorder must be a clean cache
        miss, never a silent permutation (the same guard the autotuner applies
        to its ``hyper`` schema).

        The convergence thresholds (``ep_tr`` / ``ep_vb`` / ``ep_vc``) and the
        iteration cap ``k_max`` / time limit ``t_max`` are deliberately *not*
        folded in here: they ride the :class:`AlgorithmState` pytree as runtime
        inputs, so one exported artifact serves every tolerance and
        ``max_iters`` setting. Override only if an algorithm's identity needs
        something beyond ``weights`` + ``autotuner``.
        """
        from dataclasses import fields as dc_fields

        from openscvx.utils.caching import hash_value_into

        from .state import AlgorithmState

        hasher.update(type(self).__name__.encode())
        for fld in dc_fields(self.weights):
            hash_value_into(hasher, getattr(self.weights, fld.name))
        # Schema, not values: the AlgorithmState field order fixes the exported
        # artifact's input-leaf layout, so a reorder must invalidate cached
        # artifacts rather than silently permute runtime values onto the wrong
        # avals (mirrors the hyper-schema fold in AutotuningBase._hash_into).
        for name in AlgorithmState._FIELDS:
            hasher.update(name.encode())
        self.autotuner._hash_into(hasher)

    @abstractmethod
    def citation(self) -> List[str]:
        """Return BibTeX citations for this algorithm."""
        raise NotImplementedError

__init__(weights: Weights, autotuner: AutotuningBase, k_max: int, t_max: Optional[float], ep_tr: float, ep_vb: float, ep_vc: float)

Record the SCP weights, autotuner, and convergence parameters.

Subclasses build the user-facing defaults (PTR owns them) and end their own __init__ with super().__init__(...); every parameter here is required so the ABC is not a fourth place those defaults are declared.

Parameters:

Name	Type	Description	Default
weights	Weights	SCP weights used by the algorithm and autotuner.	required
autotuner	AutotuningBase	The penalty-weight update rule (:class:AutotuningBase).	required
k_max	int	Maximum number of SCP iterations.	required
t_max	Optional[float]	Optional wall-clock time limit in seconds (solve() only; None means no limit).	required
ep_tr	float	Trust-region convergence threshold.	required
ep_vb	float	Virtual-buffer convergence threshold.	required
ep_vc	float	Virtual-control convergence threshold.	required

Source code in openscvx/algorithms/base.py

def __init__(
    self,
    weights: "Weights",
    autotuner: "AutotuningBase",
    k_max: int,
    t_max: Optional[float],
    ep_tr: float,
    ep_vb: float,
    ep_vc: float,
):
    """Record the SCP weights, autotuner, and convergence parameters.

    Subclasses build the user-facing defaults (PTR owns them) and end their
    own ``__init__`` with ``super().__init__(...)``; every parameter here is
    required so the ABC is not a fourth place those defaults are declared.

    Args:
        weights: SCP weights used by the algorithm and autotuner.
        autotuner: The penalty-weight update rule (:class:`AutotuningBase`).
        k_max: Maximum number of SCP iterations.
        t_max: Optional wall-clock time limit in seconds (``solve()`` only;
            ``None`` means no limit).
        ep_tr: Trust-region convergence threshold.
        ep_vb: Virtual-buffer convergence threshold.
        ep_vc: Virtual-control convergence threshold.
    """
    self.weights = weights
    self.autotuner = autotuner
    self.k_max = k_max
    self.t_max = t_max
    self.ep_tr = ep_tr
    self.ep_vb = ep_vb
    self.ep_vc = ep_vc

build_iteration(dis_continuous: Callable, dis_impulsive: Callable, jax_constraints: LoweredJaxConstraints, solver_callback: Callable, settings: Config) -> Callable abstractmethod

Build the JAX-pure SCP iteration body for this algorithm.

:class:~openscvx.problem.Problem assembles the discretization solvers, lowered constraints, and convex-solver callback, then asks the algorithm to fuse them into one step. The algorithm owns this because the fusion is algorithm-specific (which autotuner runs, which penalty terms are assembled); Problem stays algorithm-agnostic.

Parameters:

Name	Type	Description	Default
dis_continuous	Callable	Continuous-dynamics discretization solver, (x, u, params) -> (A_d, B_d, C_d, x_prop, V).	required
dis_impulsive	Callable	Impulsive/discrete-dynamics discretization solver, (x_nodes, u, params) -> (x_prop_plus, D_d, E_d, W).	required
jax_constraints	LoweredJaxConstraints	Lowered JAX constraints the body operates over.	required
solver_callback	Callable	The convex backend's iteration_callback, (state, SubproblemData) -> SubproblemSolution (see :class:~openscvx.solvers.ptr_solver.PTRSolver).	required
settings	Config	Problem configuration.	required

Returns:

Type	Description
Callable	The JAX-pure (state, params) -> (next_state, diagnostics)
Callable	iteration body. It must advance state.k by one per call —
Callable	both Problem.solve()'s Python loop and the lax.while_loop
Callable	behind solve_jax terminate on k reaching k_max.
Callable	class:~openscvx.problem.Problem wraps the body in jax.jit
Callable	and hands it back via :meth:initialize.

Source code in openscvx/algorithms/base.py

@abstractmethod
def build_iteration(
    self,
    dis_continuous: Callable,
    dis_impulsive: Callable,
    jax_constraints: "LoweredJaxConstraints",
    solver_callback: Callable,
    settings: "Config",
) -> Callable:
    """Build the JAX-pure SCP iteration body for this algorithm.

    :class:`~openscvx.problem.Problem` assembles the discretization
    solvers, lowered constraints, and convex-solver callback, then asks
    the algorithm to fuse them into one step. The algorithm owns this
    because the fusion is algorithm-specific (which autotuner runs, which
    penalty terms are assembled); ``Problem`` stays algorithm-agnostic.

    Args:
        dis_continuous: Continuous-dynamics discretization solver,
            ``(x, u, params) -> (A_d, B_d, C_d, x_prop, V)``.
        dis_impulsive: Impulsive/discrete-dynamics discretization solver,
            ``(x_nodes, u, params) -> (x_prop_plus, D_d, E_d, W)``.
        jax_constraints: Lowered JAX constraints the body operates over.
        solver_callback: The convex backend's ``iteration_callback``,
            ``(state, SubproblemData) -> SubproblemSolution`` (see
            :class:`~openscvx.solvers.ptr_solver.PTRSolver`).
        settings: Problem configuration.

    Returns:
        The JAX-pure ``(state, params) -> (next_state, diagnostics)``
        iteration body. It must advance ``state.k`` by one per call —
        both ``Problem.solve()``'s Python loop and the ``lax.while_loop``
        behind ``solve_jax`` terminate on ``k`` reaching ``k_max``.
        :class:`~openscvx.problem.Problem` wraps the body in ``jax.jit``
        and hands it back via :meth:`initialize`.
    """
    raise NotImplementedError

citation() -> List[str] abstractmethod

Return BibTeX citations for this algorithm.

Source code in openscvx/algorithms/base.py

@abstractmethod
def citation(self) -> List[str]:
    """Return BibTeX citations for this algorithm."""
    raise NotImplementedError

converged(state: AlgorithmState) -> jnp.ndarray

Boolean SCP convergence test from the metrics and thresholds on state.

The default — every metric below its threshold — is algorithm-agnostic: :class:AlgorithmState carries J_tr / J_vb / J_vc and the ep_* tolerances generically, so it serves any SCP algorithm. Override to change the convergence policy; the override is honored on all three solve paths (solve / solve_jax / solve_batched), since :meth:step and the lax.while_loop harness both route through it.

Must be JAX-traceable: it runs inside the lax.while_loop cond of solve_jax / solve_batched (where it is jax.vmap'd per batch element) as well as on the Python solve() path.

Under save_compiled, the predicate is baked into the exported solve_batched artifact, but the cache key sees only the algorithm's class name (:meth:_hash_into): introducing an override — a new subclass — is a clean cache miss, while editing the body of an existing override is invisible to the key and silently loads the stale artifact. Clear the solver cache (or rename the class) when iterating on a predicate with save_compiled enabled.

Source code in openscvx/algorithms/base.py

def converged(self, state: "AlgorithmState") -> jnp.ndarray:
    """Boolean SCP convergence test from the metrics and thresholds on ``state``.

    The default — every metric below its threshold — is algorithm-agnostic:
    :class:`AlgorithmState` carries ``J_tr`` / ``J_vb`` / ``J_vc`` and the
    ``ep_*`` tolerances generically, so it serves any SCP algorithm. Override
    to change the convergence policy; the override is honored on all three
    solve paths (``solve`` / ``solve_jax`` / ``solve_batched``), since
    :meth:`step` and the ``lax.while_loop`` harness both route through it.

    Must be JAX-traceable: it runs inside the ``lax.while_loop`` cond of
    ``solve_jax`` / ``solve_batched`` (where it is ``jax.vmap``'d per batch
    element) as well as on the Python ``solve()`` path.

    Under ``save_compiled``, the predicate is baked into the exported
    ``solve_batched`` artifact, but the cache key sees only the algorithm's
    class *name* (:meth:`_hash_into`): introducing an override — a new
    subclass — is a clean cache miss, while *editing the body* of an
    existing override is invisible to the key and silently loads the stale
    artifact. Clear the solver cache (or rename the class) when iterating
    on a predicate with ``save_compiled`` enabled.
    """
    return (state.J_tr < state.ep_tr) & (state.J_vb < state.ep_vb) & (state.J_vc < state.ep_vc)

get_columns(verbosity: int) -> List[Column] abstractmethod

Return the iteration-table columns to print at verbosity.

:class:~openscvx.problem.Problem calls this once per solve to lay out the progress table. Implementations typically concatenate the algorithm's own columns with the autotuner's COLUMNS and filter by each column's min_verbosity.

Source code in openscvx/algorithms/base.py

@abstractmethod
def get_columns(self, verbosity: int) -> List[Column]:
    """Return the iteration-table columns to print at ``verbosity``.

    :class:`~openscvx.problem.Problem` calls this once per solve to lay
    out the progress table. Implementations typically concatenate the
    algorithm's own columns with the autotuner's ``COLUMNS`` and filter by
    each column's ``min_verbosity``.
    """
    raise NotImplementedError

initialize(iteration_fn: Callable, emitter: Callable) -> None abstractmethod

Store the fused SCP iteration body and per-iteration infrastructure.

Parameters:

Name	Type	Description	Default
iteration_fn	Callable	The JAX-pure (state, params) -> (next_state, diagnostics) body returned by :meth:build_iteration, already wrapped in jax.jit by :class:~openscvx.problem.Problem.	required
emitter	Callable	Per-iteration diagnostics sink (printing queue / no-op).	required

Source code in openscvx/algorithms/base.py

@abstractmethod
def initialize(self, iteration_fn: Callable, emitter: Callable) -> None:
    """Store the fused SCP iteration body and per-iteration infrastructure.

    Args:
        iteration_fn: The JAX-pure ``(state, params) -> (next_state,
            diagnostics)`` body returned by :meth:`build_iteration`,
            already wrapped in ``jax.jit`` by
            :class:`~openscvx.problem.Problem`.
        emitter: Per-iteration diagnostics sink (printing queue / no-op).
    """
    raise NotImplementedError

step(state: AlgorithmState, history: AlgorithmHistory, params: dict, settings: Config) -> Tuple[AlgorithmState, bool] abstractmethod

Execute one iteration.

Parameters:

Name	Type	Description	Default
state	AlgorithmState	Current-iterate pytree.	required
history	AlgorithmHistory	CPU-side append-only iteration log.	required
params	dict	Problem parameters (may change between steps).	required
settings	Config	Configuration object (may change between steps).	required

Returns:

Type	Description
AlgorithmState	(next_state, converged) — the post-iteration state pytree
bool	and a flag indicating whether the SCP convergence criteria are
Tuple[AlgorithmState, bool]	satisfied.

Source code in openscvx/algorithms/base.py

@abstractmethod
def step(
    self,
    state: "AlgorithmState",
    history: "AlgorithmHistory",
    params: dict,
    settings: "Config",
) -> Tuple["AlgorithmState", bool]:
    """Execute one iteration.

    Args:
        state: Current-iterate pytree.
        history: CPU-side append-only iteration log.
        params: Problem parameters (may change between steps).
        settings: Configuration object (may change between steps).

    Returns:
        ``(next_state, converged)`` — the post-iteration state pytree
        and a flag indicating whether the SCP convergence criteria are
        satisfied.
    """
    raise NotImplementedError

AlgorithmHistory dataclass

Append-only iteration log, populated by the SCP loop.

Mirrors the lists that previously lived directly on AlgorithmState. Never appears on the JAX boundary; the SCP loop appends to it after each iteration via meth:record_iteration. Purely diagnostic — the subproblem solver consumes SubproblemData in-trace, not the history.

Source code in openscvx/algorithms/history.py

@dataclass
class AlgorithmHistory:
    """Append-only iteration log, populated by the SCP loop.

    Mirrors the lists that previously lived directly on ``AlgorithmState``.
    Never appears on the JAX boundary; the SCP loop appends to it after each
    iteration via :py:meth:`record_iteration`. Purely diagnostic — the
    subproblem solver consumes ``SubproblemData`` in-trace, not the history.
    """

    n_x: int
    n_u: int
    N: int
    X: List[np.ndarray] = field(default_factory=list)
    U: List[np.ndarray] = field(default_factory=list)
    discretizations: List[DiscretizationResult] = field(default_factory=list)
    VC: List[np.ndarray] = field(default_factory=list)
    TR: List[np.ndarray] = field(default_factory=list)
    lam_prox: List[np.ndarray] = field(default_factory=list)
    lam_vc: List[Union[float, np.ndarray]] = field(default_factory=list)
    lam_cost: List[Union[float, np.ndarray]] = field(default_factory=list)
    lam_vb_nodal: List[np.ndarray] = field(default_factory=list)
    lam_vb_cross: List[np.ndarray] = field(default_factory=list)
    J_nonlin: List[float] = field(default_factory=list)
    J_lin: List[float] = field(default_factory=list)
    pred_reduction: List[float] = field(default_factory=list)
    actual_reduction: List[float] = field(default_factory=list)
    acceptance_ratio: List[float] = field(default_factory=list)
    adaptive_state: List[str] = field(default_factory=list)

    @classmethod
    def from_settings(cls, settings: "Config") -> "AlgorithmHistory":
        """Construct an empty history sized to the problem."""
        return cls(
            n_x=settings.sim.n_states,
            n_u=settings.sim.n_controls,
            N=settings.sim.n,
        )

    @property
    def V_history(self) -> List[np.ndarray]:
        """View of the raw discretization matrices recorded so far."""
        return [d.V for d in self.discretizations]

    # -- Per-iteration recording -------------------------------------------

    def record_iteration(
        self,
        state: AlgorithmState,
        *,
        V: Optional[np.ndarray] = None,
        W: Optional[np.ndarray] = None,
        VC: Optional[np.ndarray] = None,
        TR: Optional[np.ndarray] = None,
        J_lin: Optional[float] = None,
        record_diagnostics: bool = True,
    ) -> Tuple[dict, np.ndarray]:
        """Append per-iteration data based on ``state.adaptive_state_code``.

        The accepted iterate (``x`` / ``u`` / weights / convergence scalars) is
        read off ``state``; the raw host-side diagnostics the SCP loop already
        synced — discretization matrices (``V`` / ``W``), virtual control
        (``VC``), trust region (``TR``), and the linearized cost (``J_lin``) —
        come in as explicit keyword arrays so the generic history never has to
        know the algorithm's diagnostics type.

        Reproduces the old ``accept_solution`` / ``reject_solution`` behavior:

        * **REJECT**: append ``lam_prox`` (the bumped weight that drives the
          next subproblem), plus the predicted/actual/acceptance diagnostics
          when ``record_diagnostics`` is set.
        * **INITIAL**: append every trajectory / weight history entry, but
          skip the predicted/actual/acceptance diagnostics (the autotuner
          didn't compute them on iter 1).
        * **ACCEPT_***: append everything, including diagnostics when
          ``record_diagnostics`` is set (the SCP loop turns this off for
          autotuners that don't compute them — :class:`ConstantProximalWeight`,
          :class:`RampProximalWeight`).

        Every device-resident leaf this method needs is pulled with one
        ``jax.device_get`` call (a single CPU<->device round trip) and the
        host-side numpy arrays are then partitioned onto the history lists.
        The SCP loop also reads the same scalars for printing + convergence,
        so the bundle is returned alongside the lam_prox array to avoid a
        second sync.

        Returns:
            ``(scalars, lam_prox_np)`` — a dict of host-side scalar floats
            (``J_tr``, ``J_vb``, ``J_vc``, ``J_nonlin``,
            ``predicted_reduction``, ``actual_reduction``,
            ``acceptance_ratio``) plus the int ``adaptive_state_code``, and
            the numpy copy of ``state.lam_prox``. The SCP emitter consumes
            these directly without re-syncing.
        """
        # Coalesce every device read into one transfer. ``jax.device_get`` on
        # a tuple lets XLA dispatch the whole bundle as a single host copy,
        # which is the entire point on tiny problems where dozens of independent
        # ``float(state.scalar)`` calls would each issue a sync.
        leaves = jax.device_get(
            (
                state.adaptive_state_code,
                state.lam_prox,
                state.J_tr,
                state.J_vb,
                state.J_vc,
                state.J_nonlin,
                state.predicted_reduction,
                state.actual_reduction,
                state.acceptance_ratio,
                state.x,
                state.u,
                state.lam_vc,
                state.lam_cost,
                state.lam_vb_nodal,
                state.lam_vb_cross,
            )
        )
        (
            asc_np,
            lam_prox_np,
            J_tr_np,
            J_vb_np,
            J_vc_np,
            J_nonlin_np,
            pred_np,
            actual_np,
            ratio_np,
            x_np,
            u_np,
            lam_vc_np,
            lam_cost_np,
            lam_vb_nodal_np,
            lam_vb_cross_np,
        ) = leaves

        adaptive_code = int(asc_np)
        code = AdaptiveStateCode(adaptive_code)

        scalars = {
            "J_tr": float(J_tr_np),
            "J_vb": float(J_vb_np),
            "J_vc": float(J_vc_np),
            "J_nonlin": float(J_nonlin_np),
            "predicted_reduction": float(pred_np),
            "actual_reduction": float(actual_np),
            "acceptance_ratio": float(ratio_np),
            "adaptive_state_code": adaptive_code,
        }

        self.lam_prox.append(np.asarray(lam_prox_np))
        self.adaptive_state.append(adaptive_state_code_to_str(code))

        if code is AdaptiveStateCode.REJECT:
            if record_diagnostics:
                self.pred_reduction.append(scalars["predicted_reduction"])
                self.actual_reduction.append(scalars["actual_reduction"])
                self.acceptance_ratio.append(scalars["acceptance_ratio"])
            return scalars, lam_prox_np

        # INITIAL and any ACCEPT_*: full record of the accepted iterate.
        self.X.append(np.asarray(x_np))
        self.U.append(np.asarray(u_np))

        if V is not None:
            if W is not None:
                self.discretizations.append(
                    DiscretizationResult.from_VW(
                        V,
                        W,
                        n_x=self.n_x,
                        n_u=self.n_u,
                        N=self.N,
                    )
                )
            else:
                self.discretizations.append(
                    DiscretizationResult.from_V(
                        V,
                        n_x=self.n_x,
                        n_u=self.n_u,
                        N=self.N,
                    )
                )
        if VC is not None:
            self.VC.append(np.asarray(VC))
        if TR is not None:
            self.TR.append(np.asarray(TR))

        self.lam_vc.append(np.asarray(lam_vc_np))
        self.lam_cost.append(np.asarray(lam_cost_np))
        self.lam_vb_nodal.append(np.asarray(lam_vb_nodal_np))
        self.lam_vb_cross.append(np.asarray(lam_vb_cross_np))

        self.J_nonlin.append(scalars["J_nonlin"])
        if J_lin is not None:
            self.J_lin.append(float(J_lin))

        # Diagnostics: only meaningful for iterations after the initial one.
        if record_diagnostics and code is not AdaptiveStateCode.INITIAL:
            self.pred_reduction.append(scalars["predicted_reduction"])
            self.actual_reduction.append(scalars["actual_reduction"])
            self.acceptance_ratio.append(scalars["acceptance_ratio"])

        return scalars, lam_prox_np

V_history: List[np.ndarray] property

View of the raw discretization matrices recorded so far.

from_settings(settings: 'Config') -> 'AlgorithmHistory' classmethod

Construct an empty history sized to the problem.

Source code in openscvx/algorithms/history.py

@classmethod
def from_settings(cls, settings: "Config") -> "AlgorithmHistory":
    """Construct an empty history sized to the problem."""
    return cls(
        n_x=settings.sim.n_states,
        n_u=settings.sim.n_controls,
        N=settings.sim.n,
    )

record_iteration(state: AlgorithmState, *, V: Optional[np.ndarray] = None, W: Optional[np.ndarray] = None, VC: Optional[np.ndarray] = None, TR: Optional[np.ndarray] = None, J_lin: Optional[float] = None, record_diagnostics: bool = True) -> Tuple[dict, np.ndarray]

Append per-iteration data based on state.adaptive_state_code.

The accepted iterate (x / u / weights / convergence scalars) is read off state; the raw host-side diagnostics the SCP loop already synced — discretization matrices (V / W), virtual control (VC), trust region (TR), and the linearized cost (J_lin) — come in as explicit keyword arrays so the generic history never has to know the algorithm's diagnostics type.

Reproduces the old accept_solution / reject_solution behavior:

• REJECT: append lam_prox (the bumped weight that drives the next subproblem), plus the predicted/actual/acceptance diagnostics when record_diagnostics is set.
• INITIAL: append every trajectory / weight history entry, but skip the predicted/actual/acceptance diagnostics (the autotuner didn't compute them on iter 1).
• ACCEPT_*: append everything, including diagnostics when record_diagnostics is set (the SCP loop turns this off for autotuners that don't compute them — :class:ConstantProximalWeight, :class:RampProximalWeight).

Every device-resident leaf this method needs is pulled with one jax.device_get call (a single CPU<->device round trip) and the host-side numpy arrays are then partitioned onto the history lists. The SCP loop also reads the same scalars for printing + convergence, so the bundle is returned alongside the lam_prox array to avoid a second sync.

Returns:

Type	Description
dict	(scalars, lam_prox_np) — a dict of host-side scalar floats
ndarray	(J_tr, J_vb, J_vc, J_nonlin,
Tuple[dict, ndarray]	predicted_reduction, actual_reduction,
Tuple[dict, ndarray]	acceptance_ratio) plus the int adaptive_state_code, and
Tuple[dict, ndarray]	the numpy copy of state.lam_prox. The SCP emitter consumes
Tuple[dict, ndarray]	these directly without re-syncing.

Source code in openscvx/algorithms/history.py

def record_iteration(
    self,
    state: AlgorithmState,
    *,
    V: Optional[np.ndarray] = None,
    W: Optional[np.ndarray] = None,
    VC: Optional[np.ndarray] = None,
    TR: Optional[np.ndarray] = None,
    J_lin: Optional[float] = None,
    record_diagnostics: bool = True,
) -> Tuple[dict, np.ndarray]:
    """Append per-iteration data based on ``state.adaptive_state_code``.

    The accepted iterate (``x`` / ``u`` / weights / convergence scalars) is
    read off ``state``; the raw host-side diagnostics the SCP loop already
    synced — discretization matrices (``V`` / ``W``), virtual control
    (``VC``), trust region (``TR``), and the linearized cost (``J_lin``) —
    come in as explicit keyword arrays so the generic history never has to
    know the algorithm's diagnostics type.

    Reproduces the old ``accept_solution`` / ``reject_solution`` behavior:

    * **REJECT**: append ``lam_prox`` (the bumped weight that drives the
      next subproblem), plus the predicted/actual/acceptance diagnostics
      when ``record_diagnostics`` is set.
    * **INITIAL**: append every trajectory / weight history entry, but
      skip the predicted/actual/acceptance diagnostics (the autotuner
      didn't compute them on iter 1).
    * **ACCEPT_***: append everything, including diagnostics when
      ``record_diagnostics`` is set (the SCP loop turns this off for
      autotuners that don't compute them — :class:`ConstantProximalWeight`,
      :class:`RampProximalWeight`).

    Every device-resident leaf this method needs is pulled with one
    ``jax.device_get`` call (a single CPU<->device round trip) and the
    host-side numpy arrays are then partitioned onto the history lists.
    The SCP loop also reads the same scalars for printing + convergence,
    so the bundle is returned alongside the lam_prox array to avoid a
    second sync.

    Returns:
        ``(scalars, lam_prox_np)`` — a dict of host-side scalar floats
        (``J_tr``, ``J_vb``, ``J_vc``, ``J_nonlin``,
        ``predicted_reduction``, ``actual_reduction``,
        ``acceptance_ratio``) plus the int ``adaptive_state_code``, and
        the numpy copy of ``state.lam_prox``. The SCP emitter consumes
        these directly without re-syncing.
    """
    # Coalesce every device read into one transfer. ``jax.device_get`` on
    # a tuple lets XLA dispatch the whole bundle as a single host copy,
    # which is the entire point on tiny problems where dozens of independent
    # ``float(state.scalar)`` calls would each issue a sync.
    leaves = jax.device_get(
        (
            state.adaptive_state_code,
            state.lam_prox,
            state.J_tr,
            state.J_vb,
            state.J_vc,
            state.J_nonlin,
            state.predicted_reduction,
            state.actual_reduction,
            state.acceptance_ratio,
            state.x,
            state.u,
            state.lam_vc,
            state.lam_cost,
            state.lam_vb_nodal,
            state.lam_vb_cross,
        )
    )
    (
        asc_np,
        lam_prox_np,
        J_tr_np,
        J_vb_np,
        J_vc_np,
        J_nonlin_np,
        pred_np,
        actual_np,
        ratio_np,
        x_np,
        u_np,
        lam_vc_np,
        lam_cost_np,
        lam_vb_nodal_np,
        lam_vb_cross_np,
    ) = leaves

    adaptive_code = int(asc_np)
    code = AdaptiveStateCode(adaptive_code)

    scalars = {
        "J_tr": float(J_tr_np),
        "J_vb": float(J_vb_np),
        "J_vc": float(J_vc_np),
        "J_nonlin": float(J_nonlin_np),
        "predicted_reduction": float(pred_np),
        "actual_reduction": float(actual_np),
        "acceptance_ratio": float(ratio_np),
        "adaptive_state_code": adaptive_code,
    }

    self.lam_prox.append(np.asarray(lam_prox_np))
    self.adaptive_state.append(adaptive_state_code_to_str(code))

    if code is AdaptiveStateCode.REJECT:
        if record_diagnostics:
            self.pred_reduction.append(scalars["predicted_reduction"])
            self.actual_reduction.append(scalars["actual_reduction"])
            self.acceptance_ratio.append(scalars["acceptance_ratio"])
        return scalars, lam_prox_np

    # INITIAL and any ACCEPT_*: full record of the accepted iterate.
    self.X.append(np.asarray(x_np))
    self.U.append(np.asarray(u_np))

    if V is not None:
        if W is not None:
            self.discretizations.append(
                DiscretizationResult.from_VW(
                    V,
                    W,
                    n_x=self.n_x,
                    n_u=self.n_u,
                    N=self.N,
                )
            )
        else:
            self.discretizations.append(
                DiscretizationResult.from_V(
                    V,
                    n_x=self.n_x,
                    n_u=self.n_u,
                    N=self.N,
                )
            )
    if VC is not None:
        self.VC.append(np.asarray(VC))
    if TR is not None:
        self.TR.append(np.asarray(TR))

    self.lam_vc.append(np.asarray(lam_vc_np))
    self.lam_cost.append(np.asarray(lam_cost_np))
    self.lam_vb_nodal.append(np.asarray(lam_vb_nodal_np))
    self.lam_vb_cross.append(np.asarray(lam_vb_cross_np))

    self.J_nonlin.append(scalars["J_nonlin"])
    if J_lin is not None:
        self.J_lin.append(float(J_lin))

    # Diagnostics: only meaningful for iterations after the initial one.
    if record_diagnostics and code is not AdaptiveStateCode.INITIAL:
        self.pred_reduction.append(scalars["predicted_reduction"])
        self.actual_reduction.append(scalars["actual_reduction"])
        self.acceptance_ratio.append(scalars["acceptance_ratio"])

    return scalars, lam_prox_np

AlgorithmState dataclass

Current-iterate pytree for SCP iterations.

Every field is a JAX array (or scalar). The state moves between iterations by way of pure functional updates: meth:replace returns a new state rather than mutating in place. Registered as a JAX pytree so the SCP body composes with jax.jit / jax.vmap / jax.grad.

Append-only iteration histories (X, U, discretizations, ...) live on :class:AlgorithmHistory, which is CPU-side and grown by the SCP loop.

Attributes:

Name	Type	Description
x	ndarray	Nodal state trajectory, shape (N, n_states).
u	ndarray	Nodal control trajectory, shape (N, n_controls).
x_prop	ndarray	Continuous-time propagation of x, shape (N-1, n_states).
x_prop_plus	ndarray	Impulsive/discrete dynamics evaluated at x_prop, shape (N, n_states). Zeros when no impulsive component is present.
lam_prox	ndarray	Trust-region weight, shape (N, n_states + n_controls).
lam_vc	ndarray	Virtual-control penalty weight, shape (N-1, n_states).
lam_cost	ndarray	Cost weight, shape (n_states,).
lam_cost_init	ndarray	Initial cost weight (weights.lam_cost broadcast to (n_states,)). Used by autotuners as the "reset" value during early iterations and on the state.k <= hyper.lam_cost_drop branch. Lives on the pytree (not closure-captured) so weight mutations between solves propagate through the JIT'd update_weights.
lam_vb_nodal	ndarray	Nodal virtual-buffer weights, shape (N, n_nodal).
lam_vb_cross	ndarray	Cross-node virtual-buffer weights, shape (n_cross,).
k	ndarray	Iteration counter (starts at 1).
J_tr	ndarray	Current trust-region cost (scalar).
J_vb	ndarray	Current virtual-buffer cost (scalar).
J_vc	ndarray	Current virtual-control cost (scalar).
J_nonlin	ndarray	Nonlinear-objective value for the current accepted iterate.
predicted_reduction	ndarray	Predicted reduction in J_nonlin for this iter.
actual_reduction	ndarray	Actual reduction in J_nonlin for this iter.
acceptance_ratio	ndarray	actual_reduction / predicted_reduction.
adaptive_state_code	ndarray	:class:AdaptiveStateCode value as int32.
x_init_pin	ndarray	Initial-state boundary condition, shape (n_states,). jnp.nan where the state is not pinned at t0 (initial_type is not "Fix"). Carried on the pytree — rather than read from settings — so the fused SCP iteration body assembles the subproblem's initial boundary rows as a pure function of state, and a future jax.vmap over problems can batch boundary conditions per element.
x_term_pin	ndarray	Terminal-state boundary condition, shape (n_states,). jnp.nan where the state is not pinned at tf (final_type is not "Fix").
ep_tr	ndarray	Convergence threshold on J_tr (scalar). Carried on the pytree — like x_init_pin / lam_cost_init — so the SCP loop reads it as a runtime input: per-solve overrides and jax.vmap sweeps need no retrace, and one exported solve_batched artifact serves every tolerance setting.
ep_vb	ndarray	Convergence threshold on J_vb (scalar).
ep_vc	ndarray	Convergence threshold on J_vc (scalar).
k_max	ndarray	SCP iteration cap (scalar, k's integer dtype). The loop runs while k <= k_max; a traced bound is valid inside lax.while_loop, so the cap is per-solve and batchable too.
hyper	HyperParams	Autotuner-declared hyperparameters — an instance of the autotuner's :class:HyperParams subclass with scalar array leaves (the empty HyperParams() when it declares none). Seeded from :attr:AutotuningBase.hyper; update_weights reads its knobs here (e.g. state.hyper.lam_cost_drop) so each is a per-solve override and a batchable sweep target like any other field.

Source code in openscvx/algorithms/state.py

@jax.tree_util.register_pytree_node_class
@dataclass(frozen=True)
class AlgorithmState:
    """Current-iterate pytree for SCP iterations.

    Every field is a JAX array (or scalar). The state moves between iterations
    by way of pure functional updates: :py:meth:`replace` returns a new state
    rather than mutating in place. Registered as a JAX pytree so the SCP body
    composes with ``jax.jit`` / ``jax.vmap`` / ``jax.grad``.

    Append-only iteration histories (X, U, discretizations, ...) live on
    :class:`AlgorithmHistory`, which is CPU-side and grown by the SCP loop.

    Attributes:
        x: Nodal state trajectory, shape ``(N, n_states)``.
        u: Nodal control trajectory, shape ``(N, n_controls)``.
        x_prop: Continuous-time propagation of ``x``, shape ``(N-1, n_states)``.
        x_prop_plus: Impulsive/discrete dynamics evaluated at ``x_prop``,
            shape ``(N, n_states)``. Zeros when no impulsive component is
            present.
        lam_prox: Trust-region weight, shape ``(N, n_states + n_controls)``.
        lam_vc: Virtual-control penalty weight, shape ``(N-1, n_states)``.
        lam_cost: Cost weight, shape ``(n_states,)``.
        lam_cost_init: Initial cost weight (``weights.lam_cost`` broadcast to
            ``(n_states,)``). Used by autotuners as the "reset" value during
            early iterations and on the ``state.k <= hyper.lam_cost_drop``
            branch. Lives on the pytree (not closure-captured) so weight
            mutations between solves propagate through the JIT'd
            update_weights.
        lam_vb_nodal: Nodal virtual-buffer weights, shape ``(N, n_nodal)``.
        lam_vb_cross: Cross-node virtual-buffer weights, shape ``(n_cross,)``.
        k: Iteration counter (starts at 1).
        J_tr: Current trust-region cost (scalar).
        J_vb: Current virtual-buffer cost (scalar).
        J_vc: Current virtual-control cost (scalar).
        J_nonlin: Nonlinear-objective value for the current accepted iterate.
        predicted_reduction: Predicted reduction in ``J_nonlin`` for this iter.
        actual_reduction: Actual reduction in ``J_nonlin`` for this iter.
        acceptance_ratio: ``actual_reduction / predicted_reduction``.
        adaptive_state_code: :class:`AdaptiveStateCode` value as ``int32``.
        x_init_pin: Initial-state boundary condition, shape ``(n_states,)``.
            ``jnp.nan`` where the state is not pinned at t0 (``initial_type``
            is not ``"Fix"``). Carried on the pytree — rather than read from
            ``settings`` — so the fused SCP iteration body assembles the
            subproblem's initial boundary rows as a pure function of state,
            and a future ``jax.vmap`` over problems can batch boundary
            conditions per element.
        x_term_pin: Terminal-state boundary condition, shape ``(n_states,)``.
            ``jnp.nan`` where the state is not pinned at tf (``final_type`` is
            not ``"Fix"``).
        ep_tr: Convergence threshold on ``J_tr`` (scalar). Carried on the
            pytree — like ``x_init_pin`` / ``lam_cost_init`` — so the SCP loop
            reads it as a runtime input: per-solve overrides and ``jax.vmap``
            sweeps need no retrace, and one exported ``solve_batched`` artifact
            serves every tolerance setting.
        ep_vb: Convergence threshold on ``J_vb`` (scalar).
        ep_vc: Convergence threshold on ``J_vc`` (scalar).
        k_max: SCP iteration cap (scalar, ``k``'s integer dtype). The loop
            runs while ``k <= k_max``; a traced bound is valid inside
            ``lax.while_loop``, so the cap is per-solve and batchable too.
        hyper: Autotuner-declared hyperparameters — an instance of the
            autotuner's :class:`HyperParams` subclass with scalar array
            leaves (the empty ``HyperParams()`` when it declares none).
            Seeded from :attr:`AutotuningBase.hyper`; ``update_weights``
            reads its knobs here (e.g. ``state.hyper.lam_cost_drop``) so
            each is a per-solve override and a batchable sweep target like
            any other field.
    """

    x: jnp.ndarray
    u: jnp.ndarray
    x_prop: jnp.ndarray
    x_prop_plus: jnp.ndarray
    lam_prox: jnp.ndarray
    lam_vc: jnp.ndarray
    lam_cost: jnp.ndarray
    lam_cost_init: jnp.ndarray
    lam_vb_nodal: jnp.ndarray
    lam_vb_cross: jnp.ndarray
    k: jnp.ndarray
    J_tr: jnp.ndarray
    J_vb: jnp.ndarray
    J_vc: jnp.ndarray
    J_nonlin: jnp.ndarray
    predicted_reduction: jnp.ndarray
    actual_reduction: jnp.ndarray
    acceptance_ratio: jnp.ndarray
    adaptive_state_code: jnp.ndarray
    x_init_pin: jnp.ndarray
    x_term_pin: jnp.ndarray
    ep_tr: jnp.ndarray
    ep_vb: jnp.ndarray
    ep_vc: jnp.ndarray
    k_max: jnp.ndarray
    hyper: HyperParams

    def replace(self, **changes) -> "AlgorithmState":
        """Return a new state with ``changes`` applied (functional update)."""
        return dc_replace(self, **changes)

    def tree_flatten(self):
        children = tuple(getattr(self, name) for name in self._FIELDS)
        return children, None

    @classmethod
    def tree_unflatten(cls, aux, children):
        return cls(**dict(zip(cls._FIELDS, children)))

    @classmethod
    def from_settings(
        cls,
        settings: "Config",
        weights: "Weights",
        *,
        ep_tr: float,
        ep_vb: float,
        ep_vc: float,
        k_max: int,
        hyper: HyperParams,
    ) -> "AlgorithmState":
        """Construct the initial iterate from configuration.

        Copies trajectory guesses, expands weights to dense arrays, and seeds
        the diagnostic scalars to zero / :py:attr:`AdaptiveStateCode.INITIAL`.
        ``x_prop`` and ``x_prop_plus`` are zero-initialized; the autotuner
        fills them when it carries the first accepted candidate onto the
        state (see :meth:`AutotuningBase.update_weights`).

        The SCP loop constants (``ep_tr`` / ``ep_vb`` / ``ep_vc`` / ``k_max``)
        and the autotuner's declared hyperparameters (*hyper*, from
        :attr:`AutotuningBase.hyper`) are snapshotted onto the pytree here —
        the caller passes them off the algorithm and its autotuner (see
        ``Problem._default_state``), and per-solve overrides land via
        ``state.replace``.
        """
        shadowed = sorted({fld.name for fld in dc_fields(hyper)} & set(cls._FIELDS))
        if shadowed:
            raise ValueError(
                f"Autotuner hyperparameter(s) {shadowed} shadow AlgorithmState "
                f"fields; rename the HyperParams fields so overrides stay "
                f"unambiguous."
            )
        n = settings.sim.n
        n_states = settings.sim.n_states
        n_controls = settings.sim.n_controls
        n_total = n_states + n_controls

        # Strong-typed dtypes so the initial state matches the autotuner's
        # outputs. JAX caches `jit` traces by argument weak/strong dtype AND
        # committed sharding; a mismatch between the seed state and the
        # post-iter-1 state would trigger an extra recompile inside the SCP
        # loop. We route every leaf through ``jax.device_put(..., device)`` to
        # produce *committed* arrays so the cache key matches what comes back
        # from the JIT'd ``update_weights``.
        #
        # Removable once the SCP loop body is a single JAX trace — at that
        # point the cache-key match is irrelevant and the strong-dtype /
        # committed-sharding routing below can collapse to plain ``jnp.asarray``.
        f = jnp.float64 if jax.config.read("jax_enable_x64") else jnp.float32
        i = jnp.int32
        device = jax.devices()[0]

        def put(arr):
            return jax.device_put(arr, device)

        def hyper_leaf(fld):
            # int knobs (e.g. lam_cost_drop) share k's dtype so comparisons
            # against the iteration counter don't promote; floats follow the
            # problem's float dtype. The annotation is the rule — validated
            # when the HyperParams subclass is defined.
            dtype = i if fld.type in (int, "int") else f
            return put(jnp.asarray(getattr(hyper, fld.name), dtype=dtype))

        lam_vc_array = np.ones((n - 1, n_states)) * weights.lam_vc
        lam_prox_array = np.ones((n, n_total)) * weights.lam_prox

        if isinstance(weights.lam_cost, np.ndarray):
            lam_cost_init = weights.lam_cost.copy()
        else:
            lam_cost_init = np.full(n_states, weights.lam_cost)

        # Boundary-condition pins: the physical value where the state is fixed
        # at t0 / tf, ``nan`` elsewhere. The subproblem only reads these at
        # ``"Fix"`` entries, so the sentinel marks "unpinned" without poisoning
        # any value the solver consumes.
        x_initial = np.asarray(settings.sim.x.initial, dtype=float).reshape(-1)
        x_final = np.asarray(settings.sim.x.final, dtype=float).reshape(-1)
        init_fixed = np.asarray(settings.sim.x.initial_type) == "Fix"
        final_fixed = np.asarray(settings.sim.x.final_type) == "Fix"
        x_init_pin = np.where(init_fixed, x_initial, np.nan)
        x_term_pin = np.where(final_fixed, x_final, np.nan)

        return cls(
            x=put(jnp.asarray(settings.sim.x.guess, dtype=f)),
            u=put(jnp.asarray(settings.sim.u.guess, dtype=f)),
            x_prop=put(jnp.zeros((n - 1, n_states), dtype=f)),
            x_prop_plus=put(jnp.zeros((n, n_states), dtype=f)),
            lam_prox=put(jnp.asarray(lam_prox_array, dtype=f)),
            lam_vc=put(jnp.asarray(lam_vc_array, dtype=f)),
            lam_cost=put(jnp.asarray(lam_cost_init, dtype=f)),
            lam_cost_init=put(jnp.asarray(lam_cost_init, dtype=f)),
            lam_vb_nodal=put(jnp.asarray(weights.lam_vb_nodal, dtype=f)),
            lam_vb_cross=put(jnp.asarray(weights.lam_vb_cross, dtype=f)),
            k=put(jnp.asarray(1, dtype=i)),
            J_tr=put(jnp.asarray(1e2, dtype=f)),
            J_vb=put(jnp.asarray(1e2, dtype=f)),
            J_vc=put(jnp.asarray(1e2, dtype=f)),
            J_nonlin=put(jnp.asarray(0.0, dtype=f)),
            predicted_reduction=put(jnp.asarray(0.0, dtype=f)),
            actual_reduction=put(jnp.asarray(0.0, dtype=f)),
            acceptance_ratio=put(jnp.asarray(0.0, dtype=f)),
            adaptive_state_code=put(jnp.asarray(int(AdaptiveStateCode.INITIAL), dtype=i)),
            x_init_pin=put(jnp.asarray(x_init_pin, dtype=f)),
            x_term_pin=put(jnp.asarray(x_term_pin, dtype=f)),
            ep_tr=put(jnp.asarray(ep_tr, dtype=f)),
            ep_vb=put(jnp.asarray(ep_vb, dtype=f)),
            ep_vc=put(jnp.asarray(ep_vc, dtype=f)),
            # k_max shares k's dtype so `k <= k_max` compares without promotion.
            k_max=put(jnp.asarray(k_max, dtype=i)),
            hyper=dc_replace(hyper, **{fld.name: hyper_leaf(fld) for fld in dc_fields(hyper)}),
        )

from_settings(settings: Config, weights: Weights, *, ep_tr: float, ep_vb: float, ep_vc: float, k_max: int, hyper: HyperParams) -> AlgorithmState classmethod

Construct the initial iterate from configuration.

Copies trajectory guesses, expands weights to dense arrays, and seeds the diagnostic scalars to zero / attr:AdaptiveStateCode.INITIAL. x_prop and x_prop_plus are zero-initialized; the autotuner fills them when it carries the first accepted candidate onto the state (see :meth:AutotuningBase.update_weights).

The SCP loop constants (ep_tr / ep_vb / ep_vc / k_max) and the autotuner's declared hyperparameters (hyper, from :attr:AutotuningBase.hyper) are snapshotted onto the pytree here — the caller passes them off the algorithm and its autotuner (see Problem._default_state), and per-solve overrides land via state.replace.

Source code in openscvx/algorithms/state.py

@classmethod
def from_settings(
    cls,
    settings: "Config",
    weights: "Weights",
    *,
    ep_tr: float,
    ep_vb: float,
    ep_vc: float,
    k_max: int,
    hyper: HyperParams,
) -> "AlgorithmState":
    """Construct the initial iterate from configuration.

    Copies trajectory guesses, expands weights to dense arrays, and seeds
    the diagnostic scalars to zero / :py:attr:`AdaptiveStateCode.INITIAL`.
    ``x_prop`` and ``x_prop_plus`` are zero-initialized; the autotuner
    fills them when it carries the first accepted candidate onto the
    state (see :meth:`AutotuningBase.update_weights`).

    The SCP loop constants (``ep_tr`` / ``ep_vb`` / ``ep_vc`` / ``k_max``)
    and the autotuner's declared hyperparameters (*hyper*, from
    :attr:`AutotuningBase.hyper`) are snapshotted onto the pytree here —
    the caller passes them off the algorithm and its autotuner (see
    ``Problem._default_state``), and per-solve overrides land via
    ``state.replace``.
    """
    shadowed = sorted({fld.name for fld in dc_fields(hyper)} & set(cls._FIELDS))
    if shadowed:
        raise ValueError(
            f"Autotuner hyperparameter(s) {shadowed} shadow AlgorithmState "
            f"fields; rename the HyperParams fields so overrides stay "
            f"unambiguous."
        )
    n = settings.sim.n
    n_states = settings.sim.n_states
    n_controls = settings.sim.n_controls
    n_total = n_states + n_controls

    # Strong-typed dtypes so the initial state matches the autotuner's
    # outputs. JAX caches `jit` traces by argument weak/strong dtype AND
    # committed sharding; a mismatch between the seed state and the
    # post-iter-1 state would trigger an extra recompile inside the SCP
    # loop. We route every leaf through ``jax.device_put(..., device)`` to
    # produce *committed* arrays so the cache key matches what comes back
    # from the JIT'd ``update_weights``.
    #
    # Removable once the SCP loop body is a single JAX trace — at that
    # point the cache-key match is irrelevant and the strong-dtype /
    # committed-sharding routing below can collapse to plain ``jnp.asarray``.
    f = jnp.float64 if jax.config.read("jax_enable_x64") else jnp.float32
    i = jnp.int32
    device = jax.devices()[0]

    def put(arr):
        return jax.device_put(arr, device)

    def hyper_leaf(fld):
        # int knobs (e.g. lam_cost_drop) share k's dtype so comparisons
        # against the iteration counter don't promote; floats follow the
        # problem's float dtype. The annotation is the rule — validated
        # when the HyperParams subclass is defined.
        dtype = i if fld.type in (int, "int") else f
        return put(jnp.asarray(getattr(hyper, fld.name), dtype=dtype))

    lam_vc_array = np.ones((n - 1, n_states)) * weights.lam_vc
    lam_prox_array = np.ones((n, n_total)) * weights.lam_prox

    if isinstance(weights.lam_cost, np.ndarray):
        lam_cost_init = weights.lam_cost.copy()
    else:
        lam_cost_init = np.full(n_states, weights.lam_cost)

    # Boundary-condition pins: the physical value where the state is fixed
    # at t0 / tf, ``nan`` elsewhere. The subproblem only reads these at
    # ``"Fix"`` entries, so the sentinel marks "unpinned" without poisoning
    # any value the solver consumes.
    x_initial = np.asarray(settings.sim.x.initial, dtype=float).reshape(-1)
    x_final = np.asarray(settings.sim.x.final, dtype=float).reshape(-1)
    init_fixed = np.asarray(settings.sim.x.initial_type) == "Fix"
    final_fixed = np.asarray(settings.sim.x.final_type) == "Fix"
    x_init_pin = np.where(init_fixed, x_initial, np.nan)
    x_term_pin = np.where(final_fixed, x_final, np.nan)

    return cls(
        x=put(jnp.asarray(settings.sim.x.guess, dtype=f)),
        u=put(jnp.asarray(settings.sim.u.guess, dtype=f)),
        x_prop=put(jnp.zeros((n - 1, n_states), dtype=f)),
        x_prop_plus=put(jnp.zeros((n, n_states), dtype=f)),
        lam_prox=put(jnp.asarray(lam_prox_array, dtype=f)),
        lam_vc=put(jnp.asarray(lam_vc_array, dtype=f)),
        lam_cost=put(jnp.asarray(lam_cost_init, dtype=f)),
        lam_cost_init=put(jnp.asarray(lam_cost_init, dtype=f)),
        lam_vb_nodal=put(jnp.asarray(weights.lam_vb_nodal, dtype=f)),
        lam_vb_cross=put(jnp.asarray(weights.lam_vb_cross, dtype=f)),
        k=put(jnp.asarray(1, dtype=i)),
        J_tr=put(jnp.asarray(1e2, dtype=f)),
        J_vb=put(jnp.asarray(1e2, dtype=f)),
        J_vc=put(jnp.asarray(1e2, dtype=f)),
        J_nonlin=put(jnp.asarray(0.0, dtype=f)),
        predicted_reduction=put(jnp.asarray(0.0, dtype=f)),
        actual_reduction=put(jnp.asarray(0.0, dtype=f)),
        acceptance_ratio=put(jnp.asarray(0.0, dtype=f)),
        adaptive_state_code=put(jnp.asarray(int(AdaptiveStateCode.INITIAL), dtype=i)),
        x_init_pin=put(jnp.asarray(x_init_pin, dtype=f)),
        x_term_pin=put(jnp.asarray(x_term_pin, dtype=f)),
        ep_tr=put(jnp.asarray(ep_tr, dtype=f)),
        ep_vb=put(jnp.asarray(ep_vb, dtype=f)),
        ep_vc=put(jnp.asarray(ep_vc, dtype=f)),
        # k_max shares k's dtype so `k <= k_max` compares without promotion.
        k_max=put(jnp.asarray(k_max, dtype=i)),
        hyper=dc_replace(hyper, **{fld.name: hyper_leaf(fld) for fld in dc_fields(hyper)}),
    )

replace(**changes) -> AlgorithmState

Return a new state with changes applied (functional update).

Source code in openscvx/algorithms/state.py

def replace(self, **changes) -> "AlgorithmState":
    """Return a new state with ``changes`` applied (functional update)."""
    return dc_replace(self, **changes)

AugmentedLagrangian

Bases: AcceptanceRatioAutotuner

Augmented Lagrangian autotuner.

Uses an acceptance-ratio :math:\rho between predicted and actual reduction in the nonlinear objective to drive a four-bucket update of the trust-region weight lam_prox (inherited from :class:~openscvx.algorithms.autotuner.acceptance_ratio.AcceptanceRatioAutotuner) and to update the virtual-control / virtual-buffer weights from constraint violations via the multiplier hook.

update_weights is a pure functional update on the :class:AlgorithmState pytree; see the base-class contract.

Source code in openscvx/algorithms/autotuner/augmented_lagrangian.py

class AugmentedLagrangian(AcceptanceRatioAutotuner):
    """Augmented Lagrangian autotuner.

    Uses an acceptance-ratio :math:`\\rho` between predicted and actual
    reduction in the nonlinear objective to drive a four-bucket update of
    the trust-region weight ``lam_prox`` (inherited from
    :class:`~openscvx.algorithms.autotuner.acceptance_ratio.AcceptanceRatioAutotuner`)
    and to update the virtual-control / virtual-buffer weights from constraint
    violations via the multiplier hook.

    ``update_weights`` is a pure functional update on the
    :class:`AlgorithmState` pytree; see the base-class contract.
    """

    def __init__(
        self,
        rho_init: float = 1.0,
        rho_max: float = 1e2,
        gamma_1: float = 2.0,
        gamma_2: float = 0.5,
        eta_0: float = 1e-2,
        eta_1: float = 1e-1,
        eta_2: float = 0.8,
        ep: float = 0.99,
        eta_lambda: float = 1e1,
        lam_vc_max: float = 1e5,
        lam_prox_min: float = 1e-3,
        lam_prox_max: float = 1e4,
        lam_cost_drop: int = -1,
        lam_cost_relax: float = 1.0,
    ):
        """Initialize Augmented Lagrangian autotuning parameters.

        Every numeric knob is a declared hyperparameter: it lives on the frozen
        :class:`AugmentedLagrangianHyper` container at ``self.hyper`` and is a
        per-solve / batched override target (see :class:`AutotuningBase`).
        Knobs may still be read and written as plain attributes
        (``autotuner.lam_prox_max = 1e6``); the attribute proxy on
        :class:`AutotuningBase` routes the access into ``hyper``.

        Args:
            rho_init: Reserved for a future penalty-growth rule; currently
                unused. Passing a non-default value warns. Defaults to 1.0.
            rho_max: Reserved for a future penalty-growth rule; currently
                unused. Passing a non-default value warns. Defaults to 1e2.
            gamma_1: Factor to increase trust region weight when ratio is low.
                Defaults to 2.0.
            gamma_2: Factor to decrease trust region weight when ratio is high.
                Defaults to 0.5.
            eta_0: Acceptance ratio threshold below which solution is rejected.
                Defaults to 1e-2.
            eta_1: Threshold above which solution is accepted with constant weight.
                Defaults to 1e-1.
            eta_2: Threshold above which solution is accepted with lower weight.
                Defaults to 0.8.
            ep: Threshold for virtual control weight update (nu > ep vs nu <= ep).
                Must lie in (0, 1). Defaults to 0.99.
            eta_lambda: Step size for virtual control weight update. Defaults to 1e1.
            lam_vc_max: Maximum virtual control penalty weight. Defaults to 1e5.
            lam_prox_min: Minimum trust region (proximal) weight. Defaults to 1e-3.
            lam_prox_max: Maximum trust region (proximal) weight. Defaults to 1e4.
            lam_cost_drop: Iteration after which cost relaxation applies
                (``state.k > lam_cost_drop``). ``-1`` relaxes from the first
                iteration; the default ``lam_cost_relax=1.0`` makes that a
                no-op. Defaults to -1.
            lam_cost_relax: Factor applied to lam_cost after lam_cost_drop.
                Defaults to 1.0.
        """
        if rho_init != 1.0 or rho_max != 1e2:
            warnings.warn(
                "AugmentedLagrangian rho_init/rho_max are reserved for a future "
                "penalty-growth rule and are currently unused — setting them has "
                "no effect on the solve.",
                UserWarning,
                stacklevel=2,
            )
        self.hyper = AugmentedLagrangianHyper(
            rho_init=rho_init,
            rho_max=rho_max,
            gamma_1=gamma_1,
            gamma_2=gamma_2,
            eta_0=eta_0,
            eta_1=eta_1,
            eta_2=eta_2,
            ep=ep,
            eta_lambda=eta_lambda,
            lam_vc_max=lam_vc_max,
            lam_prox_min=lam_prox_min,
            lam_prox_max=lam_prox_max,
            lam_cost_drop=lam_cost_drop,
            lam_cost_relax=lam_cost_relax,
        )

    # -----------------------------------------------------------------------
    # Multiplier (virtual-control / virtual-buffer) updates from violation
    # -----------------------------------------------------------------------

    def _update_multipliers(
        self,
        state: "AlgorithmState",
        candidate: "CandidateIterate",
        nodal_constraints: "LoweredJaxConstraints",
        params: dict,
        settings: Config,
        new_lam_prox: jnp.ndarray,
    ) -> dict:
        """Update the virtual-control and virtual-buffer weights from violation.

        Each weight follows the same piecewise (linear / quadratic) update keyed
        on the scaled constraint violation; the shared body gates the result by
        iterate acceptance.
        """
        candidate_x_prop = candidate.x_prop_plus[1:]
        return {
            "lam_vc": self._update_virtual_control_weights(
                state, candidate, candidate_x_prop, settings, new_lam_prox
            ),
            "lam_vb_nodal": self._update_virtual_buffer_nodal_weights(
                state, candidate, nodal_constraints, params, new_lam_prox
            ),
            "lam_vb_cross": self._update_virtual_buffer_cross_weights(
                state, candidate, nodal_constraints, params, new_lam_prox
            ),
        }

    def _update_virtual_control_weights(
        self,
        state: "AlgorithmState",
        candidate: "CandidateIterate",
        candidate_x_prop: jnp.ndarray,
        settings: Config,
        lam_prox: jnp.ndarray,
    ) -> jnp.ndarray:
        """Update virtual control penalty weights from state violation.

        Computes scaled violation ``nu = inv_S_x @ |x[1:] - x_prop|`` and
        applies a piecewise update: linear in ``nu`` when ``nu > ep``,
        quadratic otherwise. Result is clipped to ``lam_vc_max``.
        """
        nu = (settings.sim.inv_S_x @ jnp.abs(candidate.x[1:] - candidate_x_prop).T).T
        # TODO: (haynec) use per-variable lam_prox to scale VC updates proportionally
        lam_prox_scalar = jnp.max(lam_prox)
        scale = state.hyper.eta_lambda * (1.0 / (2.0 * lam_prox_scalar))
        case1 = state.lam_vc + nu * scale
        case2 = state.lam_vc + (nu**2) / state.hyper.ep * scale
        vc_new = jnp.where(nu > state.hyper.ep, case1, case2)
        return jnp.minimum(state.hyper.lam_vc_max, vc_new)

    def _update_virtual_buffer_nodal_weights(
        self,
        state: "AlgorithmState",
        candidate: "CandidateIterate",
        nodal_constraints: "LoweredJaxConstraints",
        params: dict,
        lam_prox: jnp.ndarray,
    ) -> jnp.ndarray:
        """Update virtual-buffer weights for nodal constraints.

        Evaluates each nodal constraint, computes the same piecewise update
        as for virtual control, and writes the result back into the
        ``(N, n_nodal)`` weight array via ``.at[...].set(...)``.
        """
        lam_vb_new = state.lam_vb_nodal
        lam_prox_scalar = jnp.max(lam_prox)
        scale = state.hyper.eta_lambda * (1.0 / (2.0 * lam_prox_scalar))

        for idx, constraint in enumerate(nodal_constraints.nodal):
            g = constraint.func(candidate.x, candidate.u, 0, params)
            nu = jnp.abs(g) if constraint.is_equality else jnp.maximum(0.0, g)

            if constraint.nodes is not None:
                nodes_array = jnp.asarray(constraint.nodes)
                nu_slice = nu[nodes_array]
                current = state.lam_vb_nodal[nodes_array, idx]
            else:
                nu_slice = nu
                current = state.lam_vb_nodal[:, idx]

            case1 = current + nu_slice * scale
            case2 = current + (nu_slice**2) / state.hyper.ep * scale
            updated = jnp.where(nu_slice > state.hyper.ep, case1, case2)

            if constraint.nodes is not None:
                lam_vb_new = lam_vb_new.at[nodes_array, idx].set(updated)
            else:
                lam_vb_new = lam_vb_new.at[:, idx].set(updated)

        # Intentional: the virtual-buffer weights share the virtual-control cap
        # ``lam_vc_max`` rather than a dedicated knob — one ceiling governs every
        # penalty weight in this autotuner.
        return jnp.minimum(state.hyper.lam_vc_max, lam_vb_new)

    def _update_virtual_buffer_cross_weights(
        self,
        state: "AlgorithmState",
        candidate: "CandidateIterate",
        nodal_constraints: "LoweredJaxConstraints",
        params: dict,
        lam_prox: jnp.ndarray,
    ) -> jnp.ndarray:
        """Update virtual-buffer weights for cross-node constraints."""
        lam_vb_new = state.lam_vb_cross
        lam_prox_scalar = jnp.max(lam_prox)
        scale = state.hyper.eta_lambda * (1.0 / (2.0 * lam_prox_scalar))

        for idx, constraint in enumerate(nodal_constraints.cross_node):
            g = constraint.func(candidate.x, candidate.u, params)
            nu = jnp.sum(jnp.abs(g) if constraint.is_equality else jnp.maximum(0.0, g))
            current = state.lam_vb_cross[idx]
            case1 = current + nu * scale
            case2 = current + (nu**2) / state.hyper.ep * scale
            updated = jnp.where(nu > state.hyper.ep, case1, case2)
            lam_vb_new = lam_vb_new.at[idx].set(updated)

        # Intentional: shares the virtual-control cap ``lam_vc_max`` — see
        # :py:meth:`_update_virtual_buffer_nodal_weights`.
        return jnp.minimum(state.hyper.lam_vc_max, lam_vb_new)

    def citation(self) -> List[str]:
        """Return BibTeX citations for the augmented-Lagrangian multiplier updates.

        Extends the inherited SCvx acceptance-ratio references with Luo's
        dissertation, whose penalty-weight updates this autotuner implements.
        """
        return super().citation() + [
            r"""@phdthesis{luo2025modeling,
  title={Modeling and Algorithms for Nonconvex Trajectory Generation Problems:
    From Constraint Reformulations and First-Order Proximal Methods to
    Structure-Exploiting Convex Solvers},
  author={Luo, Dayou},
  school={University of Washington},
  year={2025},
  url={https://hdl.handle.net/1773/53364}
}"""
        ]

__init__(rho_init: float = 1.0, rho_max: float = 100.0, gamma_1: float = 2.0, gamma_2: float = 0.5, eta_0: float = 0.01, eta_1: float = 0.1, eta_2: float = 0.8, ep: float = 0.99, eta_lambda: float = 10.0, lam_vc_max: float = 100000.0, lam_prox_min: float = 0.001, lam_prox_max: float = 10000.0, lam_cost_drop: int = -1, lam_cost_relax: float = 1.0)

Initialize Augmented Lagrangian autotuning parameters.

Every numeric knob is a declared hyperparameter: it lives on the frozen :class:AugmentedLagrangianHyper container at self.hyper and is a per-solve / batched override target (see :class:AutotuningBase). Knobs may still be read and written as plain attributes (autotuner.lam_prox_max = 1e6); the attribute proxy on :class:AutotuningBase routes the access into hyper.

Parameters:

Name	Type	Description	Default
rho_init	float	Reserved for a future penalty-growth rule; currently unused. Passing a non-default value warns. Defaults to 1.0.	1.0
rho_max	float	Reserved for a future penalty-growth rule; currently unused. Passing a non-default value warns. Defaults to 1e2.	100.0
gamma_1	float	Factor to increase trust region weight when ratio is low. Defaults to 2.0.	2.0
gamma_2	float	Factor to decrease trust region weight when ratio is high. Defaults to 0.5.	0.5
eta_0	float	Acceptance ratio threshold below which solution is rejected. Defaults to 1e-2.	0.01
eta_1	float	Threshold above which solution is accepted with constant weight. Defaults to 1e-1.	0.1
eta_2	float	Threshold above which solution is accepted with lower weight. Defaults to 0.8.	0.8
ep	float	Threshold for virtual control weight update (nu > ep vs nu <= ep). Must lie in (0, 1). Defaults to 0.99.	0.99
eta_lambda	float	Step size for virtual control weight update. Defaults to 1e1.	10.0
lam_vc_max	float	Maximum virtual control penalty weight. Defaults to 1e5.	100000.0
lam_prox_min	float	Minimum trust region (proximal) weight. Defaults to 1e-3.	0.001
lam_prox_max	float	Maximum trust region (proximal) weight. Defaults to 1e4.	10000.0
lam_cost_drop	int	Iteration after which cost relaxation applies (state.k > lam_cost_drop). -1 relaxes from the first iteration; the default lam_cost_relax=1.0 makes that a no-op. Defaults to -1.	-1
lam_cost_relax	float	Factor applied to lam_cost after lam_cost_drop. Defaults to 1.0.	1.0

Source code in openscvx/algorithms/autotuner/augmented_lagrangian.py

def __init__(
    self,
    rho_init: float = 1.0,
    rho_max: float = 1e2,
    gamma_1: float = 2.0,
    gamma_2: float = 0.5,
    eta_0: float = 1e-2,
    eta_1: float = 1e-1,
    eta_2: float = 0.8,
    ep: float = 0.99,
    eta_lambda: float = 1e1,
    lam_vc_max: float = 1e5,
    lam_prox_min: float = 1e-3,
    lam_prox_max: float = 1e4,
    lam_cost_drop: int = -1,
    lam_cost_relax: float = 1.0,
):
    """Initialize Augmented Lagrangian autotuning parameters.

    Every numeric knob is a declared hyperparameter: it lives on the frozen
    :class:`AugmentedLagrangianHyper` container at ``self.hyper`` and is a
    per-solve / batched override target (see :class:`AutotuningBase`).
    Knobs may still be read and written as plain attributes
    (``autotuner.lam_prox_max = 1e6``); the attribute proxy on
    :class:`AutotuningBase` routes the access into ``hyper``.

    Args:
        rho_init: Reserved for a future penalty-growth rule; currently
            unused. Passing a non-default value warns. Defaults to 1.0.
        rho_max: Reserved for a future penalty-growth rule; currently
            unused. Passing a non-default value warns. Defaults to 1e2.
        gamma_1: Factor to increase trust region weight when ratio is low.
            Defaults to 2.0.
        gamma_2: Factor to decrease trust region weight when ratio is high.
            Defaults to 0.5.
        eta_0: Acceptance ratio threshold below which solution is rejected.
            Defaults to 1e-2.
        eta_1: Threshold above which solution is accepted with constant weight.
            Defaults to 1e-1.
        eta_2: Threshold above which solution is accepted with lower weight.
            Defaults to 0.8.
        ep: Threshold for virtual control weight update (nu > ep vs nu <= ep).
            Must lie in (0, 1). Defaults to 0.99.
        eta_lambda: Step size for virtual control weight update. Defaults to 1e1.
        lam_vc_max: Maximum virtual control penalty weight. Defaults to 1e5.
        lam_prox_min: Minimum trust region (proximal) weight. Defaults to 1e-3.
        lam_prox_max: Maximum trust region (proximal) weight. Defaults to 1e4.
        lam_cost_drop: Iteration after which cost relaxation applies
            (``state.k > lam_cost_drop``). ``-1`` relaxes from the first
            iteration; the default ``lam_cost_relax=1.0`` makes that a
            no-op. Defaults to -1.
        lam_cost_relax: Factor applied to lam_cost after lam_cost_drop.
            Defaults to 1.0.
    """
    if rho_init != 1.0 or rho_max != 1e2:
        warnings.warn(
            "AugmentedLagrangian rho_init/rho_max are reserved for a future "
            "penalty-growth rule and are currently unused — setting them has "
            "no effect on the solve.",
            UserWarning,
            stacklevel=2,
        )
    self.hyper = AugmentedLagrangianHyper(
        rho_init=rho_init,
        rho_max=rho_max,
        gamma_1=gamma_1,
        gamma_2=gamma_2,
        eta_0=eta_0,
        eta_1=eta_1,
        eta_2=eta_2,
        ep=ep,
        eta_lambda=eta_lambda,
        lam_vc_max=lam_vc_max,
        lam_prox_min=lam_prox_min,
        lam_prox_max=lam_prox_max,
        lam_cost_drop=lam_cost_drop,
        lam_cost_relax=lam_cost_relax,
    )

citation() -> List[str]

Return BibTeX citations for the augmented-Lagrangian multiplier updates.

Extends the inherited SCvx acceptance-ratio references with Luo's dissertation, whose penalty-weight updates this autotuner implements.

Source code in openscvx/algorithms/autotuner/augmented_lagrangian.py

    def citation(self) -> List[str]:
        """Return BibTeX citations for the augmented-Lagrangian multiplier updates.

        Extends the inherited SCvx acceptance-ratio references with Luo's
        dissertation, whose penalty-weight updates this autotuner implements.
        """
        return super().citation() + [
            r"""@phdthesis{luo2025modeling,
  title={Modeling and Algorithms for Nonconvex Trajectory Generation Problems:
    From Constraint Reformulations and First-Order Proximal Methods to
    Structure-Exploiting Convex Solvers},
  author={Luo, Dayou},
  school={University of Washington},
  year={2025},
  url={https://hdl.handle.net/1773/53364}
}"""
        ]

AutotuningBase

Bases: ABC

Base class for autotuning strategies in SCP algorithms.

Subclasses implement meth:update_weights as a pure functional update on an :class:AlgorithmState pytree. Concretely, an autotuner must:

• Not mutate state or candidate.
• Not return strings, raise on data-dependent conditions, or append to Python lists at trace time.
• Express all branching on iterate values via jax.lax.cond / jnp.where.

The autotuner's per-iteration outcome is encoded as an :class:AdaptiveStateCode int32 on the returned state; the SCP loop converts that back to a human-readable label only on the Python-loop printing path.

The returned state IS the next iterate. The SCP loop discards everything except what update_weights returns, so accepting the subproblem's candidate means carrying its trajectory onto the returned state — x / u / x_prop / x_prop_plus from candidate — while rejecting means keeping the previous fields and adjusting only the weights. An update that never copies the candidate produces a solver that runs to k_max without the iterate ever moving, silently.

Declaring tunable hyperparameters. A numeric knob a user might sweep (a penalty ramp, a relaxation iteration) must not be read off self inside update_weights — a Python attribute is baked into the trace and invisible to every override channel. Instead, declare it on a :class:HyperParams subclass, assign an instance to self.hyper, and read it from state.hyper::

class MyHyper(HyperParams):
    ramp: float = 2.0

class MyAutotuner(AutotuningBase):
    def __init__(self, ramp: float = 2.0):
        self.hyper = MyHyper(ramp=ramp)

    def update_weights(self, state, candidate, nodal_constraints, settings, params):
        return state.replace(
            # accept the candidate: it becomes the next iterate
            x=candidate.x,
            u=candidate.u,
            x_prop=candidate.x_prop,
            x_prop_plus=candidate.x_prop_plus,
            lam_prox=state.lam_prox * state.hyper.ramp,
            adaptive_state_code=...,
        )

The declaration is the registration: the field becomes a per-solve override (solve_jax(algorithm={"ramp": ...})), a batchable sweep target (solve_batched(algorithm={"ramp": jnp.linspace(...)})), and a runtime input of the exported batched artifact — with zero core edits. Every numeric knob of the built-in autotuners is declared this way, so the only configuration left outside hyper is structural choices that select code paths (class attributes like COMPUTES_ACCEPTANCE_METRICS); they are part of the traced program, not data. For ergonomics, declared knobs are still readable and writable as bare attributes (autotuner.ramp = 3.0): the proxy below routes the access into hyper so a constructor — or a test — may keep assigning knobs by name. Assign self.hyper first in __init__; a knob assigned before hyper exists lands as a plain instance attribute the proxy cannot see.

Class Attributes

COLUMNS: List of Column specs for autotuner-specific metrics to display. Subclasses override this to add their own columns. COMPUTES_ACCEPTANCE_METRICS: Whether update_weights computes a predicted/actual reduction and the resulting acceptance ratio. The SCP loop records and prints those diagnostics only when this is True (the default); tuners that never reject an iterate — e.g. :class:~openscvx.algorithms.autotuner.constant_proximal_weight.ConstantProximalWeight and :class:~openscvx.algorithms.autotuner.ramp_proximal_weight.RampProximalWeight — set it False. hyper: The autotuner's declared hyperparameters — a :class:HyperParams instance carrying plain-Python defaults (the empty base when it declares none). Snapshotted onto AlgorithmState.hyper with array leaves (see :meth:AlgorithmState.from_settings; int fields get k's integer dtype, float fields the float dtype), which is what update_weights reads at trace time.

Source code in openscvx/algorithms/autotuner/base.py

class AutotuningBase(ABC):
    """Base class for autotuning strategies in SCP algorithms.

    Subclasses implement :py:meth:`update_weights` as a **pure functional
    update** on an :class:`AlgorithmState` pytree. Concretely, an autotuner
    must:

    * Not mutate ``state`` or ``candidate``.
    * Not return strings, raise on data-dependent conditions, or append to
      Python lists at trace time.
    * Express all branching on iterate values via ``jax.lax.cond`` /
      ``jnp.where``.

    The autotuner's per-iteration outcome is encoded as an
    :class:`AdaptiveStateCode` int32 on the returned state; the SCP loop
    converts that back to a human-readable label only on the Python-loop
    printing path.

    **The returned state IS the next iterate.** The SCP loop discards
    everything except what ``update_weights`` returns, so accepting the
    subproblem's candidate means carrying its trajectory onto the returned
    state — ``x`` / ``u`` / ``x_prop`` / ``x_prop_plus`` from ``candidate`` —
    while rejecting means keeping the previous fields and adjusting only the
    weights. An update that never copies the candidate produces a solver that
    runs to ``k_max`` without the iterate ever moving, silently.

    **Declaring tunable hyperparameters.** A numeric knob a user might sweep
    (a penalty ramp, a relaxation iteration) must not be read off ``self``
    inside ``update_weights`` — a Python attribute is baked into the trace
    and invisible to every override channel. Instead, declare it on a
    :class:`HyperParams` subclass, assign an instance to ``self.hyper``, and
    read it from ``state.hyper``::

        class MyHyper(HyperParams):
            ramp: float = 2.0

        class MyAutotuner(AutotuningBase):
            def __init__(self, ramp: float = 2.0):
                self.hyper = MyHyper(ramp=ramp)

            def update_weights(self, state, candidate, nodal_constraints, settings, params):
                return state.replace(
                    # accept the candidate: it becomes the next iterate
                    x=candidate.x,
                    u=candidate.u,
                    x_prop=candidate.x_prop,
                    x_prop_plus=candidate.x_prop_plus,
                    lam_prox=state.lam_prox * state.hyper.ramp,
                    adaptive_state_code=...,
                )

    The declaration is the registration: the field becomes a per-solve
    override (``solve_jax(algorithm={"ramp": ...})``), a batchable sweep
    target (``solve_batched(algorithm={"ramp": jnp.linspace(...)})``), and a
    runtime input of the exported batched artifact — with zero core edits.
    Every numeric knob of the built-in autotuners is declared this way, so the
    only configuration left outside ``hyper`` is structural choices that select
    code paths (class attributes like ``COMPUTES_ACCEPTANCE_METRICS``); they
    are part of the traced program, not data. For ergonomics, declared knobs
    are still readable and writable as bare attributes
    (``autotuner.ramp = 3.0``): the proxy below routes the access into
    ``hyper`` so a constructor — or a test — may keep assigning knobs by name.
    Assign ``self.hyper`` *first* in ``__init__``; a knob assigned before
    ``hyper`` exists lands as a plain instance attribute the proxy cannot see.

    Class Attributes:
        COLUMNS: List of Column specs for autotuner-specific metrics to display.
            Subclasses override this to add their own columns.
        COMPUTES_ACCEPTANCE_METRICS: Whether ``update_weights`` computes a
            predicted/actual reduction and the resulting acceptance ratio. The
            SCP loop records and prints those diagnostics only when this is
            ``True`` (the default); tuners that never reject an iterate — e.g.
            :class:`~openscvx.algorithms.autotuner.constant_proximal_weight.ConstantProximalWeight`
            and :class:`~openscvx.algorithms.autotuner.ramp_proximal_weight.RampProximalWeight`
            — set it ``False``.
        hyper: The autotuner's declared hyperparameters — a
            :class:`HyperParams` instance carrying plain-Python defaults
            (the empty base when it declares none). Snapshotted onto
            ``AlgorithmState.hyper`` with array leaves (see
            :meth:`AlgorithmState.from_settings`; ``int`` fields get ``k``'s
            integer dtype, ``float`` fields the float dtype), which is what
            ``update_weights`` reads at trace time.
    """

    COLUMNS: List[Column] = []
    COMPUTES_ACCEPTANCE_METRICS: bool = True
    hyper: HyperParams = HyperParams()

    def __getattr__(self, name: str):
        """Read a declared hyperparameter as a bare attribute.

        Promoted knobs live on the frozen ``hyper`` container, but the
        documented API and constructors still touch them by name
        (``autotuner.lam_prox_max``). ``__getattr__`` fires only on failed
        normal lookups, so real attributes (``COLUMNS``, ``hyper`` itself) are
        unaffected; a name matching a ``hyper`` field resolves there.
        """
        hyper = self.__dict__.get("hyper")
        if hyper is not None and name in {f.name for f in dc_fields(hyper)}:
            return getattr(hyper, name)
        raise AttributeError(f"{type(self).__name__!r} has no attribute {name!r}")

    def __setattr__(self, name: str, value) -> None:
        """Route a write to a declared hyperparameter into the frozen ``hyper``.

        Without this, assigning a promoted knob as a bare attribute
        (``autotuner.lam_prox_max = 1e6``) would silently shadow the ``hyper``
        field and never reach the solve. A name matching a ``hyper`` field is
        applied via :func:`dataclasses.replace`; every other attribute
        (including ``hyper`` itself) is set normally, so constructors may keep
        assigning knobs by name — they route into ``hyper``, provided
        ``self.hyper`` was assigned first (see the class docstring).
        """
        hyper = self.__dict__.get("hyper")
        if hyper is not None and name != "hyper" and name in {f.name for f in dc_fields(hyper)}:
            # Clear any instance attribute shadowing the knob (left by an
            # assignment made before ``hyper`` existed) so reads resolve
            # through ``__getattr__`` to the value just written.
            self.__dict__.pop(name, None)
            super().__setattr__("hyper", dc_replace(hyper, **{name: value}))
        else:
            super().__setattr__(name, value)

    def _hash_into(self, hasher: "hashlib._Hash") -> None:
        """Contribute the autotuner's update rule to the ``solve_batched`` cache key.

        The exported batched loop bakes in ``update_weights`` and every numeric
        parameter that steers it (penalty ramps, acceptance thresholds, weight
        clips). The default hashes the concrete class plus all instance
        attributes — sufficient because autotuner parameters are plain
        scalars. The ``hyper`` *values* are excluded: declared hyperparameters
        ride ``AlgorithmState.hyper`` as runtime inputs, so one artifact serves
        every setting of them (the same reasoning that keeps the ``ep_*``
        thresholds out of the algorithm's hash). The ``hyper`` *schema* is
        folded in, though: the ``hyper`` leaves map positionally onto the
        exported artifact's avals, so a field reorder or addition silently
        permutes runtime values against the wrong leaves. Hashing the field
        names in declaration order turns such a change into a clean cache miss
        and rebuild — this matters because :class:`LamCostRelaxHyper` moves
        ``lam_cost_drop`` / ``lam_cost_relax`` to the *front* of the field
        order. Folded in by :meth:`Algorithm._hash_into`; mirrors the symbolic
        ``_hash_into`` protocol.
        """
        from openscvx.utils.caching import hash_value_into

        hasher.update(type(self).__name__.encode())
        for name in sorted(vars(self)):
            if name == "hyper":
                continue
            hasher.update(name.encode())
            hash_value_into(hasher, getattr(self, name))
        # Schema, not values: the hyper field order/names fix the exported
        # artifact's leaf layout, so a reorder must invalidate cached artifacts
        # rather than silently permute runtime values onto the wrong avals.
        for fld in dc_fields(self.hyper):
            hasher.update(fld.name.encode())

    @staticmethod
    def _relaxed_lam_cost(state: "AlgorithmState") -> jnp.ndarray:
        """``lam_cost`` for the next iterate per the shared relaxation rule.

        When ``state.k > state.hyper.lam_cost_drop``, scale ``state.lam_cost``
        by ``state.hyper.lam_cost_relax``; otherwise reset to the algorithm's
        initial weight (carried on the pytree as ``state.lam_cost_init``,
        broadcast at :meth:`AlgorithmState.from_settings`). Both constants ride
        the pytree so per-solve overrides and ``vmap`` sweeps reach the traced
        body; the scalar ``lam_cost_relax`` preserves the user-specified
        per-state weight ratios. Shared verbatim by every built-in autotuner
        via the :class:`LamCostRelaxHyper` knobs.
        """
        return jnp.where(
            state.k > state.hyper.lam_cost_drop,
            state.lam_cost * state.hyper.lam_cost_relax,
            state.lam_cost_init,
        )

    @abstractmethod
    def update_weights(
        self,
        state: "AlgorithmState",
        candidate: "CandidateIterate",
        nodal_constraints: "LoweredJaxConstraints",
        settings: "Config",
        params: dict,
    ) -> "AlgorithmState":
        """Return the next-iterate :class:`AlgorithmState`.

        Must be JAX-traceable. See the class docstring for the contract.

        Args:
            state: Current-iterate pytree. Initial weights live on
                ``state.lam_cost_init``; autotuners read from there rather than
                from the algorithm's ``Weights`` object so the JIT'd closure
                stays cache-stable across weight mutations.
            candidate: Subproblem result (read-only here). On acceptance its
                trajectory fields must be carried onto the returned state —
                see the class docstring.
            nodal_constraints: Lowered JAX constraints.
            settings: Configuration object.
            params: Problem parameter dictionary.

        Returns:
            The next-iterate :class:`AlgorithmState` — the SCP loop uses
            nothing else, so an accepting update carries
            ``candidate.x / u / x_prop / x_prop_plus`` onto it, while a
            rejecting update keeps the previous fields and adjusts only the
            weights. Its :py:attr:`AlgorithmState.adaptive_state_code` encodes
            the autotuner's decision; the SCP loop records that into history
            and maps it to a printable label.
        """

    def citation(self) -> List[str]:
        """Return BibTeX citations for the published update rule this autotuner implements.

        Autotuners implementing a published scheme override this so that
        :meth:`openscvx.Problem.citation` can credit it alongside the
        algorithm, solver, and discretization references. Unpublished
        heuristics keep the default.

        Returns:
            List of BibTeX citation strings. Empty for uncited autotuners.
        """
        return []

citation() -> List[str]

Return BibTeX citations for the published update rule this autotuner implements.

Autotuners implementing a published scheme override this so that :meth:openscvx.Problem.citation can credit it alongside the algorithm, solver, and discretization references. Unpublished heuristics keep the default.

Returns:

Type	Description
List[str]	List of BibTeX citation strings. Empty for uncited autotuners.

Source code in openscvx/algorithms/autotuner/base.py

def citation(self) -> List[str]:
    """Return BibTeX citations for the published update rule this autotuner implements.

    Autotuners implementing a published scheme override this so that
    :meth:`openscvx.Problem.citation` can credit it alongside the
    algorithm, solver, and discretization references. Unpublished
    heuristics keep the default.

    Returns:
        List of BibTeX citation strings. Empty for uncited autotuners.
    """
    return []

update_weights(state: AlgorithmState, candidate: CandidateIterate, nodal_constraints: LoweredJaxConstraints, settings: Config, params: dict) -> AlgorithmState abstractmethod

Return the next-iterate :class:AlgorithmState.

Must be JAX-traceable. See the class docstring for the contract.

Parameters:

Name	Type	Description	Default
state	AlgorithmState	Current-iterate pytree. Initial weights live on state.lam_cost_init; autotuners read from there rather than from the algorithm's Weights object so the JIT'd closure stays cache-stable across weight mutations.	required
candidate	CandidateIterate	Subproblem result (read-only here). On acceptance its trajectory fields must be carried onto the returned state — see the class docstring.	required
nodal_constraints	LoweredJaxConstraints	Lowered JAX constraints.	required
settings	Config	Configuration object.	required
params	dict	Problem parameter dictionary.	required

Returns:

Type	Description
AlgorithmState	The next-iterate :class:AlgorithmState — the SCP loop uses
AlgorithmState	nothing else, so an accepting update carries
AlgorithmState	candidate.x / u / x_prop / x_prop_plus onto it, while a
AlgorithmState	rejecting update keeps the previous fields and adjusts only the
AlgorithmState	weights. Its attr:AlgorithmState.adaptive_state_code encodes
AlgorithmState	the autotuner's decision; the SCP loop records that into history
AlgorithmState	and maps it to a printable label.

Source code in openscvx/algorithms/autotuner/base.py

@abstractmethod
def update_weights(
    self,
    state: "AlgorithmState",
    candidate: "CandidateIterate",
    nodal_constraints: "LoweredJaxConstraints",
    settings: "Config",
    params: dict,
) -> "AlgorithmState":
    """Return the next-iterate :class:`AlgorithmState`.

    Must be JAX-traceable. See the class docstring for the contract.

    Args:
        state: Current-iterate pytree. Initial weights live on
            ``state.lam_cost_init``; autotuners read from there rather than
            from the algorithm's ``Weights`` object so the JIT'd closure
            stays cache-stable across weight mutations.
        candidate: Subproblem result (read-only here). On acceptance its
            trajectory fields must be carried onto the returned state —
            see the class docstring.
        nodal_constraints: Lowered JAX constraints.
        settings: Configuration object.
        params: Problem parameter dictionary.

    Returns:
        The next-iterate :class:`AlgorithmState` — the SCP loop uses
        nothing else, so an accepting update carries
        ``candidate.x / u / x_prop / x_prop_plus`` onto it, while a
        rejecting update keeps the previous fields and adjusts only the
        weights. Its :py:attr:`AlgorithmState.adaptive_state_code` encodes
        the autotuner's decision; the SCP loop records that into history
        and maps it to a printable label.
    """

CandidateIterate dataclass

Subproblem candidate the fused iteration hands to the autotuner.

Built once, in-trace, by the SCP iteration body (scvx/iteration.py): it carries the just-solved subproblem trajectory (x / u), its propagation (x_prop / x_prop_plus), and the linearized cost (J_lin) into meth:AutotuningBase.update_weights, which copies the accepted fields onto the next :class:AlgorithmState. It is the autotuner's input and nothing else — host-side history recording takes its arrays through meth:AlgorithmHistory.record_iteration's explicit keyword arguments rather than through this type.

Source code in openscvx/algorithms/state.py

@dataclass(frozen=True)
class CandidateIterate:
    """Subproblem candidate the fused iteration hands to the autotuner.

    Built once, in-trace, by the SCP iteration body
    (``scvx/iteration.py``): it carries the just-solved subproblem trajectory
    (``x`` / ``u``), its propagation (``x_prop`` / ``x_prop_plus``), and the
    linearized cost (``J_lin``) into :py:meth:`AutotuningBase.update_weights`,
    which copies the accepted fields onto the next :class:`AlgorithmState`. It
    is the autotuner's input and nothing else — host-side history recording
    takes its arrays through :py:meth:`AlgorithmHistory.record_iteration`'s
    explicit keyword arguments rather than through this type.
    """

    x: jnp.ndarray
    u: jnp.ndarray
    x_prop: jnp.ndarray
    x_prop_plus: jnp.ndarray
    J_lin: jnp.ndarray

ConstantProximalWeight

Bases: AutotuningBase

Hold lam_prox constant; relax lam_cost after lam_cost_drop.

Useful when you want a fixed trust-region size without adaptation.

update_weights is a pure functional update on the :class:AlgorithmState pytree; see the base-class contract.

Source code in openscvx/algorithms/autotuner/constant_proximal_weight.py

class ConstantProximalWeight(AutotuningBase):
    """Hold ``lam_prox`` constant; relax ``lam_cost`` after ``lam_cost_drop``.

    Useful when you want a fixed trust-region size without adaptation.

    ``update_weights`` is a pure functional update on the
    :class:`AlgorithmState` pytree; see the base-class contract.
    """

    COMPUTES_ACCEPTANCE_METRICS = False

    def __init__(
        self,
        lam_cost_drop: int = -1,
        lam_cost_relax: float = 1.0,
    ):
        self.hyper = ConstantProximalWeightHyper(
            lam_cost_drop=lam_cost_drop,
            lam_cost_relax=lam_cost_relax,
        )

    def update_weights(
        self,
        state: "AlgorithmState",
        candidate: "CandidateIterate",
        nodal_constraints: "LoweredJaxConstraints",
        settings: Config,
        params: dict,
    ) -> "AlgorithmState":
        """Return the next-iterate state.

        Pure functional update — see class docstring.
        """
        lam_cost_next = self._relaxed_lam_cost(state)

        return state.replace(
            x=candidate.x,
            u=candidate.u,
            x_prop=candidate.x_prop,
            x_prop_plus=candidate.x_prop_plus,
            lam_cost=lam_cost_next,
            adaptive_state_code=jnp.asarray(
                int(AdaptiveStateCode.ACCEPT_CONSTANT), dtype=jnp.int32
            ),
        )

update_weights(state: AlgorithmState, candidate: CandidateIterate, nodal_constraints: LoweredJaxConstraints, settings: Config, params: dict) -> AlgorithmState

Return the next-iterate state.

Pure functional update — see class docstring.

Source code in openscvx/algorithms/autotuner/constant_proximal_weight.py

def update_weights(
    self,
    state: "AlgorithmState",
    candidate: "CandidateIterate",
    nodal_constraints: "LoweredJaxConstraints",
    settings: Config,
    params: dict,
) -> "AlgorithmState":
    """Return the next-iterate state.

    Pure functional update — see class docstring.
    """
    lam_cost_next = self._relaxed_lam_cost(state)

    return state.replace(
        x=candidate.x,
        u=candidate.u,
        x_prop=candidate.x_prop,
        x_prop_plus=candidate.x_prop_plus,
        lam_cost=lam_cost_next,
        adaptive_state_code=jnp.asarray(
            int(AdaptiveStateCode.ACCEPT_CONSTANT), dtype=jnp.int32
        ),
    )

DiscretizationResult dataclass

Unpacked discretization data from a multi-shot discretization matrix.

The discretization solver returns a matrix V that stores multiple blocks (propagated state and linearization matrices) across nodes/time. Historically, we stored the raw V matrices and re-unpacked them repeatedly via slicing. This dataclass unpacks once and makes access trivial.

Source code in openscvx/algorithms/history.py

@dataclass(frozen=True, slots=True)
class DiscretizationResult:
    """Unpacked discretization data from a multi-shot discretization matrix.

    The discretization solver returns a matrix ``V`` that stores multiple blocks
    (propagated state and linearization matrices) across nodes/time. Historically,
    we stored the raw ``V`` matrices and re-unpacked them repeatedly via slicing.
    This dataclass unpacks once and makes access trivial.
    """

    V: np.ndarray  # raw V matrix, shape: (flattened_size, n_timesteps)
    x_prop: np.ndarray  # (N-1, n_x)
    A_d: np.ndarray  # (N-1, n_x, n_x)
    B_d: np.ndarray  # (N-1, n_x, n_u)
    C_d: np.ndarray  # (N-1, n_x, n_u)
    x_prop_plus: Optional[np.ndarray] = None  # (N, n_x), discrete dynamics on node states
    D_d: Optional[np.ndarray] = None  # (N, n_x, n_x), d(x_prop_plus)/d(x_node)
    E_d: Optional[np.ndarray] = None  # (N, n_x, n_u), d(x_prop_plus)/d(u_node)

    @classmethod
    def from_V(
        cls,
        V: np.ndarray,
        n_x: int,
        n_u: int,
        N: int,
    ) -> "DiscretizationResult":
        """Unpack the final timestep of a raw discretization matrix ``V``."""
        i1, i2 = n_x, n_x + n_x * n_x
        i3, i4 = i2 + n_x * n_u, i2 + 2 * n_x * n_u
        V_final = V[:, -1].reshape(-1, i4)
        return cls(
            V=np.asarray(V),
            x_prop=V_final[:, :i1],
            A_d=V_final[:, i1:i2].reshape(N - 1, n_x, n_x),
            B_d=V_final[:, i2:i3].reshape(N - 1, n_x, n_u),
            C_d=V_final[:, i3:i4].reshape(N - 1, n_x, n_u),
        )

    @classmethod
    def from_VW(
        cls,
        V: np.ndarray,
        W: np.ndarray,
        n_x: int,
        n_u: int,
        N: int,
    ) -> "DiscretizationResult":
        """Unpack continuous and impulsive discretization blocks from ``V`` and ``W``."""
        base = cls.from_V(V=V, n_x=n_x, n_u=n_u, N=N)

        W_arr = np.asarray(W)
        i_w = n_x + n_x * n_x + n_x * n_u
        i1 = n_x
        i2 = i1 + n_x * n_x
        i3 = i2 + n_x * n_u

        W_final = W_arr[:, -1].reshape(-1, i_w)

        return cls(
            V=base.V,
            x_prop=base.x_prop,
            A_d=base.A_d,
            B_d=base.B_d,
            C_d=base.C_d,
            x_prop_plus=W_final[:, :i1],
            D_d=W_final[:, i1:i2].reshape(W_final.shape[0], n_x, n_x),
            E_d=W_final[:, i2:i3].reshape(W_final.shape[0], n_x, n_u),
        )

from_V(V: np.ndarray, n_x: int, n_u: int, N: int) -> 'DiscretizationResult' classmethod

Unpack the final timestep of a raw discretization matrix V.

Source code in openscvx/algorithms/history.py

@classmethod
def from_V(
    cls,
    V: np.ndarray,
    n_x: int,
    n_u: int,
    N: int,
) -> "DiscretizationResult":
    """Unpack the final timestep of a raw discretization matrix ``V``."""
    i1, i2 = n_x, n_x + n_x * n_x
    i3, i4 = i2 + n_x * n_u, i2 + 2 * n_x * n_u
    V_final = V[:, -1].reshape(-1, i4)
    return cls(
        V=np.asarray(V),
        x_prop=V_final[:, :i1],
        A_d=V_final[:, i1:i2].reshape(N - 1, n_x, n_x),
        B_d=V_final[:, i2:i3].reshape(N - 1, n_x, n_u),
        C_d=V_final[:, i3:i4].reshape(N - 1, n_x, n_u),
    )

from_VW(V: np.ndarray, W: np.ndarray, n_x: int, n_u: int, N: int) -> 'DiscretizationResult' classmethod

Unpack continuous and impulsive discretization blocks from V and W.

Source code in openscvx/algorithms/history.py

@classmethod
def from_VW(
    cls,
    V: np.ndarray,
    W: np.ndarray,
    n_x: int,
    n_u: int,
    N: int,
) -> "DiscretizationResult":
    """Unpack continuous and impulsive discretization blocks from ``V`` and ``W``."""
    base = cls.from_V(V=V, n_x=n_x, n_u=n_u, N=N)

    W_arr = np.asarray(W)
    i_w = n_x + n_x * n_x + n_x * n_u
    i1 = n_x
    i2 = i1 + n_x * n_x
    i3 = i2 + n_x * n_u

    W_final = W_arr[:, -1].reshape(-1, i_w)

    return cls(
        V=base.V,
        x_prop=base.x_prop,
        A_d=base.A_d,
        B_d=base.B_d,
        C_d=base.C_d,
        x_prop_plus=W_final[:, :i1],
        D_d=W_final[:, i1:i2].reshape(W_final.shape[0], n_x, n_x),
        E_d=W_final[:, i2:i3].reshape(W_final.shape[0], n_x, n_u),
    )

HyperParams

Base for autotuner hyperparameter containers.

Subclassing is the entire integration. Declare each tunable knob as an annotated field with its default — bare annotations, no @dataclass decorator — and the base class does the rest:

• applies the frozen, keyword-only dataclass transform, so instances are immutable value objects constructed by knob name (mixin inheritance reorders fields, so positional construction would be a silent mis-binding) and updated with :func:dataclasses.replace;
• registers the subclass as a JAX pytree (and for jax.export treedef serialization), so instances ride :attr:AlgorithmState.hyper through jit / vmap / the exported batched solve_batched artifact;
• wires dtype handling off the field annotations when :meth:AlgorithmState.from_settings snapshots the instance onto the state — int fields get the iteration counter's dtype, float fields the problem float dtype. Any other annotation is rejected at class definition.

Example::

class MyHyper(HyperParams):
    ramp: float = 2.0
    drop: int = -1

The empty base is the "no declared hyperparameters" container: it flattens to zero pytree leaves.

Source code in openscvx/algorithms/hyperparams.py

class HyperParams:
    """Base for autotuner hyperparameter containers.

    Subclassing is the entire integration. Declare each tunable knob as an
    annotated field with its default — bare annotations, no ``@dataclass``
    decorator — and the base class does the rest:

    * applies the frozen, keyword-only dataclass transform, so instances are
      immutable value objects constructed by knob name (mixin inheritance
      reorders fields, so positional construction would be a silent
      mis-binding) and updated with :func:`dataclasses.replace`;
    * registers the subclass as a JAX pytree (and for ``jax.export`` treedef
      serialization), so instances ride :attr:`AlgorithmState.hyper` through
      ``jit`` / ``vmap`` / the exported batched ``solve_batched`` artifact;
    * wires dtype handling off the field annotations when
      :meth:`AlgorithmState.from_settings` snapshots the instance onto the
      state — ``int`` fields get the iteration counter's dtype, ``float``
      fields the problem float dtype. Any other annotation is rejected at
      class definition.

    Example::

        class MyHyper(HyperParams):
            ramp: float = 2.0
            drop: int = -1

    The empty base is the "no declared hyperparameters" container: it
    flattens to zero pytree leaves.
    """

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        # ``__init_subclass__`` runs before any decorator on the subclass
        # could, so the transform must happen here — which is exactly what
        # lets authors skip the dataclass and pytree mechanics entirely.
        # ``kw_only`` because mixins (e.g. ``LamCostRelaxHyper``) put inherited
        # fields ahead of a subclass's own: positional construction would bind
        # silently to the wrong knobs, and keyword-only also frees field order
        # from any inheritance constraint.
        dataclass(frozen=True, kw_only=True)(cls)
        for fld in dc_fields(cls):
            # fld.type is a string under postponed annotation evaluation
            # (``from __future__ import annotations``); accept both forms.
            if fld.type not in (int, float, "int", "float"):
                raise TypeError(
                    f"{cls.__name__}.{fld.name}: HyperParams fields must be "
                    f"annotated int or float (got {fld.type!r}) — the "
                    f"annotation decides the dtype the knob gets on "
                    f"AlgorithmState.hyper."
                )
        jax.tree_util.register_dataclass(cls)
        export.register_pytree_node_serialization(
            cls,
            serialized_name=f"{cls.__module__}.{cls.__qualname__}",
            serialize_auxdata=lambda aux: b"",
            deserialize_auxdata=lambda data: (),
        )

MultishotPropagation dataclass

Unpacked SCP multishot integration matrix V.

V shape: (n_segments * segment_size, n_substeps) where segment_size = n_x + n_x**2 + 2 * n_x * n_u.

Source code in openscvx/algorithms/history.py

@dataclass(frozen=True)
class MultishotPropagation:
    """Unpacked SCP multishot integration matrix ``V``.

    ``V`` shape: ``(n_segments * segment_size, n_substeps)`` where
    ``segment_size = n_x + n_x**2 + 2 * n_x * n_u``.
    """

    V: np.ndarray
    n_x: int
    n_u: int
    t_nodes: np.ndarray
    states: tuple = ()

    @property
    def segment_size(self) -> int:
        return segment_size(self.n_x, self.n_u)

    @property
    def n_segments(self) -> int:
        return int(self.V.shape[0] // self.segment_size)

    @property
    def n_substeps(self) -> int:
        return int(self.V.shape[1])

    def segment_states(self, seg_idx: int) -> np.ndarray:
        """Integrated states for one segment, shape ``(n_substeps, n_x)``."""
        seg_start = seg_idx * self.segment_size
        rows = self.V[seg_start : seg_start + self.n_x, :]
        return np.asarray(rows, dtype=np.float64).T

    def segments(self) -> list[np.ndarray]:
        """All segment state arrays, each shape ``(n_substeps, n_x)``."""
        return [self.segment_states(seg_idx) for seg_idx in range(self.n_segments)]

    def chronological(self) -> tuple[np.ndarray, np.ndarray]:
        """Time-ordered stitched full states and times.

        Skips duplicated segment-boundary samples (``j0 = 0`` on seg 0, else ``1``).

        Returns:
            ``states`` — ``(n_samples, n_x)``
            ``t`` — ``(n_samples,)``, linearly interpolated within each segment
        """
        t_nodes = np.asarray(self.t_nodes, dtype=np.float64).ravel()
        n_sub = self.n_substeps
        state_rows: list[np.ndarray] = []
        t_rows: list[float] = []
        for seg_idx in range(self.n_segments):
            t0, t1 = float(t_nodes[seg_idx]), float(t_nodes[seg_idx + 1])
            j0 = 0 if seg_idx == 0 else 1
            for j in range(j0, n_sub):
                alpha = j / (n_sub - 1) if n_sub > 1 else 0.0
                state_rows.append(self.segment_states(seg_idx)[j])
                t_rows.append((1.0 - alpha) * t0 + alpha * t1)
        if not state_rows:
            return np.empty((0, self.n_x), dtype=np.float64), np.empty(0, dtype=np.float64)
        return np.stack(state_rows, axis=0), np.asarray(t_rows, dtype=np.float64)

    def state(self, name_or_state: StateLike) -> tuple[np.ndarray, np.ndarray]:
        """Chronological trajectory for one symbolic state.

        ``name_or_state`` may be a :class:`~openscvx.symbolic.expr.state.State`
        instance or a string name matching ``State.name`` in ``self.states``.
        """
        state_slice = _resolve_state_slice(name_or_state, self.states)
        return self.slice_states(state_slice)

    def slice_states(self, state_slice: slice) -> tuple[np.ndarray, np.ndarray]:
        """``chronological()`` then apply ``state_slice`` to the state dimension."""
        states, t = self.chronological()
        return states[:, state_slice], t

chronological() -> tuple[np.ndarray, np.ndarray]

Time-ordered stitched full states and times.

Skips duplicated segment-boundary samples (j0 = 0 on seg 0, else 1).

Returns:

Type	Description
ndarray	states — (n_samples, n_x)
ndarray	t — (n_samples,), linearly interpolated within each segment

Source code in openscvx/algorithms/history.py

def chronological(self) -> tuple[np.ndarray, np.ndarray]:
    """Time-ordered stitched full states and times.

    Skips duplicated segment-boundary samples (``j0 = 0`` on seg 0, else ``1``).

    Returns:
        ``states`` — ``(n_samples, n_x)``
        ``t`` — ``(n_samples,)``, linearly interpolated within each segment
    """
    t_nodes = np.asarray(self.t_nodes, dtype=np.float64).ravel()
    n_sub = self.n_substeps
    state_rows: list[np.ndarray] = []
    t_rows: list[float] = []
    for seg_idx in range(self.n_segments):
        t0, t1 = float(t_nodes[seg_idx]), float(t_nodes[seg_idx + 1])
        j0 = 0 if seg_idx == 0 else 1
        for j in range(j0, n_sub):
            alpha = j / (n_sub - 1) if n_sub > 1 else 0.0
            state_rows.append(self.segment_states(seg_idx)[j])
            t_rows.append((1.0 - alpha) * t0 + alpha * t1)
    if not state_rows:
        return np.empty((0, self.n_x), dtype=np.float64), np.empty(0, dtype=np.float64)
    return np.stack(state_rows, axis=0), np.asarray(t_rows, dtype=np.float64)

segment_states(seg_idx: int) -> np.ndarray

Integrated states for one segment, shape (n_substeps, n_x).

Source code in openscvx/algorithms/history.py

def segment_states(self, seg_idx: int) -> np.ndarray:
    """Integrated states for one segment, shape ``(n_substeps, n_x)``."""
    seg_start = seg_idx * self.segment_size
    rows = self.V[seg_start : seg_start + self.n_x, :]
    return np.asarray(rows, dtype=np.float64).T

segments() -> list[np.ndarray]

All segment state arrays, each shape (n_substeps, n_x).

Source code in openscvx/algorithms/history.py

def segments(self) -> list[np.ndarray]:
    """All segment state arrays, each shape ``(n_substeps, n_x)``."""
    return [self.segment_states(seg_idx) for seg_idx in range(self.n_segments)]

slice_states(state_slice: slice) -> tuple[np.ndarray, np.ndarray]

chronological() then apply state_slice to the state dimension.

Source code in openscvx/algorithms/history.py

def slice_states(self, state_slice: slice) -> tuple[np.ndarray, np.ndarray]:
    """``chronological()`` then apply ``state_slice`` to the state dimension."""
    states, t = self.chronological()
    return states[:, state_slice], t

state(name_or_state: StateLike) -> tuple[np.ndarray, np.ndarray]

Chronological trajectory for one symbolic state.

name_or_state may be a :class:~openscvx.symbolic.expr.state.State instance or a string name matching State.name in self.states.

Source code in openscvx/algorithms/history.py

def state(self, name_or_state: StateLike) -> tuple[np.ndarray, np.ndarray]:
    """Chronological trajectory for one symbolic state.

    ``name_or_state`` may be a :class:`~openscvx.symbolic.expr.state.State`
    instance or a string name matching ``State.name`` in ``self.states``.
    """
    state_slice = _resolve_state_slice(name_or_state, self.states)
    return self.slice_states(state_slice)

OptimizationResults dataclass

Structured container for optimization results from the Successive Convexification (SCP) solver.

This class provides a type-safe and organized way to store and access optimization results, replacing the previous dictionary-based approach. It includes core optimization data, iteration history for convergence analysis, post-processing results, and flexible storage for plotting and application-specific data.

Attributes:

Name	Type	Description
converged	bool	Whether the optimization successfully converged.
t_final	float	Final time of the optimized trajectory.
x_guess	ndarray	Optimized state trajectory at discretization nodes, shape (N, n_states).
u_guess	ndarray	Optimized control trajectory at discretization nodes, shape (N, n_controls).
nodes	dict[str, ndarray]	Dictionary mapping state/control names to arrays at optimization nodes. Includes both user-defined and augmented variables.
trajectory	dict[str, ndarray]	Dictionary mapping state/control names to arrays along the propagated trajectory. Added by post_process().
x_history	list[ndarray]	State trajectories from each SCP iteration.
u_history	list[ndarray]	Control trajectories from each SCP iteration.
discretization_history	list[ndarray]	Time discretization from each iteration.
J_tr_history	list[ndarray]	Trust region cost history.
J_vb_history	list[ndarray]	Virtual buffer cost history.
J_vc_history	list[ndarray]	Virtual control cost history.
t_full	Optional[ndarray]	Full time grid for interpolated trajectory. Added by propagate_trajectory_results.
x_full	Optional[ndarray]	Interpolated state trajectory on full time grid. Added by propagate_trajectory_results.
u_full	Optional[ndarray]	Interpolated control trajectory on full time grid. Added by propagate_trajectory_results.
cost	Optional[float]	Total cost of the optimized trajectory. Added by propagate_trajectory_results.
ctcs_violation	Optional[ndarray]	Continuous-time constraint violations. Added by propagate_trajectory_results.
parameters	dict[str, ndarray]	Parameter values the solve used, recorded by :meth:~openscvx.problem.Problem.solve and :meth:~openscvx.problem.Problem.solve_batched (batched solves store every value with a leading (B,) axis, shared values replicated). Results record what happened: post-processing propagates with these values, and the Problem's parameter dict only seeds the next solve. Empty when no snapshot was recorded (e.g. solve_jax, whose callers manage parameters themselves).
plotting_data	dict[str, Any]	Flexible storage for plotting and application data.

JAX pytree registration

OptimizationResults is a registered JAX pytree, so the output of :meth:~openscvx.problem.Problem.solve_jax composes with jax.vmap / jax.jit / jax.grad. The pytree children are converged, t_final, nodes, trajectory, X, U, and every *_history field — the data the user actually consumes. Post-process fields (t_full, x_full, u_full, cost, ctcs_violation), parameters, plotting_data, and the internal _states / _controls metadata are stashed in the treedef's aux instead, so a batched solve_jax doesn't force callers to handle None leaves and post_process() outputs stay outside the vmap surface (post-process per element after the batched solve).

The two construction paths produce different histories:

• :meth:from_history (solve()) populates the per-iteration X / U / *_history lists from the Python loop's :class:~openscvx.algorithms.history.AlgorithmHistory.
• :meth:from_final_state (solve_jax()) wraps only the final iterate: X = [state.x] and U = [state.u] (single-element lists so result.x / result.u continue to return the final iterate), and every *_history field is empty — per-iteration history requires Python-side list growth that can't run inside lax.while_loop.

For Developers

The metadata={"npz": ...} parameter on each field below is a built-in feature of dataclasses.field. It attaches a read-only mapping to the field definition, not to instances. Instances still only have the normal attributes (result.X, result.converged, etc.).

The save() / load() methods iterate over field metadata to determine how to serialize each field, so adding a new field only requires tagging it here, no separate field-name lists to maintain.

Fields without an "npz" metadata entry are skipped during serialization (e.g. _states, _controls).

Source code in openscvx/algorithms/optimization_results.py

@jax.tree_util.register_pytree_node_class
@dataclass
class OptimizationResults:
    """
    Structured container for optimization results from the Successive Convexification (SCP) solver.

    This class provides a type-safe and organized way to store and access optimization results,
    replacing the previous dictionary-based approach. It includes core optimization data,
    iteration history for convergence analysis, post-processing results, and flexible
    storage for plotting and application-specific data.

    Attributes:
        converged (bool): Whether the optimization successfully converged.
        t_final (float): Final time of the optimized trajectory.
        x_guess (np.ndarray): Optimized state trajectory at discretization nodes,
            shape (N, n_states).
        u_guess (np.ndarray): Optimized control trajectory at discretization nodes,
            shape (N, n_controls).
        nodes (dict[str, np.ndarray]): Dictionary mapping state/control names to arrays
            at optimization nodes. Includes both user-defined and augmented variables.
        trajectory (dict[str, np.ndarray]): Dictionary mapping state/control names to arrays
            along the propagated trajectory. Added by post_process().
        x_history (list[np.ndarray]): State trajectories from each SCP iteration.
        u_history (list[np.ndarray]): Control trajectories from each SCP iteration.
        discretization_history (list[np.ndarray]): Time discretization from each iteration.
        J_tr_history (list[np.ndarray]): Trust region cost history.
        J_vb_history (list[np.ndarray]): Virtual buffer cost history.
        J_vc_history (list[np.ndarray]): Virtual control cost history.
        t_full (Optional[np.ndarray]): Full time grid for interpolated trajectory.
            Added by propagate_trajectory_results.
        x_full (Optional[np.ndarray]): Interpolated state trajectory on full time grid.
            Added by propagate_trajectory_results.
        u_full (Optional[np.ndarray]): Interpolated control trajectory on full time grid.
            Added by propagate_trajectory_results.
        cost (Optional[float]): Total cost of the optimized trajectory.
            Added by propagate_trajectory_results.
        ctcs_violation (Optional[np.ndarray]): Continuous-time constraint violations.
            Added by propagate_trajectory_results.
        parameters (dict[str, np.ndarray]): Parameter values the solve used, recorded
            by :meth:`~openscvx.problem.Problem.solve` and
            :meth:`~openscvx.problem.Problem.solve_batched` (batched solves store every
            value with a leading ``(B,)`` axis, shared values replicated). Results
            record what happened: post-processing propagates with these values, and
            the Problem's parameter dict only seeds the next solve. Empty when no
            snapshot was recorded (e.g. ``solve_jax``, whose callers manage
            parameters themselves).
        plotting_data (dict[str, Any]): Flexible storage for plotting and application data.

    !!! note "JAX pytree registration"
        ``OptimizationResults`` is a registered JAX pytree, so the output of
        :meth:`~openscvx.problem.Problem.solve_jax` composes with
        ``jax.vmap`` / ``jax.jit`` / ``jax.grad``. The pytree *children* are
        ``converged``, ``t_final``, ``nodes``, ``trajectory``, ``X``, ``U``,
        and every ``*_history`` field — the data the user actually consumes.
        Post-process fields (``t_full``, ``x_full``, ``u_full``, ``cost``,
        ``ctcs_violation``), ``parameters``, ``plotting_data``, and the
        internal ``_states`` / ``_controls`` metadata are stashed in the
        treedef's *aux* instead, so
        a batched ``solve_jax`` doesn't force callers to handle ``None``
        leaves and ``post_process()`` outputs stay outside the vmap surface
        (post-process per element after the batched solve).

        The two construction paths produce different histories:

        * :meth:`from_history` (``solve()``) populates the per-iteration
          ``X`` / ``U`` / ``*_history`` lists from the Python loop's
          :class:`~openscvx.algorithms.history.AlgorithmHistory`.
        * :meth:`from_final_state` (``solve_jax()``) wraps only the final
          iterate: ``X = [state.x]`` and ``U = [state.u]`` (single-element
          lists so ``result.x`` / ``result.u`` continue to return the final
          iterate), and every ``*_history`` field is empty — per-iteration
          history requires Python-side list growth that can't run inside
          ``lax.while_loop``.

    !!! note "For Developers"
        The ``metadata={"npz": ...}`` parameter on each field below is a built-in feature of
        `dataclasses.field`.  It attaches a read-only mapping to the **field definition**,
        *not* to instances. Instances still only have the normal attributes (``result.X``,
        ``result.converged``, etc.).

        The ``save()`` / ``load()`` methods iterate over field metadata to
        determine how to serialize each field, so adding a new field only
        requires tagging it here, no separate field-name lists to maintain.

        Fields *without* an ``"npz"`` metadata entry are skipped during
        serialization (e.g. ``_states``, ``_controls``).
    """

    _SCALAR = "scalar"
    _ARRAY_LIST = "array_list"
    _FLOAT_LIST = "float_list"
    _DICT = "dict"
    _OPT_ARRAY = "optional_array"
    _OPT_SCALAR = "optional_scalar"

    # Pytree children: fields whose contents flow through ``jax.vmap`` /
    # ``jax.jit`` / ``jax.grad``. Excludes ``_states`` / ``_controls`` (Python
    # symbolic metadata), ``plotting_data`` (user scratch), and the
    # post-process fields (``None`` until ``post_process()`` runs and not part
    # of the ``solve_jax`` vmap surface — see the class docstring).
    _PYTREE_CHILDREN = (
        "converged",
        "t_final",
        "nodes",
        "trajectory",
        "X",
        "U",
        "discretization_history",
        "J_tr_history",
        "J_vb_history",
        "J_vc_history",
        "TR_history",
        "VC_history",
        "lam_prox_history",
        "actual_reduction_history",
        "pred_reduction_history",
        "acceptance_ratio_history",
    )

    # Core optimization results
    converged: bool = field(metadata={"npz": "scalar"})
    t_final: float = field(metadata={"npz": "scalar"})

    # Dictionary-based access to states and controls
    nodes: dict[str, np.ndarray] = field(default_factory=dict, metadata={"npz": "dict"})
    trajectory: dict[str, np.ndarray] = field(default_factory=dict, metadata={"npz": "dict"})

    # Internal metadata for dictionary construction (not serialized)
    _states: list = field(default_factory=list, repr=False)
    _controls: list = field(default_factory=list, repr=False)

    # History of SCP iterations (single source of truth)
    X: list[np.ndarray] = field(default_factory=list, metadata={"npz": "array_list"})
    U: list[np.ndarray] = field(default_factory=list, metadata={"npz": "array_list"})
    discretization_history: list[np.ndarray] = field(
        default_factory=list, metadata={"npz": "array_list"}
    )
    J_tr_history: list[np.ndarray] = field(default_factory=list, metadata={"npz": "array_list"})
    J_vb_history: list[np.ndarray] = field(default_factory=list, metadata={"npz": "array_list"})
    J_vc_history: list[np.ndarray] = field(default_factory=list, metadata={"npz": "array_list"})
    TR_history: list[np.ndarray] = field(default_factory=list, metadata={"npz": "array_list"})
    VC_history: list[np.ndarray] = field(default_factory=list, metadata={"npz": "array_list"})

    # Convergence histories
    lam_prox_history: list[np.ndarray] = field(default_factory=list, metadata={"npz": "array_list"})
    actual_reduction_history: list[float] = field(
        default_factory=list, metadata={"npz": "float_list"}
    )
    pred_reduction_history: list[float] = field(
        default_factory=list, metadata={"npz": "float_list"}
    )
    acceptance_ratio_history: list[float] = field(
        default_factory=list, metadata={"npz": "float_list"}
    )

    @property
    def x(self) -> np.ndarray:
        """Optimal state trajectory at discretization nodes.

        Returns the final converged solution from the SCP iteration history.

        Returns:
            State trajectory array, shape (N, n_states)
        """
        return self.X[-1]

    @property
    def u(self) -> np.ndarray:
        """Optimal control trajectory at discretization nodes.

        Returns the final converged solution from the SCP iteration history.

        Returns:
            Control trajectory array, shape (N, n_controls)
        """
        return self.U[-1]

    def multishot_propagation(
        self,
        iteration: int = -1,
        *,
        t_nodes: Optional[np.ndarray] = None,
    ) -> Optional["MultishotPropagation"]:
        """Return unpacked multishot data for SCP iteration ``iteration`` (default: final).

        Returns ``None`` when ``discretization_history`` is empty (e.g. ``solve_jax``).

        Attaches ``self._states`` so ``.state("q")`` works without passing slices manually.
        """
        from openscvx.algorithms.history import unpack_multishot_V

        dh = self.discretization_history
        if not dh:
            return None
        V = np.asarray(dh[iteration], dtype=np.float64)
        n_x, n_u = self.x.shape[1], self.u.shape[1]
        if t_nodes is None:
            t_nodes_raw = self.nodes.get("time")
            if t_nodes_raw is None:
                t_final = float(self.t_final)
                t_nodes = np.linspace(0.0, t_final, self.x.shape[0])
            else:
                t_nodes = np.asarray(t_nodes_raw, dtype=np.float64).ravel()
        return unpack_multishot_V(
            V, n_x=n_x, n_u=n_u, t_nodes=t_nodes, states=tuple(self._states or ())
        )

    # Post-processing results (added by propagate_trajectory_results)
    t_full: Optional[np.ndarray] = field(default=None, metadata={"npz": "optional_array"})
    x_full: Optional[np.ndarray] = field(default=None, metadata={"npz": "optional_array"})
    u_full: Optional[np.ndarray] = field(default=None, metadata={"npz": "optional_array"})
    cost: Optional[float] = field(default=None, metadata={"npz": "optional_scalar"})
    ctcs_violation: Optional[np.ndarray] = field(default=None, metadata={"npz": "optional_array"})

    # Parameter values the solve used (recorded by Problem.solve() /
    # Problem.solve_batched(); empty when nothing was recorded)
    parameters: dict[str, np.ndarray] = field(default_factory=dict, metadata={"npz": "dict"})

    # Additional plotting/application data (added by user)
    plotting_data: dict[str, Any] = field(default_factory=dict, metadata={"npz": "dict"})

    def __post_init__(self):
        """Initialize the results object."""
        pass

    # ------------------------------------------------------------------
    # JAX pytree registration
    # ------------------------------------------------------------------

    def tree_flatten(self):
        """Split into JAX-traceable children and host-side aux.

        See :attr:`_PYTREE_CHILDREN` for the children list. Aux carries the
        symbolic metadata, scratch dict, and post-process fields verbatim so
        the round-trip ``tree_unflatten(*tree_flatten(self))`` preserves the
        instance.
        """
        children = tuple(getattr(self, name) for name in self._PYTREE_CHILDREN)
        aux = (
            self._states,
            self._controls,
            self.t_full,
            self.x_full,
            self.u_full,
            self.cost,
            self.ctcs_violation,
            self.parameters,
            self.plotting_data,
        )
        return children, aux

    @classmethod
    def tree_unflatten(cls, aux, children):
        kwargs = dict(zip(cls._PYTREE_CHILDREN, children))
        (
            _states,
            _controls,
            t_full,
            x_full,
            u_full,
            cost,
            ctcs_violation,
            parameters,
            plotting_data,
        ) = aux
        return cls(
            **kwargs,
            _states=_states,
            _controls=_controls,
            t_full=t_full,
            x_full=x_full,
            u_full=u_full,
            cost=cost,
            ctcs_violation=ctcs_violation,
            parameters=parameters,
            plotting_data=plotting_data,
        )

    # ------------------------------------------------------------------
    # Constructors
    # ------------------------------------------------------------------

    @classmethod
    def from_final_state(
        cls,
        state: "AlgorithmState",
        *,
        problem: "Problem",
    ) -> "OptimizationResults":
        """JAX-pure construction from the final SCP iterate.

        Used by :meth:`~openscvx.problem.Problem.solve_jax`. Every leaf is a
        ``jnp.ndarray`` derived from ``state`` via static slicing, so the
        result composes with ``jax.vmap`` / ``jax.jit`` / ``jax.grad`` over a
        batched ``state``. ``X = [state.x]`` and ``U = [state.u]`` are
        single-element lists so ``result.x`` / ``result.u`` continue to
        return the final iterate; every ``*_history`` field is empty — the
        Python-loop ``.solve()`` is the path that populates them.
        """
        settings = problem.settings
        symbolic = problem.symbolic

        has_impulsive_controls = any(
            control.parameterization == "impulsive" for control in symbolic.controls
        )

        nodes_dict: dict[str, jnp.ndarray] = {}
        for sym_state in symbolic.states:
            state_nodes = state.x[..., sym_state._slice]
            if has_impulsive_controls:
                bc = jnp.asarray(settings.sim.x.initial[sym_state._slice])
                state_nodes = state_nodes.at[..., 0, :].set(bc)
            nodes_dict[sym_state.name] = state_nodes

        for control in symbolic.controls:
            nodes_dict[control.name] = state.u[..., control._slice]

        # Match the host-path ``t_final`` shape: a 1-vector along the time
        # slice of the last node (vmap'd to ``(B, 1)``).
        t_final = state.x[..., -1, :][..., settings.sim.time_slice]

        # Convergence flag through the algorithm's own predicate, so this JAX
        # results path agrees with ``solve()`` / the ``lax.while_loop``. The
        # per-element thresholds ride the state pytree, so a batched solve with
        # a tolerance sweep reports convergence against each element's ep_*.
        converged = problem.algorithm.converged(state)

        return cls(
            converged=converged,
            t_final=t_final,
            nodes=nodes_dict,
            trajectory={},
            _states=symbolic.states_prop,
            _controls=symbolic.controls,
            X=[state.x],
            U=[state.u],
            discretization_history=[],
            J_tr_history=[],
            J_vb_history=[],
            J_vc_history=[],
            TR_history=[],
            VC_history=[],
            lam_prox_history=[],
            actual_reduction_history=[],
            pred_reduction_history=[],
            acceptance_ratio_history=[],
        )

    @classmethod
    def from_history(
        cls,
        history: "AlgorithmHistory",
        final_state: "AlgorithmState",
        *,
        problem: "Problem",
        converged: bool,
    ) -> "OptimizationResults":
        """Host-side construction from the Python-loop iteration history.

        Used by :meth:`~openscvx.problem.Problem.solve` to package the full
        per-iteration log produced by the Python ``while`` loop. Symmetric
        with :meth:`from_final_state` — same return type, populated history.
        """
        settings = problem.settings
        symbolic = problem.symbolic

        has_impulsive_controls = any(
            control.parameterization == "impulsive" for control in symbolic.controls
        )

        x_arr = np.asarray(final_state.x)
        u_arr = np.asarray(final_state.u)

        nodes_dict: dict[str, np.ndarray] = {}
        for sym_state in symbolic.states:
            state_nodes = x_arr[:, sym_state._slice].copy()
            if has_impulsive_controls:
                state_nodes[0] = settings.sim.x.initial[sym_state._slice]
            nodes_dict[sym_state.name] = state_nodes

        for control in symbolic.controls:
            nodes_dict[control.name] = u_arr[:, control._slice]

        return cls(
            converged=converged,
            t_final=x_arr[:, settings.sim.time_slice][-1],
            nodes=nodes_dict,
            trajectory={},
            _states=symbolic.states_prop,
            _controls=symbolic.controls,
            X=history.X,
            U=history.U,
            discretization_history=history.V_history,
            J_tr_history=float(final_state.J_tr),
            J_vb_history=float(final_state.J_vb),
            J_vc_history=float(final_state.J_vc),
            TR_history=history.TR,
            VC_history=history.VC,
            lam_prox_history=history.lam_prox.copy(),
            actual_reduction_history=history.actual_reduction.copy(),
            pred_reduction_history=history.pred_reduction.copy(),
            acceptance_ratio_history=history.acceptance_ratio.copy(),
        )

    def update_plotting_data(self, **kwargs: Any) -> None:
        """
        Update the plotting data with additional information.

        Args:
            **kwargs: Key-value pairs to add to plotting_data
        """
        self.plotting_data.update(kwargs)

    def get(self, key: str, default: Any = None) -> Any:
        """
        Get a value from the results, similar to dict.get().

        Args:
            key: The key to look up
            default: Default value if key is not found

        Returns:
            The value associated with the key, or default if not found
        """
        # Check if it's a direct attribute
        if hasattr(self, key):
            return getattr(self, key)

        # Check if it's in plotting_data
        if key in self.plotting_data:
            return self.plotting_data[key]

        return default

    def __getitem__(self, key: str) -> Any:
        """
        Allow dictionary-style access to results.

        Args:
            key: The key to look up

        Returns:
            The value associated with the key

        Raises:
            KeyError: If key is not found
        """
        # Check if it's a direct attribute
        if hasattr(self, key):
            return getattr(self, key)

        # Check if it's in plotting_data
        if key in self.plotting_data:
            return self.plotting_data[key]

        raise KeyError(f"Key '{key}' not found in results")

    def __setitem__(self, key: str, value: Any):
        """
        Allow dictionary-style assignment to results.

        Args:
            key: The key to set
            value: The value to assign
        """
        # Check if it's a direct attribute
        if hasattr(self, key):
            setattr(self, key, value)
        else:
            # Store in plotting_data
            self.plotting_data[key] = value

    def __contains__(self, key: str) -> bool:
        """
        Check if a key exists in the results.

        Args:
            key: The key to check

        Returns:
            True if key exists, False otherwise
        """
        return hasattr(self, key) or key in self.plotting_data

    def update(self, other: dict[str, Any]):
        """
        Update the results with additional data from a dictionary.

        Args:
            other: Dictionary containing additional data
        """
        for key, value in other.items():
            self[key] = value

    def to_dict(self) -> dict[str, Any]:
        """
        Convert the results to a dictionary for backward compatibility.

        Returns:
            Dictionary representation of the results
        """
        result_dict = {}

        # Add all direct attributes
        for attr_name in self.__dataclass_fields__:
            if attr_name != "plotting_data":
                result_dict[attr_name] = getattr(self, attr_name)

        # Add plotting data
        result_dict.update(self.plotting_data)

        return result_dict

    # ------------------------------------------------------------------
    # Serialization
    # ------------------------------------------------------------------

    def save(self, path: Union[str, Path]) -> None:
        """Save results to a compressed ``.npz`` file.

        Serialization behaviour is driven by the ``"npz"`` key in each
        dataclass field's ``metadata``.  Fields without that key are
        skipped.

        Args:
            path: Output file path.  ``.npz`` is appended automatically
                by numpy if not already present.
        """
        data: dict[str, np.ndarray] = {}

        for name, f in self.__dataclass_fields__.items():
            tag = f.metadata.get("npz")
            if tag is None:
                continue
            val = getattr(self, name)

            if tag == self._SCALAR:
                data[name] = np.array(val)

            elif tag == self._ARRAY_LIST:
                data[name] = np.stack([np.asarray(a) for a in val]) if val else np.array([])

            elif tag == self._FLOAT_LIST:
                data[name] = np.array(val, dtype=float)

            elif tag == self._DICT:
                for k, v in val.items():
                    try:
                        arr = np.asarray(v)
                        if arr.dtype.kind != "O":
                            data[f"{name}/{k}"] = arr
                    except (TypeError, ValueError):
                        pass

            elif tag == self._OPT_ARRAY:
                if val is not None:
                    data[name] = np.asarray(val)

            elif tag == self._OPT_SCALAR:
                if val is not None:
                    data[name] = np.array(val)

        np.savez_compressed(str(path), **data)

    @classmethod
    def load(cls, path: Union[str, Path]) -> "OptimizationResults":
        """Load results from a ``.npz`` file previously created by :meth:`save`.

        Args:
            path: Path to the ``.npz`` file.  If the path has no suffix,
                ``.npz`` is appended automatically.

        Returns:
            Reconstructed :class:`OptimizationResults`.
        """
        path = Path(path)
        if not path.suffix:
            path = path.with_suffix(".npz")

        archive = np.load(str(path), allow_pickle=False)
        keys = set(archive.files)

        def _unstack(name: str) -> list:
            if name not in keys:
                return []
            arr = archive[name]
            return [arr[i] for i in range(arr.shape[0])] if arr.size else []

        def _prefixed_dict(prefix: str) -> dict:
            p = prefix + "/"
            return {k[len(p) :]: archive[k] for k in keys if k.startswith(p)}

        kwargs: dict[str, Any] = {}
        deferred: dict[str, tuple] = {}  # fields to set after construction

        for name, f in cls.__dataclass_fields__.items():
            tag = f.metadata.get("npz")
            if tag is None:
                continue

            if tag == cls._SCALAR:
                kwargs[name] = archive[name].item()

            elif tag == cls._ARRAY_LIST:
                kwargs[name] = _unstack(name)

            elif tag == cls._FLOAT_LIST:
                kwargs[name] = archive[name].tolist() if name in keys else []

            elif tag == cls._DICT:
                kwargs[name] = _prefixed_dict(name)

            elif tag == cls._OPT_ARRAY:
                deferred[name] = archive[name] if name in keys else None

            elif tag == cls._OPT_SCALAR:
                deferred[name] = float(archive[name]) if name in keys else None

        result = cls(**kwargs)
        for name, val in deferred.items():
            setattr(result, name, val)
        return result

u: np.ndarray property

Optimal control trajectory at discretization nodes.

Returns the final converged solution from the SCP iteration history.

Returns:

Type	Description
ndarray	Control trajectory array, shape (N, n_controls)

x: np.ndarray property

Optimal state trajectory at discretization nodes.

Returns the final converged solution from the SCP iteration history.

Returns:

Type	Description
ndarray	State trajectory array, shape (N, n_states)

from_final_state(state: AlgorithmState, *, problem: Problem) -> OptimizationResults classmethod

JAX-pure construction from the final SCP iterate.

Used by :meth:~openscvx.problem.Problem.solve_jax. Every leaf is a jnp.ndarray derived from state via static slicing, so the result composes with jax.vmap / jax.jit / jax.grad over a batched state. X = [state.x] and U = [state.u] are single-element lists so result.x / result.u continue to return the final iterate; every *_history field is empty — the Python-loop .solve() is the path that populates them.

Source code in openscvx/algorithms/optimization_results.py

@classmethod
def from_final_state(
    cls,
    state: "AlgorithmState",
    *,
    problem: "Problem",
) -> "OptimizationResults":
    """JAX-pure construction from the final SCP iterate.

    Used by :meth:`~openscvx.problem.Problem.solve_jax`. Every leaf is a
    ``jnp.ndarray`` derived from ``state`` via static slicing, so the
    result composes with ``jax.vmap`` / ``jax.jit`` / ``jax.grad`` over a
    batched ``state``. ``X = [state.x]`` and ``U = [state.u]`` are
    single-element lists so ``result.x`` / ``result.u`` continue to
    return the final iterate; every ``*_history`` field is empty — the
    Python-loop ``.solve()`` is the path that populates them.
    """
    settings = problem.settings
    symbolic = problem.symbolic

    has_impulsive_controls = any(
        control.parameterization == "impulsive" for control in symbolic.controls
    )

    nodes_dict: dict[str, jnp.ndarray] = {}
    for sym_state in symbolic.states:
        state_nodes = state.x[..., sym_state._slice]
        if has_impulsive_controls:
            bc = jnp.asarray(settings.sim.x.initial[sym_state._slice])
            state_nodes = state_nodes.at[..., 0, :].set(bc)
        nodes_dict[sym_state.name] = state_nodes

    for control in symbolic.controls:
        nodes_dict[control.name] = state.u[..., control._slice]

    # Match the host-path ``t_final`` shape: a 1-vector along the time
    # slice of the last node (vmap'd to ``(B, 1)``).
    t_final = state.x[..., -1, :][..., settings.sim.time_slice]

    # Convergence flag through the algorithm's own predicate, so this JAX
    # results path agrees with ``solve()`` / the ``lax.while_loop``. The
    # per-element thresholds ride the state pytree, so a batched solve with
    # a tolerance sweep reports convergence against each element's ep_*.
    converged = problem.algorithm.converged(state)

    return cls(
        converged=converged,
        t_final=t_final,
        nodes=nodes_dict,
        trajectory={},
        _states=symbolic.states_prop,
        _controls=symbolic.controls,
        X=[state.x],
        U=[state.u],
        discretization_history=[],
        J_tr_history=[],
        J_vb_history=[],
        J_vc_history=[],
        TR_history=[],
        VC_history=[],
        lam_prox_history=[],
        actual_reduction_history=[],
        pred_reduction_history=[],
        acceptance_ratio_history=[],
    )

from_history(history: AlgorithmHistory, final_state: AlgorithmState, *, problem: Problem, converged: bool) -> OptimizationResults classmethod

Host-side construction from the Python-loop iteration history.

Used by :meth:~openscvx.problem.Problem.solve to package the full per-iteration log produced by the Python while loop. Symmetric with :meth:from_final_state — same return type, populated history.

Source code in openscvx/algorithms/optimization_results.py

@classmethod
def from_history(
    cls,
    history: "AlgorithmHistory",
    final_state: "AlgorithmState",
    *,
    problem: "Problem",
    converged: bool,
) -> "OptimizationResults":
    """Host-side construction from the Python-loop iteration history.

    Used by :meth:`~openscvx.problem.Problem.solve` to package the full
    per-iteration log produced by the Python ``while`` loop. Symmetric
    with :meth:`from_final_state` — same return type, populated history.
    """
    settings = problem.settings
    symbolic = problem.symbolic

    has_impulsive_controls = any(
        control.parameterization == "impulsive" for control in symbolic.controls
    )

    x_arr = np.asarray(final_state.x)
    u_arr = np.asarray(final_state.u)

    nodes_dict: dict[str, np.ndarray] = {}
    for sym_state in symbolic.states:
        state_nodes = x_arr[:, sym_state._slice].copy()
        if has_impulsive_controls:
            state_nodes[0] = settings.sim.x.initial[sym_state._slice]
        nodes_dict[sym_state.name] = state_nodes

    for control in symbolic.controls:
        nodes_dict[control.name] = u_arr[:, control._slice]

    return cls(
        converged=converged,
        t_final=x_arr[:, settings.sim.time_slice][-1],
        nodes=nodes_dict,
        trajectory={},
        _states=symbolic.states_prop,
        _controls=symbolic.controls,
        X=history.X,
        U=history.U,
        discretization_history=history.V_history,
        J_tr_history=float(final_state.J_tr),
        J_vb_history=float(final_state.J_vb),
        J_vc_history=float(final_state.J_vc),
        TR_history=history.TR,
        VC_history=history.VC,
        lam_prox_history=history.lam_prox.copy(),
        actual_reduction_history=history.actual_reduction.copy(),
        pred_reduction_history=history.pred_reduction.copy(),
        acceptance_ratio_history=history.acceptance_ratio.copy(),
    )

get(key: str, default: Any = None) -> Any

Get a value from the results, similar to dict.get().

Parameters:

Name	Type	Description	Default
key	str	The key to look up	required
default	Any	Default value if key is not found	None

Returns:

Type	Description
Any	The value associated with the key, or default if not found

Source code in openscvx/algorithms/optimization_results.py

def get(self, key: str, default: Any = None) -> Any:
    """
    Get a value from the results, similar to dict.get().

    Args:
        key: The key to look up
        default: Default value if key is not found

    Returns:
        The value associated with the key, or default if not found
    """
    # Check if it's a direct attribute
    if hasattr(self, key):
        return getattr(self, key)

    # Check if it's in plotting_data
    if key in self.plotting_data:
        return self.plotting_data[key]

    return default

load(path: Union[str, Path]) -> OptimizationResults classmethod

Load results from a .npz file previously created by :meth:save.

Parameters:

Name	Type	Description	Default
path	Union[str, Path]	Path to the .npz file. If the path has no suffix, .npz is appended automatically.	required

Returns:

Name	Type	Description
Reconstructed	OptimizationResults	class:OptimizationResults.

Source code in openscvx/algorithms/optimization_results.py

@classmethod
def load(cls, path: Union[str, Path]) -> "OptimizationResults":
    """Load results from a ``.npz`` file previously created by :meth:`save`.

    Args:
        path: Path to the ``.npz`` file.  If the path has no suffix,
            ``.npz`` is appended automatically.

    Returns:
        Reconstructed :class:`OptimizationResults`.
    """
    path = Path(path)
    if not path.suffix:
        path = path.with_suffix(".npz")

    archive = np.load(str(path), allow_pickle=False)
    keys = set(archive.files)

    def _unstack(name: str) -> list:
        if name not in keys:
            return []
        arr = archive[name]
        return [arr[i] for i in range(arr.shape[0])] if arr.size else []

    def _prefixed_dict(prefix: str) -> dict:
        p = prefix + "/"
        return {k[len(p) :]: archive[k] for k in keys if k.startswith(p)}

    kwargs: dict[str, Any] = {}
    deferred: dict[str, tuple] = {}  # fields to set after construction

    for name, f in cls.__dataclass_fields__.items():
        tag = f.metadata.get("npz")
        if tag is None:
            continue

        if tag == cls._SCALAR:
            kwargs[name] = archive[name].item()

        elif tag == cls._ARRAY_LIST:
            kwargs[name] = _unstack(name)

        elif tag == cls._FLOAT_LIST:
            kwargs[name] = archive[name].tolist() if name in keys else []

        elif tag == cls._DICT:
            kwargs[name] = _prefixed_dict(name)

        elif tag == cls._OPT_ARRAY:
            deferred[name] = archive[name] if name in keys else None

        elif tag == cls._OPT_SCALAR:
            deferred[name] = float(archive[name]) if name in keys else None

    result = cls(**kwargs)
    for name, val in deferred.items():
        setattr(result, name, val)
    return result

multishot_propagation(iteration: int = -1, *, t_nodes: Optional[np.ndarray] = None) -> Optional[MultishotPropagation]

Return unpacked multishot data for SCP iteration iteration (default: final).

Returns None when discretization_history is empty (e.g. solve_jax).

Attaches self._states so .state("q") works without passing slices manually.

Source code in openscvx/algorithms/optimization_results.py

def multishot_propagation(
    self,
    iteration: int = -1,
    *,
    t_nodes: Optional[np.ndarray] = None,
) -> Optional["MultishotPropagation"]:
    """Return unpacked multishot data for SCP iteration ``iteration`` (default: final).

    Returns ``None`` when ``discretization_history`` is empty (e.g. ``solve_jax``).

    Attaches ``self._states`` so ``.state("q")`` works without passing slices manually.
    """
    from openscvx.algorithms.history import unpack_multishot_V

    dh = self.discretization_history
    if not dh:
        return None
    V = np.asarray(dh[iteration], dtype=np.float64)
    n_x, n_u = self.x.shape[1], self.u.shape[1]
    if t_nodes is None:
        t_nodes_raw = self.nodes.get("time")
        if t_nodes_raw is None:
            t_final = float(self.t_final)
            t_nodes = np.linspace(0.0, t_final, self.x.shape[0])
        else:
            t_nodes = np.asarray(t_nodes_raw, dtype=np.float64).ravel()
    return unpack_multishot_V(
        V, n_x=n_x, n_u=n_u, t_nodes=t_nodes, states=tuple(self._states or ())
    )

save(path: Union[str, Path]) -> None

Save results to a compressed .npz file.

Serialization behaviour is driven by the "npz" key in each dataclass field's metadata. Fields without that key are skipped.

Parameters:

Name	Type	Description	Default
path	Union[str, Path]	Output file path. .npz is appended automatically by numpy if not already present.	required

Source code in openscvx/algorithms/optimization_results.py

def save(self, path: Union[str, Path]) -> None:
    """Save results to a compressed ``.npz`` file.

    Serialization behaviour is driven by the ``"npz"`` key in each
    dataclass field's ``metadata``.  Fields without that key are
    skipped.

    Args:
        path: Output file path.  ``.npz`` is appended automatically
            by numpy if not already present.
    """
    data: dict[str, np.ndarray] = {}

    for name, f in self.__dataclass_fields__.items():
        tag = f.metadata.get("npz")
        if tag is None:
            continue
        val = getattr(self, name)

        if tag == self._SCALAR:
            data[name] = np.array(val)

        elif tag == self._ARRAY_LIST:
            data[name] = np.stack([np.asarray(a) for a in val]) if val else np.array([])

        elif tag == self._FLOAT_LIST:
            data[name] = np.array(val, dtype=float)

        elif tag == self._DICT:
            for k, v in val.items():
                try:
                    arr = np.asarray(v)
                    if arr.dtype.kind != "O":
                        data[f"{name}/{k}"] = arr
                except (TypeError, ValueError):
                    pass

        elif tag == self._OPT_ARRAY:
            if val is not None:
                data[name] = np.asarray(val)

        elif tag == self._OPT_SCALAR:
            if val is not None:
                data[name] = np.array(val)

    np.savez_compressed(str(path), **data)

to_dict() -> dict[str, Any]

Convert the results to a dictionary for backward compatibility.

Returns:

Type	Description
dict[str, Any]	Dictionary representation of the results

Source code in openscvx/algorithms/optimization_results.py

def to_dict(self) -> dict[str, Any]:
    """
    Convert the results to a dictionary for backward compatibility.

    Returns:
        Dictionary representation of the results
    """
    result_dict = {}

    # Add all direct attributes
    for attr_name in self.__dataclass_fields__:
        if attr_name != "plotting_data":
            result_dict[attr_name] = getattr(self, attr_name)

    # Add plotting data
    result_dict.update(self.plotting_data)

    return result_dict

tree_flatten()

Split into JAX-traceable children and host-side aux.

See :attr:_PYTREE_CHILDREN for the children list. Aux carries the symbolic metadata, scratch dict, and post-process fields verbatim so the round-trip tree_unflatten(*tree_flatten(self)) preserves the instance.

Source code in openscvx/algorithms/optimization_results.py

def tree_flatten(self):
    """Split into JAX-traceable children and host-side aux.

    See :attr:`_PYTREE_CHILDREN` for the children list. Aux carries the
    symbolic metadata, scratch dict, and post-process fields verbatim so
    the round-trip ``tree_unflatten(*tree_flatten(self))`` preserves the
    instance.
    """
    children = tuple(getattr(self, name) for name in self._PYTREE_CHILDREN)
    aux = (
        self._states,
        self._controls,
        self.t_full,
        self.x_full,
        self.u_full,
        self.cost,
        self.ctcs_violation,
        self.parameters,
        self.plotting_data,
    )
    return children, aux

update(other: dict[str, Any])

Update the results with additional data from a dictionary.

Parameters:

Name	Type	Description	Default
other	dict[str, Any]	Dictionary containing additional data	required

Source code in openscvx/algorithms/optimization_results.py

def update(self, other: dict[str, Any]):
    """
    Update the results with additional data from a dictionary.

    Args:
        other: Dictionary containing additional data
    """
    for key, value in other.items():
        self[key] = value

update_plotting_data(**kwargs: Any) -> None

Update the plotting data with additional information.

Parameters:

Name	Type	Description	Default
**kwargs	Any	Key-value pairs to add to plotting_data	{}

Source code in openscvx/algorithms/optimization_results.py

def update_plotting_data(self, **kwargs: Any) -> None:
    """
    Update the plotting data with additional information.

    Args:
        **kwargs: Key-value pairs to add to plotting_data
    """
    self.plotting_data.update(kwargs)

PenalizedTrustRegion

Bases: Algorithm

Penalized Trust Region (PTR) successive convexification algorithm.

PTR solves non-convex trajectory optimization problems through iterative convex approximation. Each subproblem balances competing cost terms:

• Trust region penalty: Discourages large deviations from the previous iterate, keeping the solution within the region where linearization is valid.
• Virtual control: Relaxes dynamics constraints, penalized to drive defects toward zero as the algorithm converges.
• Virtual buffer: Relaxes non-convex constraints, similarly penalized to enforce feasibility at convergence.
• Problem objective and other terms: The user-defined cost (e.g., minimum fuel, minimum time) and any additional penalty terms.

The interplay between these terms guides the optimization: the trust region anchors the solution near the linearization point while virtual terms allow temporary constraint violations that shrink over iterations.

Example

Using PTR with a Problem::

from openscvx.algorithms import PenalizedTrustRegion

problem = Problem(dynamics, constraints, states, controls, N, time)
problem.initialize()
result = problem.solve()

Source code in openscvx/algorithms/scvx/penalized_trust_region.py

class PenalizedTrustRegion(Algorithm):
    """Penalized Trust Region (PTR) successive convexification algorithm.

    PTR solves non-convex trajectory optimization problems through iterative
    convex approximation. Each subproblem balances competing cost terms:

    - **Trust region penalty**: Discourages large deviations from the previous
      iterate, keeping the solution within the region where linearization is valid.
    - **Virtual control**: Relaxes dynamics constraints, penalized to drive
      defects toward zero as the algorithm converges.
    - **Virtual buffer**: Relaxes non-convex constraints, similarly penalized
      to enforce feasibility at convergence.
    - **Problem objective and other terms**: The user-defined cost (e.g., minimum
      fuel, minimum time) and any additional penalty terms.

    The interplay between these terms guides the optimization: the trust region
    anchors the solution near the linearization point while virtual terms allow
    temporary constraint violations that shrink over iterations.

    Example:
        Using PTR with a Problem::

            from openscvx.algorithms import PenalizedTrustRegion

            problem = Problem(dynamics, constraints, states, controls, N, time)
            problem.initialize()
            result = problem.solve()
    """

    # Base columns emitted by PTR algorithm (before autotuner columns)
    BASE_COLUMNS: List[Column] = [
        Column("iter", "Iter", 4, "{:4d}"),
        Column("subprop_time", "Step (ms)", 10, "{:6.2f}", min_verbosity=Verbosity.STANDARD),
        Column("cost", "Cost", 8, "{: .1e}"),
        Column("J_tr", "J_tr", 8, "{: .1e}", color_J_tr, Verbosity.STANDARD),
        Column("J_vb", "J_vb", 8, "{: .1e}", color_J_vb, Verbosity.STANDARD),
        Column("J_vc", "J_vc", 8, "{: .1e}", color_J_vc, Verbosity.STANDARD),
    ]

    # Columns that always appear last (after autotuner columns)
    TAIL_COLUMNS: List[Column] = [
        Column("prob_stat", "Cvx Status", 11, "{}", color_prob_stat),
    ]

    def __init__(
        self,
        autotuner: "AutotuningBase" = None,
        k_max: int = 200,
        t_max: float | None = None,
        lam_prox: Union[float, Dict[str, Union[float, list]]] = 1e-1,
        lam_vc: Union[float, Dict[str, Union[float, list]]] = 1e0,
        lam_cost: Union[float, Dict[str, float]] = 1e-2,
        lam_vb: float = 0.0,
        ep_tr: float = 1e-4,
        ep_vb: float = 1e-4,
        ep_vc: float = 1e-8,
        states: List["State"] = None,
        controls: List["Control"] = None,
    ):
        """Initialize PTR with algorithm parameters and optional autotuner."""
        # Compiled infrastructure (set by initialize()). ``_iteration_fn`` is
        # the fused JAX-pure SCP body; ``step()`` is a thin Python wrapper that
        # calls it, records history, and emits.
        self._iteration_fn: Callable | None = None
        self._emitter: Callable | None = None

        # Store states/controls for later re-resolution.
        self._states: List["State"] = states
        self._controls: List["Control"] = controls

        super().__init__(
            # SCP weights (grouped dataclass, dict inputs expanded to arrays).
            weights=Weights.build(
                lam_prox=lam_prox,
                lam_vc=lam_vc,
                lam_cost=lam_cost,
                lam_vb=lam_vb,
                states=states,
                controls=controls,
            ),
            autotuner=autotuner if autotuner is not None else AugmentedLagrangian(),
            k_max=k_max,
            t_max=t_max,
            ep_tr=ep_tr,
            ep_vb=ep_vb,
            ep_vc=ep_vc,
        )

    @property
    def lam_prox(self) -> Union[float, np.ndarray]:
        return self.weights.lam_prox

    @lam_prox.setter
    def lam_prox(self, value: Union[float, Dict[str, Union[float, list]]]) -> None:
        self.weights.lam_prox = Weights.resolve_lam_prox(value, self._states, self._controls)

    @property
    def lam_vc(self) -> Union[float, np.ndarray]:
        return self.weights.lam_vc

    @lam_vc.setter
    def lam_vc(self, value: Union[float, Dict[str, Union[float, list]]]) -> None:
        self.weights.lam_vc = Weights.resolve_lam_vc(value, self._states)

    @property
    def lam_cost(self) -> Union[float, np.ndarray]:
        return self.weights.lam_cost

    @lam_cost.setter
    def lam_cost(self, value: Union[float, Dict[str, float]]) -> None:
        self.weights.lam_cost = Weights.resolve_lam_cost(value, self._states)

    @property
    def lam_vb(self) -> float:
        return self.weights.lam_vb

    @lam_vb.setter
    def lam_vb(self, value: float) -> None:
        self.weights.lam_vb = float(value)

    def build_iteration(
        self,
        dis_continuous: Callable,
        dis_impulsive: Callable,
        jax_constraints: "LoweredJaxConstraints",
        solver_callback: Callable,
        settings: Config,
    ) -> Callable:
        """Fuse the discretizers, constraints, and solver into the PTR step.

        Thin wrapper around
        :func:`~openscvx.algorithms.scvx.iteration.make_scp_iteration`,
        threading this algorithm's :attr:`autotuner` into the fused body.
        """
        return make_scp_iteration(
            dis_continuous=dis_continuous,
            dis_impulsive=dis_impulsive,
            jax_constraints=jax_constraints,
            solver_callback=solver_callback,
            autotuner=self.autotuner,
            settings=settings,
        )

    def get_columns(self, verbosity: int = Verbosity.STANDARD) -> List[Column]:
        """Get the columns to display for iteration output."""
        all_columns = self.BASE_COLUMNS + self.autotuner.COLUMNS + self.TAIL_COLUMNS
        return [col for col in all_columns if col.min_verbosity <= verbosity]

    def initialize(self, iteration_fn: Callable, emitter: Callable) -> None:
        """Store the fused SCP iteration body and per-iteration infrastructure.

        ``iteration_fn`` is built and JIT-warmed by :meth:`Problem.initialize`;
        :meth:`step` is a thin Python wrapper that calls it, records history,
        and emits diagnostics. The boundary conditions the subproblem pins are
        carried on :class:`AlgorithmState` (``x_init_pin`` / ``x_term_pin``), so
        there is no solver state to prime here.
        """
        self._iteration_fn = iteration_fn
        self._emitter = emitter

    def step(
        self,
        state: AlgorithmState,
        history: AlgorithmHistory,
        params: dict,
        settings: Config,
    ) -> Tuple[AlgorithmState, bool]:
        """Execute one PTR iteration and return ``(next_state, converged)``.

        Calls the fused JAX iteration body, records the per-iteration
        diagnostics it returns into ``history``, emits progress, and reports
        convergence from the metrics on the next state.
        """
        if self._iteration_fn is None:
            raise RuntimeError(
                "PenalizedTrustRegion.step() called before initialize(). "
                "Call initialize() first to set up the iteration body."
            )

        iter_index = int(state.k)

        t0 = time.time()
        next_state, diag = self._iteration_fn(state, params)
        jax.block_until_ready((next_state, diag))
        step_time = time.time() - t0

        # Fail loudly on a bad subproblem solve before it becomes the next
        # linearization point. Two gates: a non-OPTIMAL status from the backend
        # (e.g. QPAX divergence, CVXPy infeasibility) and a non-finite iterate
        # (NaN/Inf from anywhere in the body). Either would silently corrupt the
        # rest of the SCP run.
        if int(diag.status) != int(StatusCode.OPTIMAL):
            raise RuntimeError(
                f"Convex subproblem did not solve to optimality "
                f"(status={status_code_to_str(int(diag.status))!r}). The backend failed "
                f"to converge or the subproblem is infeasible — adjust solver tolerances, "
                f"rescale the problem, or use float_dtype='float64'."
            )
        if not bool(np.all(np.isfinite(np.asarray(next_state.x)))):
            raise RuntimeError(
                "Subproblem solve produced a non-finite iterate (NaN/Inf in the state)."
            )

        use_full_metrics = self.autotuner.COMPUTES_ACCEPTANCE_METRICS
        scalars, lam_prox_np = history.record_iteration(
            next_state,
            V=np.asarray(diag.V),
            W=np.asarray(diag.W),
            VC=np.asarray(diag.VC),
            TR=np.asarray(diag.TR),
            J_lin=float(diag.J_lin),
            record_diagnostics=use_full_metrics,
        )

        emission_data = {
            "iter": iter_index,
            "subprop_time": step_time * 1000.0,
            "J_tr": scalars["J_tr"],
            "J_vb": scalars["J_vb"],
            "J_vc": scalars["J_vc"],
            "cost": float(diag.cost),
            # TODO: (haynec) log per-variable lam_prox detail (e.g. min/max range)
            "lam_prox": float(np.max(lam_prox_np)),
            "prob_stat": status_code_to_str(int(diag.status)),
            "adaptive_state": adaptive_state_code_to_str(scalars["adaptive_state_code"]),
            "ep_tr": self.ep_tr,
            "ep_vb": self.ep_vb,
            "ep_vc": self.ep_vc,
        }

        if use_full_metrics:
            emission_data.update(
                {
                    "J_nonlin": scalars["J_nonlin"],
                    "J_lin": float(diag.J_lin),
                    "pred_reduction": scalars["predicted_reduction"],
                    "actual_reduction": scalars["actual_reduction"],
                    "acceptance_ratio": scalars["acceptance_ratio"],
                }
            )

        self._emitter(emission_data)

        # Same convergence test the lax.while_loop path uses, reading the
        # thresholds off the state pytree (synced from this algorithm's
        # ep_* attributes by Problem._sync_scp_constants at solve() time).
        return next_state, bool(self.converged(next_state))

    def citation(self) -> List[str]:
        """Return BibTeX citations for the PTR algorithm."""
        return [
            r"""@article{drusvyatskiy2018error,
  title={Error bounds, quadratic growth, and linear convergence of proximal methods},
  author={Drusvyatskiy, Dmitriy and Lewis, Adrian S},
  journal={Mathematics of operations research},
  volume={43},
  number={3},
  pages={919--948},
  year={2018},
  publisher={INFORMS}
}""",
            r"""@article{szmuk2020successive,
  title={Successive convexification for real-time six-degree-of-freedom powered descent guidance
    with state-triggered constraints},
  author={Szmuk, Michael and Reynolds, Taylor P and A{\c{c}}{\i}kme{\c{s}}e, Beh{\c{c}}et},
  journal={Journal of Guidance, Control, and Dynamics},
  volume={43},
  number={8},
  pages={1399--1413},
  year={2020},
  publisher={American Institute of Aeronautics and Astronautics}
}""",
            r"""@article{reynolds2020dual,
  title={Dual quaternion-based powered descent guidance with state-triggered constraints},
  author={Reynolds, Taylor P and Szmuk, Michael and Malyuta, Danylo and Mesbahi, Mehran and
    A{\c{c}}{\i}kme{\c{s}}e, Beh{\c{c}}et and Carson III, John M},
  journal={Journal of Guidance, Control, and Dynamics},
  volume={43},
  number={9},
  pages={1584--1599},
  year={2020},
  publisher={American Institute of Aeronautics and Astronautics}
}""",
        ]

__init__(autotuner: AutotuningBase = None, k_max: int = 200, t_max: float | None = None, lam_prox: Union[float, Dict[str, Union[float, list]]] = 0.1, lam_vc: Union[float, Dict[str, Union[float, list]]] = 1.0, lam_cost: Union[float, Dict[str, float]] = 0.01, lam_vb: float = 0.0, ep_tr: float = 0.0001, ep_vb: float = 0.0001, ep_vc: float = 1e-08, states: List[State] = None, controls: List[Control] = None)

Initialize PTR with algorithm parameters and optional autotuner.

Source code in openscvx/algorithms/scvx/penalized_trust_region.py

def __init__(
    self,
    autotuner: "AutotuningBase" = None,
    k_max: int = 200,
    t_max: float | None = None,
    lam_prox: Union[float, Dict[str, Union[float, list]]] = 1e-1,
    lam_vc: Union[float, Dict[str, Union[float, list]]] = 1e0,
    lam_cost: Union[float, Dict[str, float]] = 1e-2,
    lam_vb: float = 0.0,
    ep_tr: float = 1e-4,
    ep_vb: float = 1e-4,
    ep_vc: float = 1e-8,
    states: List["State"] = None,
    controls: List["Control"] = None,
):
    """Initialize PTR with algorithm parameters and optional autotuner."""
    # Compiled infrastructure (set by initialize()). ``_iteration_fn`` is
    # the fused JAX-pure SCP body; ``step()`` is a thin Python wrapper that
    # calls it, records history, and emits.
    self._iteration_fn: Callable | None = None
    self._emitter: Callable | None = None

    # Store states/controls for later re-resolution.
    self._states: List["State"] = states
    self._controls: List["Control"] = controls

    super().__init__(
        # SCP weights (grouped dataclass, dict inputs expanded to arrays).
        weights=Weights.build(
            lam_prox=lam_prox,
            lam_vc=lam_vc,
            lam_cost=lam_cost,
            lam_vb=lam_vb,
            states=states,
            controls=controls,
        ),
        autotuner=autotuner if autotuner is not None else AugmentedLagrangian(),
        k_max=k_max,
        t_max=t_max,
        ep_tr=ep_tr,
        ep_vb=ep_vb,
        ep_vc=ep_vc,
    )

build_iteration(dis_continuous: Callable, dis_impulsive: Callable, jax_constraints: LoweredJaxConstraints, solver_callback: Callable, settings: Config) -> Callable

Fuse the discretizers, constraints, and solver into the PTR step.

Thin wrapper around :func:~openscvx.algorithms.scvx.iteration.make_scp_iteration, threading this algorithm's :attr:autotuner into the fused body.

Source code in openscvx/algorithms/scvx/penalized_trust_region.py

def build_iteration(
    self,
    dis_continuous: Callable,
    dis_impulsive: Callable,
    jax_constraints: "LoweredJaxConstraints",
    solver_callback: Callable,
    settings: Config,
) -> Callable:
    """Fuse the discretizers, constraints, and solver into the PTR step.

    Thin wrapper around
    :func:`~openscvx.algorithms.scvx.iteration.make_scp_iteration`,
    threading this algorithm's :attr:`autotuner` into the fused body.
    """
    return make_scp_iteration(
        dis_continuous=dis_continuous,
        dis_impulsive=dis_impulsive,
        jax_constraints=jax_constraints,
        solver_callback=solver_callback,
        autotuner=self.autotuner,
        settings=settings,
    )

citation() -> List[str]

Return BibTeX citations for the PTR algorithm.

Source code in openscvx/algorithms/scvx/penalized_trust_region.py

    def citation(self) -> List[str]:
        """Return BibTeX citations for the PTR algorithm."""
        return [
            r"""@article{drusvyatskiy2018error,
  title={Error bounds, quadratic growth, and linear convergence of proximal methods},
  author={Drusvyatskiy, Dmitriy and Lewis, Adrian S},
  journal={Mathematics of operations research},
  volume={43},
  number={3},
  pages={919--948},
  year={2018},
  publisher={INFORMS}
}""",
            r"""@article{szmuk2020successive,
  title={Successive convexification for real-time six-degree-of-freedom powered descent guidance
    with state-triggered constraints},
  author={Szmuk, Michael and Reynolds, Taylor P and A{\c{c}}{\i}kme{\c{s}}e, Beh{\c{c}}et},
  journal={Journal of Guidance, Control, and Dynamics},
  volume={43},
  number={8},
  pages={1399--1413},
  year={2020},
  publisher={American Institute of Aeronautics and Astronautics}
}""",
            r"""@article{reynolds2020dual,
  title={Dual quaternion-based powered descent guidance with state-triggered constraints},
  author={Reynolds, Taylor P and Szmuk, Michael and Malyuta, Danylo and Mesbahi, Mehran and
    A{\c{c}}{\i}kme{\c{s}}e, Beh{\c{c}}et and Carson III, John M},
  journal={Journal of Guidance, Control, and Dynamics},
  volume={43},
  number={9},
  pages={1584--1599},
  year={2020},
  publisher={American Institute of Aeronautics and Astronautics}
}""",
        ]

get_columns(verbosity: int = Verbosity.STANDARD) -> List[Column]

Get the columns to display for iteration output.

Source code in openscvx/algorithms/scvx/penalized_trust_region.py

def get_columns(self, verbosity: int = Verbosity.STANDARD) -> List[Column]:
    """Get the columns to display for iteration output."""
    all_columns = self.BASE_COLUMNS + self.autotuner.COLUMNS + self.TAIL_COLUMNS
    return [col for col in all_columns if col.min_verbosity <= verbosity]

initialize(iteration_fn: Callable, emitter: Callable) -> None

Store the fused SCP iteration body and per-iteration infrastructure.

iteration_fn is built and JIT-warmed by :meth:Problem.initialize; :meth:step is a thin Python wrapper that calls it, records history, and emits diagnostics. The boundary conditions the subproblem pins are carried on :class:AlgorithmState (x_init_pin / x_term_pin), so there is no solver state to prime here.

Source code in openscvx/algorithms/scvx/penalized_trust_region.py

def initialize(self, iteration_fn: Callable, emitter: Callable) -> None:
    """Store the fused SCP iteration body and per-iteration infrastructure.

    ``iteration_fn`` is built and JIT-warmed by :meth:`Problem.initialize`;
    :meth:`step` is a thin Python wrapper that calls it, records history,
    and emits diagnostics. The boundary conditions the subproblem pins are
    carried on :class:`AlgorithmState` (``x_init_pin`` / ``x_term_pin``), so
    there is no solver state to prime here.
    """
    self._iteration_fn = iteration_fn
    self._emitter = emitter

step(state: AlgorithmState, history: AlgorithmHistory, params: dict, settings: Config) -> Tuple[AlgorithmState, bool]

Execute one PTR iteration and return (next_state, converged).

Calls the fused JAX iteration body, records the per-iteration diagnostics it returns into history, emits progress, and reports convergence from the metrics on the next state.

Source code in openscvx/algorithms/scvx/penalized_trust_region.py

def step(
    self,
    state: AlgorithmState,
    history: AlgorithmHistory,
    params: dict,
    settings: Config,
) -> Tuple[AlgorithmState, bool]:
    """Execute one PTR iteration and return ``(next_state, converged)``.

    Calls the fused JAX iteration body, records the per-iteration
    diagnostics it returns into ``history``, emits progress, and reports
    convergence from the metrics on the next state.
    """
    if self._iteration_fn is None:
        raise RuntimeError(
            "PenalizedTrustRegion.step() called before initialize(). "
            "Call initialize() first to set up the iteration body."
        )

    iter_index = int(state.k)

    t0 = time.time()
    next_state, diag = self._iteration_fn(state, params)
    jax.block_until_ready((next_state, diag))
    step_time = time.time() - t0

    # Fail loudly on a bad subproblem solve before it becomes the next
    # linearization point. Two gates: a non-OPTIMAL status from the backend
    # (e.g. QPAX divergence, CVXPy infeasibility) and a non-finite iterate
    # (NaN/Inf from anywhere in the body). Either would silently corrupt the
    # rest of the SCP run.
    if int(diag.status) != int(StatusCode.OPTIMAL):
        raise RuntimeError(
            f"Convex subproblem did not solve to optimality "
            f"(status={status_code_to_str(int(diag.status))!r}). The backend failed "
            f"to converge or the subproblem is infeasible — adjust solver tolerances, "
            f"rescale the problem, or use float_dtype='float64'."
        )
    if not bool(np.all(np.isfinite(np.asarray(next_state.x)))):
        raise RuntimeError(
            "Subproblem solve produced a non-finite iterate (NaN/Inf in the state)."
        )

    use_full_metrics = self.autotuner.COMPUTES_ACCEPTANCE_METRICS
    scalars, lam_prox_np = history.record_iteration(
        next_state,
        V=np.asarray(diag.V),
        W=np.asarray(diag.W),
        VC=np.asarray(diag.VC),
        TR=np.asarray(diag.TR),
        J_lin=float(diag.J_lin),
        record_diagnostics=use_full_metrics,
    )

    emission_data = {
        "iter": iter_index,
        "subprop_time": step_time * 1000.0,
        "J_tr": scalars["J_tr"],
        "J_vb": scalars["J_vb"],
        "J_vc": scalars["J_vc"],
        "cost": float(diag.cost),
        # TODO: (haynec) log per-variable lam_prox detail (e.g. min/max range)
        "lam_prox": float(np.max(lam_prox_np)),
        "prob_stat": status_code_to_str(int(diag.status)),
        "adaptive_state": adaptive_state_code_to_str(scalars["adaptive_state_code"]),
        "ep_tr": self.ep_tr,
        "ep_vb": self.ep_vb,
        "ep_vc": self.ep_vc,
    }

    if use_full_metrics:
        emission_data.update(
            {
                "J_nonlin": scalars["J_nonlin"],
                "J_lin": float(diag.J_lin),
                "pred_reduction": scalars["predicted_reduction"],
                "actual_reduction": scalars["actual_reduction"],
                "acceptance_ratio": scalars["acceptance_ratio"],
            }
        )

    self._emitter(emission_data)

    # Same convergence test the lax.while_loop path uses, reading the
    # thresholds off the state pytree (synced from this algorithm's
    # ep_* attributes by Problem._sync_scp_constants at solve() time).
    return next_state, bool(self.converged(next_state))

PenalizedTrustRegionConfig

Bases: BaseModel

Validates PenalizedTrustRegion configuration from dict input.

The autotuner field accepts:

• None — defaults to :class:AugmentedLagrangian.
• A string — class name only, default parameters.
• A dict — class name via "type" key plus overrides.
• An instance — already-constructed autotuner (pass-through).

Source code in openscvx/algorithms/__init__.py

class PenalizedTrustRegionConfig(BaseModel):
    """Validates PenalizedTrustRegion configuration from dict input.

    The ``autotuner`` field accepts:

    * ``None`` — defaults to :class:`AugmentedLagrangian`.
    * A **string** — class name only, default parameters.
    * A **dict** — class name via ``"type"`` key plus overrides.
    * An **instance** — already-constructed autotuner (pass-through).
    """

    autotuner: Optional[Union[AutotunerConfig, AutotuningBase]] = None

    @field_validator("autotuner", mode="before")
    @classmethod
    def _wrap_bare_string(cls, v: Any) -> Any:
        if isinstance(v, str):
            return {"type": v}
        return v

    k_max: int = 200
    t_max: Optional[float] = None
    lam_prox: Union[float, Dict[str, Any]] = 1e-1
    lam_vc: Union[float, Dict[str, Any]] = 1e0
    lam_cost: Union[float, Dict[str, Any]] = 1e-2
    lam_vb: float = 0.0
    ep_tr: float = 1e-4
    ep_vb: float = 1e-4
    ep_vc: float = 1e-8

    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)

    def to_algorithm(
        self,
        states: Optional[List[Any]] = None,
        controls: Optional[List[Any]] = None,
    ) -> PenalizedTrustRegion:
        at = self.autotuner
        if at is None:
            autotuner = None
        elif isinstance(at, AutotuningBase):
            autotuner = at
        else:
            autotuner = at.build()
        kwargs = self.model_dump(exclude={"autotuner"}, exclude_unset=True)
        return PenalizedTrustRegion(
            autotuner=autotuner,
            states=states,
            controls=controls,
            **kwargs,
        )

RampProximalWeight

Bases: AutotuningBase

Ramp lam_prox toward lam_prox_max then hold.

update_weights is a pure functional update on the :class:AlgorithmState pytree; see the base-class contract.

Source code in openscvx/algorithms/autotuner/ramp_proximal_weight.py

class RampProximalWeight(AutotuningBase):
    """Ramp ``lam_prox`` toward ``lam_prox_max`` then hold.

    ``update_weights`` is a pure functional update on the
    :class:`AlgorithmState` pytree; see the base-class contract.
    """

    COMPUTES_ACCEPTANCE_METRICS = False

    def __init__(
        self,
        ramp_factor: float = 1.0,
        lam_prox_max: float = 1e3,
        lam_cost_drop: int = -1,
        lam_cost_relax: float = 1.0,
    ):
        self.hyper = RampProximalWeightHyper(
            ramp_factor=ramp_factor,
            lam_prox_max=lam_prox_max,
            lam_cost_drop=lam_cost_drop,
            lam_cost_relax=lam_cost_relax,
        )

    def update_weights(
        self,
        state: "AlgorithmState",
        candidate: "CandidateIterate",
        nodal_constraints: "LoweredJaxConstraints",
        settings: Config,
        params: dict,
    ) -> "AlgorithmState":
        """Return the next-iterate state.

        Pure functional update — see class docstring.
        """
        lam_cost_next = self._relaxed_lam_cost(state)

        was_at_max = jnp.all(state.lam_prox >= state.hyper.lam_prox_max)
        new_lam_prox = jnp.minimum(
            state.lam_prox * state.hyper.ramp_factor, state.hyper.lam_prox_max
        )

        code = jnp.where(
            was_at_max,
            jnp.int32(AdaptiveStateCode.ACCEPT_CONSTANT),
            jnp.int32(AdaptiveStateCode.ACCEPT_HIGHER),
        )

        return state.replace(
            x=candidate.x,
            u=candidate.u,
            x_prop=candidate.x_prop,
            x_prop_plus=candidate.x_prop_plus,
            lam_prox=new_lam_prox,
            lam_cost=lam_cost_next,
            adaptive_state_code=code,
        )

update_weights(state: AlgorithmState, candidate: CandidateIterate, nodal_constraints: LoweredJaxConstraints, settings: Config, params: dict) -> AlgorithmState

Return the next-iterate state.

Pure functional update — see class docstring.

Source code in openscvx/algorithms/autotuner/ramp_proximal_weight.py

def update_weights(
    self,
    state: "AlgorithmState",
    candidate: "CandidateIterate",
    nodal_constraints: "LoweredJaxConstraints",
    settings: Config,
    params: dict,
) -> "AlgorithmState":
    """Return the next-iterate state.

    Pure functional update — see class docstring.
    """
    lam_cost_next = self._relaxed_lam_cost(state)

    was_at_max = jnp.all(state.lam_prox >= state.hyper.lam_prox_max)
    new_lam_prox = jnp.minimum(
        state.lam_prox * state.hyper.ramp_factor, state.hyper.lam_prox_max
    )

    code = jnp.where(
        was_at_max,
        jnp.int32(AdaptiveStateCode.ACCEPT_CONSTANT),
        jnp.int32(AdaptiveStateCode.ACCEPT_HIGHER),
    )

    return state.replace(
        x=candidate.x,
        u=candidate.u,
        x_prop=candidate.x_prop,
        x_prop_plus=candidate.x_prop_plus,
        lam_prox=new_lam_prox,
        lam_cost=lam_cost_next,
        adaptive_state_code=code,
    )

Weights dataclass

Penalty weights for the SCP objective function.

Each SCP subproblem minimizes a weighted sum of the original cost, trust region penalty, virtual control penalty, and virtual buffer penalty. This dataclass holds the weights for each term.

Weights can be accessed directly (weights.lam_prox) or through the algorithm's convenience properties (algorithm.lam_prox). During SCP iteration, the autotuner may mutate these fields; those changes are tracked in :class:AlgorithmState weight histories.

Use :meth:build to construct from user-friendly inputs (floats or {name: weight} dicts). Use :meth:build_vb_arrays to populate lam_vb_nodal and lam_vb_cross once symbolic constraints are available.

Attributes:

Name	Type	Description
lam_prox	Union[float, ndarray]	Trust region (proximal) weight. Scalar (uniform) or array of shape (n_x + n_u,) for per-variable weighting, or (N, n_x + n_u) for per-node-per-variable weighting.
lam_vc	Union[float, ndarray]	Virtual control penalty weight. Scalar (uniform) or array of shape (n_x,) for per-state weighting, or (N-1, n_x) for per-node-per-state weighting.
lam_cost	Union[float, ndarray]	Cost weight. Scalar (uniform across all minimize/maximize states) or array of shape (n_x,) for per-state weighting.
lam_vb	float	Global virtual buffer penalty weight. Scalar default applied to every constraint unless overridden via .weight() on individual constraints.
lam_vb_nodal	Optional[ndarray]	Per-node weights for nodal constraints, shape (N, n_nodal). Built by :meth:build_vb_arrays.
lam_vb_cross	Optional[ndarray]	Weights for cross-node constraints, shape (n_cross,). Built by :meth:build_vb_arrays.

Source code in openscvx/algorithms/weights.py

@dataclass
class Weights:
    """Penalty weights for the SCP objective function.

    Each SCP subproblem minimizes a weighted sum of the original cost,
    trust region penalty, virtual control penalty, and virtual buffer
    penalty. This dataclass holds the weights for each term.

    Weights can be accessed directly (``weights.lam_prox``) or through
    the algorithm's convenience properties (``algorithm.lam_prox``).
    During SCP iteration, the autotuner may mutate these fields; those
    changes are tracked in :class:`AlgorithmState` weight histories.

    Use :meth:`build` to construct from user-friendly inputs (floats or
    ``{name: weight}`` dicts). Use :meth:`build_vb_arrays` to populate
    ``lam_vb_nodal`` and ``lam_vb_cross`` once symbolic constraints are
    available.

    Attributes:
        lam_prox: Trust region (proximal) weight. Scalar (uniform) or
            array of shape ``(n_x + n_u,)`` for per-variable weighting,
            or ``(N, n_x + n_u)`` for per-node-per-variable weighting.
        lam_vc: Virtual control penalty weight. Scalar (uniform) or
            array of shape ``(n_x,)`` for per-state weighting, or
            ``(N-1, n_x)`` for per-node-per-state weighting.
        lam_cost: Cost weight. Scalar (uniform across all
            minimize/maximize states) or array of shape ``(n_x,)`` for
            per-state weighting.
        lam_vb: Global virtual buffer penalty weight. Scalar default
            applied to every constraint unless overridden via
            ``.weight()`` on individual constraints.
        lam_vb_nodal: Per-node weights for nodal constraints, shape
            ``(N, n_nodal)``. Built by :meth:`build_vb_arrays`.
        lam_vb_cross: Weights for cross-node constraints, shape
            ``(n_cross,)``. Built by :meth:`build_vb_arrays`.
    """

    lam_prox: Union[float, np.ndarray] = 1e-1
    lam_vc: Union[float, np.ndarray] = 1e0
    lam_cost: Union[float, np.ndarray] = 1e-2
    lam_vb: float = 0.0
    lam_vb_nodal: Optional[np.ndarray] = None
    lam_vb_cross: Optional[np.ndarray] = None

    def __post_init__(self):
        # Coerce lists/lists-of-lists to numpy arrays.
        if isinstance(self.lam_prox, (list, tuple)):
            self.lam_prox = np.asarray(self.lam_prox, dtype=float)
        if isinstance(self.lam_vc, (list, tuple)):
            self.lam_vc = np.asarray(self.lam_vc, dtype=float)
        if isinstance(self.lam_cost, (list, tuple)):
            self.lam_cost = np.asarray(self.lam_cost, dtype=float)

    @classmethod
    def build(
        cls,
        lam_prox: Union[float, Dict[str, Union[float, list, np.ndarray]]] = 1e-1,
        lam_vc: Union[float, Dict[str, Union[float, list, np.ndarray]]] = 1e0,
        lam_cost: Union[float, Dict[str, float]] = 1e-2,
        lam_vb: float = 0.0,
        states: Optional[List["State"]] = None,
        controls: Optional[List["Control"]] = None,
    ) -> "Weights":
        """Construct Weights from user-friendly inputs.

        Accepts floats (applied uniformly) or dicts mapping state/control
        names to per-variable weights. Dict inputs are expanded to dense
        arrays via the ``resolve_lam_*`` methods.

        Args:
            lam_prox: Trust region weight. Float or ``{name: weight}`` dict.
                Dict requires *states* and *controls*.
            lam_vc: Virtual control weight. Float or ``{state_name: weight}``
                dict. Dict requires *states*.
            lam_cost: Cost weight. Float or ``{state_name: weight}`` dict.
                Dict requires *states*.
            lam_vb: Virtual buffer default weight (scalar).
            states: Symbolic State objects (required when any weight is a dict).
            controls: Symbolic Control objects (required when *lam_prox* is a dict).

        Returns:
            A new Weights instance with resolved numeric values.

        Raises:
            ValueError: If a dict weight is given without the required
                states/controls.
        """
        return cls(
            lam_prox=cls.resolve_lam_prox(lam_prox, states, controls),
            lam_vc=cls.resolve_lam_vc(lam_vc, states),
            lam_cost=cls.resolve_lam_cost(lam_cost, states),
            lam_vb=float(lam_vb),
        )

    @staticmethod
    def resolve_lam_prox(
        lam_prox: Union[float, Dict[str, Union[float, list, np.ndarray]]],
        states: Optional[List["State"]] = None,
        controls: Optional[List["Control"]] = None,
    ) -> Union[float, np.ndarray]:
        """Resolve a ``lam_prox`` spec to a numeric value.

        If *lam_prox* is a float it is returned as-is. If it is a dict
        mapping state/control names to weights, it is expanded to a dense
        array using each variable's ``_slice``. Variables not in the dict
        default to ``1.0``.

        Dict values may be scalars, 1-D arrays (per-component), or 2-D
        arrays of shape ``(K, n_components)`` for per-node-per-component
        weighting. All 2-D entries must agree on *K*.

        Args:
            lam_prox: Scalar weight or ``{name: weight}`` dict.
            states: Symbolic State objects (required when *lam_prox* is a dict).
            controls: Symbolic Control objects (required when *lam_prox* is a dict).

        Returns:
            float or np.ndarray of shape ``(n_states + n_controls,)`` or
            ``(K, n_states + n_controls)``.

        Raises:
            ValueError: If *lam_prox* is a dict and *states*/*controls* is ``None``,
                or if dict contains unknown names, or if 2-D entries disagree on K.
        """
        if not isinstance(lam_prox, dict):
            return lam_prox

        if states is None or controls is None:
            raise ValueError(
                "lam_prox was specified as a dict but states and/or "
                "controls were not provided. Pass both so the dict can "
                "be expanded to a per-variable weight array."
            )

        n_states = sum(s.shape[0] if len(s.shape) > 0 else 1 for s in states)
        n_controls = sum(c.shape[0] if len(c.shape) > 0 else 1 for c in controls)
        n_total = n_states + n_controls

        valid_state_names = {s.name for s in states}
        valid_control_names = {c.name for c in controls}
        valid_names = valid_state_names | valid_control_names
        unknown = set(lam_prox.keys()) - valid_names
        if unknown:
            raise ValueError(
                f"lam_prox dict contains unknown name(s): {unknown}. "
                f"Valid names: {sorted(valid_names)}"
            )

        # Build a unified list of (name, n_components, slice_in_output).
        # States occupy columns [0, n_states), controls occupy [n_states, n_total).
        variables: list = []
        for s in states:
            nc = s.shape[0] if len(s.shape) > 0 else 1
            variables.append((s.name, nc, s._slice))
        for c in controls:
            nc = c.shape[0] if len(c.shape) > 0 else 1
            out_slice = slice(n_states + c._slice.start, n_states + c._slice.stop)
            variables.append((c.name, nc, out_slice))

        # First pass: determine if any entry is 2-D and infer K.
        n_nodes: Optional[int] = None
        for name, n_comp, _ in variables:
            if name not in lam_prox:
                continue
            val = np.asarray(lam_prox[name], dtype=float)
            if val.ndim == 2:
                if n_nodes is None:
                    n_nodes = val.shape[0]
                elif val.shape[0] != n_nodes:
                    raise ValueError(
                        f"lam_prox['{name}'] has {val.shape[0]} rows, but a "
                        f"previous entry had {n_nodes} rows. All 2-D entries "
                        f"must have the same number of rows (n_nodes)."
                    )

        # Build the output array.
        if n_nodes is not None:
            lam_arr = np.ones((n_nodes, n_total))
        else:
            lam_arr = np.ones(n_total)

        for name, n_comp, out_slice in variables:
            if name not in lam_prox:
                continue
            val = np.asarray(lam_prox[name], dtype=float)

            if val.ndim == 0:
                lam_arr[..., out_slice] = float(val)
            elif val.ndim == 1:
                if val.shape[0] != n_comp:
                    raise ValueError(
                        f"lam_prox['{name}'] has length {val.shape[0]}, "
                        f"expected scalar or length {n_comp}"
                    )
                lam_arr[..., out_slice] = val
            elif val.ndim == 2:
                if val.shape[1] != n_comp:
                    raise ValueError(
                        f"lam_prox['{name}'] has {val.shape[1]} columns, expected {n_comp}"
                    )
                lam_arr[:, out_slice] = val
            else:
                raise ValueError(
                    f"lam_prox['{name}'] has {val.ndim} dimensions, expected scalar, 1-D, or 2-D"
                )

        return lam_arr

    @staticmethod
    def resolve_lam_vc(
        lam_vc: Union[float, Dict[str, Union[float, list, np.ndarray]]],
        states: Optional[List["State"]] = None,
    ) -> Union[float, np.ndarray]:
        """Resolve a ``lam_vc`` spec to a numeric value.

        If *lam_vc* is a float it is returned as-is. If it is a dict
        mapping state names to weights, it is expanded to a dense array
        using each state's ``_slice``. States not in the dict default
        to ``1.0``.

        Dict values may be scalars, 1-D arrays (per-component), or 2-D
        arrays of shape ``(K, n_components)`` for per-node-per-component
        weighting. All 2-D entries must agree on *K*.

        Args:
            lam_vc: Scalar weight or ``{state_name: weight}`` dict.
            states: Symbolic State objects (required when *lam_vc* is a dict).

        Returns:
            float or np.ndarray of shape ``(n_states,)`` or ``(K, n_states)``.

        Raises:
            ValueError: If *lam_vc* is a dict and *states* is ``None``,
                or if dict contains unknown names, or if 2-D entries disagree on K.
        """
        if not isinstance(lam_vc, dict):
            return lam_vc

        if states is None:
            raise ValueError(
                "lam_vc was specified as a dict but no states were "
                "provided. Pass states so the dict can be expanded to "
                "a per-state weight array."
            )

        n_states = sum(s.shape[0] if len(s.shape) > 0 else 1 for s in states)

        valid_names = {s.name for s in states}
        unknown = set(lam_vc.keys()) - valid_names
        if unknown:
            raise ValueError(
                f"lam_vc dict contains unknown state name(s): {unknown}. "
                f"Valid state names: {sorted(valid_names)}"
            )

        # First pass: determine if any entry is 2-D and infer K.
        n_nodes: Optional[int] = None
        for state in states:
            if state.name not in lam_vc:
                continue
            val = np.asarray(lam_vc[state.name], dtype=float)
            if val.ndim == 2:
                if n_nodes is None:
                    n_nodes = val.shape[0]
                elif val.shape[0] != n_nodes:
                    raise ValueError(
                        f"lam_vc['{state.name}'] has {val.shape[0]} rows, but a "
                        f"previous entry had {n_nodes} rows. All 2-D entries "
                        f"must have the same number of rows (n_nodes-1)."
                    )

        # Build the output array.
        if n_nodes is not None:
            lam_arr = np.ones((n_nodes, n_states))
        else:
            lam_arr = np.ones(n_states)

        for state in states:
            if state.name not in lam_vc:
                continue
            val = np.asarray(lam_vc[state.name], dtype=float)
            n_components = state.shape[0] if len(state.shape) > 0 else 1

            if val.ndim == 0:
                lam_arr[..., state._slice] = float(val)
            elif val.ndim == 1:
                if val.shape[0] != n_components:
                    raise ValueError(
                        f"lam_vc['{state.name}'] has length {val.shape[0]}, "
                        f"expected scalar or length {n_components}"
                    )
                lam_arr[..., state._slice] = val
            elif val.ndim == 2:
                if val.shape[1] != n_components:
                    raise ValueError(
                        f"lam_vc['{state.name}'] has {val.shape[1]} columns, "
                        f"expected {n_components}"
                    )
                lam_arr[:, state._slice] = val
            else:
                raise ValueError(
                    f"lam_vc['{state.name}'] has {val.ndim} dimensions, "
                    f"expected scalar, 1-D, or 2-D"
                )

        return lam_arr

    @staticmethod
    def resolve_lam_cost(
        lam_cost: Union[float, Dict[str, Union[float, list, np.ndarray]]],
        states: Optional[List["State"]] = None,
    ) -> Union[float, np.ndarray]:
        """Resolve a ``lam_cost`` spec to a numeric value.

        If *lam_cost* is a float it is returned as-is. If it is a dict
        mapping state names to weights, it is expanded to a dense array
        of shape ``(n_states,)`` using each state's ``_slice``. States
        without a minimize/maximize objective receive weight 0. States
        **with** an objective **must** appear in the dict.

        Dict values may be scalars (broadcast to every component) or
        arrays matching the state's shape for per-component weighting,
        e.g. ``{"position": [0, 0, 1e-6]}``.

        Args:
            lam_cost: Scalar weight or ``{state_name: weight}`` dict.
            states: Symbolic State objects (required when *lam_cost* is a dict).

        Returns:
            float or np.ndarray of shape ``(n_states,)``.

        Raises:
            ValueError: If *lam_cost* is a dict and *states* is ``None``,
                or if dict contains unknown names, or if dict is missing
                entries for states with minimize/maximize objectives.
        """
        if not isinstance(lam_cost, dict):
            return lam_cost

        if states is None:
            raise ValueError(
                "lam_cost was specified as a dict but no states were "
                "provided. Pass states so the dict can be expanded to "
                "a per-state weight array."
            )

        n_states = sum(s.shape[0] if len(s.shape) > 0 else 1 for s in states)
        lam_arr = np.zeros(n_states)

        valid_names = {s.name for s in states}
        unknown = set(lam_cost.keys()) - valid_names
        if unknown:
            raise ValueError(
                f"lam_cost dict contains unknown state name(s): {unknown}. "
                f"Valid state names: {sorted(valid_names)}"
            )

        # Identify states that have minimize/maximize objectives.
        cost_states: Set[str] = set()
        for state in states:
            if state.initial_type is not None:
                for t in state.initial_type:
                    if t in ("Minimize", "Maximize"):
                        cost_states.add(state.name)
                        break
            if state.final_type is not None:
                for t in state.final_type:
                    if t in ("Minimize", "Maximize"):
                        cost_states.add(state.name)
                        break

        # Check that all cost states are in the dict
        missing = cost_states - set(lam_cost.keys())
        if missing:
            raise ValueError(
                f"lam_cost dict is missing weight(s) for state(s) with "
                f"minimize/maximize objectives: {missing}. All states with "
                f"cost terms must have a weight in the dict."
            )

        # Fill the array.
        for state in states:
            if state.name in lam_cost:
                val = np.asarray(lam_cost[state.name], dtype=float)
                n_components = state.shape[0] if len(state.shape) > 0 else 1
                if val.ndim > 0 and val.shape[0] != n_components:
                    raise ValueError(
                        f"lam_cost['{state.name}'] has length {val.shape[0]}, "
                        f"expected scalar or length {n_components}"
                    )
                lam_arr[state._slice] = val

        return lam_arr

    def build_vb_arrays(
        self,
        N: int,
        nodal_constraints: list,
        cross_node_constraints: list,
        n_byof_nodal: int = 0,
        n_byof_cross: int = 0,
    ) -> None:
        """Build per-constraint virtual buffer weight arrays.

        Inspects each symbolic constraint's shape (to account for vector
        decomposition) and ``.weight()`` overrides, then populates
        ``self.lam_vb_nodal`` and ``self.lam_vb_cross``.

        Args:
            N: Number of trajectory nodes.
            nodal_constraints: Symbolic ``NodalConstraint`` objects (post-
                preprocessing, pre-lowering).
            cross_node_constraints: Symbolic ``CrossNodeConstraint`` objects.
            n_byof_nodal: Number of byof nodal constraints (each adds one
                column with the default weight).
            n_byof_cross: Number of byof cross-node constraints (each adds
                one entry with the default weight).
        """
        default_vb = float(self.lam_vb)

        # Count decomposed nodal constraints (vector → multiple scalars).
        # Vector constraints are decomposed element-wise during lowering
        # (see decompose_vector_nodal_constraints), so each element gets its
        # own column.  We mirror that here via check_shape() to ensure the
        # array dimensions match the post-decomposition constraint count.
        n_nodal = 0
        for nc in nodal_constraints:
            shape = nc.constraint.lhs.check_shape()
            n_nodal += int(np.prod(shape)) if len(shape) > 0 else 1

        # Byof constraints are scalar (one column each), added after symbolic.
        n_nodal += n_byof_nodal
        n_cross = len(cross_node_constraints) + n_byof_cross

        # max(..., 1) avoids size-0 CVXPy parameters.
        n_nodal_param = max(n_nodal, 1)
        n_cross_param = max(n_cross, 1)

        lam_vb_nodal = np.full((N, n_nodal_param), default_vb)
        lam_vb_cross = np.full(n_cross_param, default_vb)

        # Apply per-constraint .weight() overrides for nodal constraints.
        col = 0
        for nc in nodal_constraints:
            shape = nc.constraint.lhs.check_shape()
            n_elem = int(np.prod(shape)) if len(shape) > 0 else 1

            w = nc._lam_vb
            if w is not None:
                nodes = nc.nodes if nc.nodes is not None else list(range(N))
                if isinstance(w, (int, float)):
                    lam_vb_nodal[nodes, col : col + n_elem] = float(w)
                elif isinstance(w, np.ndarray):
                    if w.ndim == 1:
                        # (n_elem,) — broadcast across nodes
                        for i in range(n_elem):
                            val = float(w[0]) if len(w) == 1 else float(w[i])
                            lam_vb_nodal[nodes, col + i] = val
                    elif w.ndim == 2:
                        # (n_nodes, n_elem) — per-node-per-element
                        for i in range(n_elem):
                            c_i = 0 if w.shape[1] == 1 else i
                            lam_vb_nodal[nodes, col + i] = w[:, c_i]

            col += n_elem

        # Apply per-constraint .weight() overrides for cross-node constraints.
        for idx, cc in enumerate(cross_node_constraints):
            if cc._lam_vb is not None:
                lam_vb_cross[idx] = float(cc._lam_vb)

        self.lam_vb_nodal = lam_vb_nodal
        self.lam_vb_cross = lam_vb_cross

build(lam_prox: Union[float, Dict[str, Union[float, list, np.ndarray]]] = 0.1, lam_vc: Union[float, Dict[str, Union[float, list, np.ndarray]]] = 1.0, lam_cost: Union[float, Dict[str, float]] = 0.01, lam_vb: float = 0.0, states: Optional[List[State]] = None, controls: Optional[List[Control]] = None) -> Weights classmethod

Construct Weights from user-friendly inputs.

Accepts floats (applied uniformly) or dicts mapping state/control names to per-variable weights. Dict inputs are expanded to dense arrays via the resolve_lam_* methods.

Parameters:

Name	Type	Description	Default
lam_prox	Union[float, Dict[str, Union[float, list, ndarray]]]	Trust region weight. Float or {name: weight} dict. Dict requires states and controls.	0.1
lam_vc	Union[float, Dict[str, Union[float, list, ndarray]]]	Virtual control weight. Float or {state_name: weight} dict. Dict requires states.	1.0
lam_cost	Union[float, Dict[str, float]]	Cost weight. Float or {state_name: weight} dict. Dict requires states.	0.01
lam_vb	float	Virtual buffer default weight (scalar).	0.0
states	Optional[List[State]]	Symbolic State objects (required when any weight is a dict).	None
controls	Optional[List[Control]]	Symbolic Control objects (required when lam_prox is a dict).	None

Returns:

Type	Description
Weights	A new Weights instance with resolved numeric values.

Raises:

Type	Description
ValueError	If a dict weight is given without the required states/controls.

Source code in openscvx/algorithms/weights.py

@classmethod
def build(
    cls,
    lam_prox: Union[float, Dict[str, Union[float, list, np.ndarray]]] = 1e-1,
    lam_vc: Union[float, Dict[str, Union[float, list, np.ndarray]]] = 1e0,
    lam_cost: Union[float, Dict[str, float]] = 1e-2,
    lam_vb: float = 0.0,
    states: Optional[List["State"]] = None,
    controls: Optional[List["Control"]] = None,
) -> "Weights":
    """Construct Weights from user-friendly inputs.

    Accepts floats (applied uniformly) or dicts mapping state/control
    names to per-variable weights. Dict inputs are expanded to dense
    arrays via the ``resolve_lam_*`` methods.

    Args:
        lam_prox: Trust region weight. Float or ``{name: weight}`` dict.
            Dict requires *states* and *controls*.
        lam_vc: Virtual control weight. Float or ``{state_name: weight}``
            dict. Dict requires *states*.
        lam_cost: Cost weight. Float or ``{state_name: weight}`` dict.
            Dict requires *states*.
        lam_vb: Virtual buffer default weight (scalar).
        states: Symbolic State objects (required when any weight is a dict).
        controls: Symbolic Control objects (required when *lam_prox* is a dict).

    Returns:
        A new Weights instance with resolved numeric values.

    Raises:
        ValueError: If a dict weight is given without the required
            states/controls.
    """
    return cls(
        lam_prox=cls.resolve_lam_prox(lam_prox, states, controls),
        lam_vc=cls.resolve_lam_vc(lam_vc, states),
        lam_cost=cls.resolve_lam_cost(lam_cost, states),
        lam_vb=float(lam_vb),
    )

build_vb_arrays(N: int, nodal_constraints: list, cross_node_constraints: list, n_byof_nodal: int = 0, n_byof_cross: int = 0) -> None

Build per-constraint virtual buffer weight arrays.

Inspects each symbolic constraint's shape (to account for vector decomposition) and .weight() overrides, then populates self.lam_vb_nodal and self.lam_vb_cross.

Parameters:

Name	Type	Description	Default
N	int	Number of trajectory nodes.	required
nodal_constraints	list	Symbolic NodalConstraint objects (post- preprocessing, pre-lowering).	required
cross_node_constraints	list	Symbolic CrossNodeConstraint objects.	required
n_byof_nodal	int	Number of byof nodal constraints (each adds one column with the default weight).	0
n_byof_cross	int	Number of byof cross-node constraints (each adds one entry with the default weight).	0

Source code in openscvx/algorithms/weights.py

def build_vb_arrays(
    self,
    N: int,
    nodal_constraints: list,
    cross_node_constraints: list,
    n_byof_nodal: int = 0,
    n_byof_cross: int = 0,
) -> None:
    """Build per-constraint virtual buffer weight arrays.

    Inspects each symbolic constraint's shape (to account for vector
    decomposition) and ``.weight()`` overrides, then populates
    ``self.lam_vb_nodal`` and ``self.lam_vb_cross``.

    Args:
        N: Number of trajectory nodes.
        nodal_constraints: Symbolic ``NodalConstraint`` objects (post-
            preprocessing, pre-lowering).
        cross_node_constraints: Symbolic ``CrossNodeConstraint`` objects.
        n_byof_nodal: Number of byof nodal constraints (each adds one
            column with the default weight).
        n_byof_cross: Number of byof cross-node constraints (each adds
            one entry with the default weight).
    """
    default_vb = float(self.lam_vb)

    # Count decomposed nodal constraints (vector → multiple scalars).
    # Vector constraints are decomposed element-wise during lowering
    # (see decompose_vector_nodal_constraints), so each element gets its
    # own column.  We mirror that here via check_shape() to ensure the
    # array dimensions match the post-decomposition constraint count.
    n_nodal = 0
    for nc in nodal_constraints:
        shape = nc.constraint.lhs.check_shape()
        n_nodal += int(np.prod(shape)) if len(shape) > 0 else 1

    # Byof constraints are scalar (one column each), added after symbolic.
    n_nodal += n_byof_nodal
    n_cross = len(cross_node_constraints) + n_byof_cross

    # max(..., 1) avoids size-0 CVXPy parameters.
    n_nodal_param = max(n_nodal, 1)
    n_cross_param = max(n_cross, 1)

    lam_vb_nodal = np.full((N, n_nodal_param), default_vb)
    lam_vb_cross = np.full(n_cross_param, default_vb)

    # Apply per-constraint .weight() overrides for nodal constraints.
    col = 0
    for nc in nodal_constraints:
        shape = nc.constraint.lhs.check_shape()
        n_elem = int(np.prod(shape)) if len(shape) > 0 else 1

        w = nc._lam_vb
        if w is not None:
            nodes = nc.nodes if nc.nodes is not None else list(range(N))
            if isinstance(w, (int, float)):
                lam_vb_nodal[nodes, col : col + n_elem] = float(w)
            elif isinstance(w, np.ndarray):
                if w.ndim == 1:
                    # (n_elem,) — broadcast across nodes
                    for i in range(n_elem):
                        val = float(w[0]) if len(w) == 1 else float(w[i])
                        lam_vb_nodal[nodes, col + i] = val
                elif w.ndim == 2:
                    # (n_nodes, n_elem) — per-node-per-element
                    for i in range(n_elem):
                        c_i = 0 if w.shape[1] == 1 else i
                        lam_vb_nodal[nodes, col + i] = w[:, c_i]

        col += n_elem

    # Apply per-constraint .weight() overrides for cross-node constraints.
    for idx, cc in enumerate(cross_node_constraints):
        if cc._lam_vb is not None:
            lam_vb_cross[idx] = float(cc._lam_vb)

    self.lam_vb_nodal = lam_vb_nodal
    self.lam_vb_cross = lam_vb_cross

resolve_lam_cost(lam_cost: Union[float, Dict[str, Union[float, list, np.ndarray]]], states: Optional[List[State]] = None) -> Union[float, np.ndarray] staticmethod

Resolve a lam_cost spec to a numeric value.

If lam_cost is a float it is returned as-is. If it is a dict mapping state names to weights, it is expanded to a dense array of shape (n_states,) using each state's _slice. States without a minimize/maximize objective receive weight 0. States with an objective must appear in the dict.

Dict values may be scalars (broadcast to every component) or arrays matching the state's shape for per-component weighting, e.g. {"position": [0, 0, 1e-6]}.

Parameters:

Name	Type	Description	Default
lam_cost	Union[float, Dict[str, Union[float, list, ndarray]]]	Scalar weight or {state_name: weight} dict.	required
states	Optional[List[State]]	Symbolic State objects (required when lam_cost is a dict).	None

Returns:

Type	Description
Union[float, ndarray]	float or np.ndarray of shape (n_states,).

Raises:

Type	Description
ValueError	If lam_cost is a dict and states is None, or if dict contains unknown names, or if dict is missing entries for states with minimize/maximize objectives.

Source code in openscvx/algorithms/weights.py

@staticmethod
def resolve_lam_cost(
    lam_cost: Union[float, Dict[str, Union[float, list, np.ndarray]]],
    states: Optional[List["State"]] = None,
) -> Union[float, np.ndarray]:
    """Resolve a ``lam_cost`` spec to a numeric value.

    If *lam_cost* is a float it is returned as-is. If it is a dict
    mapping state names to weights, it is expanded to a dense array
    of shape ``(n_states,)`` using each state's ``_slice``. States
    without a minimize/maximize objective receive weight 0. States
    **with** an objective **must** appear in the dict.

    Dict values may be scalars (broadcast to every component) or
    arrays matching the state's shape for per-component weighting,
    e.g. ``{"position": [0, 0, 1e-6]}``.

    Args:
        lam_cost: Scalar weight or ``{state_name: weight}`` dict.
        states: Symbolic State objects (required when *lam_cost* is a dict).

    Returns:
        float or np.ndarray of shape ``(n_states,)``.

    Raises:
        ValueError: If *lam_cost* is a dict and *states* is ``None``,
            or if dict contains unknown names, or if dict is missing
            entries for states with minimize/maximize objectives.
    """
    if not isinstance(lam_cost, dict):
        return lam_cost

    if states is None:
        raise ValueError(
            "lam_cost was specified as a dict but no states were "
            "provided. Pass states so the dict can be expanded to "
            "a per-state weight array."
        )

    n_states = sum(s.shape[0] if len(s.shape) > 0 else 1 for s in states)
    lam_arr = np.zeros(n_states)

    valid_names = {s.name for s in states}
    unknown = set(lam_cost.keys()) - valid_names
    if unknown:
        raise ValueError(
            f"lam_cost dict contains unknown state name(s): {unknown}. "
            f"Valid state names: {sorted(valid_names)}"
        )

    # Identify states that have minimize/maximize objectives.
    cost_states: Set[str] = set()
    for state in states:
        if state.initial_type is not None:
            for t in state.initial_type:
                if t in ("Minimize", "Maximize"):
                    cost_states.add(state.name)
                    break
        if state.final_type is not None:
            for t in state.final_type:
                if t in ("Minimize", "Maximize"):
                    cost_states.add(state.name)
                    break

    # Check that all cost states are in the dict
    missing = cost_states - set(lam_cost.keys())
    if missing:
        raise ValueError(
            f"lam_cost dict is missing weight(s) for state(s) with "
            f"minimize/maximize objectives: {missing}. All states with "
            f"cost terms must have a weight in the dict."
        )

    # Fill the array.
    for state in states:
        if state.name in lam_cost:
            val = np.asarray(lam_cost[state.name], dtype=float)
            n_components = state.shape[0] if len(state.shape) > 0 else 1
            if val.ndim > 0 and val.shape[0] != n_components:
                raise ValueError(
                    f"lam_cost['{state.name}'] has length {val.shape[0]}, "
                    f"expected scalar or length {n_components}"
                )
            lam_arr[state._slice] = val

    return lam_arr

resolve_lam_prox(lam_prox: Union[float, Dict[str, Union[float, list, np.ndarray]]], states: Optional[List[State]] = None, controls: Optional[List[Control]] = None) -> Union[float, np.ndarray] staticmethod

Resolve a lam_prox spec to a numeric value.

If lam_prox is a float it is returned as-is. If it is a dict mapping state/control names to weights, it is expanded to a dense array using each variable's _slice. Variables not in the dict default to 1.0.

Dict values may be scalars, 1-D arrays (per-component), or 2-D arrays of shape (K, n_components) for per-node-per-component weighting. All 2-D entries must agree on K.

Parameters:

Name	Type	Description	Default
lam_prox	Union[float, Dict[str, Union[float, list, ndarray]]]	Scalar weight or {name: weight} dict.	required
states	Optional[List[State]]	Symbolic State objects (required when lam_prox is a dict).	None
controls	Optional[List[Control]]	Symbolic Control objects (required when lam_prox is a dict).	None

Returns:

Type	Description
Union[float, ndarray]	float or np.ndarray of shape (n_states + n_controls,) or
Union[float, ndarray]	(K, n_states + n_controls).

Raises:

Type	Description
ValueError	If lam_prox is a dict and states/controls is None, or if dict contains unknown names, or if 2-D entries disagree on K.

Source code in openscvx/algorithms/weights.py

@staticmethod
def resolve_lam_prox(
    lam_prox: Union[float, Dict[str, Union[float, list, np.ndarray]]],
    states: Optional[List["State"]] = None,
    controls: Optional[List["Control"]] = None,
) -> Union[float, np.ndarray]:
    """Resolve a ``lam_prox`` spec to a numeric value.

    If *lam_prox* is a float it is returned as-is. If it is a dict
    mapping state/control names to weights, it is expanded to a dense
    array using each variable's ``_slice``. Variables not in the dict
    default to ``1.0``.

    Dict values may be scalars, 1-D arrays (per-component), or 2-D
    arrays of shape ``(K, n_components)`` for per-node-per-component
    weighting. All 2-D entries must agree on *K*.

    Args:
        lam_prox: Scalar weight or ``{name: weight}`` dict.
        states: Symbolic State objects (required when *lam_prox* is a dict).
        controls: Symbolic Control objects (required when *lam_prox* is a dict).

    Returns:
        float or np.ndarray of shape ``(n_states + n_controls,)`` or
        ``(K, n_states + n_controls)``.

    Raises:
        ValueError: If *lam_prox* is a dict and *states*/*controls* is ``None``,
            or if dict contains unknown names, or if 2-D entries disagree on K.
    """
    if not isinstance(lam_prox, dict):
        return lam_prox

    if states is None or controls is None:
        raise ValueError(
            "lam_prox was specified as a dict but states and/or "
            "controls were not provided. Pass both so the dict can "
            "be expanded to a per-variable weight array."
        )

    n_states = sum(s.shape[0] if len(s.shape) > 0 else 1 for s in states)
    n_controls = sum(c.shape[0] if len(c.shape) > 0 else 1 for c in controls)
    n_total = n_states + n_controls

    valid_state_names = {s.name for s in states}
    valid_control_names = {c.name for c in controls}
    valid_names = valid_state_names | valid_control_names
    unknown = set(lam_prox.keys()) - valid_names
    if unknown:
        raise ValueError(
            f"lam_prox dict contains unknown name(s): {unknown}. "
            f"Valid names: {sorted(valid_names)}"
        )

    # Build a unified list of (name, n_components, slice_in_output).
    # States occupy columns [0, n_states), controls occupy [n_states, n_total).
    variables: list = []
    for s in states:
        nc = s.shape[0] if len(s.shape) > 0 else 1
        variables.append((s.name, nc, s._slice))
    for c in controls:
        nc = c.shape[0] if len(c.shape) > 0 else 1
        out_slice = slice(n_states + c._slice.start, n_states + c._slice.stop)
        variables.append((c.name, nc, out_slice))

    # First pass: determine if any entry is 2-D and infer K.
    n_nodes: Optional[int] = None
    for name, n_comp, _ in variables:
        if name not in lam_prox:
            continue
        val = np.asarray(lam_prox[name], dtype=float)
        if val.ndim == 2:
            if n_nodes is None:
                n_nodes = val.shape[0]
            elif val.shape[0] != n_nodes:
                raise ValueError(
                    f"lam_prox['{name}'] has {val.shape[0]} rows, but a "
                    f"previous entry had {n_nodes} rows. All 2-D entries "
                    f"must have the same number of rows (n_nodes)."
                )

    # Build the output array.
    if n_nodes is not None:
        lam_arr = np.ones((n_nodes, n_total))
    else:
        lam_arr = np.ones(n_total)

    for name, n_comp, out_slice in variables:
        if name not in lam_prox:
            continue
        val = np.asarray(lam_prox[name], dtype=float)

        if val.ndim == 0:
            lam_arr[..., out_slice] = float(val)
        elif val.ndim == 1:
            if val.shape[0] != n_comp:
                raise ValueError(
                    f"lam_prox['{name}'] has length {val.shape[0]}, "
                    f"expected scalar or length {n_comp}"
                )
            lam_arr[..., out_slice] = val
        elif val.ndim == 2:
            if val.shape[1] != n_comp:
                raise ValueError(
                    f"lam_prox['{name}'] has {val.shape[1]} columns, expected {n_comp}"
                )
            lam_arr[:, out_slice] = val
        else:
            raise ValueError(
                f"lam_prox['{name}'] has {val.ndim} dimensions, expected scalar, 1-D, or 2-D"
            )

    return lam_arr

resolve_lam_vc(lam_vc: Union[float, Dict[str, Union[float, list, np.ndarray]]], states: Optional[List[State]] = None) -> Union[float, np.ndarray] staticmethod

Resolve a lam_vc spec to a numeric value.

If lam_vc is a float it is returned as-is. If it is a dict mapping state names to weights, it is expanded to a dense array using each state's _slice. States not in the dict default to 1.0.

Dict values may be scalars, 1-D arrays (per-component), or 2-D arrays of shape (K, n_components) for per-node-per-component weighting. All 2-D entries must agree on K.

Parameters:

Name	Type	Description	Default
lam_vc	Union[float, Dict[str, Union[float, list, ndarray]]]	Scalar weight or {state_name: weight} dict.	required
states	Optional[List[State]]	Symbolic State objects (required when lam_vc is a dict).	None

Returns:

Type	Description
Union[float, ndarray]	float or np.ndarray of shape (n_states,) or (K, n_states).

Raises:

Type	Description
ValueError	If lam_vc is a dict and states is None, or if dict contains unknown names, or if 2-D entries disagree on K.

Source code in openscvx/algorithms/weights.py

@staticmethod
def resolve_lam_vc(
    lam_vc: Union[float, Dict[str, Union[float, list, np.ndarray]]],
    states: Optional[List["State"]] = None,
) -> Union[float, np.ndarray]:
    """Resolve a ``lam_vc`` spec to a numeric value.

    If *lam_vc* is a float it is returned as-is. If it is a dict
    mapping state names to weights, it is expanded to a dense array
    using each state's ``_slice``. States not in the dict default
    to ``1.0``.

    Dict values may be scalars, 1-D arrays (per-component), or 2-D
    arrays of shape ``(K, n_components)`` for per-node-per-component
    weighting. All 2-D entries must agree on *K*.

    Args:
        lam_vc: Scalar weight or ``{state_name: weight}`` dict.
        states: Symbolic State objects (required when *lam_vc* is a dict).

    Returns:
        float or np.ndarray of shape ``(n_states,)`` or ``(K, n_states)``.

    Raises:
        ValueError: If *lam_vc* is a dict and *states* is ``None``,
            or if dict contains unknown names, or if 2-D entries disagree on K.
    """
    if not isinstance(lam_vc, dict):
        return lam_vc

    if states is None:
        raise ValueError(
            "lam_vc was specified as a dict but no states were "
            "provided. Pass states so the dict can be expanded to "
            "a per-state weight array."
        )

    n_states = sum(s.shape[0] if len(s.shape) > 0 else 1 for s in states)

    valid_names = {s.name for s in states}
    unknown = set(lam_vc.keys()) - valid_names
    if unknown:
        raise ValueError(
            f"lam_vc dict contains unknown state name(s): {unknown}. "
            f"Valid state names: {sorted(valid_names)}"
        )

    # First pass: determine if any entry is 2-D and infer K.
    n_nodes: Optional[int] = None
    for state in states:
        if state.name not in lam_vc:
            continue
        val = np.asarray(lam_vc[state.name], dtype=float)
        if val.ndim == 2:
            if n_nodes is None:
                n_nodes = val.shape[0]
            elif val.shape[0] != n_nodes:
                raise ValueError(
                    f"lam_vc['{state.name}'] has {val.shape[0]} rows, but a "
                    f"previous entry had {n_nodes} rows. All 2-D entries "
                    f"must have the same number of rows (n_nodes-1)."
                )

    # Build the output array.
    if n_nodes is not None:
        lam_arr = np.ones((n_nodes, n_states))
    else:
        lam_arr = np.ones(n_states)

    for state in states:
        if state.name not in lam_vc:
            continue
        val = np.asarray(lam_vc[state.name], dtype=float)
        n_components = state.shape[0] if len(state.shape) > 0 else 1

        if val.ndim == 0:
            lam_arr[..., state._slice] = float(val)
        elif val.ndim == 1:
            if val.shape[0] != n_components:
                raise ValueError(
                    f"lam_vc['{state.name}'] has length {val.shape[0]}, "
                    f"expected scalar or length {n_components}"
                )
            lam_arr[..., state._slice] = val
        elif val.ndim == 2:
            if val.shape[1] != n_components:
                raise ValueError(
                    f"lam_vc['{state.name}'] has {val.shape[1]} columns, "
                    f"expected {n_components}"
                )
            lam_arr[:, state._slice] = val
        else:
            raise ValueError(
                f"lam_vc['{state.name}'] has {val.ndim} dimensions, "
                f"expected scalar, 1-D, or 2-D"
            )

    return lam_arr

adaptive_state_code_to_str(code: Union[int, jnp.ndarray, np.ndarray]) -> str

Map an :class:AdaptiveStateCode value (int or 0-d array) to its label.

Source code in openscvx/algorithms/state.py

def adaptive_state_code_to_str(code: Union[int, jnp.ndarray, np.ndarray]) -> str:
    """Map an :class:`AdaptiveStateCode` value (int or 0-d array) to its label."""
    return _ADAPTIVE_STATE_NAMES[AdaptiveStateCode(int(code))]

unpack_multishot_V(V: np.ndarray, *, n_x: int, n_u: int, t_nodes: np.ndarray, states: Sequence = ()) -> MultishotPropagation

Validate V layout and return :class:MultishotPropagation.

Raises:

Type	Description
ValueError	invalid V shape, empty matrix, or len(t_nodes) != n_segments + 1.

Source code in openscvx/algorithms/history.py

def unpack_multishot_V(
    V: np.ndarray,
    *,
    n_x: int,
    n_u: int,
    t_nodes: np.ndarray,
    states: Sequence = (),
) -> MultishotPropagation:
    """Validate ``V`` layout and return :class:`MultishotPropagation`.

    Raises:
        ValueError: invalid ``V`` shape, empty matrix, or ``len(t_nodes) != n_segments + 1``.
    """
    V = np.asarray(V, dtype=np.float64)
    if V.size == 0:
        raise ValueError("multishot V matrix is empty")
    seg_size = segment_size(n_x, n_u)
    if seg_size <= 0:
        raise ValueError(f"invalid multishot dimensions n_x={n_x}, n_u={n_u}")
    n_rows, n_sub = V.shape
    if n_sub < 1:
        raise ValueError("multishot V must have at least one substep column")
    if n_rows % seg_size != 0:
        raise ValueError(
            f"multishot V row count {n_rows} is not divisible by segment_size {seg_size}"
        )
    n_segments = n_rows // seg_size
    t_nodes_arr = np.asarray(t_nodes, dtype=np.float64).ravel()
    if len(t_nodes_arr) != n_segments + 1:
        raise ValueError(
            f"multishot t_nodes length {len(t_nodes_arr)} != n_segments + 1 "
            f"({n_segments + 1}); pass optimization node times from "
            "results.nodes['time'] or an explicit array."
        )
    return MultishotPropagation(
        V=V,
        n_x=n_x,
        n_u=n_u,
        t_nodes=t_nodes_arr,
        states=tuple(states),
    )