openscvx 0.5.3.dev33__tar.gz → 0.5.3.dev34__tar.gz

{openscvx-0.5.3.dev33/openscvx.egg-info → openscvx-0.5.3.dev34}/PKG-INFO

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

{openscvx-0.5.3.dev33 → openscvx-0.5.3.dev34}/openscvx/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.5.3.dev33'
-__version_tuple__ = version_tuple = (0, 5, 3, 'dev33')
+__version__ = version = '0.5.3.dev34'
+__version_tuple__ = version_tuple = (0, 5, 3, 'dev34')
 
-__commit_id__ = commit_id = 'g3d6f6768a'
+__commit_id__ = commit_id = 'g264808b5c'

{openscvx-0.5.3.dev33 → openscvx-0.5.3.dev34}/openscvx/algorithms/optimization_results.py

@@ -50,6 +50,14 @@ class OptimizationResults:
             Added by propagate_trajectory_results.
         ctcs_violation (Optional[np.ndarray]): Continuous-time constraint violations.
             Added by propagate_trajectory_results.
+        parameters (dict[str, np.ndarray]): Parameter values the solve used, recorded
+            by :meth:`~openscvx.problem.Problem.solve` and
+            :meth:`~openscvx.problem.Problem.solve_batched` (batched solves store every
+            value with a leading ``(B,)`` axis, shared values replicated). Results
+            record what happened: post-processing propagates with these values, and
+            the Problem's parameter dict only seeds the next solve. Empty when no
+            snapshot was recorded (e.g. ``solve_jax``, whose callers manage
+            parameters themselves).
         plotting_data (dict[str, Any]): Flexible storage for plotting and application data.
 
     !!! note "JAX pytree registration"
@@ -59,8 +67,9 @@ class OptimizationResults:
         ``converged``, ``t_final``, ``nodes``, ``trajectory``, ``X``, ``U``,
         and every ``*_history`` field — the data the user actually consumes.
         Post-process fields (``t_full``, ``x_full``, ``u_full``, ``cost``,
-        ``ctcs_violation``), ``plotting_data``, and the internal ``_states`` /
-        ``_controls`` metadata are stashed in the treedef's *aux* instead, so
+        ``ctcs_violation``), ``parameters``, ``plotting_data``, and the
+        internal ``_states`` / ``_controls`` metadata are stashed in the
+        treedef's *aux* instead, so
         a batched ``solve_jax`` doesn't force callers to handle ``None``
         leaves and ``post_process()`` outputs stay outside the vmap surface
         (post-process per element after the batched solve).
@@ -217,6 +226,10 @@ class OptimizationResults:
     cost: Optional[float] = field(default=None, metadata={"npz": "optional_scalar"})
     ctcs_violation: Optional[np.ndarray] = field(default=None, metadata={"npz": "optional_array"})
 
+    # Parameter values the solve used (recorded by Problem.solve() /
+    # Problem.solve_batched(); empty when nothing was recorded)
+    parameters: dict[str, np.ndarray] = field(default_factory=dict, metadata={"npz": "dict"})
+
     # Additional plotting/application data (added by user)
     plotting_data: dict[str, Any] = field(default_factory=dict, metadata={"npz": "dict"})
 
@@ -245,6 +258,7 @@ class OptimizationResults:
             self.u_full,
             self.cost,
             self.ctcs_violation,
+            self.parameters,
             self.plotting_data,
         )
         return children, aux
@@ -260,6 +274,7 @@ class OptimizationResults:
             u_full,
             cost,
             ctcs_violation,
+            parameters,
             plotting_data,
         ) = aux
         return cls(
@@ -271,6 +286,7 @@ class OptimizationResults:
             u_full=u_full,
             cost=cost,
             ctcs_violation=ctcs_violation,
+            parameters=parameters,
             plotting_data=plotting_data,
         )
 

{openscvx-0.5.3.dev33 → openscvx-0.5.3.dev34}/openscvx/problem.py

@@ -23,6 +23,7 @@ import os
 import queue
 import threading
 import time
+import warnings
 from dataclasses import fields as dc_fields
 from dataclasses import replace as dc_replace
 from typing import Any, Dict, List, Optional, Tuple, Union
@@ -672,6 +673,7 @@ class Problem:
         # Final solution state (saved after successful solve)
         self._solution: Optional[AlgorithmState] = None
         self._solution_history: Optional[AlgorithmHistory] = None
+        self._solution_parameters: Optional[Dict[str, np.ndarray]] = None
 
         # SCP algorithm (resolved from `algorithm` parameter above)
 
@@ -1287,6 +1289,7 @@ class Problem:
         # Reset solution
         self._solution = None
         self._solution_history = None
+        self._solution_parameters = None
 
         # Reset timing
         self.timing_solve = None
@@ -1355,7 +1358,10 @@ class Problem:
 
         Returns:
             OptimizationResults with trajectory and convergence info
-                (call post_process() for full propagation)
+                (call post_process() for full propagation). The parameter
+                values used for the solve are recorded on
+                ``results.parameters``; :meth:`post_process` propagates
+                with them.
         """
         # Sync parameters, boundary conditions, and SCP constants before solving
         self._sync_parameters()
@@ -1408,16 +1414,22 @@ class Problem:
 
         profiling.profiling_end(pr, "solve")
 
-        # Snapshot state + history for post_process() / plotting reuse.
+        # Snapshot state, history, and parameters for post_process() / plotting
+        # reuse.
         self._solution = self._state
         self._solution_history = copy.deepcopy(self._history)
+        self._solution_parameters = {
+            name: np.array(value) for name, value in self._parameters.items()
+        }
 
         timed_out = t_max is not None and self.timing_solve >= t_max
-        return self._format_result(
+        result = self._format_result(
             self._state,
             self._history,
             int(self._state.k) <= k_max and not timed_out,
         )
+        result.parameters = self._solution_parameters
+        return result
 
     def _default_state(self) -> AlgorithmState:
         """Fresh :class:`AlgorithmState` from settings and the algorithm's SCP constants.
@@ -1665,6 +1677,11 @@ class Problem:
             parameters: Problem parameters dict for the current solve. Use
                 this rather than mutating ``self.parameters`` if the
                 parameters are traced. ``None`` reuses ``self._parameters``.
+                Unlike :meth:`solve` / :meth:`solve_batched`, the merged
+                values are *not* recorded on the returned results — under a
+                caller's ``jit`` / ``vmap`` they are tracers, which must not
+                leak into the pytree's aux — so callers that post-process
+                manage parameter bookkeeping themselves.
             x_guess: Initial state trajectory guess, shape ``(N, n_states)``.
                 Replaces ``state.x``; ``None`` uses the default from
                 settings (``State.guess`` at ``initialize()``).
@@ -1846,7 +1863,12 @@ class Problem:
 
         Returns:
             :class:`OptimizationResults` pytree with a leading ``B`` axis on
-            every leaf and empty per-iteration histories.
+            every leaf and empty per-iteration histories. The merged
+            parameter values used for the solve are recorded on
+            ``results.parameters``, post-broadcast — every value carries the
+            leading ``(B,)`` axis, shared values replicated — and
+            :meth:`post_process_batched` propagates each element with its
+            own values.
         """
         if self._iteration_fn_jit_inner is None:
             raise ValueError(
@@ -1953,7 +1975,16 @@ class Problem:
             final_states = batched_solve(states, params_for_solve)
         jax.block_until_ready(final_states)
         self.timing_solve = time.time() - t_0_solve
-        return OptimizationResults.from_final_state(final_states, problem=self)
+        results = OptimizationResults.from_final_state(final_states, problem=self)
+        # Record the as-used parameter values with a leading (B,) axis on
+        # every leaf, so post_process_batched() can index element b directly.
+        results.parameters = {
+            name: np.array(np.moveaxis(np.asarray(value), param_axes[name], 0))
+            if param_axes[name] is not None
+            else np.broadcast_to(np.asarray(value), (B,) + np.shape(value)).copy()
+            for name, value in params_for_solve.items()
+        }
+        return results
 
     def post_process(self) -> OptimizationResults:
         """Propagate solution through full nonlinear dynamics for high-fidelity trajectory.
@@ -1961,6 +1992,10 @@ class Problem:
         Integrates the converged SCP solution through the nonlinear dynamics to
         produce x_full, u_full, and t_full. Call after solve() for final results.
 
+        Propagation uses the parameter values recorded at solve time
+        (``results.parameters``), so mutating ``self.parameters`` between
+        :meth:`solve` and here only seeds the next solve.
+
         Returns:
             OptimizationResults with propagated trajectory fields
 
@@ -1980,9 +2015,11 @@ class Problem:
             int(self._solution.k) <= self._algorithm.k_max,
         )
 
+        result.parameters = self._solution_parameters
+
         t_0_post = time.time()
         result = propagate_trajectory_results(
-            self._parameters,
+            self._solution_parameters,
             self.settings,
             result,
             self._propagation_solver,
@@ -2022,9 +2059,11 @@ class Problem:
         - ``results.cost``    — ``(B,)``
         - ``results.ctcs_violation`` — ``(B, ...)`` or ``None``
 
-        Propagation runs sequentially over the batch in a Python loop;
-        all elements share ``self._parameters`` for any parameter values not
-        embedded in the trajectory state vector.
+        Propagation runs sequentially over the batch in a Python loop; each
+        element propagates with its own parameter values from the snapshot
+        :meth:`solve_batched` recorded on ``results.parameters``. Results
+        without a snapshot (e.g. loaded from an older save) fall back to
+        ``self.parameters`` with a warning.
 
         Args:
             results: Batched :class:`~openscvx.algorithms.OptimizationResults`
@@ -2068,11 +2107,23 @@ class Problem:
         costs: List[float] = []
         ctcs_violations = []
 
+        # Empty for results saved before the snapshot existed — those fall
+        # back to the Problem's current values, the pre-snapshot behaviour.
+        snapshot = getattr(results, "parameters", None) or {}
+        if not snapshot and self._parameters:
+            warnings.warn(
+                "results carry no record of the parameter values they were solved "
+                "with; propagating every element with the Problem's current "
+                "parameters, which may differ from the solved values.",
+                UserWarning,
+                stacklevel=2,
+            )
+
         t_0_post = time.time()
         for b in range(B):
             single = _make_single_result(results, b)
             prop = propagate_trajectory_results(
-                self._parameters,
+                dict(self._parameters, **{k: np.asarray(v)[b] for k, v in snapshot.items()}),
                 self.settings,
                 single,
                 self._propagation_solver,

{openscvx-0.5.3.dev33 → openscvx-0.5.3.dev34/openscvx.egg-info}/PKG-INFO

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

{openscvx-0.5.3.dev33 → openscvx-0.5.3.dev34}/openscvx.egg-info/SOURCES.txt

@@ -355,6 +355,7 @@ tests/test_multishot_propagation.py
 tests/test_optimization_results.py
 tests/test_plotting.py
 tests/test_propagation.py
+tests/test_results_parameters.py
 tests/test_solve_batched_brachistochrone.py
 tests/test_solve_batched_cvxpy_export_errors.py
 tests/test_solve_batched_export_roundtrip.py

{openscvx-0.5.3.dev33 → openscvx-0.5.3.dev34}/openscvx.egg-info/scm_file_list.json

@@ -220,6 +220,7 @@
     "tests/test_propagation.py",
     "tests/hohmann_analytical.py",
     "tests/test_expert.py",
+    "tests/test_results_parameters.py",
     "tests/test_solve_batched_inference.py",
     "tests/test_solve_jax_bare_brachistochrone.py",
     "tests/test_cvxpygen_optional.py",

openscvx-0.5.3.dev34/openscvx.egg-info/scm_version.json

@@ -0,0 +1,8 @@
+{
+  "tag": "0.5.2",
+  "distance": 34,
+  "node": "g264808b5c8dfd13d52b35c4d96c885b010058203",
+  "dirty": false,
+  "branch": "main",
+  "node_date": "2026-07-06"
+}

{openscvx-0.5.3.dev33 → openscvx-0.5.3.dev34}/tests/test_optimization_results.py

@@ -1,5 +1,6 @@
 """Tests for OptimizationResults save/load round-trip."""
 
+import jax
 import numpy as np
 import pytest
 
@@ -184,3 +185,29 @@ def test_plotting_data_skips_non_array(tmp_path):
     loaded = OptimizationResults.load(path)
     np.testing.assert_array_equal(loaded.plotting_data["good"], np.array([1, 2, 3]))
     assert "bad_func" not in loaded.plotting_data
+
+
+def test_save_load_roundtrip_parameters(tmp_path):
+    """The recorded parameter snapshot survives a save/load round-trip."""
+    original = _make_result()
+    original.parameters = {"gain": np.array([2.0]), "center": np.array([1.0, 2.0, 3.0])}
+    path = tmp_path / "with_parameters.npz"
+    original.save(path)
+
+    loaded = OptimizationResults.load(path)
+    assert set(loaded.parameters) == set(original.parameters)
+    for key, value in original.parameters.items():
+        np.testing.assert_array_equal(loaded.parameters[key], value)
+
+
+def test_parameters_snapshot_is_pytree_aux_not_child():
+    """The parameter snapshot rides the treedef's aux, leaving the children unchanged."""
+    result = _make_result()
+    n_leaves = len(jax.tree_util.tree_leaves(result))
+
+    result.parameters = {"gain": np.array([2.0])}
+    leaves, treedef = jax.tree_util.tree_flatten(result)
+    assert len(leaves) == n_leaves  # no new pytree children
+
+    restored = jax.tree_util.tree_unflatten(treedef, leaves)
+    np.testing.assert_array_equal(restored.parameters["gain"], np.array([2.0]))

openscvx-0.5.3.dev34/tests/test_results_parameters.py

@@ -0,0 +1,93 @@
+"""Solved parameter values are recorded on results and drive post-processing.
+
+``Problem.solve`` and ``Problem.solve_batched`` snapshot the merged, as-used
+parameter values onto ``results.parameters``; ``post_process`` /
+``post_process_batched`` propagate with the recorded values, so mutating
+``problem.parameters`` between solve and post-process — or batching a
+parameter that enters the dynamics — can no longer corrupt the dense
+trajectory. The fixture's ``gain`` parameter scales the position kinematics,
+which makes a wrong propagation parameter directly visible: gain 0 freezes
+the dense positions, gain 1 moves them.
+"""
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from tests.test_solve_batched_brachistochrone import _build_brachistochrone_with_params
+
+# === solve() -> post_process() ===
+
+
+def test_post_process_propagates_with_solved_parameters():
+    prob = _build_brachistochrone_with_params("cvxpy", n=8, k_max=20)
+    prob.initialize()
+    result = prob.solve()
+
+    # The solve records the as-used values as plain numpy.
+    np.testing.assert_array_equal(result.parameters["gain"], np.array([1.0]))
+
+    # Mutating the Problem's parameters afterwards only seeds the next solve;
+    # post-processing propagates with the recorded gain of 1.0. Under gain 0
+    # the position kinematics would be identically zero.
+    prob.parameters["gain"] = np.array([0.0])
+    result = prob.post_process()
+
+    np.testing.assert_array_equal(result.parameters["gain"], np.array([1.0]))
+    position = np.asarray(result.trajectory["position"])  # (n_times, 2)
+    assert np.ptp(position[:, 0]) > 5.0  # moved from [0, 10] toward [10, 5]
+
+    jax.clear_caches()
+
+
+# === solve_batched() -> post_process_batched() ===
+
+
+def test_post_process_batched_propagates_each_element_with_its_own_parameters():
+    pytest.importorskip("qpax")
+
+    prob = _build_brachistochrone_with_params("qpax", n=8, k_max=5)
+    prob.initialize()
+
+    gains = jnp.array([[0.0], [1.0]])  # (B, 1) vs declared (1,) -> batched
+    results = prob.solve_batched(parameters={"gain": gains})
+
+    # Post-broadcast snapshot: every value carries the leading (B,) axis,
+    # the shared gravity replicated.
+    np.testing.assert_allclose(np.asarray(results.parameters["gain"]), np.asarray(gains))
+    np.testing.assert_allclose(
+        np.asarray(results.parameters["gravity"]), np.broadcast_to([9.81], (2, 1))
+    )
+
+    results = prob.post_process_batched(results)
+    position = np.asarray(results.trajectory["position"])  # (B, n_times, 2)
+
+    # Element 0 solved with gain 0 and must propagate with gain 0: its
+    # position kinematics are identically zero, so the dense trajectory
+    # stays at the initial point. Element 1 (gain 1) moves.
+    assert np.abs(position[0] - position[0, 0]).max() < 1e-6
+    assert np.ptp(position[1, :, 0]) > 1.0
+
+    jax.clear_caches()
+
+
+def test_post_process_batched_without_snapshot_falls_back_to_problem_parameters():
+    pytest.importorskip("qpax")
+
+    prob = _build_brachistochrone_with_params("qpax", n=8, k_max=5)
+    prob.initialize()
+
+    results = prob.solve_batched(parameters={"gain": jnp.array([[0.0], [1.0]])})
+    del results.parameters  # a results object that predates the snapshot
+
+    with pytest.warns(UserWarning, match="no record of the parameter values"):
+        results = prob.post_process_batched(results)
+    position = np.asarray(results.trajectory["position"])
+
+    # Every element falls back to the Problem's own gain of 1.0: even the
+    # element solved with gain 0 moves during propagation.
+    for b in range(2):
+        assert np.ptp(position[b, :, 0]) > 0.5
+
+    jax.clear_caches()

openscvx-0.5.3.dev33/openscvx.egg-info/scm_version.json

@@ -1,8 +0,0 @@
-{
-  "tag": "0.5.2",
-  "distance": 33,
-  "node": "g3d6f6768a9d011c99ab12c85b6695002b0606580",
-  "dirty": false,
-  "branch": "main",
-  "node_date": "2026-07-05"
-}