openscvx 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

openscvx/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.2'
-__version_tuple__ = version_tuple = (0, 1, 2)
+__version__ = version = '0.1.3'
+__version_tuple__ = version_tuple = (0, 1, 3)

openscvx/config.py

@@ -2,6 +2,8 @@ import numpy as np
 from dataclasses import dataclass, field
 from typing import Dict, List
 
+from openscvx.constraints.boundary import BoundaryConstraint
+
 
 def get_affine_scaling_matrices(n, minimum, maximum):
     S = np.diag(np.maximum(np.ones(n), abs(minimum - maximum) / 2))
@@ -120,7 +122,8 @@ class PropagationConfig:
 class SimConfig:
     x_bar: np.ndarray
     u_bar: np.ndarray
-    initial_state: np.ndarray
+    initial_state: BoundaryConstraint
+    initial_state_prop: BoundaryConstraint
     final_state: np.ndarray
     max_state: np.ndarray
     min_state: np.ndarray
@@ -128,9 +131,11 @@ class SimConfig:
     min_control: np.ndarray
     total_time: float
     idx_x_true: slice
+    idx_x_true_prop: slice
     idx_u_true: slice
     idx_t: slice
     idx_y: slice
+    idx_y_prop: slice
     idx_s: slice
     ctcs_node_intervals: list = None
     constraints_ctcs: List[callable] = field(
@@ -138,6 +143,7 @@ class SimConfig:
     )  # TODO (norrisg): clean this up, consider moving to dedicated `constraints` dataclass
     constraints_nodal: List[callable] = field(default_factory=list)
     n_states: int = None
+    n_states_prop: int = None
     n_controls: int = None
     S_x: np.ndarray = None
     inv_S_x: np.ndarray = None

openscvx/post_processing.py

@@ -1,4 +1,5 @@
 import numpy as np
+import jax.numpy as jnp
 
 from openscvx.propagation import s_to_t, t_to_tau, simulate_nonlinear_time
 from openscvx.config import Config
@@ -14,9 +15,13 @@ def propagate_trajectory_results(params: Config, result: dict, propagation_solve
 
     tau_vals, u_full = t_to_tau(u, t_full, u, t, params)
 
-    x_full = simulate_nonlinear_time(x[0], u, tau_vals, t, params, propagation_solver)
+    # Match free values from initial state to the initial value from the result
+    mask = jnp.array([t == "Free" for t in params.sim.initial_state_prop.types], dtype=bool)
+    params.sim.initial_state_prop.value = jnp.where(mask, x[0], params.sim.initial_state_prop.value)
 
-    print("Total CTCS Constraint Violation:", x_full[-1, params.sim.idx_y])
+    x_full = simulate_nonlinear_time(params.sim.initial_state_prop.value, u, tau_vals, t, params, propagation_solver)
+
+    print("Total CTCS Constraint Violation:", x_full[-1, params.sim.idx_y_prop])
     i = 0
     cost = np.zeros_like(x[-1, i])
     for type in params.sim.initial_state.type:
@@ -28,6 +33,16 @@ def propagate_trajectory_results(params: Config, result: dict, propagation_solve
         if type == "Minimize":
             cost += x[-1, i]
         i += 1
+    i=0
+    for type in params.sim.initial_state.type:
+        if type == "Maximize":
+            cost -= x[0, i]
+        i += 1
+    i = 0
+    for type in params.sim.final_state.type:
+        if type == "Maximize":
+            cost -= x[-1, i]
+        i += 1
     print("Cost: ", cost)
 
     more_result = dict(t_full=t_full, x_full=x_full, u_full=u_full)

openscvx/trajoptproblem.py

@@ -1,5 +1,5 @@
 import jax.numpy as jnp
-from typing import List, Union
+from typing import List, Union, Optional
 import queue
 import threading
 import time
@@ -23,7 +23,7 @@ from openscvx.augmentation.ctcs import sort_ctcs_constraints
 from openscvx.constraints.violation import get_g_funcs, CTCSViolation
 from openscvx.discretization import get_discretization_solver
 from openscvx.propagation import get_propagation_solver
-from openscvx.constraints.boundary import BoundaryConstraint
+from openscvx.constraints.boundary import BoundaryConstraint, boundary
 from openscvx.constraints.ctcs import CTCSConstraint
 from openscvx.constraints.nodal import NodalConstraint
 from openscvx.ptr import PTR_init, PTR_main
@@ -49,17 +49,24 @@ class TrajOptProblem:
         x_min: jnp.ndarray,
         u_max: jnp.ndarray,
         u_min: jnp.ndarray,
-        scp: ScpConfig = None,
-        dis: DiscretizationConfig = None,
-        prp: PropagationConfig = None,
-        sim: SimConfig = None,
-        dev: DevConfig = None,
-        cvx: ConvexSolverConfig = None,
+        dynamics_prop: callable = None,
+        initial_state_prop: BoundaryConstraint = None,
+        scp: Optional[ScpConfig] = None,
+        dis: Optional[DiscretizationConfig] = None,
+        prp: Optional[PropagationConfig] = None,
+        sim: Optional[SimConfig] = None,
+        dev: Optional[DevConfig] = None,
+        cvx: Optional[ConvexSolverConfig] = None,
         licq_min=0.0,
         licq_max=1e-4,
         time_dilation_factor_min=0.3,
         time_dilation_factor_max=3.0,
     ):
+        if dynamics_prop is None:
+            dynamics_prop = dynamics
+        
+        if initial_state_prop is None:
+            initial_state_prop = initial_state
 
         # TODO (norrisg) move this into some augmentation function, if we want to make this be executed after the init (i.e. within problem.initialize) need to rethink how problem is defined
         constraints_ctcs = []
@@ -81,11 +88,15 @@ class TrajOptProblem:
         constraints_ctcs, node_intervals, num_augmented_states = sort_ctcs_constraints(constraints_ctcs, N)
 
         # Index tracking
-        idx_x_true = slice(0, len(x_max))
+        idx_x_true = slice(0, len(initial_state.value))
+        idx_x_true_prop = slice(0, len(initial_state_prop.value))
         idx_u_true = slice(0, len(u_max))
         idx_constraint_violation = slice(
             idx_x_true.stop, idx_x_true.stop + num_augmented_states
         )
+        idx_constraint_violation_prop = slice(
+            idx_x_true_prop.stop, idx_x_true_prop.stop + num_augmented_states
+        )
 
         idx_time_dilation = slice(idx_u_true.stop, idx_u_true.stop + 1)
 
@@ -106,6 +117,11 @@ class TrajOptProblem:
             [u_guess, np.full((u_guess.shape[0], 1), time_init)]
         )
 
+        initial_state_prop_values = np.hstack([initial_state_prop.value, np.repeat(licq_min, num_augmented_states)])
+        initial_state_prop_types = np.hstack([initial_state_prop.type, ["Fix"] * num_augmented_states])
+        initial_state_prop = boundary(initial_state_prop_values)
+        initial_state_prop.types = initial_state_prop_types
+
         if dis is None:
             dis = DiscretizationConfig()
 
@@ -114,17 +130,21 @@ class TrajOptProblem:
                 x_bar=x_bar_augmented,
                 u_bar=u_bar_augmented,
                 initial_state=initial_state,
+                initial_state_prop=initial_state_prop,
                 final_state=final_state,
                 max_state=x_max_augmented,
                 min_state=x_min_augmented,
                 max_control=u_max_augmented,
                 min_control=u_min_augmented,
                 total_time=time_init,
-                n_states=len(x_max),
+                n_states=len(initial_state.value),
+                n_states_prop=len(initial_state_prop.value),
                 idx_x_true=idx_x_true,
+                idx_x_true_prop=idx_x_true_prop,
                 idx_u_true=idx_u_true,
                 idx_t=idx_time,
                 idx_y=idx_constraint_violation,
+                idx_y_prop=idx_constraint_violation_prop,
                 idx_s=idx_time_dilation,
                 ctcs_node_intervals=node_intervals,
             )
@@ -162,6 +182,7 @@ class TrajOptProblem:
 
         ctcs_violation_funcs = get_g_funcs(constraints_ctcs)
         self.dynamics_augmented = build_augmented_dynamics(dynamics, ctcs_violation_funcs, idx_x_true, idx_u_true)
+        self.dynamics_augmented_prop = build_augmented_dynamics(dynamics_prop, ctcs_violation_funcs, idx_x_true_prop, idx_u_true)
 
         self.params = Config(
             sim=sim,
@@ -215,6 +236,9 @@ class TrajOptProblem:
         self.dynamics_augmented.A = jax.jit(jax.vmap(self.dynamics_augmented.A, in_axes=(0, 0, 0)))
         self.dynamics_augmented.B = jax.jit(jax.vmap(self.dynamics_augmented.B, in_axes=(0, 0, 0)))
 
+
+        self.dynamics_augmented_prop.f = jax.vmap(self.dynamics_augmented_prop.f)
+
         for constraint in self.params.sim.constraints_nodal:
             if not constraint.convex:
                 # TODO: (haynec) switch to AOT instead of JIT
@@ -224,7 +248,7 @@ class TrajOptProblem:
 
         # Generate solvers and optimal control problem
         self.discretization_solver = get_discretization_solver(self.dynamics_augmented, self.params)
-        self.propagation_solver = get_propagation_solver(self.dynamics_augmented.f, self.params)
+        self.propagation_solver = get_propagation_solver(self.dynamics_augmented_prop.f, self.params)
         self.optimal_control_problem = OptimalControlProblem(self.params)
 
         # Initialize the PTR loop
@@ -248,7 +272,7 @@ class TrajOptProblem:
         self.propagation_solver = (
             jax.jit(self.propagation_solver)
             .lower(
-                np.ones((self.params.sim.n_states)),
+                np.ones((self.params.sim.n_states_prop)),
                 (0.0, 0.0),
                 np.ones((1, self.params.sim.n_controls)),
                 np.ones((1, self.params.sim.n_controls)),

{openscvx-0.1.2.dist-info → openscvx-0.1.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.1.2
+Version: 0.1.3
 Summary: A general Python-based successive convexification implementation which uses a JAX backend.
 Home-page: https://haynec.github.io/openscvx/
 Author: Chris Hayner and Griffin Norris

{openscvx-0.1.2.dist-info → openscvx-0.1.3.dist-info}/RECORD

@@ -1,16 +1,16 @@
 openscvx/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-openscvx/_version.py,sha256=bSmADqydH8nBu-J4lG8UVuR7hnU_zcwhnSav2oQ0W0A,511
-openscvx/config.py,sha256=8Cl5O0ekf9MGNDTEeMHsp1C4XvY9NfJQkxd80lvnafM,10296
+openscvx/_version.py,sha256=NIzzV8ZM0W-CSLuEs1weG4zPrn_-8yr1AwwI1iuS6yo,511
+openscvx/config.py,sha256=lTdvJtU-Xiwq55zGtW3QvmldZj9Ne0xyKsSDSIKwqiI,10488
 openscvx/discretization.py,sha256=YF3mEeyYHgyTWQVNQsqpi1Mv72zDLyNfaMJSWqxj34c,4745
 openscvx/dynamics.py,sha256=X9sPpxUGGbdsnvQzgyrb_939N9ctBSsWVyI1eXtOKpc,1118
 openscvx/integrators.py,sha256=msIS-1Ehj-9TJLHfoCMs3vdyZ8NXz-TM0RII6aqRf4E,3821
 openscvx/io.py,sha256=fOvNWQWAegcN1gejeToaNbXenP5H5bAifNU8edJvdk4,4127
 openscvx/ocp.py,sha256=L_509EQiMsI6s5gBYlYyxKaHEzzRdpo-XAMjliCU3Rc,7544
 openscvx/plotting.py,sha256=fCvWJV4qWMhVyJlh18s12S_5xhj6EviF-_FuP0tWjx4,31207
-openscvx/post_processing.py,sha256=TP1gi4TVlDS2HHpdqaIPCqfM5o4w7a7RCMU3Pu3czHw,1024
+openscvx/post_processing.py,sha256=t3fUeDfA2PZq8S-WVnxDyN0xIgFuvpDP-5wbGQuIOVY,1618
 openscvx/propagation.py,sha256=XNezQnAM-NXb9L7aHUgKQOBn0CNUPeGGDL3_BbGoODU,3758
 openscvx/ptr.py,sha256=itDTR6RQUphnU226jaeRaAKuia-6v8U3MqAdw5-BYOk,5268
-openscvx/trajoptproblem.py,sha256=3yufy-egU7m0NV834TH8csY1HJqM90Is7VYw0gQe3pk,11996
+openscvx/trajoptproblem.py,sha256=hcju3UB3iZ5peqhxh6NvDQYnsArnabR0mGUAKZ3sIaw,13360
 openscvx/utils.py,sha256=zmkKyto8Jowe_RAdOe8K0w6gzOu4JfxmX1RUL-3OFlY,2408
 openscvx/augmentation/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 openscvx/augmentation/ctcs.py,sha256=m1jdALXSqHq3WD6lCBAUI7FR0Sfs8aCYr66h0EwE4z4,1707
@@ -20,8 +20,8 @@ openscvx/constraints/boundary.py,sha256=yEhEnkKJ5f8NUeTksigEJjgBeE_YyuG_PJb_DWxg
 openscvx/constraints/ctcs.py,sha256=V763033aV82nAu7y4653KsAs11A7RpUysR_oUcnLfko,2572
 openscvx/constraints/nodal.py,sha256=YCS0cwUurA2OTQcHBb1EQqLxNt_w3MX8Nj8FH3GYClo,1726
 openscvx/constraints/violation.py,sha256=aIdDhHd-UndT0XB2QeuwLBKSNSAUWVkha_GeHOw9cQg,2362
-openscvx-0.1.2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-openscvx-0.1.2.dist-info/METADATA,sha256=MDHeKrpE_3FKRiQD5fVKVzNBWerOvcY0vfapGSRTlbk,6911
-openscvx-0.1.2.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
-openscvx-0.1.2.dist-info/top_level.txt,sha256=nUT4Ybefzh40H8tVXqc1RzKESy_MAowElb-CIvAbd4Q,9
-openscvx-0.1.2.dist-info/RECORD,,
+openscvx-0.1.3.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+openscvx-0.1.3.dist-info/METADATA,sha256=CVr7tDDV75qDlnuEuxf-IL0Arxkb1RXwv0Xj9oVxf_s,6911
+openscvx-0.1.3.dist-info/WHEEL,sha256=QZxptf4Y1BKFRCEDxD4h2V0mBFQOVFLFEpvxHmIs52A,91
+openscvx-0.1.3.dist-info/top_level.txt,sha256=nUT4Ybefzh40H8tVXqc1RzKESy_MAowElb-CIvAbd4Q,9
+openscvx-0.1.3.dist-info/RECORD,,

{openscvx-0.1.2.dist-info → openscvx-0.1.3.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.4.0)
+Generator: setuptools (80.6.0)
 Root-Is-Purelib: true
 Tag: py3-none-any