openscvx 2.dev1__py3-none-any.whl → 2.dev2__py3-none-any.whl

openscvx/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '2.dev1'
-__version_tuple__ = version_tuple = (2, 'dev1')
+__version__ = version = '2.dev2'
+__version_tuple__ = version_tuple = (2, 'dev2')
 
 __commit_id__ = commit_id = None

openscvx/algorithms/penalized_trust_region.py

@@ -6,8 +6,9 @@ optimization problems through iterative convex approximation.
 
 import time
 import warnings
-from typing import TYPE_CHECKING, Callable, Dict, List, Union
+from typing import TYPE_CHECKING, Callable, Dict, List, Tuple, Union
 
+import jax
 import numpy as np
 import numpy.linalg as la
 
@@ -188,6 +189,11 @@ class PenalizedTrustRegion(Algorithm):
             return solver.call(*args)
         return solver(*args)
 
+    @staticmethod
+    def _block_until_ready_outputs(outputs: Tuple[object, ...]) -> None:
+        """Finish any pending XLA work from discretization exports (warm-up helper)."""
+        jax.block_until_ready(outputs)
+
     def _recover_prior_node_from_initial(
         self,
         settings: Config,
@@ -286,7 +292,8 @@ class PenalizedTrustRegion(Algorithm):
         """Initialize PTR algorithm.
 
         Stores compiled infrastructure and performs a warm-start solve to
-        initialize DPP and JAX jacobians.
+        initialize DPP and JAX jacobians. Also runs post-subproblem discretization
+        on the throwaway CVX solution so XLA/export caches match the first ``step()``.
 
         Args:
             solver: Convex subproblem solver (e.g., CVXPySolver)
@@ -333,7 +340,29 @@ class PenalizedTrustRegion(Algorithm):
                 params,
             )
             init_state.add_impulsive_discretization(W_multi_shoot.__array__())
-        _ = self._subproblem(params, init_state, settings)
+        (x_sol, u_sol, *_) = self._subproblem(params, init_state, settings)
+
+        # Prime the same exported discretization calls used after every subproblem in
+        # step() (candidate trajectory). initialize() previously only discretized the
+        # initial guess, so the first step() in solve() could still hit an XLA cache_miss
+        # on post-CVX (x_sol, u_sol). Running that path here moves compilation into init.
+        cont_out = self._invoke_solver(
+            self._discretization_solver, x_sol, u_sol.astype(float), params
+        )
+        x_prop_c = cont_out[3]
+        u_candidate = u_sol.astype(float)
+        x0_prior_c = self._recover_prior_node_from_initial(settings, x_sol[0])
+        x_nodes_prior_c = np.vstack((x0_prior_c, np.asarray(x_prop_c)))
+        if self._discretization_solver_impulsive is not None:
+            imp_out = self._invoke_solver(
+                self._discretization_solver_impulsive,
+                x_nodes_prior_c,
+                u_candidate,
+                params,
+            )
+            self._block_until_ready_outputs(cont_out + imp_out)
+        else:
+            self._block_until_ready_outputs(cont_out)
 
     def step(
         self,

openscvx/symbolic/hashing.py

@@ -9,16 +9,29 @@ the same structure, the cached compiled artifacts can be reused.
 """
 
 import hashlib
+from importlib.metadata import PackageNotFoundError, version
 from typing import TYPE_CHECKING
 
 import numpy as np
 
-from openscvx._version import __version__
-
 if TYPE_CHECKING:
     from openscvx.symbolic.problem import SymbolicProblem
 
 
+def _openscvx_version() -> str:
+    """Package version for cache invalidation; works without generated _version.py."""
+    try:
+        return version("openscvx")
+    except PackageNotFoundError:
+        pass
+    try:
+        from openscvx._version import __version__ as v
+
+        return v
+    except ImportError:
+        return "0.0.0"
+
+
 def hash_symbolic_problem(problem: "SymbolicProblem") -> str:
     """Compute a structural hash of a symbolic optimization problem.
 
@@ -47,7 +60,7 @@ def hash_symbolic_problem(problem: "SymbolicProblem") -> str:
     hasher = hashlib.sha256()
 
     # Include library version to invalidate cache on version changes
-    hasher.update(f"openscvx:{__version__}:".encode())
+    hasher.update(f"openscvx:{_openscvx_version()}:".encode())
 
     # Hash the dynamics
     hasher.update(b"dynamics:")

{openscvx-2.dev1.dist-info → openscvx-2.dev2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 2.dev1
+Version: 2.dev2
 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-2.dev1.dist-info → openscvx-2.dev2.dist-info}/RECORD

@@ -1,6 +1,6 @@
 openscvx/__init__.py,sha256=S-pt3p_BI5aOdOYPLrxKOshOeLEBXmqPoqZfROqq730,3542
 openscvx/__main__.py,sha256=Hwm7mtVg3tLdvoUPkpcQv8KF3wxl72PNLBp9axFu8GY,2991
-openscvx/_version.py,sha256=4w1vpaqG6FPfOy5YFO_ugNNX5OrGGKO6fMH-0ZI5wFs,523
+openscvx/_version.py,sha256=U-q9bAeU59rRgUryJHtruyMJsxMPd7Yq_xsfqRREV4A,523
 openscvx/config.py,sha256=qfDDYoCe6WqJglKsx5b2W48YOglXenKr-PVRRdCFhYE,9898
 openscvx/loader.py,sha256=FvKLkkXd4ihd5FqLFF8Cd9VnPbPwTV_azBRnEipi28c,7654
 openscvx/problem.py,sha256=g-DKsdpZ5rOxEoMVOUj_QwB1L1edN0rgtRn8Z-TjgTI,47177
@@ -9,7 +9,7 @@ openscvx/algorithms/augmented_lagrangian.py,sha256=liEHtqONbpmw7CZJJCkAluPOEbFsl
 openscvx/algorithms/base.py,sha256=JsaVfS-hyHeGU1GUuQV9i-EpKQ1HBE0RHABIBGrEonM,30205
 openscvx/algorithms/constant_proximal_weight.py,sha256=UopFJHFHaFRTUnpeBEc_wYsSDvmxS9-bSKpI8JGYQII,2846
 openscvx/algorithms/optimization_results.py,sha256=Fk7tiW8IuMZa8Ri9-SHq9kmw3BPxXl1mc0yXy11w_-c,13416
-openscvx/algorithms/penalized_trust_region.py,sha256=ncM3BczFGIPRBrWXcMp_K-ZbJ1IkMJu7qxzPrmV2ilY,28687
+openscvx/algorithms/penalized_trust_region.py,sha256=2UDZMYshpn0Z52JFjqPji5ececD0SGjZqhSrAT1DbD0,30164
 openscvx/algorithms/ramp_proximal_weight.py,sha256=19A-oerRMJjajQpHxeUeXdyTV0_Onr9cw1ttlwbY5_w,3383
 openscvx/algorithms/weights.py,sha256=DQV9NOR_-59lISv6d_el3cpQ-EcFasSGb1oZr4ma-yA,20682
 openscvx/discretization/__init__.py,sha256=lDMVvJMuaP3s78IqwQ2N-BmPPp7_vnll-TX5cU3ezlo,2799
@@ -61,7 +61,7 @@ openscvx/symbolic/__init__.py,sha256=aBwrKg6OnmbDL9pma9k4G009QWCOjM2Lmt_6MZPF7I4
 openscvx/symbolic/augmentation.py,sha256=qZWpcFenJawRM_VIRZmoJtYOLErx0jmb8eLvSf_xCvI,29555
 openscvx/symbolic/builder.py,sha256=ivkPY9tecg6KEgguEuOHknh_FRabWXlCUjXrnWuiwe0,24631
 openscvx/symbolic/constraint_set.py,sha256=HtLKmXvji9B47l1jzd5fEJSAFAKX93vyzDwux1xxayM,3476
-openscvx/symbolic/hashing.py,sha256=u9JJQpAskuC8mPntwdmWIdEtoL0BF96a-zbqkEXJX78,4320
+openscvx/symbolic/hashing.py,sha256=OLSwuqFal9tKKtMxLAM9HBGb1bo5vzo-DG7srSy21lI,4686
 openscvx/symbolic/lower.py,sha256=zGMVRduOqzC-u3cU3Rw1y61EG-86waHdc-hP0XDh5io,31790
 openscvx/symbolic/preprocessing.py,sha256=Sqy_i5ntGWXdxcJCvql_034ZWoYU-kBRPEX34wgB344,36318
 openscvx/symbolic/problem.py,sha256=TblK3Z0orLaI_6_uPBVZP4SMxTSFXnUfNb0SICR-vFg,4853
@@ -137,9 +137,9 @@ openscvx/utils/caching.py,sha256=BPkT_IbmYT-i-BZ-himdWUc_4oBcwXWJxeUMwQWnSNc,934
 openscvx/utils/printing.py,sha256=zl3IxnKhwITqB5dK0Ru2IlORPqB61Y3DgsTryqMNu9M,13360
 openscvx/utils/profiling.py,sha256=k2x-i0CpG_kRe6dNcNBGu-ylrOtQw4B4C1UaOTjUMfU,1678
 openscvx/utils/utils.py,sha256=M25RHE_7DSr3Reaca0xCXnDSY9KHuqYvXdh5m1ZotEc,3047
-openscvx-2.dev1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-openscvx-2.dev1.dist-info/METADATA,sha256=pYPpRpTVKMuApseklvwDpBvNFFl_znVuCJw81FQC8_E,10662
-openscvx-2.dev1.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
-openscvx-2.dev1.dist-info/entry_points.txt,sha256=1Oqek8Sy28hmAZFgZXDxFXYVf56YLYWlHjhh9RYJ7wE,52
-openscvx-2.dev1.dist-info/top_level.txt,sha256=nUT4Ybefzh40H8tVXqc1RzKESy_MAowElb-CIvAbd4Q,9
-openscvx-2.dev1.dist-info/RECORD,,
+openscvx-2.dev2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+openscvx-2.dev2.dist-info/METADATA,sha256=ZS2YI58BaQKZfK5RsfxJXH8KjaHNtrYDhHk0iJzlNGU,10662
+openscvx-2.dev2.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+openscvx-2.dev2.dist-info/entry_points.txt,sha256=1Oqek8Sy28hmAZFgZXDxFXYVf56YLYWlHjhh9RYJ7wE,52
+openscvx-2.dev2.dist-info/top_level.txt,sha256=nUT4Ybefzh40H8tVXqc1RzKESy_MAowElb-CIvAbd4Q,9
+openscvx-2.dev2.dist-info/RECORD,,