openscvx 0.5.3.dev24__tar.gz → 0.5.3.dev26__tar.gz

{openscvx-0.5.3.dev24/openscvx.egg-info → openscvx-0.5.3.dev26}/PKG-INFO

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

{openscvx-0.5.3.dev24 → openscvx-0.5.3.dev26}/docs/UnderTheHood/batching_jit_grad.md

@@ -161,15 +161,27 @@ The batch runs until its slowest element converges or exhausts its own cap;
 finished elements are frozen, not re-iterated, so each element matches the
 corresponding single `solve_jax` run.
 
-**New autotuners batch for free.** An autotuner's tunable knobs follow the
-same derivation: declare them as fields on a `HyperParams` subclass — bare
-annotations, no decorator; subclassing applies the frozen-dataclass transform
-and the JAX pytree registration, and the annotation picks the dtype (`int`
-gets the iteration counter's, `float` the problem's) — assign an instance to
-`self.hyper`, and read `state.hyper.<field>` inside `update_weights` (see the
-`AutotuningBase` docstring). The declaration is the registration — each field
-immediately works as a `solve_jax` override, a `solve_batched` sweep target,
-and a runtime input of the exported artifact, with zero edits anywhere else:
+**Every built-in autotuner knob batches for free.** All numeric knobs of the
+built-in autotuners are declared hyperparameters — the trust-region factors
+(`gamma_1` / `gamma_2`), the acceptance thresholds (`eta_0` / `eta_1` /
+`eta_2`), the weight clips, the cost-relaxation schedule — so any of them is a
+`solve_jax` override and a `solve_batched` sweep target by name:
+
+```python
+# sweep the trust-region growth factor across a batch, one cached artifact
+results = problem.solve_batched(algorithm={"gamma_1": jnp.linspace(1.5, 3.0, B)})
+```
+
+**New autotuners batch the same way.** A custom autotuner's tunable knobs
+follow the same derivation: declare them as fields on a `HyperParams` subclass
+— bare annotations, no decorator; subclassing applies the frozen-dataclass
+transform and the JAX pytree registration, and the annotation picks the dtype
+(`int` gets the iteration counter's, `float` the problem's) — assign an
+instance to `self.hyper`, and read `state.hyper.<field>` inside
+`update_weights` (see the `AutotuningBase` docstring). The declaration is the
+registration — each field immediately works as a `solve_jax` override, a
+`solve_batched` sweep target, and a runtime input of the exported artifact,
+with zero edits anywhere else:
 
 ```python
 from openscvx.algorithms import AutotuningBase, HyperParams

openscvx-0.5.3.dev26/docs/UnderTheHood/custom_algorithms_autotuners.md

@@ -0,0 +1,223 @@
+# Custom Algorithms and Autotuners
+
+OpenSCvx solves through a pluggable SCP loop: an **algorithm** owns the
+iteration (how a subproblem is built, when the loop is converged) and an
+**autotuner** owns the penalty-weight update applied each step. Both are
+ordinary Python classes behind small abstract bases — `Algorithm` and
+`AutotuningBase`, importable from `openscvx.algorithms` — so swapping either
+out is a subclass, not a fork.
+
+Most users never touch these: the default
+`PenalizedTrustRegion` + `AugmentedLagrangian` pair covers the common case.
+Reach for this page when you want a different penalty schedule, a different
+convergence rule, or a whole new SCP variant — and you want it to keep
+working under `solve()`, `solve_jax()`, and `solve_batched()` alike.
+
+Both bases live under `openscvx.algorithms`:
+
+```python
+from openscvx.algorithms import Algorithm, AutotuningBase, HyperParams
+```
+
+## A custom autotuner
+
+An autotuner is a **pure functional update on the iterate**. Its one required
+method, `update_weights`, takes the current `AlgorithmState`, the just-solved
+subproblem `candidate`, the constraints, settings, and parameters, and returns
+the *next* `AlgorithmState`:
+
+```python
+import jax.numpy as jnp
+from openscvx.algorithms import AdaptiveStateCode, AutotuningBase, HyperParams
+
+
+class MyHyper(HyperParams):
+    ramp: float = 2.0
+
+
+class MyAutotuner(AutotuningBase):
+    COMPUTES_ACCEPTANCE_METRICS = False  # this tuner never rejects an iterate
+
+    def __init__(self, ramp: float = 2.0):
+        self.hyper = MyHyper(ramp=ramp)  # assign self.hyper FIRST
+
+    def update_weights(self, state, candidate, nodal_constraints, settings, params):
+        return state.replace(
+            # accept the candidate: its trajectory becomes the next iterate
+            x=candidate.x,
+            u=candidate.u,
+            x_prop=candidate.x_prop,
+            x_prop_plus=candidate.x_prop_plus,
+            # ramp the proximal weight every iteration
+            lam_prox=state.lam_prox * state.hyper.ramp,
+            # jnp.int32, not bare asarray: the code must keep the state field's
+            # strong int32 dtype or the lax.while_loop carry check fails.
+            adaptive_state_code=jnp.int32(AdaptiveStateCode.ACCEPT_HIGHER),
+        )
+```
+
+Three rules carry the whole contract.
+
+### The returned state *is* the next iterate
+
+The SCP loop discards everything except what `update_weights` returns.
+Accepting the subproblem's candidate therefore 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.
+
+!!! warning "The silent no-op"
+    An update that never copies the candidate produces a solver that runs to
+    `k_max` without the iterate ever moving — no error, just a flat history.
+    If a custom tuner seems to "not converge," check that the accept branch
+    copies all four trajectory fields.
+
+Express any accept/reject branching with `jax.lax.cond` or `jnp.where`, not a
+Python `if` on iterate values: `update_weights` is traced (it runs inside the
+`lax.while_loop` of `solve_jax` / `solve_batched`, vmapped per batch element),
+so it must not mutate `state` or `candidate`, raise on data-dependent
+conditions, or return strings.
+
+### Declare tunable knobs as `HyperParams`, not `self` attributes
+
+A numeric knob a user might sweep (a ramp factor, a relaxation iteration) must
+**not** be read off `self` inside `update_weights` — a plain Python attribute
+is baked into the trace and invisible to every override channel. Declare it on
+a `HyperParams` subclass, assign an instance to `self.hyper`, and read it from
+`state.hyper`.
+
+The declaration *is* the registration. With nothing else, the field becomes:
+
+* a **per-solve override** — `problem.solve_jax(algorithm={"ramp": 1.5})`;
+* a **batchable sweep target** —
+  `problem.solve_batched(algorithm={"ramp": jnp.linspace(1.1, 2.0, 8)})`;
+* a **runtime input** of the exported `solve_batched` artifact, so changing it
+  needs no recompile.
+
+!!! danger "Assign `self.hyper` first"
+    Declared knobs are also readable and writable as bare attributes
+    (`autotuner.ramp = 3.0`) for ergonomics — a proxy on `AutotuningBase`
+    routes the access into `hyper`. That proxy only works once `self.hyper`
+    exists, so assign it as the **first** line of `__init__`. A knob assigned
+    before `hyper` exists lands as a plain instance attribute the proxy cannot
+    see, and your override channels silently break.
+
+### Acceptance metrics are opt-in
+
+`COMPUTES_ACCEPTANCE_METRICS` (default `True`) tells the loop whether
+`update_weights` produces a predicted/actual reduction and an acceptance
+ratio. The loop records and prints those diagnostics only when it is `True`.
+A tuner that never rejects an iterate — like the built-in
+`ConstantProximalWeight` and `RampProximalWeight` — sets it `False`.
+
+To add columns to the iteration table, set the class attribute `COLUMNS` to a
+list of `openscvx.utils.printing.Column` specs; the algorithm concatenates
+them into its own table via `get_columns`.
+
+## A custom algorithm
+
+An `Algorithm` owns the SCP iteration end to end: it *builds* the JAX-pure
+iteration body, *stores* it, and *drives* it one step at a time. All iteration
+state lives on `AlgorithmState` (JAX-traceable) and `AlgorithmHistory`
+(Python-side) threaded explicitly through `step()` — avoid storing mutable
+iteration state on `self`.
+
+The surface is six methods. `converged` has a working default; the other five
+are abstract.
+
+### `__init__` — end with `super().__init__`
+
+The base records the weights, autotuner, and convergence parameters. Every
+parameter is required, so the ABC is not yet another place defaults are
+declared — your subclass owns the user-facing defaults and forwards them:
+
+```python
+from openscvx.algorithms import Algorithm
+from openscvx.algorithms.weights import Weights
+from openscvx.algorithms import AugmentedLagrangian
+
+
+class MyAlgorithm(Algorithm):
+    def __init__(self, autotuner=None, k_max=200, **weight_kwargs):
+        super().__init__(
+            weights=Weights.build(**weight_kwargs),
+            autotuner=autotuner if autotuner is not None else AugmentedLagrangian(),
+            k_max=k_max,
+            t_max=None,
+            ep_tr=1e-4,
+            ep_vb=1e-4,
+            ep_vc=1e-8,
+        )
+```
+
+### `build_iteration` — fuse the step
+
+`Problem` assembles the discretization solvers, lowered constraints, and
+convex-solver callback, then asks the algorithm to fuse them into one JAX-pure
+`(state, params) -> (next_state, diagnostics)` body. The algorithm owns this
+because the fusion is algorithm-specific (which autotuner runs, which penalty
+terms are assembled). PTR's implementation is a thin wrapper around
+`make_scp_iteration` (`openscvx/algorithms/scvx/iteration.py`) that threads its
+`autotuner` into the fused body. The returned body must advance `state.k` by
+one per call — every solve path terminates on `k` reaching `k_max`.
+
+### `initialize` and `step` — store, then drive
+
+`Problem.initialize()` builds and JIT-warms the body, then hands it back via
+`initialize(iteration_fn, emitter)` for the algorithm to store. `step()` is the
+Python-loop driver: it calls the stored body, records the per-iteration
+diagnostics into `history`, emits progress, and returns
+`(next_state, converged)`.
+
+### `converged` — the one-method convergence hook
+
+`converged(state)` is the only piece of convergence policy that is
+algorithm-specific; it has a concrete default, so override it only to change
+the rule:
+
+```python
+def converged(self, state):
+    return state.J_vc < state.ep_vc   # custom: virtual control alone
+```
+
+The default — every metric below its threshold — is algorithm-agnostic, since
+`AlgorithmState` carries `J_tr` / `J_vb` / `J_vc` and the `ep_*` tolerances
+generically.
+
+!!! note "Honored on every solve path"
+    The override is honored on **all three** solve paths. `step()` routes
+    through it on the Python `solve()` path, and the `lax.while_loop` harness
+    (`openscvx/algorithms/loop.py`) takes it as the loop predicate for
+    `solve_jax` / `solve_batched`, where it is vmapped per batch element. It
+    must therefore be JAX-traceable. An override implies a subclass, hence a
+    distinct class name that the export cache key already folds in — so a
+    changed predicate invalidates stale batched artifacts on its own.
+
+By contrast, `t_max` (the wall-clock time limit) is honored **only** on the
+Python `solve()` path; the JAX loop terminates on `k_max` and `converged`
+alone, since there is no wall-clock probe inside a trace.
+
+### `get_columns` and `citation`
+
+`get_columns(verbosity)` returns the iteration-table columns to print —
+implementations typically concatenate the algorithm's own columns with the
+autotuner's `COLUMNS` and filter by each column's `min_verbosity`.
+`citation()` returns a list of BibTeX strings for the algorithm.
+
+## Using a custom class
+
+Pass either to the `Problem` constructor — the algorithm directly, an
+autotuner through the algorithm:
+
+```python
+import openscvx as ox
+
+problem = ox.Problem(
+    ...,
+    algorithm=MyAlgorithm(autotuner=MyAutotuner(ramp=1.5)),
+)
+```
+
+The default `PenalizedTrustRegion` lives at
+`openscvx/algorithms/scvx/penalized_trust_region.py` and is the reference
+implementation to read alongside this guide.

{openscvx-0.5.3.dev24 → openscvx-0.5.3.dev26}/mkdocs.yml

@@ -198,4 +198,5 @@ nav:
     - Under the Hood:
       - 'UnderTheHood/vectorization_and_vmapping.md'
       - 'UnderTheHood/lowering_architecture.md'
-      - 'UnderTheHood/batching_jit_grad.md'
+      - 'UnderTheHood/batching_jit_grad.md'
+      - 'UnderTheHood/custom_algorithms_autotuners.md'

{openscvx-0.5.3.dev24 → openscvx-0.5.3.dev26}/openscvx/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.5.3.dev24'
-__version_tuple__ = version_tuple = (0, 5, 3, 'dev24')
+__version__ = version = '0.5.3.dev26'
+__version_tuple__ = version_tuple = (0, 5, 3, 'dev26')
 
-__commit_id__ = commit_id = 'gba72d9b80'
+__commit_id__ = commit_id = 'g053a32537'

{openscvx-0.5.3.dev24 → openscvx-0.5.3.dev26}/openscvx/algorithms/__init__.py

@@ -5,9 +5,11 @@ 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. Immutable components (solver,
-discretization_solver, jax_constraints, etc.) are stored during ``initialize()``;
-mutable configuration (params, settings) is passed per-step.
+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:
 
@@ -37,6 +39,8 @@ from typing import Annotated, Any, Dict, List, Optional, Union
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 from .autotuner import (
+    AcceptanceRatioAutotuner,
+    AcceptanceRatioHyper,
     AdaptiveProximalWeight,
     AdaptiveProximalWeightSpec,
     AugmentedLagrangian,
@@ -46,18 +50,18 @@ from .autotuner import (
     RampProximalWeight,
     RampProximalWeightSpec,
 )
-from .base import (
+from .autotuner.base import AutotuningBase
+from .base import Algorithm
+from .history import AlgorithmHistory, DiscretizationResult
+from .hyperparams import HyperParams
+from .optimization_results import OptimizationResults
+from .scvx import PenalizedTrustRegion
+from .state import (
     AdaptiveStateCode,
-    Algorithm,
-    AlgorithmHistory,
     AlgorithmState,
-    AutotuningBase,
-    DiscretizationResult,
-    HyperParams,
+    CandidateIterate,
     adaptive_state_code_to_str,
 )
-from .optimization_results import OptimizationResults
-from .scvx import PenalizedTrustRegion
 from .weights import Weights
 
 # ---------------------------------------------------------------------------
@@ -138,6 +142,7 @@ __all__ = [
     "AlgorithmHistory",
     "AlgorithmState",
     "AdaptiveStateCode",
+    "CandidateIterate",
     "adaptive_state_code_to_str",
     "DiscretizationResult",
     "Weights",
@@ -147,6 +152,8 @@ __all__ = [
     "PenalizedTrustRegion",
     "AutotuningBase",
     "HyperParams",
+    "AcceptanceRatioAutotuner",
+    "AcceptanceRatioHyper",
     "AugmentedLagrangian",
     "AdaptiveProximalWeight",
     "ConstantProximalWeight",

{openscvx-0.5.3.dev24 → openscvx-0.5.3.dev26}/openscvx/algorithms/autotuner/__init__.py

@@ -1,11 +1,14 @@
 """SCP weight autotuning strategies."""
 
+from .acceptance_ratio import AcceptanceRatioAutotuner, AcceptanceRatioHyper
 from .adaptive_proximal_weight import AdaptiveProximalWeight, AdaptiveProximalWeightSpec
 from .augmented_lagrangian import AugmentedLagrangian, AugmentedLagrangianSpec
 from .constant_proximal_weight import ConstantProximalWeight, ConstantProximalWeightSpec
 from .ramp_proximal_weight import RampProximalWeight, RampProximalWeightSpec
 
 __all__ = [
+    "AcceptanceRatioAutotuner",
+    "AcceptanceRatioHyper",
     "AdaptiveProximalWeight",
     "AdaptiveProximalWeightSpec",
     "AugmentedLagrangian",

openscvx-0.5.3.dev26/openscvx/algorithms/autotuner/acceptance_ratio.py

@@ -0,0 +1,234 @@
+"""Acceptance-ratio autotuner family.
+
+The acceptance-ratio autotuners share one decision rule: an acceptance ratio
+:math:`\\rho` between the predicted and actual reduction of the nonlinear
+objective drives a four-bucket update of the trust-region weight ``lam_prox``
+(``REJECT`` / ``ACCEPT_HIGHER`` / ``ACCEPT_CONSTANT`` / ``ACCEPT_LOWER``). What
+differs between members is the *multiplier* rule — how the virtual-control and
+virtual-buffer weights (``lam_vc`` / ``lam_vb_*``) respond once an iterate is
+accepted.
+
+:class:`AcceptanceRatioAutotuner` owns the shared decision body and exposes a
+single :meth:`AcceptanceRatioAutotuner._update_multipliers` hook for the
+multiplier rule; subclasses supply only that hook.
+:class:`AdaptiveProximalWeight` is the hook-default (multipliers held
+constant); :class:`AugmentedLagrangian` overrides it with the
+constraint-violation update.
+"""
+
+from typing import TYPE_CHECKING, List
+
+import jax
+import jax.numpy as jnp
+
+from openscvx.config import Config
+from openscvx.utils.printing import (
+    Column,
+    Verbosity,
+    color_acceptance_ratio,
+    color_adaptive_state,
+)
+
+from ..penalty import calculate_nonlinear_penalty
+from ..state import AdaptiveStateCode
+from .base import AutotuningBase, LamCostRelaxHyper
+
+if TYPE_CHECKING:
+    from openscvx.lowered import LoweredJaxConstraints
+
+    from ..state import AlgorithmState, CandidateIterate
+
+
+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
+
+
+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 _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)

openscvx-0.5.3.dev26/openscvx/algorithms/autotuner/adaptive_proximal_weight.py

@@ -0,0 +1,74 @@
+"""Adaptive proximal-weight autotuner.
+
+Same acceptance-ratio decision logic as the rest of the family for
+``lam_prox``, but ``lam_vc`` and ``lam_vb_*`` are held constant at their
+current values — the default multiplier hook.
+"""
+
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict
+
+from .acceptance_ratio import AcceptanceRatioAutotuner, AcceptanceRatioHyper
+
+
+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,
+        )
+
+
+# =============================================================================
+# Pydantic spec for dict / YAML validation
+# =============================================================================
+
+
+class AdaptiveProximalWeightSpec(BaseModel):
+    """Validates AdaptiveProximalWeight configuration from dict/YAML input."""
+
+    type: Literal["AdaptiveProximalWeight"] = "AdaptiveProximalWeight"
+    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
+
+    model_config = ConfigDict(extra="forbid")
+
+    def build(self) -> AdaptiveProximalWeight:
+        return AdaptiveProximalWeight(**self.model_dump(exclude={"type"}, exclude_unset=True))