openscvx 0.4.1.dev144__tar.gz → 0.4.1.dev146__tar.gz

{openscvx-0.4.1.dev144/openscvx.egg-info → openscvx-0.4.1.dev146}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.4.1.dev144
+Version: 0.4.1.dev146
 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.dev144 → openscvx-0.4.1.dev146}/openscvx/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.4.1.dev144'
-__version_tuple__ = version_tuple = (0, 4, 1, 'dev144')
+__version__ = version = '0.4.1.dev146'
+__version_tuple__ = version_tuple = (0, 4, 1, 'dev146')
 
-__commit_id__ = commit_id = 'g137ef78b8'
+__commit_id__ = commit_id = 'g66b4191a1'

{openscvx-0.4.1.dev144 → openscvx-0.4.1.dev146}/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.dev144 → openscvx-0.4.1.dev146}/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.dev144 → openscvx-0.4.1.dev146}/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
 
 
@@ -101,6 +102,207 @@ def _expand_lam_cost_dict(
     return lam_arr
 
 
+def _expand_lam_vc_dict(
+    lam_vc_dict: Dict[str, Union[float, list, np.ndarray]],
+    states: List["State"],
+) -> np.ndarray:
+    """Expand a ``{state_name: weight}`` dict to a per-state VC weight array.
+
+    Maps user-provided per-state virtual control weights to a dense array
+    using each state's ``_slice``.  States not present in the dict default
+    to ``1.0`` (neutral weight — dynamics violations are still penalised at
+    the base rate).
+
+    Per-state 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 state supplies a 2-D value the output is 2-D with first
+      dimension *K* (all 2-D entries must agree on *K*).
+
+    Args:
+        lam_vc_dict: Mapping from state names to VC weights.
+        states: List of State objects (must already have ``_slice`` assigned).
+
+    Returns:
+        np.ndarray of shape ``(n_states,)`` when all entries are scalar/1-D,
+        or ``(K, n_states)`` when any entry is 2-D.
+
+    Raises:
+        ValueError: If the dict contains unknown state 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)
+
+    valid_names = {s.name for s in states}
+    unknown = set(lam_vc_dict.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_dict:
+            continue
+        val = np.asarray(lam_vc_dict[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_dict:
+            continue
+        val = np.asarray(lam_vc_dict[state.name], dtype=float)
+        n_components = state.shape[0] if len(state.shape) > 0 else 1
+
+        if val.ndim == 0:
+            # Scalar — broadcast to all components (and all nodes if 2-D).
+            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, expected {n_components}"
+                )
+            lam_arr[:, state._slice] = val
+        else:
+            raise ValueError(
+                f"lam_vc['{state.name}'] has {val.ndim} dimensions, expected scalar, 1-D, or 2-D"
+            )
+
+    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.
@@ -120,8 +322,13 @@ class Weights:
     drift.
 
     Attributes:
-        lam_prox: Trust region (proximal) weight (normalized).
-        lam_vc: Virtual control penalty 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.
         lam_cost: Cost weight per state (normalized). Scalar or array of
             shape ``(n_states,)`` for per-state weighting.
         lam_vb: Global virtual buffer penalty weight (normalized). Scalar
@@ -135,8 +342,8 @@ class Weights:
             :meth:`set_vb_arrays`.
     """
 
-    lam_prox: float = 1e0
-    lam_vc: float = 1e1
+    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
     lam_vb_nodal: Optional[np.ndarray] = None
@@ -144,12 +351,20 @@ 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_vc = self.lam_vc
+        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
+        )
         self._raw_lam_cost = (
             self.lam_cost.copy() if isinstance(self.lam_cost, np.ndarray) else self.lam_cost
         )
@@ -164,6 +379,16 @@ 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)
+            else self._raw_lam_vc
+        )
         raw_cost_max = (
             float(np.max(self._raw_lam_cost))
             if isinstance(self._raw_lam_cost, np.ndarray)
@@ -176,7 +401,7 @@ class Weights:
             )
         else:
             raw_vb_max = self._raw_lam_vb
-        scale = max(self._raw_lam_prox, self._raw_lam_vc, 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
@@ -509,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)
 
@@ -715,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().")
@@ -787,16 +1012,22 @@ class AlgorithmState:
 
         Args:
             settings: Configuration object containing initial guesses and SCP parameters
-            weights: Normalized initial weights from the algorithm. The scalar
-                ``lam_vc`` is expanded to an ``(N-1, n_states)`` array here.
+            weights: Normalized initial weights from the algorithm.
+                ``lam_vc`` is expanded to an ``(N-1, n_states)`` array here
+                (scalar or per-state values are broadcast).
 
         Returns:
             Fresh AlgorithmState initialized from settings with copied arrays
         """
         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()
@@ -825,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],
         )
 
 
@@ -883,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]],
@@ -915,6 +1181,39 @@ class Algorithm(ABC):
             return _expand_lam_cost_dict(lam_cost, states)
         return lam_cost
 
+    @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, *states* must be provided so the
+        dict can be expanded to a per-state array via
+        :func:`_expand_lam_vc_dict`.
+
+        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``.
+        """
+        if isinstance(lam_vc, dict):
+            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."
+                )
+            return _expand_lam_vc_dict(lam_vc, states)
+        return lam_vc
+
     def _resolve_lam_vb(
         self,
         N: int,

{openscvx-0.4.1.dev144 → openscvx-0.4.1.dev146}/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"}
     )