class MoreauPTRSolver(PTRSolver):
"""JAX-native conic backend for the PTR convex subproblem.
Assembles each SCP subproblem as a sparse conic program and dispatches to
``moreau.jax.Solver``. See the module docstring for the SOC epigraph
encoding, warm-start policy, and float64 advice.
Compared to :class:`QPAXPTRSolver`, this backend uses fewer decision
variables and constraint rows for the same PTR physics (SOC epigraphs for
``|nu|`` instead of two nonneg slack rows each), warm-starts between SCP
iterations on successful solves, and natively handles user
:class:`~openscvx.solvers.cones.SOCConstraint` convex constraints.
The JAX-pure entry point :meth:`iteration_callback` calls the jitted
kernel on the :class:`moreau.jax.Solver` built in :meth:`initialize`,
as ``(P_data, A_data, q, b) -> (JaxSolution, JaxSolveInfo)``, so the
backend composes with ``jax.jit`` and ``jax.vmap``. Moreau 0.4 removed
the public ``moreau.jax.solver`` factory. ``Solver.solve`` returns only
the solution and stores status on ``solver.info``, which is not batched
under ``vmap``. The kernel has no warm-start argument, so every call is
a cold start: both
:meth:`~openscvx.problem.Problem.solve` and
:meth:`~openscvx.problem.Problem.solve_jax` drive
``iteration_callback`` for the SCP body, so neither thread Moreau's
``_warm_start`` carry across SCP iterations. The OO warm-start path
still mutates ``_warm_start`` on each :meth:`solve` call (the legacy
NumPy entry point, no longer on the SCP hot path); restoring
warm-start to the SCP loop requires routing ``(x, z, s)`` through
:class:`AlgorithmState` as a real pytree carry — see the
Future Extensions in ``plans/jax-pure-solve.md``.
Scope:
Supported — state/control box, dynamics linearization (continuous
and impulsive), boundary Fix, uniform time grid, linearized nodal
nonconvex, CTCS LICQ rows, user ``.convex()`` affine equality,
affine inequality, and second-order-cone constraints.
Not supported — cross-node constraints (raises
:class:`NotImplementedError` with a "use CVXPyPTRSolver" pointer).
Differentiability hook for future work:
``moreau.jax.Solver`` differentiates through the solve via implicit
differentiation on the KKT conditions. Once the surrounding SCP
pipeline stays in ``jit``, ``jax.grad`` / ``jax.vmap`` can reach
through a full SCvx solve.
Class attributes:
SUPPORTED_CONE_TYPES: Accepts
:class:`~openscvx.solvers.cones.ZeroConeConstraint`,
:class:`~openscvx.solvers.cones.NonnegConeConstraint`, and
:class:`~openscvx.solvers.cones.SOCConstraint`.
Args:
max_iter: Maximum number of IPM iterations forwarded to
:class:`moreau.Settings`. Defaults to ``200``.
verbose: Whether Moreau prints per-iteration diagnostics.
Forwarded to :class:`moreau.Settings`. Defaults to ``False``.
device: Compute device for Moreau's JAX kernels. One of
``"auto"``, ``"cpu"``, or ``"cuda"``. Forwarded to
:class:`moreau.Settings`. Defaults to ``"auto"``.
tol_gap_abs: Absolute duality-gap tolerance forwarded to
:class:`moreau.IPMSettings`. ``None`` uses Moreau's default.
tol_feas: Primal/dual feasibility tolerance forwarded to
:class:`moreau.IPMSettings`. ``None`` uses Moreau's default.
solver_args: Additional keyword arguments forwarded verbatim to
:class:`moreau.Settings`. Use for settings not covered by the
named params above. ``solver_args["ipm_settings"]`` may be a
dict or a :class:`moreau.IPMSettings` object; if it is a dict,
``tol_gap_abs`` / ``tol_feas`` are merged into it. Raises
``ValueError`` at construction time if any top-level key or
IPM tolerance overlaps with a named param.
Attributes:
layout: :class:`_ConicLayout` describing flat decision-vector slot
ranges. Populated by :meth:`create_variables`.
"""
SUPPORTED_CONE_TYPES = frozenset({ZeroConeConstraint, NonnegConeConstraint, SOCConstraint})
def __init__(
self,
*,
max_iter: Optional[int] = None,
verbose: Optional[bool] = None,
device: Optional[str] = None,
tol_gap_abs: Optional[float] = None,
tol_feas: Optional[float] = None,
solver_args: Optional[Dict] = None,
):
if not _MOREAU_AVAILABLE:
raise ImportError(
"MoreauPTRSolver requires the `moreau` package. "
"Install it with: pip install openscvx[moreau]"
)
# None sentinels at the signature surface let us distinguish
# "user explicitly passed this" from "default applied" — only the
# former should collide with the same key inside solver_args.
_named = {
k: v
for k, v in (
("max_iter", max_iter),
("verbose", verbose),
("device", device),
)
if v is not None
}
_extra = dict(solver_args) if solver_args else {}
_overlap = _named.keys() & _extra.keys()
if _overlap:
raise ValueError(
f"Moreau settings {sorted(_overlap)} appear as both named arguments "
"and inside solver_args; use one or the other."
)
merged = {**_named, **_extra}
merged.setdefault("max_iter", 200)
merged.setdefault("verbose", False)
merged.setdefault("device", "auto")
_ipm = {
k: v for k, v in [("tol_gap_abs", tol_gap_abs), ("tol_feas", tol_feas)] if v is not None
}
if _ipm:
existing_ipm = merged.get("ipm_settings", {})
if not isinstance(existing_ipm, dict):
raise ValueError(
"Cannot combine tol_gap_abs / tol_feas named arguments with an "
"ipm_settings object in solver_args; use one form or the other."
)
ipm_overlap = _ipm.keys() & existing_ipm.keys()
if ipm_overlap:
raise ValueError(
f"Moreau IPM settings {sorted(ipm_overlap)} appear as both named "
"arguments and inside solver_args['ipm_settings']; use one or the other."
)
merged["ipm_settings"] = {**_ipm, **existing_ipm}
self.solver_args = merged
self.layout: Optional[_ConicLayout] = None
self._S_x: Optional[np.ndarray] = None
self._c_x: Optional[np.ndarray] = None
self._S_u: Optional[np.ndarray] = None
self._c_u: Optional[np.ndarray] = None
self._S_x_diag: Optional[np.ndarray] = None
self._S_u_diag: Optional[np.ndarray] = None
self._inv_S_x_diag: Optional[np.ndarray] = None
self._inv_S_u_diag: Optional[np.ndarray] = None
self._settings: Optional["Config"] = None
self._jax_constraints: Optional["LoweredJaxConstraints"] = None
# moreau.jax.Solver constructed at initialize().
self._moreau: Optional[_MoreauJaxSolver] = None
# Fixed CSR structure stored after initialize().
self._P_diag_slots: Optional[np.ndarray] = None
self._P_n_dx: int = 0 # N*n_x, for indexing into P_data
self._coo_rows: Optional[np.ndarray] = None
self._coo_cols: Optional[np.ndarray] = None
self._n_con: int = 0
# Populated by PTRSolver.lower_convex_constraints.
self._impulsive_pins: List[Tuple[List[int], slice]] = []
self._user_cone_constraints: List = []
self._parameters_dict: dict = {}
# Per-iteration data, set by update_* methods.
self._dyn: dict = {}
self._cons: dict = {}
self._pen: dict = {}
self._x_init: Optional[np.ndarray] = None
self._x_term: Optional[np.ndarray] = None
# Warm-start carry; reset to None when problem structure changes.
self._warm_start = None
# Last solve status from moreau (populated by solve()).
self._last_status: moreau.SolverStatus = moreau.SolverStatus.Unsolved
# ------------------------------------------------------------------
# Setup
# ------------------------------------------------------------------
def create_variables(
self,
N: int,
x_unified: "UnifiedState",
u_unified: "UnifiedControl",
jax_constraints: "LoweredJaxConstraints",
dynamics_sparsity: Optional[tuple] = None,
constraint_sparsity: Optional[list] = None,
) -> None:
"""Compute scaling matrices and the static conic decision-vector layout.
Sparsity hints are accepted for interface symmetry with
:class:`openscvx.solvers.cvxpy_ptr_solver.CVXPyPTRSolver` but ignored;
numeric values are filled into a fixed CSR structure at solve time.
Args:
N: Number of discretization nodes.
x_unified: Unified state bounds and scaling metadata.
u_unified: Unified control bounds and scaling metadata.
jax_constraints: Lowered JAX constraints (nodal structure only).
dynamics_sparsity: Ignored.
constraint_sparsity: Ignored.
"""
del dynamics_sparsity, constraint_sparsity
n_x = len(x_unified.max)
n_u = len(u_unified.max)
S_x, c_x = self._scaling(x_unified)
S_u, c_u = self._scaling(u_unified)
self._S_x = S_x
self._c_x = c_x
self._S_u = S_u
self._c_u = c_u
self._S_x_diag = np.diag(S_x)
self._S_u_diag = np.diag(S_u)
self._inv_S_x_diag = 1.0 / self._S_x_diag
self._inv_S_u_diag = 1.0 / self._S_u_diag
# JAX-side scaling diagonals for ``iteration_callback``. The dtype
# tracks the global ``jax_enable_x64`` setting so the assembled
# ``(P_data, A_data, q, b)`` lands in the dtype Moreau's functional
# solver expects.
f = jnp.float64 if jax.config.read("jax_enable_x64") else jnp.float32
self._S_x_diag_j = jnp.asarray(self._S_x_diag, dtype=f)
self._S_u_diag_j = jnp.asarray(self._S_u_diag, dtype=f)
self._inv_S_x_diag_j = jnp.asarray(self._inv_S_x_diag, dtype=f)
self._inv_S_u_diag_j = jnp.asarray(self._inv_S_u_diag, dtype=f)
self._c_x_j = jnp.asarray(self._c_x, dtype=f)
self._c_u_j = jnp.asarray(self._c_u, dtype=f)
self._jax_dtype = f
if any(c.is_equality for c in jax_constraints.nodal):
raise NotImplementedError(
"MoreauPTRSolver does not support L1-penalized equality "
"constraints. Use CVXPyPTRSolver."
)
self.layout = _ConicLayout(N=N, n_x=n_x, n_u=n_u, n_nodal=len(jax_constraints.nodal))
self._jax_constraints = jax_constraints
self._x_unified = x_unified
def initialize(self, lowered: "LoweredProblem", settings: "Config") -> None:
"""Build the static conic structure and construct ``moreau.jax.Solver``.
Validates the supported constraint subset, enumerates all A-matrix
nonzero positions (fixed CSR structure), builds a :class:`moreau.Cones`
spec, and constructs :class:`moreau.jax.Solver`. Numeric values in
``A``, ``b``, ``P``, and ``q`` are filled at each :meth:`solve` call.
Args:
lowered: Lowered problem with constraint structure.
settings: Problem configuration.
Raises:
RuntimeError: If :meth:`create_variables` was not called first.
NotImplementedError: If cross-node constraints are present.
"""
if self.layout is None:
raise RuntimeError(
"MoreauPTRSolver.initialize() called before create_variables(). "
"Call create_variables() first."
)
if lowered.jax_constraints.cross_node:
raise NotImplementedError(
"MoreauPTRSolver does not yet support cross-node constraints "
f"({len(lowered.jax_constraints.cross_node)} defined). "
"Use CVXPyPTRSolver."
)
self._settings = settings
# Reset warm-start on new initialization (problem structure may change).
self._warm_start = None
# Enumerate all (row, col) pairs in the A matrix (structural pass).
coo_rows, coo_cols, n_eq, n_nn, soc_dims = self._structural_pass(settings)
n_soc = sum(soc_dims)
self._n_con = n_eq + n_nn + n_soc
self._coo_rows = np.asarray(coo_rows, dtype=np.int32)
self._coo_cols = np.asarray(coo_cols, dtype=np.int32)
# Build CSR for A from the structural COO.
n_coo = len(coo_rows)
if n_coo > 0:
A_struct = sp.csr_matrix(
(np.ones(n_coo), (coo_rows, coo_cols)),
shape=(self._n_con, self.layout.n_z),
)
A_struct.sort_indices()
else:
A_struct = sp.csr_matrix((self._n_con, self.layout.n_z))
A_indptr = np.asarray(A_struct.indptr, dtype=np.int32)
A_indices = np.asarray(A_struct.indices, dtype=np.int32)
self._A_indptr = A_indptr
self._A_indices = A_indices
# Precompute the COO→CSR permutation so the JAX assembly can emit
# ``coo_vals`` in COO traversal order (matching ``_structural_pass``)
# and then ``A_data = coo_vals[csr_to_coo_perm]`` matches what
# ``solve_qp_primal``/Moreau's ``solve_fn`` expects under the fixed
# ``A_col_indices``. NumPy ``solve()`` keeps using scipy's sort path;
# this permutation is consumed only by ``_assemble_conic_jax``.
if n_coo > 0:
A_pos = sp.csr_matrix(
(np.arange(n_coo, dtype=np.int32), (coo_rows, coo_cols)),
shape=(self._n_con, self.layout.n_z),
)
A_pos.sort_indices()
self._csr_to_coo_perm = np.asarray(A_pos.data, dtype=np.int32)
else:
self._csr_to_coo_perm = np.zeros((0,), dtype=np.int32)
# Build CSR for P (diagonal on dx/du slots only).
L = self.layout
N, n_x = L.N, L.n_x
dx_cols = np.arange(L.sl_dx.start, L.sl_dx.stop, dtype=np.int32)
du_cols = np.arange(L.sl_du.start, L.sl_du.stop, dtype=np.int32)
self._P_diag_slots = np.concatenate([dx_cols, du_cols])
self._P_n_dx = N * n_x # boundary between dx and du in P_data
# P is diagonal: for each slot col, the single nonzero is P[col, col].
P_indptr = np.zeros(L.n_z + 1, dtype=np.int32)
for col in self._P_diag_slots:
P_indptr[col + 1] = 1
P_indptr = np.cumsum(P_indptr)
P_indices = self._P_diag_slots # diagonal → col index == row index
# Construct moreau.Cones spec.
cones = moreau.Cones(
num_zero_cones=n_eq,
num_nonneg_cones=n_nn,
so_cone_dims=soc_dims,
)
# Build moreau.Settings from solver_args.
moreau_settings = _build_moreau_settings(self.solver_args)
self._moreau = _MoreauJaxSolver(
n=L.n_z,
m=self._n_con,
P_row_offsets=jnp.array(P_indptr),
P_col_indices=jnp.array(P_indices),
A_row_offsets=jnp.array(A_indptr),
A_col_indices=jnp.array(A_indices),
cones=cones,
settings=moreau_settings,
)
# ------------------------------------------------------------------
# Per-iteration update hooks
# ------------------------------------------------------------------
def update_dynamics_linearization(
self,
x_bar: np.ndarray,
u_bar: np.ndarray,
A_d: np.ndarray,
B_d: np.ndarray,
C_d: np.ndarray,
x_prop: np.ndarray,
x_prop_plus: Optional[np.ndarray] = None,
D_d: Optional[np.ndarray] = None,
E_d: Optional[np.ndarray] = None,
) -> None:
A_eff = np.asarray(A_d, dtype=float)
B_eff = np.asarray(B_d, dtype=float)
C_eff = np.asarray(C_d, dtype=float)
# Absorb the impulsive state Jacobian into the continuous step
# matrices so the assembly row keeps a single A_d·dx + B_d·du term,
# matching the recipe in CVXPyPTRSolver.update_dynamics_linearization
# at openscvx/solvers/cvxpy_ptr_solver.py:636-654.
if D_d is not None:
D_arr = np.asarray(D_d, dtype=float)
if D_arr.ndim == 3 and D_arr.shape[0] == A_eff.shape[0] + 1:
D_steps = D_arr[1:]
elif D_arr.ndim == 3 and D_arr.shape[0] == A_eff.shape[0]:
D_steps = D_arr
else:
raise ValueError(
"Unexpected D_d shape for dynamics update: "
f"{D_arr.shape}, expected "
f"{(A_eff.shape[0] + 1, A_eff.shape[1], A_eff.shape[2])} "
f"or {(A_eff.shape[0], A_eff.shape[1], A_eff.shape[2])}."
)
A_eff = np.einsum("kij,kjl->kil", D_steps, A_eff)
B_eff = np.einsum("kij,kjl->kil", D_steps, B_eff)
C_eff = np.einsum("kij,kjl->kil", D_steps, C_eff)
self._dyn = {
"x_bar": np.asarray(x_bar, dtype=float),
"u_bar": np.asarray(u_bar, dtype=float),
"A_d": A_eff,
"B_d": B_eff,
"C_d": C_eff,
"x_prop": np.asarray(x_prop, dtype=float),
"x_prop_plus": (
np.asarray(x_prop_plus, dtype=float) if x_prop_plus is not None else None
),
"E_d": np.asarray(E_d, dtype=float) if E_d is not None else None,
}
def update_constraint_linearizations(
self,
nodal: Optional[List[dict]] = None,
cross_node: Optional[List[dict]] = None,
) -> None:
if cross_node:
raise NotImplementedError(
"MoreauPTRSolver received cross-node linearization data; "
"cross-node constraints are not supported. Use CVXPyPTRSolver."
)
self._cons = {
"nodal": [
{
"g": np.asarray(d["g"], dtype=float),
"grad_g_x": np.asarray(d["grad_g_x"], dtype=float),
"grad_g_u": np.asarray(d["grad_g_u"], dtype=float),
}
for d in (nodal or [])
],
}
def update_penalties(
self,
lam_prox: np.ndarray,
lam_cost: Union[float, np.ndarray],
lam_vc: np.ndarray,
lam_vb_nodal: np.ndarray,
lam_vb_cross: np.ndarray,
) -> None:
del lam_vb_cross # cross-node rejected at initialize()
self._pen = {
"lam_prox": np.asarray(lam_prox, dtype=float),
"lam_cost": np.asarray(lam_cost, dtype=float),
"lam_vc": np.asarray(lam_vc, dtype=float),
"lam_vb_nodal": np.asarray(lam_vb_nodal, dtype=float),
}
def update_boundary_conditions(
self,
x_init: Optional[np.ndarray] = None,
x_term: Optional[np.ndarray] = None,
) -> None:
if x_init is not None:
self._x_init = np.asarray(x_init, dtype=float)
if x_term is not None:
self._x_term = np.asarray(x_term, dtype=float)
# ------------------------------------------------------------------
# Structural pass (called once at initialize)
# ------------------------------------------------------------------
def _structural_pass(
self,
settings: "Config",
) -> Tuple[List[int], List[int], int, int, List[int]]:
"""Enumerate all A-matrix ``(row, col)`` nonzero positions.
Iterates through every constraint row in the exact order used by
:meth:`_assemble_conic`, recording which columns are structurally
nonzero. Only the positions (not the values) matter here; values come
from per-iteration data at solve time.
Args:
settings: Problem configuration (bounds, CTCS groups, time grid).
Returns:
tuple: ``(coo_rows, coo_cols, n_eq, n_nn, soc_dims)`` where
``coo_rows`` / ``coo_cols`` list every A-matrix nonzero in
traversal order, ``n_eq`` / ``n_nn`` count zero- and nonneg-cone
rows, and ``soc_dims`` lists SOC dimensions (2 per ``|nu[k,j]|``).
"""
L = self.layout
N, n_x, n_u = L.N, L.n_x, L.n_u
jax_constraints = self._jax_constraints
coo_rows: List[int] = []
coo_cols: List[int] = []
row = 0
def add(r: int, c: int) -> None:
coo_rows.append(r)
coo_cols.append(c)
# ================================================================
# ZERO CONE (equality, s = 0 → Az = b)
# ================================================================
# State error definitions: x[k,j] − dx[k,j] = rhs (N·n_x rows)
for k in range(N):
for j in range(n_x):
add(row, L.x_idx(k, j))
add(row, L.dx_idx(k, j))
row += 1
# Control error definitions: u[k,j] − du[k,j] = rhs (N·n_u rows)
for k in range(N):
for j in range(n_u):
add(row, L.u_idx(k, j))
add(row, L.du_idx(k, j))
row += 1
# Dynamics (FOH, continuous PTR) ((N-1)·n_x rows)
for k in range(1, N):
kp = k - 1
for i in range(n_x):
add(row, L.x_idx(k, i))
for j in range(n_x):
add(row, L.dx_idx(kp, j))
for j in range(n_u):
add(row, L.du_idx(kp, j))
add(row, L.du_idx(k, j))
add(row, L.nu_idx(kp, i))
row += 1
# Linearized nodal constraints: grad·dx + grad·du − nu_vb = −g
for c_idx, constraint in enumerate(jax_constraints.nodal):
for node in constraint.nodes:
for j in range(n_x):
add(row, L.dx_idx(node, j))
for j in range(n_u):
add(row, L.du_idx(node, j))
add(row, L.nu_vb_idx(c_idx, node))
row += 1
# Fix boundary conditions (initial / terminal). Under impulsive
# control the initial Fix row couples the post-impulse state at
# node 0 to du[0, slice_imp]; declare those columns structurally
# nonzero here so the static CSR pattern accommodates the
# numerical values emitted in _assemble_conic.
slice_imp = settings.sim.u.slice_impulsive
has_impulsive = slice_imp.stop > slice_imp.start
for i in range(settings.sim.true_state_slice.start, settings.sim.true_state_slice.stop):
if settings.sim.x.initial_type[i] == "Fix":
add(row, L.x_idx(0, i))
if has_impulsive:
for j in range(slice_imp.start, slice_imp.stop):
add(row, L.du_idx(0, j))
row += 1
if settings.sim.x.final_type[i] == "Fix":
add(row, L.x_idx(N - 1, i))
row += 1
# Impulsive zero-pin equalities: u[node, j] = const, one row per
# (node, j) absorbed by lower_convex_constraints.
for nodes, ctrl_slice in self._impulsive_pins:
for node in nodes:
for j in range(ctrl_slice.start, ctrl_slice.stop):
add(row, L.u_idx(node, j))
row += 1
# User ZeroConeConstraint rows (hard equality, one row per scalar output).
for cc in self._user_cone_constraints:
if isinstance(cc, ZeroConeConstraint):
for k in cc.nodes:
for _ in range(cc.m):
for j in range(n_x):
add(row, L.dx_idx(k, j))
for j in range(n_u):
add(row, L.du_idx(k, j))
row += 1
# Uniform time grid: u[k,j] − u[k-1,j] = 0
if settings.sim._uniform_time_grid:
td = settings.sim.time_dilation_slice
for k in range(1, N):
for j in range(td.start, td.stop):
add(row, L.u_idx(k, j))
add(row, L.u_idx(k - 1, j))
row += 1
# CTCS initial: S_x[idx]·x[0,idx] = −c_x[idx] (one per CTCS group)
for idx in range(settings.sim.ctcs_slice.start, settings.sim.ctcs_slice.stop):
add(row, L.x_idx(0, idx))
row += 1
n_eq = row
# ================================================================
# NONNEG CONE (s = b − Az ≥ 0)
# ================================================================
# State box upper: x[k,j] ≤ x_max_scaled[j]
for k in range(N):
for j in range(n_x):
add(row, L.x_idx(k, j))
row += 1
# State box lower: x[k,j] ≥ x_min_scaled[j]
for k in range(N):
for j in range(n_x):
add(row, L.x_idx(k, j))
row += 1
# Control box upper
for k in range(N):
for j in range(n_u):
add(row, L.u_idx(k, j))
row += 1
# Control box lower
for k in range(N):
for j in range(n_u):
add(row, L.u_idx(k, j))
row += 1
# CTCS difference rows (two per node interval per group)
for idx, nodes in zip(
range(settings.sim.ctcs_slice.start, settings.sim.ctcs_slice.stop),
settings.sim.ctcs_node_intervals,
):
start_i = 1 if nodes[0] == 0 else nodes[0]
for i in range(start_i, nodes[1]):
# +direction
add(row, L.x_idx(i, idx))
add(row, L.x_idx(i - 1, idx))
row += 1
# −direction
add(row, L.x_idx(i, idx))
add(row, L.x_idx(i - 1, idx))
row += 1
# Pos epigraph: t_vb[c,k] ≥ nu_vb[c,k]
for c_idx in range(L.n_nodal):
for k in range(N):
add(row, L.nu_vb_idx(c_idx, k))
add(row, L.t_vb_idx(c_idx, k))
row += 1
# Pos epigraph non-negativity: t_vb[c,k] ≥ 0
for c_idx in range(L.n_nodal):
for k in range(N):
add(row, L.t_vb_idx(c_idx, k))
row += 1
# User NonnegConeConstraint rows.
for cc in self._user_cone_constraints:
if isinstance(cc, NonnegConeConstraint):
for k in cc.nodes:
for _ in range(cc.m):
for j in range(n_x):
add(row, L.dx_idx(k, j))
for j in range(n_u):
add(row, L.du_idx(k, j))
row += 1
n_nn = row - n_eq
# ================================================================
# SOC rows: for each scalar |nu[k,j]| ≤ t_vc[k,j], SOC(2)
# ================================================================
soc_dims: List[int] = []
for k in range(N - 1):
for j in range(n_x):
# s[0] = t_vc (row 0 of this SOC block)
add(row, L.t_vc_idx(k, j))
row += 1
# s[1] = nu (row 1 of this SOC block)
add(row, L.nu_idx(k, j))
row += 1
soc_dims.append(2)
# User SOCConstraint rows: each node contributes SOC(m_arg + 1).
# Row layout: [bound, arg[0], arg[1], ..., arg[m_arg-1]].
# Moreau convention s = b - Az ∈ K_soc: ‖s[1:]‖ ≤ s[0].
# Since arg and bound are affine in (x, u), columns are dx[k, :] and du[k, :].
for cc in self._user_cone_constraints:
if isinstance(cc, SOCConstraint):
soc_dim = cc.m_arg + 1
for k in cc.nodes:
for _ in range(soc_dim):
for j in range(n_x):
add(row, L.dx_idx(k, j))
for j in range(n_u):
add(row, L.du_idx(k, j))
row += 1
soc_dims.append(soc_dim)
return coo_rows, coo_cols, n_eq, n_nn, soc_dims
# ------------------------------------------------------------------
# Assembly (called each SCP iteration)
# ------------------------------------------------------------------
def _assemble_conic(
self,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Fill ``P_data``, COO A-values, ``q``, and ``b`` from per-iteration data.
Iterates through every row in the same order as :meth:`_structural_pass`
so the emitted value list matches ``self._coo_rows`` /
``self._coo_cols``. scipy builds the same sorted CSR each call, so
``A_csr.data`` aligns with the fixed ``A_indices`` from
:meth:`initialize`.
Returns:
tuple: ``(P_data, coo_vals, q, b)`` as 1-D NumPy float arrays.
Raises:
RuntimeError: If per-iteration update hooks were not called.
"""
if not (self._dyn and self._cons and self._pen):
raise RuntimeError(
"MoreauPTRSolver.solve() requires update_dynamics_linearization, "
"update_constraint_linearizations, and update_penalties to all "
"have been called this iteration."
)
L = self.layout
settings = self._settings
N, n_x, n_u = L.N, L.n_x, L.n_u
inv_S_x = self._inv_S_x_diag
inv_S_u = self._inv_S_u_diag
S_x = self._S_x_diag
S_u = self._S_u_diag
c_x = self._c_x
c_u = self._c_u
slice_imp = settings.sim.u.slice_impulsive
has_impulsive = slice_imp.stop > slice_imp.start
lam_prox = self._pen["lam_prox"] # (N, n_x + n_u)
lam_cost = self._pen["lam_cost"] # scalar or (n_x,)
lam_vc = self._pen["lam_vc"] # (N-1, n_x)
lam_vb_nodal = self._pen["lam_vb_nodal"] # (N, n_nodal or 1)
x_bar = self._dyn["x_bar"] # (N, n_x)
u_bar = self._dyn["u_bar"] # (N, n_u)
A_d = self._dyn["A_d"] # (N-1, n_x, n_x)
B_d = self._dyn["B_d"] # (N-1, n_x, n_u)
C_d = self._dyn["C_d"] # (N-1, n_x, n_u)
x_prop = self._dyn["x_prop"] # (N-1, n_x)
E_d_arr = self._dyn["E_d"] # (N, n_x, n_u) or None
x_prop_plus_arr = self._dyn["x_prop_plus"] # (N, n_x) or None
lam_cost_arr = np.broadcast_to(lam_cost, (settings.sim.n_states,))
# ---- P_data (diagonal trust-region weights on dx/du) ----
# P_diag_slots = [dx_0..dx_{N·n_x-1}, du_0..du_{N·n_u-1}]
P_data = np.empty(len(self._P_diag_slots))
n_dx = self._P_n_dx # N * n_x
for k in range(N):
for j in range(n_x):
P_data[k * n_x + j] = 2.0 * lam_prox[k, j] + _P_DIAG_EPS
for k in range(N):
for j in range(n_u):
P_data[n_dx + k * n_u + j] = 2.0 * lam_prox[k, n_x + j] + _P_DIAG_EPS
# ---- q (linear cost vector) ----
q = np.zeros(L.n_z)
for i in range(settings.sim.true_state_slice.start, settings.sim.true_state_slice.stop):
init_t = settings.sim.x.initial_type[i]
final_t = settings.sim.x.final_type[i]
if init_t == "Minimize":
q[L.x_idx(0, i)] += lam_cost_arr[i]
elif init_t == "Maximize":
q[L.x_idx(0, i)] -= lam_cost_arr[i]
if final_t == "Minimize":
q[L.x_idx(N - 1, i)] += lam_cost_arr[i]
elif final_t == "Maximize":
q[L.x_idx(N - 1, i)] -= lam_cost_arr[i]
for k in range(N - 1):
for j in range(n_x):
q[L.t_vc_idx(k, j)] += lam_vc[k, j]
for c_idx in range(L.n_nodal):
for k in range(N):
q[L.t_vb_idx(c_idx, k)] += lam_vb_nodal[k, c_idx]
# ---- A values and b (in same row order as _structural_pass) ----
# Convention: s = b − Az.
# Zero cone: s = 0 → A has same signs as "Az = b" formulation.
# Nonneg: s ≥ 0 → b − Az ≥ 0.
# SOC: s ∈ K → (b − Az) ∈ K_soc.
coo_vals: List[float] = []
b_list: List[float] = []
def emit(a_coeffs: List[float], rhs: float) -> None:
"""Append one constraint row's A coefficients and RHS."""
coo_vals.extend(a_coeffs)
b_list.append(rhs)
# ================================================================
# ZERO CONE rows
# ================================================================
# State error defs: x[k,j] − dx[k,j] = inv_S_x[j]·(x_bar[k,j] − c_x[j])
for k in range(N):
for j in range(n_x):
emit([1.0, -1.0], inv_S_x[j] * (x_bar[k, j] - c_x[j]))
# Control error defs: u[k,j] − du[k,j] = inv_S_u[j]·(u_bar[k,j] − c_u[j])
for k in range(N):
for j in range(n_u):
emit([1.0, -1.0], inv_S_u[j] * (u_bar[k, j] - c_u[j]))
# Dynamics (continuous FOH, with optional impulsive coupling at node k):
# x[k] − A_blk·dx[k-1] − B_blk·du[k-1] − C_blk·du[k]
# − E_blk·du[k][slice_imp] − nu[k-1]
# = inv_S_x·(x_prop_plus[k] − c_x) if has_impulsive
# = inv_S_x·(x_prop[k-1] − c_x) otherwise
# Mirrors CVXPyPTRSolver.constraints at cvxpy_ptr_solver.py:506-530.
for k in range(1, N):
kp = k - 1
A_blk = (inv_S_x[:, None] * A_d[kp]) * S_x[None, :]
B_blk = (inv_S_x[:, None] * B_d[kp]) * S_u[None, :]
C_blk = (inv_S_x[:, None] * C_d[kp]) * S_u[None, :]
if has_impulsive:
E_blk = (inv_S_x[:, None] * E_d_arr[k]) * S_u[None, :]
rhs_k = inv_S_x * (x_prop_plus_arr[k] - c_x)
else:
E_blk = None
rhs_k = inv_S_x * (x_prop[kp] - c_x)
for i in range(n_x):
# Coefficients in the same col order as _structural_pass added them,
# then sorted by scipy within the row — values line up correctly.
coeffs: List[float] = [1.0] # x[k, i]
for j in range(n_x):
coeffs.append(-A_blk[i, j]) # dx[kp, j]
for j in range(n_u):
coeffs.append(-B_blk[i, j]) # du[kp, j]
c_kj = -C_blk[i, j]
if has_impulsive and slice_imp.start <= j < slice_imp.stop:
c_kj -= E_blk[i, j]
coeffs.append(c_kj) # du[k, j]
coeffs.append(-1.0) # nu[kp, i]
emit(coeffs, rhs_k[i])
# Linearized nodal constraints:
# grad_x·dx[k] + grad_u·du[k] − nu_vb[c,k] = −g[c,k]
jax_constraints = self._jax_constraints
for c_idx, constraint in enumerate(jax_constraints.nodal):
data = self._cons["nodal"][c_idx]
g = data["g"] # (N,)
grad_x = data["grad_g_x"] # (N, n_x)
grad_u = data["grad_g_u"] # (N, n_u)
for node in constraint.nodes:
coeffs = []
for j in range(n_x):
coeffs.append(grad_x[node, j]) # dx[node, j]
for j in range(n_u):
coeffs.append(grad_u[node, j]) # du[node, j]
coeffs.append(-1.0) # nu_vb[c, node]
emit(coeffs, -g[node])
# Fix boundary conditions. Under impulsive control the initial Fix
# row couples x[0, i] to du[0, slice_imp] via the linearized impulse
# Jacobian (CVXPy reference: cvxpy_ptr_solver.py:484-495). Emit
# coefficients in ascending column-index order so they match the
# CSR sort applied to the structural pass.
for i in range(settings.sim.true_state_slice.start, settings.sim.true_state_slice.stop):
if settings.sim.x.initial_type[i] == "Fix":
if has_impulsive:
coeffs = [S_x[i]]
for j in range(slice_imp.start, slice_imp.stop):
coeffs.append(-E_d_arr[0, i, j] * S_u[j])
emit(coeffs, x_prop_plus_arr[0, i] - c_x[i])
else:
if self._x_init is None:
raise RuntimeError(
f"Fix initial condition on state {i} requires x_init; "
"call update_boundary_conditions() before solve()."
)
emit([S_x[i]], self._x_init[i] - c_x[i])
if settings.sim.x.final_type[i] == "Fix":
if self._x_term is None:
raise RuntimeError(
f"Fix final condition on state {i} requires x_term; "
"call update_boundary_conditions() before solve()."
)
emit([S_x[i]], self._x_term[i] - c_x[i])
# Impulsive zero-pin equalities: u[node, j] = −inv_S_u[j]·c_u[j].
# Mirrors CVXPy's lowering of ``u_nonscaled[node][slice_imp] == 0``.
for nodes, ctrl_slice in self._impulsive_pins:
for node in nodes:
for j in range(ctrl_slice.start, ctrl_slice.stop):
emit([1.0], -inv_S_u[j] * c_u[j])
# User ZeroConeConstraint rows.
# Linearize: (J_x·S_x)·dx_scaled + (J_u·S_u)·du_scaled = -f0.
# Jacobians are w.r.t. physical x/u; dx_scaled = Δx_phys / S_x, so
# the coefficient for dx_scaled[k, j] is J_x_phys[i, j] * S_x[j].
params = self._parameters_dict
for cc in self._user_cone_constraints:
if isinstance(cc, ZeroConeConstraint):
for k in cc.nodes:
x_k = x_bar[k]
u_k = u_bar[k]
f0 = np.asarray(cc.jax_fn(x_k, u_k, k, params), dtype=float).reshape(cc.m)
J_x = (
np.asarray(
jax.jacfwd(cc.jax_fn, argnums=0)(x_k, u_k, k, params), dtype=float
).reshape(cc.m, n_x)
* S_x[None, :]
)
J_u = (
np.asarray(
jax.jacfwd(cc.jax_fn, argnums=1)(x_k, u_k, k, params), dtype=float
).reshape(cc.m, n_u)
* S_u[None, :]
)
rhs = -f0
for i in range(cc.m):
coeffs = list(J_x[i]) + list(J_u[i])
emit(coeffs, float(rhs[i]))
# Uniform time grid: u[k,j] − u[k-1,j] = 0
if settings.sim._uniform_time_grid:
td = settings.sim.time_dilation_slice
for k in range(1, N):
for j in range(td.start, td.stop):
emit([1.0, -1.0], 0.0)
# CTCS initial: S_x[idx]·x[0,idx] = −c_x[idx]
for idx in range(settings.sim.ctcs_slice.start, settings.sim.ctcs_slice.stop):
emit([S_x[idx]], -c_x[idx])
# ================================================================
# NONNEG CONE rows (s = b − Az ≥ 0)
# ================================================================
x_max_sc = inv_S_x * (np.asarray(settings.sim.x.max, dtype=float) - c_x)
x_min_sc = inv_S_x * (np.asarray(settings.sim.x.min, dtype=float) - c_x)
u_max_sc = inv_S_u * (np.asarray(settings.sim.u.max, dtype=float) - c_u)
u_min_sc = inv_S_u * (np.asarray(settings.sim.u.min, dtype=float) - c_u)
# State box upper: s = x_max_sc[j] − x[k,j] ≥ 0
for k in range(N):
for j in range(n_x):
emit([1.0], x_max_sc[j])
# State box lower: s = x[k,j] − x_min_sc[j] ≥ 0 → b=−x_min_sc, A=−1
for k in range(N):
for j in range(n_x):
emit([-1.0], -x_min_sc[j])
# Control box upper
for k in range(N):
for j in range(n_u):
emit([1.0], u_max_sc[j])
# Control box lower
for k in range(N):
for j in range(n_u):
emit([-1.0], -u_min_sc[j])
# CTCS difference rows
for idx, nodes in zip(
range(settings.sim.ctcs_slice.start, settings.sim.ctcs_slice.stop),
settings.sim.ctcs_node_intervals,
):
start_i = 1 if nodes[0] == 0 else nodes[0]
x_max_un = float(settings.sim.x.max[idx])
scale = S_x[idx]
for i in range(start_i, nodes[1]):
# Row for +direction: s = x_max − scale·(x[i]−x[i-1]) ≥ 0
# structural cols: x[i,idx] then x[i-1,idx]; scipy sorts by col.
# x[i-1,idx] < x[i,idx] in column index, so after sort:
# data = [−scale (at x[i-1]), +scale (at x[i])]
# → s = x_max − (−scale·x[i-1] + scale·x[i]) = x_max − scale·(x[i]−x[i-1]) ✓
emit([scale, -scale], x_max_un) # col order: x[i], x[i-1]
# Row for −direction: s = x_max − scale·(x[i-1]−x[i]) ≥ 0
emit([-scale, scale], x_max_un) # col order: x[i], x[i-1]
# Pos epigraph: t_vb[c,k] − nu_vb[c,k] ≥ 0
# structural col order: nu_vb first, t_vb second (nu_vb < t_vb in layout)
# scipy sort: nu_vb first (smaller idx) → coefficients below match that order
for c_idx in range(L.n_nodal):
for k in range(N):
# [nu_vb coeff, t_vb coeff] after sorting by col
emit([1.0, -1.0], 0.0) # s = 0 − (nu_vb − t_vb) = t_vb − nu_vb ≥ 0
# Pos epigraph non-negativity: t_vb[c,k] ≥ 0
for c_idx in range(L.n_nodal):
for k in range(N):
emit([-1.0], 0.0) # s = 0 − (−t_vb) = t_vb ≥ 0
# User NonnegConeConstraint rows.
# s = b − Az ≥ 0; A = J·S (scaled Jacobian), b = -f0.
# Jacobians are w.r.t. physical x/u; multiply by S_x/S_u to get
# coefficients for the scaled decision variables dx_scaled / du_scaled.
for cc in self._user_cone_constraints:
if isinstance(cc, NonnegConeConstraint):
for k in cc.nodes:
x_k = x_bar[k]
u_k = u_bar[k]
f0 = np.asarray(cc.jax_fn(x_k, u_k, k, params), dtype=float).reshape(cc.m)
J_x = (
np.asarray(
jax.jacfwd(cc.jax_fn, argnums=0)(x_k, u_k, k, params), dtype=float
).reshape(cc.m, n_x)
* S_x[None, :]
)
J_u = (
np.asarray(
jax.jacfwd(cc.jax_fn, argnums=1)(x_k, u_k, k, params), dtype=float
).reshape(cc.m, n_u)
* S_u[None, :]
)
rhs = -f0
for i in range(cc.m):
coeffs = list(J_x[i]) + list(J_u[i])
emit(coeffs, float(rhs[i]))
# ================================================================
# SOC rows: SOC(2) for |nu[k,j]| ≤ t_vc[k,j]
# s = b − Az: row 0 has only t_vc (coeff −1), row 1 has only nu (coeff −1)
# s[0] = 0 − (−t_vc) = t_vc; s[1] = 0 − (−nu) = nu
# ‖s[1:]‖ = |nu| ≤ s[0] = t_vc ✓
# ================================================================
for k in range(N - 1):
for j in range(n_x):
emit([-1.0], 0.0) # row for t_vc
emit([-1.0], 0.0) # row for nu
# User SOCConstraint rows: SOC(m_arg + 1) per node.
# Layout: [bound_row, arg_row_0, …, arg_row_{m_arg-1}].
# Moreau convention: s = b − Az ∈ K_soc, ‖s[1:]‖ ≤ s[0].
# Linearized: s[0] = bound0 + J_b·Δ, s[1+i] = arg0[i] + J_a[i]·Δ.
# s = b − A·z → A = −J·S, b = val. Jacobians scaled by S_x/S_u to
# convert from physical-space to scaled-decision-variable coefficients.
for cc in self._user_cone_constraints:
if isinstance(cc, SOCConstraint):
for k in cc.nodes:
x_k = x_bar[k]
u_k = u_bar[k]
# Bound function.
bound0 = float(np.asarray(cc.bound_fn(x_k, u_k, k, params)))
Jb_x = (
np.asarray(
jax.jacfwd(cc.bound_fn, argnums=0)(x_k, u_k, k, params), dtype=float
).reshape(n_x)
* S_x
)
Jb_u = (
np.asarray(
jax.jacfwd(cc.bound_fn, argnums=1)(x_k, u_k, k, params), dtype=float
).reshape(n_u)
* S_u
)
# Arg function.
arg0 = np.asarray(cc.arg_fn(x_k, u_k, k, params), dtype=float).reshape(cc.m_arg)
Ja_x = (
np.asarray(
jax.jacfwd(cc.arg_fn, argnums=0)(x_k, u_k, k, params), dtype=float
).reshape(cc.m_arg, n_x)
* S_x[None, :]
)
Ja_u = (
np.asarray(
jax.jacfwd(cc.arg_fn, argnums=1)(x_k, u_k, k, params), dtype=float
).reshape(cc.m_arg, n_u)
* S_u[None, :]
)
# Bound row: A = −J_b·S, b = bound0 → s[0] = bound0 + J_b·Δ.
emit(list(-Jb_x) + list(-Jb_u), bound0)
# Arg rows: A = −J_a[i]·S, b = arg0[i] → s[1+i] = arg0[i] + J_a[i]·Δ.
for i in range(cc.m_arg):
emit(list(-Ja_x[i]) + list(-Ja_u[i]), float(arg0[i]))
return (
P_data,
np.asarray(coo_vals, dtype=float),
q,
np.asarray(b_list, dtype=float),
)
# ------------------------------------------------------------------
# Solve / unpack
# ------------------------------------------------------------------
def solve(self) -> PTRSolveResult:
"""Assemble the conic subproblem and dispatch to ``moreau.jax.Solver``.
Updates the warm-start carry only when the solve status is
:attr:`moreau.SolverStatus.Solved` or
:attr:`moreau.SolverStatus.AlmostSolved`.
Returns:
PTRSolveResult: Unscaled trajectories and solver status for this
SCP iteration.
"""
P_data, coo_vals, q, b = self._assemble_conic()
# Rebuild CSR from the same COO col/row arrays (with real values).
# scipy sorts col indices within each row identically to how the
# structural CSR was built at initialize(), so A_csr.data aligns with
# the fixed A_col_indices passed to moreau.
A_csr = sp.csr_matrix(
(coo_vals, (self._coo_rows, self._coo_cols)),
shape=(self._n_con, self.layout.n_z),
)
A_csr.sort_indices()
P_j = jnp.asarray(P_data)
A_j = jnp.asarray(A_csr.data)
q_j = jnp.asarray(q)
b_j = jnp.asarray(b)
solution = self._moreau.solve(P_j, A_j, q_j, b_j, warm_start=self._warm_start)
self._last_status = moreau.SolverStatus(int(np.asarray(self._moreau.info.status)))
if _moreau_solve_ok(self._last_status):
self._warm_start = solution.to_warm_start()
z = np.asarray(solution.x)
return self._unpack(z)
def iteration_callback(self):
"""JAX-pure ``(state, SubproblemData) -> SubproblemSolution``.
Composes :meth:`_assemble_conic_jax` + Moreau's jitted kernel
+ :meth:`_build_solution_jax` into a single ``@jax.jit``-decorated
closure.
The kernel (``moreau.jax.Solver._solve_raw``) returns
``(JaxSolution, JaxSolveInfo)`` and does **not** accept a warm-start
(see the 2026-05-17 Decision Log entry in
``plans/solver-iteration-callbacks.md``). Moreau 0.4 removed the
public ``moreau.jax.solver`` factory; ``Solver.solve`` drops status
onto ``solver.info``, which does not batch under ``vmap``. Every call
is a cold start; ``state`` is accepted only for cross-backend
signature uniformity. Both :meth:`~openscvx.problem.Problem.solve`
and :meth:`~openscvx.problem.Problem.solve_jax` route through this
callback for the SCP body, so the warm-start carry isn't threaded
on either path — threading it would require an
:class:`AlgorithmState.moreau_carry` field (future extension in
``plans/jax-pure-solve.md``).
The returned callable takes ``(state, data)``: ``state`` is the
:class:`AlgorithmState` pytree, accepted for cross-backend signature
uniformity but unused (the kernel takes no warm-start);
``data`` is the :class:`SubproblemData` pytree carrying the
per-iteration linearization arrays, penalty weights, and boundary
conditions.
"""
if self.layout is None or self._settings is None:
raise RuntimeError(
"MoreauPTRSolver.iteration_callback() requires initialize() to "
"have been called first."
)
# Public ``Solver.solve`` returns only the solution. The jitted kernel
# it calls still returns ``(JaxSolution, JaxSolveInfo)``, which is
# what ``vmap`` can carry as a pytree.
solve_fn = self._moreau._solve_raw
csr_to_coo_perm = jnp.asarray(self._csr_to_coo_perm)
def step(state, data: SubproblemData) -> SubproblemSolution:
del state # The jitted kernel takes no warm-start.
P_data, coo_vals, q, b = self._assemble_conic_jax(data)
A_data = coo_vals[csr_to_coo_perm]
sol, info = solve_fn(P_data, A_data, q, b)
return self._build_solution_jax(sol, info, data)
return jax.jit(step)
# ------------------------------------------------------------------
# JAX-pure assembly / unpack (used by iteration_callback)
# ------------------------------------------------------------------
def _assemble_conic_jax(
self,
data: SubproblemData,
) -> Tuple["jnp.ndarray", "jnp.ndarray", "jnp.ndarray", "jnp.ndarray"]:
"""JAX-pure counterpart of :meth:`_assemble_conic`.
Returns ``(P_data, coo_vals, q, b)`` where ``coo_vals`` is emitted in
the same COO traversal order as :meth:`_structural_pass` — the
caller applies the precomputed ``csr_to_coo_perm`` to obtain the
``A_data`` array Moreau's fixed CSR layout expects.
Mirrors ``_assemble_conic`` row-block by row-block; structural
metadata (layout, BC types, impulsive pins, CTCS intervals) is
Python-side and closes over the function.
"""
L = self.layout
sim = self._settings.sim
N, n_x, n_u = L.N, L.n_x, L.n_u
inv_S_x = self._inv_S_x_diag_j
inv_S_u = self._inv_S_u_diag_j
S_x = self._S_x_diag_j
S_u = self._S_u_diag_j
c_x = self._c_x_j
c_u = self._c_u_j
f = self._jax_dtype
slice_imp = sim.u.slice_impulsive
has_impulsive = slice_imp.stop > slice_imp.start
eps = jnp.asarray(_P_DIAG_EPS, dtype=f)
lam_prox = data.lam_prox
lam_cost = data.lam_cost
lam_vc = data.lam_vc
lam_vb_nodal = data.lam_vb_nodal
# ---------- P_data (diagonal on dx/du slots, in that order) ----------
P_dx = (2.0 * lam_prox[:, :n_x] + eps).reshape(-1)
P_du = (2.0 * lam_prox[:, n_x:] + eps).reshape(-1)
P_data = jnp.concatenate([P_dx, P_du])
# ---------- q (linear cost) ----------
init_signs_np = np.zeros(sim.n_states, dtype=float)
final_signs_np = np.zeros(sim.n_states, dtype=float)
for i in range(sim.true_state_slice.start, sim.true_state_slice.stop):
init_t = sim.x.initial_type[i]
final_t = sim.x.final_type[i]
if init_t == "Minimize":
init_signs_np[i] = 1.0
elif init_t == "Maximize":
init_signs_np[i] = -1.0
if final_t == "Minimize":
final_signs_np[i] = 1.0
elif final_t == "Maximize":
final_signs_np[i] = -1.0
init_signs = jnp.asarray(init_signs_np, dtype=f)
final_signs = jnp.asarray(final_signs_np, dtype=f)
lam_cost_arr = jnp.broadcast_to(jnp.asarray(lam_cost, dtype=f), (sim.n_states,))
q = jnp.zeros((L.n_z,), dtype=f)
x0_start = L.sl_x.start
xN_start = L.sl_x.start + (N - 1) * n_x
q = q.at[x0_start : x0_start + n_x].add(init_signs * lam_cost_arr)
q = q.at[xN_start : xN_start + n_x].add(final_signs * lam_cost_arr)
# L1 penalty on t_vc (epigraph of |nu|).
q = q.at[L.sl_t_vc].add(lam_vc.reshape(-1))
# Pos epigraph t_vb.
for c_idx in range(L.n_nodal):
q = q.at[L.sl_t_vb[c_idx]].add(lam_vb_nodal[:, c_idx])
# ---------- coo_vals & b in COO traversal order ----------
coo_blocks: List[jnp.ndarray] = []
b_blocks: List[jnp.ndarray] = []
# State error defs (N*n_x rows, 2 coeffs per row).
state_err_coeffs = jnp.tile(jnp.asarray([1.0, -1.0], dtype=f), N * n_x)
coo_blocks.append(state_err_coeffs)
b_blocks.append((inv_S_x * (data.x_bar - c_x[None, :])).reshape(-1))
# Control error defs (N*n_u rows, 2 coeffs per row).
ctrl_err_coeffs = jnp.tile(jnp.asarray([1.0, -1.0], dtype=f), N * n_u)
coo_blocks.append(ctrl_err_coeffs)
b_blocks.append((inv_S_u * (data.u_bar - c_u[None, :])).reshape(-1))
# Dynamics rows ((N-1) * n_x), n_per_row = 2 + n_x + 2*n_u.
# Match _assemble_conic's emission order:
# [1.0 (x term)] + [-A_blk[i, :n_x]] +
# interleaved [-B_blk[i, j], -C_blk_with_E[i, j] for j in 0..n_u-1] +
# [-1.0 (nu term)]
if has_impulsive:
D_steps = data.D_d[1:]
A_eff = jnp.einsum("kij,kjl->kil", D_steps, data.A_d)
B_eff = jnp.einsum("kij,kjl->kil", D_steps, data.B_d)
C_eff = jnp.einsum("kij,kjl->kil", D_steps, data.C_d)
else:
A_eff = data.A_d
B_eff = data.B_d
C_eff = data.C_d
A_blk = (inv_S_x[None, :, None] * A_eff) * S_x[None, None, :]
B_blk = (inv_S_x[None, :, None] * B_eff) * S_u[None, None, :]
C_blk = (inv_S_x[None, :, None] * C_eff) * S_u[None, None, :]
C_with_E = C_blk
if has_impulsive:
E_blk = (inv_S_x[None, :, None] * data.E_d[1:]) * S_u[None, None, :]
imp_slice = slice(slice_imp.start, slice_imp.stop)
C_with_E = C_with_E.at[:, :, imp_slice].add(E_blk[:, :, imp_slice])
n_per_row = 2 + n_x + 2 * n_u
dyn_vals = jnp.zeros((N - 1, n_x, n_per_row), dtype=f)
dyn_vals = dyn_vals.at[:, :, 0].set(1.0) # x[k, i]
dyn_vals = dyn_vals.at[:, :, 1 : 1 + n_x].set(-A_blk) # dx[kp, :]
for j in range(n_u):
base = 1 + n_x + 2 * j
dyn_vals = dyn_vals.at[:, :, base].set(-B_blk[:, :, j]) # du[kp, j]
dyn_vals = dyn_vals.at[:, :, base + 1].set(-C_with_E[:, :, j]) # du[k, j]
dyn_vals = dyn_vals.at[:, :, -1].set(-1.0) # nu[kp, i]
coo_blocks.append(dyn_vals.reshape(-1))
if has_impulsive:
rhs_dyn = inv_S_x[None, :] * (data.x_prop_plus[1:] - c_x[None, :])
else:
rhs_dyn = inv_S_x[None, :] * (data.x_prop - c_x[None, :])
b_blocks.append(rhs_dyn.reshape(-1))
# Linearized nodal constraints (variable rows per constraint).
for c_idx, constraint in enumerate(self._jax_constraints.nodal):
nodes_tuple = tuple(constraint.nodes)
if not nodes_tuple:
continue
nodes_arr = jnp.asarray(nodes_tuple)
n_nodes = len(nodes_tuple)
grad_x_vals = data.nodal_grad_x[nodes_arr, c_idx] # (n_nodes, n_x)
grad_u_vals = data.nodal_grad_u[nodes_arr, c_idx] # (n_nodes, n_u)
neg_one = jnp.full((n_nodes, 1), -1.0, dtype=f)
nodal_vals = jnp.concatenate([grad_x_vals, grad_u_vals, neg_one], axis=1)
coo_blocks.append(nodal_vals.reshape(-1))
b_blocks.append(-data.nodal_g[nodes_arr, c_idx])
# Boundary Fix rows — emit per-state, init then final, matching
# _assemble_conic. Under impulsive control the initial Fix row also
# carries du[0, slice_imp] coefficients.
imp_j_list = list(range(slice_imp.start, slice_imp.stop)) if has_impulsive else []
for i in range(sim.true_state_slice.start, sim.true_state_slice.stop):
init_t = sim.x.initial_type[i]
final_t = sim.x.final_type[i]
if init_t == "Fix":
if has_impulsive:
coeffs = [S_x[i]] + [-data.E_d[0, i, j] * S_u[j] for j in imp_j_list]
rhs_val = data.x_prop_plus[0, i] - c_x[i]
else:
coeffs = [S_x[i]]
rhs_val = data.x_init[i] - c_x[i]
coo_blocks.append(jnp.stack(coeffs))
b_blocks.append(rhs_val[None])
if final_t == "Fix":
coo_blocks.append(jnp.stack([S_x[i]]))
b_blocks.append((data.x_term[i] - c_x[i])[None])
# Impulsive zero-pin equalities.
for nodes, ctrl_slice in self._impulsive_pins:
for _node in nodes:
for j in range(ctrl_slice.start, ctrl_slice.stop):
coo_blocks.append(jnp.asarray([1.0], dtype=f))
b_blocks.append((-inv_S_u[j] * c_u[j])[None])
# User ZeroConeConstraint rows (hard equality).
# Coefficients scaled by S_x/S_u: dx_scaled = Δx_phys / S_x, so the
# coefficient for dx_scaled[k, j] is J_x_phys[i, j] * S_x[j].
params_user = data.params
for cc in self._user_cone_constraints:
if isinstance(cc, ZeroConeConstraint):
for k in cc.nodes:
x_k = data.x_bar[k]
u_k = data.u_bar[k]
f0 = jnp.reshape(cc.jax_fn(x_k, u_k, k, params_user), (cc.m,))
J_x = (
jnp.reshape(
jax.jacfwd(cc.jax_fn, argnums=0)(x_k, u_k, k, params_user), (cc.m, n_x)
)
* S_x[None, :]
)
J_u = (
jnp.reshape(
jax.jacfwd(cc.jax_fn, argnums=1)(x_k, u_k, k, params_user), (cc.m, n_u)
)
* S_u[None, :]
)
row_vals = jnp.concatenate([J_x, J_u], axis=1) # (cc.m, n_x+n_u)
coo_blocks.append(row_vals.reshape(-1))
b_blocks.append(-f0)
# Uniform time grid: u[k, j] - u[k-1, j] = 0.
if sim._uniform_time_grid:
td = sim.time_dilation_slice
n_tg_rows = (N - 1) * (td.stop - td.start)
tg_coeffs = jnp.tile(jnp.asarray([1.0, -1.0], dtype=f), n_tg_rows)
coo_blocks.append(tg_coeffs)
b_blocks.append(jnp.zeros((n_tg_rows,), dtype=f))
# CTCS initial: S_x[idx] * x[0, idx] = -c_x[idx].
ctcs_start = sim.ctcs_slice.start
ctcs_stop = sim.ctcs_slice.stop
if ctcs_stop > ctcs_start:
ctcs_idx_arr = jnp.arange(ctcs_start, ctcs_stop)
coo_blocks.append(S_x[ctcs_idx_arr])
b_blocks.append(-c_x[ctcs_idx_arr])
# State / control box (nonneg cone). All-constant coefficients.
x_max_sc = inv_S_x * (jnp.asarray(sim.x.max, dtype=f) - c_x)
x_min_sc = inv_S_x * (jnp.asarray(sim.x.min, dtype=f) - c_x)
u_max_sc = inv_S_u * (jnp.asarray(sim.u.max, dtype=f) - c_u)
u_min_sc = inv_S_u * (jnp.asarray(sim.u.min, dtype=f) - c_u)
# State box upper: N*n_x rows of coeff +1, RHS x_max_sc tiled.
coo_blocks.append(jnp.ones((N * n_x,), dtype=f))
b_blocks.append(jnp.tile(x_max_sc, N))
# State box lower: N*n_x rows of coeff -1, RHS -x_min_sc tiled.
coo_blocks.append(jnp.full((N * n_x,), -1.0, dtype=f))
b_blocks.append(jnp.tile(-x_min_sc, N))
# Control box upper
coo_blocks.append(jnp.ones((N * n_u,), dtype=f))
b_blocks.append(jnp.tile(u_max_sc, N))
# Control box lower
coo_blocks.append(jnp.full((N * n_u,), -1.0, dtype=f))
b_blocks.append(jnp.tile(-u_min_sc, N))
# CTCS difference rows.
for idx, nodes in zip(
range(ctcs_start, ctcs_stop),
sim.ctcs_node_intervals,
):
start_i = 1 if nodes[0] == 0 else nodes[0]
scale = S_x[idx]
x_max_un = float(sim.x.max[idx])
for _ in range(start_i, nodes[1]):
# +direction: emission order [x[i], x[i-1]] = [scale, -scale]
coo_blocks.append(jnp.stack([scale, -scale]))
b_blocks.append(jnp.asarray([x_max_un], dtype=f))
# -direction: [-scale, scale]
coo_blocks.append(jnp.stack([-scale, scale]))
b_blocks.append(jnp.asarray([x_max_un], dtype=f))
# Pos epigraph: t_vb[c, k] >= nu_vb[c, k].
# Emission order matches _assemble_conic: [+1 (nu_vb), -1 (t_vb)].
for _c_idx in range(L.n_nodal):
coeffs_block = jnp.tile(jnp.asarray([1.0, -1.0], dtype=f), N)
coo_blocks.append(coeffs_block)
b_blocks.append(jnp.zeros((N,), dtype=f))
# Pos epigraph non-negativity: t_vb[c, k] >= 0.
for _c_idx in range(L.n_nodal):
coo_blocks.append(jnp.full((N,), -1.0, dtype=f))
b_blocks.append(jnp.zeros((N,), dtype=f))
# User NonnegConeConstraint rows.
# Jacobians scaled by S_x/S_u for the dx_scaled / du_scaled convention.
for cc in self._user_cone_constraints:
if isinstance(cc, NonnegConeConstraint):
for k in cc.nodes:
x_k = data.x_bar[k]
u_k = data.u_bar[k]
f0 = jnp.reshape(cc.jax_fn(x_k, u_k, k, params_user), (cc.m,))
J_x = (
jnp.reshape(
jax.jacfwd(cc.jax_fn, argnums=0)(x_k, u_k, k, params_user), (cc.m, n_x)
)
* S_x[None, :]
)
J_u = (
jnp.reshape(
jax.jacfwd(cc.jax_fn, argnums=1)(x_k, u_k, k, params_user), (cc.m, n_u)
)
* S_u[None, :]
)
row_vals = jnp.concatenate([J_x, J_u], axis=1)
coo_blocks.append(row_vals.reshape(-1))
b_blocks.append(-f0)
# SOC: 2 rows per |nu| pair, coeffs [-1] each (t_vc row, then nu row).
n_soc = 2 * (N - 1) * n_x
soc_coeffs = jnp.full((n_soc,), -1.0, dtype=f)
coo_blocks.append(soc_coeffs)
b_blocks.append(jnp.zeros((n_soc,), dtype=f))
# User SOCConstraint rows: SOC(m_arg + 1) per node.
# Bound row: A = −J_b·S, b = bound0 → s[0] = b − A·z = bound0 + J_b·Δ.
# Arg rows: A = −J_a[i]·S, b = arg0[i].
# Jacobians scaled by S_x/S_u to convert physical Jacobians to
# coefficients for scaled decision variables dx_scaled / du_scaled.
for cc in self._user_cone_constraints:
if isinstance(cc, SOCConstraint):
for k in cc.nodes:
x_k = data.x_bar[k]
u_k = data.u_bar[k]
bound0 = jnp.reshape(cc.bound_fn(x_k, u_k, k, params_user), ())
Jb_x = (
jnp.reshape(
jax.jacfwd(cc.bound_fn, argnums=0)(x_k, u_k, k, params_user), (n_x,)
)
* S_x
)
Jb_u = (
jnp.reshape(
jax.jacfwd(cc.bound_fn, argnums=1)(x_k, u_k, k, params_user), (n_u,)
)
* S_u
)
arg0 = jnp.reshape(cc.arg_fn(x_k, u_k, k, params_user), (cc.m_arg,))
Ja_x = (
jnp.reshape(
jax.jacfwd(cc.arg_fn, argnums=0)(x_k, u_k, k, params_user),
(cc.m_arg, n_x),
)
* S_x[None, :]
)
Ja_u = (
jnp.reshape(
jax.jacfwd(cc.arg_fn, argnums=1)(x_k, u_k, k, params_user),
(cc.m_arg, n_u),
)
* S_u[None, :]
)
# Bound row A coeffs (negated, concatenated x then u).
bound_A = jnp.concatenate([-Jb_x, -Jb_u]) # (n_x + n_u,)
coo_blocks.append(bound_A)
b_blocks.append(bound0[None])
# Arg rows A coeffs (negated).
arg_A = jnp.concatenate([-Ja_x, -Ja_u], axis=1) # (m_arg, n_x+n_u)
coo_blocks.append(arg_A.reshape(-1))
b_blocks.append(arg0)
coo_vals = jnp.concatenate(coo_blocks) if coo_blocks else jnp.zeros((0,), dtype=f)
b_vec = jnp.concatenate(b_blocks) if b_blocks else jnp.zeros((0,), dtype=f)
return P_data, coo_vals, q, b_vec
def _build_solution_jax(
self,
sol,
info,
data: SubproblemData,
) -> SubproblemSolution:
"""Unpack ``(JaxSolution, JaxSolveInfo)`` into a :class:`SubproblemSolution`.
Mirrors :meth:`_unpack`: scaled-to-physical trajectories, virtual
controls in their natural orientation, ``nu_vb`` stacked to
``(N, n_nodal)``. ``status_code`` is derived from ``info.status`` via
:func:`_moreau_status_to_code`.
"""
L = self.layout
N, n_x, n_u = L.N, L.n_x, L.n_u
f = self._jax_dtype
z = sol.x
x_scaled = z[L.sl_x].reshape(N, n_x)
u_scaled = z[L.sl_u].reshape(N, n_u)
x = x_scaled * self._S_x_diag_j[None, :] + self._c_x_j[None, :]
u = u_scaled * self._S_u_diag_j[None, :] + self._c_u_j[None, :]
nu = z[L.sl_nu].reshape(N - 1, n_x)
if L.n_nodal:
nu_vb = jnp.stack([z[sl] for sl in L.sl_nu_vb], axis=-1)
else:
nu_vb = jnp.zeros((N, 0), dtype=f)
cost = self._reconstruct_cost_jax(z, data)
status_code = _moreau_status_to_code(info.status)
return SubproblemSolution(
x=x,
u=u,
nu=nu,
nu_vb=nu_vb,
nu_vb_cross=jnp.zeros((0,), dtype=f),
cost=cost,
status_code=status_code,
)
def _reconstruct_cost_jax(
self,
z: "jnp.ndarray",
data: SubproblemData,
) -> "jnp.ndarray":
"""JAX-pure mirror of :meth:`_reconstruct_cost`."""
L = self.layout
sim = self._settings.sim
N, n_x = L.N, L.n_x
f = self._jax_dtype
lam_prox = data.lam_prox
lam_cost_arr = jnp.broadcast_to(jnp.asarray(data.lam_cost, dtype=f), (sim.n_states,))
lam_vc = data.lam_vc
lam_vb_nodal = data.lam_vb_nodal
dx = z[L.sl_dx].reshape(N, n_x)
du = z[L.sl_du].reshape(N, L.n_u)
x_scaled = z[L.sl_x].reshape(N, n_x)
t_vc = z[L.sl_t_vc].reshape(N - 1, n_x)
cost = jnp.sum(lam_prox[:, :n_x] * dx**2) + jnp.sum(lam_prox[:, n_x:] * du**2)
for i in range(sim.true_state_slice.start, sim.true_state_slice.stop):
init_t = sim.x.initial_type[i]
final_t = sim.x.final_type[i]
if init_t == "Minimize":
cost = cost + lam_cost_arr[i] * x_scaled[0, i]
elif init_t == "Maximize":
cost = cost - lam_cost_arr[i] * x_scaled[0, i]
if final_t == "Minimize":
cost = cost + lam_cost_arr[i] * x_scaled[-1, i]
elif final_t == "Maximize":
cost = cost - lam_cost_arr[i] * x_scaled[-1, i]
cost = cost + jnp.sum(lam_vc * t_vc)
for c_idx in range(L.n_nodal):
t_vb = z[L.sl_t_vb[c_idx]]
cost = cost + jnp.dot(lam_vb_nodal[:, c_idx], t_vb)
return cost.astype(f)
def _unpack(self, z: np.ndarray) -> PTRSolveResult:
"""Map the flat solution vector into a :class:`PTRSolveResult`.
Args:
z: Primal solution from Moreau, length ``layout.n_z``.
Returns:
PTRSolveResult: Unscaled ``x`` / ``u`` trajectories, slacks, cost,
and ``"optimal"`` or ``"infeasible"`` status.
"""
L = self.layout
N, n_x, n_u = L.N, L.n_x, L.n_u
x_scaled = z[L.sl_x].reshape(N, n_x)
u_scaled = z[L.sl_u].reshape(N, n_u)
x = x_scaled * self._S_x_diag[None, :] + self._c_x[None, :]
u = u_scaled * self._S_u_diag[None, :] + self._c_u[None, :]
nu = z[L.sl_nu].reshape(N - 1, n_x)
nu_vb = [np.asarray(z[sl]) for sl in L.sl_nu_vb]
nu_vb_cross: List[float] = []
is_ok = _moreau_solve_ok(self._last_status)
status = "optimal" if is_ok else "infeasible"
cost = self._reconstruct_cost(z)
return PTRSolveResult(
x=x,
u=u,
nu=nu,
nu_vb=nu_vb,
nu_vb_cross=nu_vb_cross,
cost=cost,
status=status,
)
def _reconstruct_cost(self, z: np.ndarray) -> float:
"""Recompute the PTR objective at ``z`` from stored penalty weights.
Uses epigraph variables ``t_vc`` / ``t_vb`` (equal to ``|nu|`` /
``max(nu_vb, 0)`` at optimality) instead of absolute values or
positive parts.
Args:
z: Primal solution vector.
Returns:
Scalar PTR cost at ``z``.
"""
L = self.layout
settings = self._settings
N, n_x = L.N, L.n_x
lam_prox = self._pen["lam_prox"]
lam_cost_arr = np.broadcast_to(self._pen["lam_cost"], (settings.sim.n_states,))
lam_vc = self._pen["lam_vc"]
lam_vb_nodal = self._pen["lam_vb_nodal"]
dx = z[L.sl_dx].reshape(N, n_x)
du = z[L.sl_du].reshape(N, L.n_u)
x = z[L.sl_x].reshape(N, n_x)
t_vc = z[L.sl_t_vc].reshape(N - 1, n_x)
cost = float(np.sum(lam_prox[:, :n_x] * dx**2) + np.sum(lam_prox[:, n_x:] * du**2))
for i in range(settings.sim.true_state_slice.start, settings.sim.true_state_slice.stop):
init_t = settings.sim.x.initial_type[i]
final_t = settings.sim.x.final_type[i]
if init_t == "Minimize":
cost += float(lam_cost_arr[i] * x[0, i])
elif init_t == "Maximize":
cost -= float(lam_cost_arr[i] * x[0, i])
if final_t == "Minimize":
cost += float(lam_cost_arr[i] * x[-1, i])
elif final_t == "Maximize":
cost -= float(lam_cost_arr[i] * x[-1, i])
cost += float(np.sum(lam_vc * t_vc))
for c_idx in range(L.n_nodal):
t_vb = z[L.sl_t_vb[c_idx]]
cost += float(np.dot(lam_vb_nodal[:, c_idx], t_vb))
return cost
# ------------------------------------------------------------------
# Misc API
# ------------------------------------------------------------------
def get_stats(self) -> dict:
"""Return conic problem dimensions for the diagnostics summary box.
Returns:
dict: Keys ``n_variables``, ``n_parameters`` (always 0), and
``n_constraints``.
"""
if self.layout is None:
return {"n_variables": 0, "n_parameters": 0, "n_constraints": 0}
return {
"n_variables": self.layout.n_z,
"n_parameters": 0,
"n_constraints": self._n_con if self._n_con else -1,
}
def citation(self) -> List[str]:
"""Return BibTeX citation entries for Moreau.
Returns:
list[str]: BibTeX strings for inclusion in a bibliography.
"""
return [
r"""@software{moreau2024,
title={Moreau: GPU-Accelerated Conic Optimization Solver},
author={{Moreau Contributors}},
url={https://docs.moreau.so},
year={2024}
}"""
]