openscvx 0.4.1.dev145__tar.gz → 0.4.1.dev147__tar.gz

{openscvx-0.4.1.dev145/openscvx.egg-info → openscvx-0.4.1.dev147}/PKG-INFO

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

{openscvx-0.4.1.dev145 → openscvx-0.4.1.dev147}/openscvx/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.4.1.dev145'
-__version_tuple__ = version_tuple = (0, 4, 1, 'dev145')
+__version__ = version = '0.4.1.dev147'
+__version_tuple__ = version_tuple = (0, 4, 1, 'dev147')
 
-__commit_id__ = commit_id = 'g8b0169a22'
+__commit_id__ = commit_id = 'g50fc80214'

{openscvx-0.4.1.dev145 → openscvx-0.4.1.dev147}/openscvx/algorithms/__init__.py

@@ -140,7 +140,9 @@ def _resolve_autotuner(val: Any) -> Any:
         raise TypeError(f"Invalid autotuner keyword argument: {e}. Valid keys: {valid}") from None
 
 
-def _resolve_algorithm(kwargs: dict, states: list = None) -> "PenalizedTrustRegion":
+def _resolve_algorithm(
+    kwargs: dict, states: list = None, controls: list = None
+) -> "PenalizedTrustRegion":
     """Build a :class:`PenalizedTrustRegion` from a user-supplied dict.
 
     Supports a nested ``autotuner`` key that is resolved via
@@ -149,7 +151,11 @@ def _resolve_algorithm(kwargs: dict, states: list = None) -> "PenalizedTrustRegi
     Args:
         kwargs: Algorithm keyword arguments (e.g. ``lam_cost``, ``autotuner``).
         states: Symbolic State objects, forwarded to the algorithm constructor
-            so that dict-valued ``lam_cost`` can be expanded immediately.
+            so that dict-valued ``lam_cost`` / ``lam_prox`` can be expanded
+            immediately.
+        controls: Symbolic Control objects, forwarded to the algorithm
+            constructor so that dict-valued ``lam_prox`` can be expanded
+            immediately.
     """
     kwargs = dict(kwargs)  # copy to avoid mutating the caller's dict
 
@@ -157,9 +163,11 @@ def _resolve_algorithm(kwargs: dict, states: list = None) -> "PenalizedTrustRegi
     if "autotuner" in kwargs:
         kwargs["autotuner"] = _resolve_autotuner(kwargs["autotuner"])
 
-    # Forward states so dict lam_cost can be expanded eagerly
+    # Forward states/controls so dict weights can be expanded eagerly
     if states is not None:
         kwargs.setdefault("states", states)
+    if controls is not None:
+        kwargs.setdefault("controls", controls)
 
     try:
         return PenalizedTrustRegion(**kwargs)

{openscvx-0.4.1.dev145 → openscvx-0.4.1.dev147}/openscvx/algorithms/augmented_lagrangian.py

@@ -116,7 +116,7 @@ class AugmentedLagrangian(AutotuningBase):
         candidate_x_prop: np.ndarray,
         settings: Config,
         lam_vc: np.ndarray,
-        lam_prox: float,
+        lam_prox: np.ndarray,
     ) -> np.ndarray:
         """Update virtual control penalty weights from state violation.
 
@@ -126,7 +126,9 @@ class AugmentedLagrangian(AutotuningBase):
         """
         nu = (settings.sim.inv_S_x @ abs(candidate.x[1:] - candidate_x_prop).T).T
         mask = nu > self.ep
-        scale = self.eta_lambda * (1 / (2 * lam_prox))
+        # TODO: (haynec) use per-variable lam_prox to scale VC updates proportionally
+        lam_prox_scalar = float(np.max(lam_prox))
+        scale = self.eta_lambda * (1 / (2 * lam_prox_scalar))
         case1 = lam_vc + nu * scale
         case2 = lam_vc + (nu**2) / self.ep * scale
         vc_new = np.where(mask, case1, case2)
@@ -222,12 +224,12 @@ class AugmentedLagrangian(AutotuningBase):
 
             if rho < self.eta_0:
                 # Reject Solution and higher weight
-                lam_prox_k1 = min(self.lam_prox_max, self.gamma_1 * lam_prox_k)
+                lam_prox_k1 = np.minimum(self.lam_prox_max, self.gamma_1 * lam_prox_k)
                 state.lam_prox_history.append(lam_prox_k1)
                 adaptive_state = "Reject Higher"
             elif rho >= self.eta_0 and rho < self.eta_1:
                 # Accept Solution with heigher weight
-                lam_prox_k1 = min(self.lam_prox_max, self.gamma_1 * lam_prox_k)
+                lam_prox_k1 = np.minimum(self.lam_prox_max, self.gamma_1 * lam_prox_k)
                 state.lam_prox_history.append(lam_prox_k1)
 
                 # Update virtual control weight matrix
@@ -255,7 +257,7 @@ class AugmentedLagrangian(AutotuningBase):
                 adaptive_state = "Accept Constant"
             else:
                 # Accept Solution with lower weight
-                lam_prox_k1 = max(self.lam_prox_min, self.gamma_2 * lam_prox_k)
+                lam_prox_k1 = np.maximum(self.lam_prox_min, self.gamma_2 * lam_prox_k)
                 state.lam_prox_history.append(lam_prox_k1)
                 # Update virtual control weight matrix
                 candidate.lam_vc = self._update_virtual_control_weights(

{openscvx-0.4.1.dev145 → openscvx-0.4.1.dev147}/openscvx/algorithms/base.py

@@ -17,6 +17,7 @@ if TYPE_CHECKING:
     from openscvx.config import Config
     from openscvx.lowered.jax_constraints import LoweredJaxConstraints
     from openscvx.solvers import ConvexSolver
+    from openscvx.symbolic.expr.control import Control
     from openscvx.symbolic.expr.state import State
 
 
@@ -195,6 +196,113 @@ def _expand_lam_vc_dict(
     return lam_arr
 
 
+def _expand_lam_prox_dict(
+    lam_prox_dict: Dict[str, Union[float, list, np.ndarray]],
+    states: List["State"],
+    controls: List["Control"],
+) -> np.ndarray:
+    """Expand a ``{name: weight}`` dict to a per-variable proximal weight array.
+
+    Maps user-provided per-state/per-control trust region weights to a dense
+    array using each variable's ``_slice``.  Variables not present in the dict
+    default to ``1.0``.
+
+    Per-variable values may be:
+
+    * **scalar** — broadcast to every component and every node.
+    * **1-D array** of length ``n_components`` — per-component weight,
+      same for every node.
+    * **2-D array** of shape ``(K, n_components)`` — per-node-per-component.
+      When any variable supplies a 2-D value the output is 2-D with first
+      dimension *K* (all 2-D entries must agree on *K*).
+
+    Args:
+        lam_prox_dict: Mapping from state/control names to proximal weights.
+        states: List of State objects (must already have ``_slice`` assigned).
+        controls: List of Control objects (must already have ``_slice`` assigned).
+
+    Returns:
+        np.ndarray of shape ``(n_states + n_controls,)`` when all entries are
+        scalar/1-D, or ``(K, n_states + n_controls)`` when any entry is 2-D.
+
+    Raises:
+        ValueError: If the dict contains unknown names or 2-D entries disagree
+            on the number of nodes.
+    """
+    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_dict.keys()) - valid_names
+    if unknown:
+        raise ValueError(
+            f"lam_prox dict contains unknown name(s): {unknown}. 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_dict:
+            continue
+        val = np.asarray(lam_prox_dict[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_dict:
+            continue
+        val = np.asarray(lam_prox_dict[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
+
+
 @dataclass
 class Weights:
     """Normalized SCP weights used internally by the algorithm and autotuner.
@@ -214,7 +322,10 @@ class Weights:
     drift.
 
     Attributes:
-        lam_prox: Trust region (proximal) weight (normalized).
+        lam_prox: Trust region (proximal) weight (normalized). Scalar or
+            array of shape ``(n_states + n_controls,)`` or
+            ``(N, n_states + n_controls)`` for per-variable / per-node
+            weighting.
         lam_vc: Virtual control penalty weight (normalized). Scalar or
             array of shape ``(n_states,)`` or ``(n_nodes-1, n_states)``
             for per-state / per-node weighting.
@@ -231,7 +342,7 @@ class Weights:
             :meth:`set_vb_arrays`.
     """
 
-    lam_prox: float = 1e0
+    lam_prox: Union[float, np.ndarray] = 1e0
     lam_vc: Union[float, np.ndarray] = 1e1
     lam_cost: Union[float, np.ndarray] = 1e-1
     lam_vb: float = 0.0
@@ -240,13 +351,17 @@ class Weights:
 
     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)
 
         # Snapshot the user-specified values so normalize() is idempotent.
-        self._raw_lam_prox = self.lam_prox
+        self._raw_lam_prox = (
+            self.lam_prox.copy() if isinstance(self.lam_prox, np.ndarray) else self.lam_prox
+        )
         self._raw_lam_vc = (
             self.lam_vc.copy() if isinstance(self.lam_vc, np.ndarray) else self.lam_vc
         )
@@ -264,6 +379,11 @@ class Weights:
         making this method idempotent and safe to call after updating
         any individual raw weight.
         """
+        raw_prox_max = (
+            float(np.max(self._raw_lam_prox))
+            if isinstance(self._raw_lam_prox, np.ndarray)
+            else self._raw_lam_prox
+        )
         raw_vc_max = (
             float(np.max(self._raw_lam_vc))
             if isinstance(self._raw_lam_vc, np.ndarray)
@@ -281,7 +401,7 @@ class Weights:
             )
         else:
             raw_vb_max = self._raw_lam_vb
-        scale = max(self._raw_lam_prox, raw_vc_max, raw_cost_max, raw_vb_max)
+        scale = max(raw_prox_max, raw_vc_max, raw_cost_max, raw_vb_max)
         if scale > 0:
             self.lam_prox = self._raw_lam_prox / scale
             self.lam_vc = self._raw_lam_vc / scale
@@ -614,7 +734,7 @@ class AlgorithmState:
     lam_cost_history: List[Union[float, np.ndarray]] = field(default_factory=list)
     lam_vb_nodal_history: List[np.ndarray] = field(default_factory=list)
     lam_vb_cross_history: List[np.ndarray] = field(default_factory=list)
-    lam_prox_history: List[float] = field(default_factory=list)
+    lam_prox_history: List[np.ndarray] = field(default_factory=list)
     x_full: List[np.ndarray] = field(default_factory=list)
     x_prop_full: List[np.ndarray] = field(default_factory=list)
 
@@ -820,11 +940,11 @@ class AlgorithmState:
         return self.discretizations[index].E_d
 
     @property
-    def lam_prox(self) -> float:
+    def lam_prox(self) -> np.ndarray:
         """Get current trust region weight.
 
         Returns:
-            Current trust region weight (latest entry in lam_prox_history)
+            Array of shape ``(N, n_states + n_controls)``.
         """
         if not self.lam_prox_history:
             raise ValueError("lam_prox_history is empty. Initialize state using from_settings().")
@@ -901,8 +1021,13 @@ class AlgorithmState:
         """
         n = settings.sim.n
         n_states = settings.sim.n_states
+        n_controls = settings.sim.n_controls
+        n_total = n_states + n_controls
         lam_vc_array = np.ones((n - 1, n_states)) * weights.lam_vc
 
+        # Expand lam_prox to (N, n_states + n_controls) array
+        lam_prox_array = np.ones((n, n_total)) * weights.lam_prox
+
         # Expand scalar lam_cost to per-state array
         if isinstance(weights.lam_cost, np.ndarray):
             lam_cost_init = weights.lam_cost.copy()
@@ -931,7 +1056,7 @@ class AlgorithmState:
             lam_cost_history=[lam_cost_init],
             lam_vb_nodal_history=[weights.lam_vb_nodal.copy()],
             lam_vb_cross_history=[weights.lam_vb_cross.copy()],
-            lam_prox_history=[weights.lam_prox],
+            lam_prox_history=[lam_prox_array],
         )
 
 
@@ -989,6 +1114,41 @@ class Algorithm(ABC):
     #: Maximum number of SCP iterations. Subclasses must set this in ``__init__``.
     k_max: int
 
+    @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, *states* and *controls*
+        must be provided so the dict can be expanded to a per-variable
+        array via :func:`_expand_lam_prox_dict`.
+
+        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``.
+        """
+        if isinstance(lam_prox, dict):
+            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."
+                )
+            return _expand_lam_prox_dict(lam_prox, states, controls)
+        return lam_prox
+
     @staticmethod
     def _resolve_lam_cost(
         lam_cost: Union[float, Dict[str, float]],

{openscvx-0.4.1.dev145 → openscvx-0.4.1.dev147}/openscvx/algorithms/optimization_results.py

@@ -91,7 +91,7 @@ class OptimizationResults:
     VC_history: list[np.ndarray] = field(default_factory=list, metadata={"npz": "array_list"})
 
     # Convergence histories
-    lam_prox_history: list[float] = field(default_factory=list, metadata={"npz": "float_list"})
+    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"}
     )

{openscvx-0.4.1.dev145 → openscvx-0.4.1.dev147}/openscvx/algorithms/penalized_trust_region.py

@@ -29,6 +29,7 @@ from .ramp_proximal_weight import RampProximalWeight
 if TYPE_CHECKING:
     from openscvx.lowered import LoweredJaxConstraints
     from openscvx.solvers import ConvexSolver
+    from openscvx.symbolic.expr.control import Control
     from openscvx.symbolic.expr.state import State
 
     from .base import AutotuningBase
@@ -85,7 +86,7 @@ class PenalizedTrustRegion(Algorithm):
         self,
         autotuner: "AutotuningBase" = None,
         k_max: int = 200,
-        lam_prox: float = 1e0,
+        lam_prox: Union[float, Dict[str, Union[float, list]]] = 1e0,
         lam_vc: Union[float, Dict[str, Union[float, list]]] = 1e1,
         lam_cost: Union[float, Dict[str, float]] = 1e-1,
         lam_vb: float = 0.0,
@@ -93,6 +94,7 @@ class PenalizedTrustRegion(Algorithm):
         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.
 
@@ -100,7 +102,14 @@ class PenalizedTrustRegion(Algorithm):
             autotuner: Weight adaptation strategy. Defaults to
                 :class:`AugmentedLagrangian` when ``None``.
             k_max: Maximum SCP iterations. Defaults to 200.
-            lam_prox: Trust region (proximal) weight. Defaults to 1.0.
+            lam_prox: Trust region (proximal) weight. Either a float
+                (applied uniformly to all states and controls) or a dict
+                mapping state/control names to weights, e.g.
+                ``{"velocity": 1e0, "thrust": 5e-1}``.  Dict values may
+                be scalars, 1-D arrays for per-component weighting, or
+                2-D arrays of shape ``(n_nodes, n_components)`` for
+                per-node-per-component weighting.  Variables not in the
+                dict default to ``1.0``. Defaults to 1.0.
             lam_vc: Virtual control penalty weight. Either a float
                 (applied uniformly to all states) or a dict mapping state
                 names to per-state weights, e.g.
@@ -119,8 +128,12 @@ class PenalizedTrustRegion(Algorithm):
             ep_tr: Trust region convergence tolerance. Defaults to 1e-4.
             ep_vb: Virtual buffer convergence tolerance. Defaults to 1e-4.
             ep_vc: Virtual control convergence tolerance. Defaults to 1e-8.
-            states: Symbolic State objects (required when *lam_cost* is a
-                dict). Normally provided automatically by :class:`Problem`.
+            states: Symbolic State objects (required when *lam_cost* or
+                *lam_prox* is a dict). Normally provided automatically by
+                :class:`Problem`.
+            controls: Symbolic Control objects (required when *lam_prox*
+                is a dict). Normally provided automatically by
+                :class:`Problem`.
         """
         # Compiled infrastructure (set by initialize())
         self._solver: "ConvexSolver" = None
@@ -134,17 +147,20 @@ class PenalizedTrustRegion(Algorithm):
             autotuner if autotuner is not None else AugmentedLagrangian()
         )
 
-        # Store states for later re-resolution (e.g. user sets lam_cost to a
-        # new dict via the property setter after construction).
+        # Store states/controls for later re-resolution (e.g. user sets
+        # lam_cost or lam_prox to a new dict via the property setter after
+        # construction).
         self._states: List["State"] = states
+        self._controls: List["Control"] = controls
 
-        # Resolve dict lam_vc / lam_cost → ndarray (requires states)
+        # Resolve dict lam_prox / lam_vc / lam_cost → ndarray (requires states/controls)
+        resolved_lam_prox = self._resolve_lam_prox(lam_prox, states, controls)
         resolved_lam_vc = self._resolve_lam_vc(lam_vc, states)
         resolved_lam_cost = self._resolve_lam_cost(lam_cost, states)
 
         # SCP weights (grouped dataclass, normalized for numerical conditioning)
         self.weights = Weights(
-            lam_prox=lam_prox,
+            lam_prox=resolved_lam_prox,
             lam_vc=resolved_lam_vc,
             lam_cost=resolved_lam_cost,
             lam_vb=lam_vb,
@@ -187,11 +203,12 @@ class PenalizedTrustRegion(Algorithm):
         return x0_prior.reshape(1, -1)
 
     @property
-    def lam_prox(self) -> float:
+    def lam_prox(self) -> Union[float, np.ndarray]:
         """Trust region (proximal) weight.
 
         This is the user-specified value before normalization. Setting this
-        property triggers automatic re-normalization of all weights.
+        property triggers automatic re-normalization of all weights.  May be
+        a scalar or array for per-variable / per-node weighting.
 
         !!! note
             The autotuner may modify the normalized weight in
@@ -201,8 +218,9 @@ class PenalizedTrustRegion(Algorithm):
         return self.weights._raw_lam_prox
 
     @lam_prox.setter
-    def lam_prox(self, value: float) -> None:
-        self.weights._raw_lam_prox = value
+    def lam_prox(self, value: Union[float, Dict[str, Union[float, list]]]) -> None:
+        resolved = self._resolve_lam_prox(value, self._states, self._controls)
+        self.weights._raw_lam_prox = resolved
         self.weights.normalize()
 
     @property
@@ -465,7 +483,8 @@ class PenalizedTrustRegion(Algorithm):
             "J_vb": state.J_vb,
             "J_vc": state.J_vc,
             "cost": cost[-1],
-            "lam_prox": state.lam_prox,
+            # TODO: (haynec) log per-variable lam_prox detail (e.g. min/max range)
+            "lam_prox": float(np.max(state.lam_prox)),
             "prob_stat": prob_stat,
             "adaptive_state": adaptive_state,
             "ep_tr": self.ep_tr,

{openscvx-0.4.1.dev145 → openscvx-0.4.1.dev147}/openscvx/algorithms/ramp_proximal_weight.py

@@ -2,6 +2,8 @@
 
 from typing import TYPE_CHECKING
 
+import numpy as np
+
 from openscvx.config import Config
 
 from .base import AutotuningBase
@@ -61,10 +63,10 @@ class RampProximalWeight(AutotuningBase):
             candidate.lam_cost = weights.lam_cost
 
         # Check if we're already at max before updating
-        was_at_max = state.lam_prox >= self.lam_prox_max
+        was_at_max = np.all(state.lam_prox >= self.lam_prox_max)
 
         # Calculate and append new value
-        new_lam_prox = min(state.lam_prox * self.ramp_factor, self.lam_prox_max)
+        new_lam_prox = np.minimum(state.lam_prox * self.ramp_factor, self.lam_prox_max)
         state.lam_prox_history.append(new_lam_prox)
 
         # If we were already at max, or if we just reached it and it's staying constant

{openscvx-0.4.1.dev145 → openscvx-0.4.1.dev147}/openscvx/lowered/cvxpy_variables.py

@@ -28,7 +28,7 @@ class CVXPyVariables:
         - Scaled expressions: CVXPy expressions for scaled state/control at each node
 
     Attributes:
-        lam_prox: Trust region weight parameter (scalar, nonneg)
+        lam_prox: Trust region weight parameter (N x (n_states+n_controls), nonneg)
         lam_cost: Cost function weight parameter (n_states, nonneg)
         lam_vc: Virtual control penalty weights (N-1 x n_states, nonneg)
         lam_vb_nodal: Virtual buffer penalty weights for nodal constraints (N x n_nodal, nonneg)

{openscvx-0.4.1.dev145 → openscvx-0.4.1.dev147}/openscvx/plotting/scp_iteration.py

@@ -339,12 +339,13 @@ def plot_scp_convergence_histories(result: OptimizationResults) -> go.Figure:
     iterations_reduction = np.arange(len(result.actual_reduction_history))
     iterations_ratio = np.arange(len(result.acceptance_ratio_history))
 
-    # Plot 1: Trust region weight history
+    # Plot 1: Trust region weight history (plot max across all variables)
     if len(result.lam_prox_history) > 0:
+        lam_prox_max_per_iter = [float(np.max(lp)) for lp in result.lam_prox_history]
         fig.add_trace(
             go.Scatter(
                 x=iterations_lam_prox,
-                y=result.lam_prox_history,
+                y=lam_prox_max_per_iter,
                 mode="lines+markers",
                 name="lam_prox",
                 line={"color": "cyan", "width": 2},

{openscvx-0.4.1.dev145 → openscvx-0.4.1.dev147}/openscvx/problem.py

@@ -18,7 +18,7 @@ import os
 import queue
 import threading
 import time
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import jax
 
@@ -83,8 +83,8 @@ class Problem:
         dynamics_prop: Optional[dict] = None,
         states_prop: Optional[List[State]] = None,
         algebraic_prop: Optional[dict] = None,
-        licq_min: float = 0.0,
-        licq_max: float = 1e-4,
+        licq_min: Union[float, Dict[int, float]] = 0.0,
+        licq_max: Union[float, Dict[int, float]] = 1e-4,
         algorithm: Optional[Union[Algorithm, dict]] = None,
         discretizer: Optional[Union[Discretizer, dict]] = None,
         solver: Optional[Union[ConvexSolver, dict]] = None,
@@ -114,8 +114,12 @@ class Problem:
                 Only specify additional states beyond optimization states. Used with dynamics_prop.
             algebraic_prop (dict, optional): Dictionary mapping names to symbolic expressions
                 for outputs evaluated (not integrated) during propagation.
-            licq_min (float): Minimum LICQ constraint value. Defaults to 0.0.
-            licq_max (float): Maximum LICQ constraint value. Defaults to 1e-4.
+            licq_min: Minimum LICQ constraint value. Defaults to 0.0.
+                Either a scalar (applied to all CTCS groups) or a dict
+                mapping CTCS group ``idx`` to per-group bounds.
+            licq_max: Maximum LICQ constraint value. Defaults to 1e-4.
+                Either a scalar (applied to all CTCS groups) or a dict
+                mapping CTCS group ``idx`` to per-group bounds.
             algorithm: SCP algorithm configuration. Accepts:
 
                 - ``None`` — uses ``PenalizedTrustRegion()`` with defaults.
@@ -272,11 +276,15 @@ class Problem:
         self._byof = byof
 
         # Resolve algorithm: None → default PTR, dict → PTR(**dict), instance → use directly
-        # Pass symbolic states so dict-valued lam_cost can be expanded eagerly.
+        # Pass symbolic states/controls so dict-valued weights can be expanded eagerly.
         if algorithm is None:
-            self._algorithm = PenalizedTrustRegion(states=self.symbolic.states)
+            self._algorithm = PenalizedTrustRegion(
+                states=self.symbolic.states, controls=self.symbolic.controls
+            )
         elif isinstance(algorithm, dict):
-            self._algorithm = _resolve_algorithm(algorithm, states=self.symbolic.states)
+            self._algorithm = _resolve_algorithm(
+                algorithm, states=self.symbolic.states, controls=self.symbolic.controls
+            )
         else:
             if not isinstance(algorithm, Algorithm):
                 raise TypeError(

{openscvx-0.4.1.dev145 → openscvx-0.4.1.dev147}/openscvx/solvers/ptr_solver.py

@@ -375,9 +375,10 @@ class PTRSolver(ConvexSolver):
             if settings.sim.x.final_type[i] == "Maximize":
                 cost -= lam_cost[i] * x[-1][i]
 
-        # Trust Region Cost
+        # Trust Region Cost (per-variable weighting)
         cost += sum(
-            lam_prox * cp.sum_squares(cp.hstack((dx[i], du[i]))) for i in range(settings.sim.n)
+            cp.sum(cp.multiply(lam_prox[i], cp.square(cp.hstack((dx[i], du[i])))))
+            for i in range(settings.sim.n)
         )
 
         # Virtual Control Slack
@@ -724,7 +725,7 @@ class PTRSolver(ConvexSolver):
 
     def update_penalties(
         self,
-        lam_prox: float,
+        lam_prox: np.ndarray,
         lam_cost: Union[float, np.ndarray],
         lam_vc: np.ndarray,
         lam_vb_nodal: np.ndarray,
@@ -736,7 +737,7 @@ class PTRSolver(ConvexSolver):
         PTR convex subproblem.
 
         Args:
-            lam_prox: Trust region weight (penalizes deviation from linearization point)
+            lam_prox: Trust region weights, shape ``(N, n_states + n_controls)``.
             lam_cost: Cost function weight. Scalar or array of shape
                 ``(n_states,)`` for per-state weighting.
             lam_vc: Virtual control penalty weights, shape (N-1, n_states)