openscvx 0.3.2.dev170__py3-none-any.whl

openscvx/config.py

@@ -0,0 +1,437 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+
+from openscvx.lowered.unified import UnifiedControl, UnifiedState
+
+
+def get_affine_scaling_matrices(n, minimum, maximum):
+    S = np.diag(np.maximum(np.ones(n), abs(minimum - maximum) / 2))
+    c = (maximum + minimum) / 2
+    return S, c
+
+
+@dataclass
+class DiscretizationConfig:
+    def __init__(
+        self,
+        dis_type: str = "FOH",
+        custom_integrator: bool = False,
+        solver: str = "Tsit5",
+        args: Optional[dict] = None,
+        atol: float = 1e-3,
+        rtol: float = 1e-6,
+    ):
+        """
+        Configuration class for discretization settings.
+
+        This class defines the parameters required for discretizing system dynamics.
+
+        Main arguments:
+        These are the arguments most commonly used day-to-day.
+
+        Args:
+            dis_type (str): The type of discretization to use (e.g., "FOH" for
+                First-Order Hold). Defaults to "FOH".
+            custom_integrator (bool): This enables our custom fixed-step RK45
+                algorithm. This tends to be faster than Diffrax but unless you're
+                going for speed, it's recommended to stick with Diffrax for
+                robustness and other solver options. Defaults to False.
+            solver (str): Not used if custom_integrator is enabled. Any choice of
+                solver in Diffrax is valid, please refer here,
+                [How to Choose a Solver](https://docs.kidger.site/diffrax/usage/
+                how-to-choose-a-solver/). Defaults to "Tsit5".
+
+        Other arguments:
+        These arguments are less frequently used, and for most purposes you
+        shouldn't need to understand these.
+
+        Args:
+            args (Dict): Additional arguments to pass to the solver which can be
+                found [here](https://docs.kidger.site/diffrax/api/diffeqsolve/).
+                Defaults to an empty dictionary.
+            atol (float): Absolute tolerance for the solver. Defaults to 1e-3.
+            rtol (float): Relative tolerance for the solver. Defaults to 1e-6.
+        """
+        self.dis_type = dis_type
+        self.custom_integrator = custom_integrator
+        self.solver = solver
+        self.args = args if args is not None else {}
+        self.atol = atol
+        self.rtol = rtol
+
+
+@dataclass
+class DevConfig:
+    def __init__(self, profiling: bool = False, debug: bool = False, printing: bool = True):
+        """
+        Configuration class for development settings.
+
+        This class defines the parameters used for development and debugging
+        purposes.
+
+        Main arguments:
+        These are the arguments most commonly used day-to-day.
+
+        Args:
+            profiling (bool): Whether to enable profiling for performance
+                analysis. Defaults to False.
+            debug (bool): Disables all precompilation so you can place
+                breakpoints and inspect values. Defaults to False.
+            printing (bool): Whether to enable printing during development.
+                Defaults to True.
+        """
+        self.profiling = profiling
+        self.debug = debug
+        self.printing = printing
+
+
+@dataclass
+class ConvexSolverConfig:
+    def __init__(
+        self,
+        solver: str = "QOCO",
+        solver_args: Optional[dict] = None,
+        cvxpygen: bool = False,
+        cvxpygen_override: bool = False,
+    ):
+        """
+        Configuration class for convex solver settings.
+
+        This class defines the parameters required for configuring a convex solver.
+
+        These are the arguments most commonly used day-to-day. Generally I have
+        found [QOCO](https://qoco-org.github.io/qoco/index.html) to be the most
+        performant of the CVXPY solvers for these types of problems (I do have a
+        bias as the author is from my group) and can handle up to SOCP's.
+        [CLARABEL](https://clarabel.org/stable/) is also a great option with
+        feasibility checking and can handle a few more problem types.
+        [CVXPYGen](https://github.com/cvxgrp/cvxpygen) is also great if your
+        problem isn't too large. I have found qocogen to be the most performant
+        of the CVXPYGen solvers.
+
+        Args:
+            solver (str): The name of the CVXPY solver to use. A list of options
+                can be found [here](https://www.cvxpy.org/tutorial/solvers/
+                index.html). Defaults to "QOCO".
+            solver_args (dict, optional): Ensure you are using the correct
+                arguments for your solver as they are not all common. Additional
+                arguments to configure the solver, such as tolerances. Defaults
+                to {"abstol": 1e-6, "reltol": 1e-9}.
+            cvxpygen (bool): Whether to enable CVXPY code generation for the
+                solver. Defaults to False.
+        """
+        if solver_args is None:
+            solver_args = {"abstol": 1e-06, "reltol": 1e-09, "enforce_dpp": True}
+        self.solver = solver
+        self.solver_args = (
+            solver_args if solver_args is not None else {"abstol": 1e-6, "reltol": 1e-9}
+        )
+        self.cvxpygen = cvxpygen
+        self.cvxpygen_override = cvxpygen_override
+
+
+@dataclass
+class PropagationConfig:
+    def __init__(
+        self,
+        inter_sample: int = 30,
+        dt: float = 0.01,
+        solver: str = "Dopri8",
+        max_tau_len: int = 1000,
+        args: Optional[dict] = None,
+        atol: float = 1e-3,
+        rtol: float = 1e-6,
+    ):
+        """
+        Configuration class for propagation settings.
+
+        This class defines the parameters required for propagating the nonlinear
+        system dynamics using the optimal control sequence.
+
+        Main arguments:
+        These are the arguments most commonly used day-to-day.
+
+        Other arguments:
+        The solver should likely not be changed as it is a high accuracy 8th-order
+        Runge-Kutta method.
+
+        Args:
+            inter_sample (int): How dense the propagation within multishot
+                discretization should be. Defaults to 30.
+            dt (float): The time step for propagation. Defaults to 0.1.
+            solver (str): The numerical solver to use for propagation
+                (e.g., "Dopri8"). Defaults to "Dopri8".
+            max_tau_len (int): The maximum length of the time vector for
+                propagation. Defaults to 1000.
+            args (Dict, optional): Additional arguments to pass to the solver.
+                Defaults to an empty dictionary.
+            atol (float): Absolute tolerance for the solver. Defaults to 1e-3.
+            rtol (float): Relative tolerance for the solver. Defaults to 1e-6.
+        """
+        self.inter_sample = inter_sample
+        self.dt = dt
+        self.solver = solver
+        self.max_tau_len = max_tau_len
+        self.args = args if args is not None else {}
+        self.atol = atol
+        self.rtol = rtol
+
+
+@dataclass(init=False)
+class SimConfig:
+    # No class-level field declarations
+
+    def __init__(
+        self,
+        x: UnifiedState,
+        x_prop: UnifiedState,
+        u: UnifiedControl,
+        total_time: float,
+        save_compiled: bool = False,
+        ctcs_node_intervals: Optional[list] = None,
+        n_states: Optional[int] = None,
+        n_states_prop: Optional[int] = None,
+        n_controls: Optional[int] = None,
+    ):
+        """
+        Configuration class for simulation settings.
+
+        This class defines the parameters required for simulating a trajectory
+        optimization problem.
+
+        Main arguments:
+        These are the arguments most commonly used day-to-day.
+
+        Args:
+            x (State): State object, must have .min and .max attributes for bounds.
+            x_prop (State): Propagation state object, must have .min and .max
+                attributes for bounds.
+            u (Control): Control object, must have .min and .max attributes for
+                bounds.
+            total_time (float): The total simulation time.
+            idx_x_true (slice): Slice for true state indices.
+            idx_x_true_prop (slice): Slice for true propagation state indices.
+            idx_u_true (slice): Slice for true control indices.
+            idx_t (slice): Slice for time index.
+            idx_y (slice): Slice for constraint violation indices.
+            idx_y_prop (slice): Slice for propagation constraint violation
+                indices.
+            idx_s (slice): Slice for time dilation index.
+            save_compiled (bool): If True, save and reuse compiled solver
+                functions. Defaults to False.
+            ctcs_node_intervals (list, optional): Node intervals for CTCS
+                constraints.
+            n_states (int, optional): The number of state variables. Defaults to
+                `None` (inferred from x.max).
+            n_states_prop (int, optional): The number of propagation state
+                variables. Defaults to `None` (inferred from x_prop.max).
+            n_controls (int, optional): The number of control variables. Defaults
+                to `None` (inferred from u.max).
+
+        Note:
+            You can specify custom scaling for specific states/controls using
+            the `scaling_min` and `scaling_max` attributes on State, Control, and Time objects.
+            If not set, the default min/max bounds will be used for scaling.
+        """
+        # Assign core arguments to self
+        self.x = x
+        self.x_prop = x_prop
+        self.u = u
+        self.total_time = total_time
+        self.save_compiled = save_compiled
+        self.ctcs_node_intervals = ctcs_node_intervals
+        self.n_states = n_states
+        self.n_states_prop = n_states_prop
+        self.n_controls = n_controls
+
+        # Call post init logic
+        self.__post_init__()
+
+    def __post_init__(self):
+        self.n_states = len(self.x.max)
+        self.n_controls = len(self.u.max)
+
+        # State scaling
+        # Use scaling_min/max if provided, otherwise use regular min/max
+        min_x = np.array(self.x.min, dtype=float)
+        max_x = np.array(self.x.max, dtype=float)
+
+        # UnifiedState now always provides full-size scaling arrays when any state has scaling
+        if self.x.scaling_min is not None:
+            lower_x = np.array(self.x.scaling_min, dtype=float)
+        else:
+            lower_x = min_x
+
+        if self.x.scaling_max is not None:
+            upper_x = np.array(self.x.scaling_max, dtype=float)
+        else:
+            upper_x = max_x
+
+        S_x, c_x = get_affine_scaling_matrices(self.n_states, lower_x, upper_x)
+        self.S_x = S_x
+        self.c_x = c_x
+        self.inv_S_x = np.diag(1 / np.diag(self.S_x))
+
+        # Control scaling
+        # Use scaling_min/max if provided, otherwise use regular min/max
+        min_u = np.array(self.u.min, dtype=float)
+        max_u = np.array(self.u.max, dtype=float)
+
+        # UnifiedControl now always provides full-size scaling arrays when any control has scaling
+        if self.u.scaling_min is not None:
+            lower_u = np.array(self.u.scaling_min, dtype=float)
+        else:
+            lower_u = min_u
+
+        if self.u.scaling_max is not None:
+            upper_u = np.array(self.u.scaling_max, dtype=float)
+        else:
+            upper_u = max_u
+
+        S_u, c_u = get_affine_scaling_matrices(self.n_controls, lower_u, upper_u)
+        self.S_u = S_u
+        self.c_u = c_u
+        self.inv_S_u = np.diag(1 / np.diag(self.S_u))
+
+    # Properties for accessing slices from unified objects
+    @property
+    def time_slice(self):
+        """Slice for accessing time in the state vector."""
+        return self.x.time_slice
+
+    @property
+    def ctcs_slice(self):
+        """Slice for accessing CTCS augmented states."""
+        return self.x.ctcs_slice
+
+    @property
+    def ctcs_slice_prop(self):
+        """Slice for accessing CTCS augmented states in propagation."""
+        return self.x_prop.ctcs_slice
+
+    @property
+    def time_dilation_slice(self):
+        """Slice for accessing time dilation in the control vector."""
+        return self.u.time_dilation_slice
+
+    @property
+    def true_state_slice(self):
+        """Slice for accessing true (non-augmented) states."""
+        return self.x._true_slice
+
+    @property
+    def true_state_slice_prop(self):
+        """Slice for accessing true (non-augmented) propagation states."""
+        return self.x_prop._true_slice
+
+    @property
+    def true_control_slice(self):
+        """Slice for accessing true (non-augmented) controls."""
+        return self.u._true_slice
+
+
+@dataclass
+class ScpConfig:
+    def __init__(
+        self,
+        n: Optional[int] = None,
+        k_max: int = 200,
+        w_tr: float = 1.0,
+        lam_vc: float = 1.0,
+        ep_tr: float = 1e-4,
+        ep_vb: float = 1e-4,
+        ep_vc: float = 1e-8,
+        lam_cost: float = 0.0,
+        lam_vb: float = 0.0,
+        uniform_time_grid: bool = False,
+        cost_drop: int = -1,
+        cost_relax: float = 1.0,
+        w_tr_adapt: float = 1.0,
+        w_tr_max: Optional[float] = None,
+        w_tr_max_scaling_factor: Optional[float] = None,
+    ):
+        """
+        Configuration class for Sequential Convex Programming (SCP).
+
+        This class defines the parameters used to configure the SCP solver. You
+        will very likely need to modify the weights for your problem. Please
+        refer to my guide [here](https://haynec.github.io/openscvx/
+        hyperparameter_tuning) for more information.
+
+        Attributes:
+            n (int): The number of discretization nodes. Defaults to `None`.
+            k_max (int): The maximum number of SCP iterations. Defaults to 200.
+            w_tr (float): The trust region weight. Defaults to 1.0.
+            lam_vc (float): The penalty weight for virtual control. Defaults to 1.0.
+            ep_tr (float): The trust region convergence tolerance. Defaults to 1e-4.
+            ep_vb (float): The boundary constraint convergence tolerance.
+                Defaults to 1e-4.
+            ep_vc (float): The virtual constraint convergence tolerance.
+                Defaults to 1e-8.
+            lam_cost (float): The weight for original cost. Defaults to 0.0.
+            lam_vb (float): The weight for virtual buffer. This is only used if
+                there are nonconvex nodal constraints present. Defaults to 0.0.
+            uniform_time_grid (bool): Whether to use a uniform time grid.
+                Defaults to `False`.
+            cost_drop (int): The number of iterations to allow for cost
+                stagnation before termination. Defaults to -1 (disabled).
+            cost_relax (float): The relaxation factor for cost reduction.
+                Defaults to 1.0.
+            w_tr_adapt (float): The adaptation factor for the trust region
+                weight. Defaults to 1.0.
+            w_tr_max (float): The maximum allowable trust region weight.
+                Defaults to `None`.
+            w_tr_max_scaling_factor (float): The scaling factor for the maximum
+                trust region weight. Defaults to `None`.
+        """
+        self.n = n
+        self.k_max = k_max
+        self.w_tr = w_tr
+        self.lam_vc = lam_vc
+        self.ep_tr = ep_tr
+        self.ep_vb = ep_vb
+        self.ep_vc = ep_vc
+        self.lam_cost = lam_cost
+        self.lam_vb = lam_vb
+        self.uniform_time_grid = uniform_time_grid
+        self.cost_drop = cost_drop
+        self.cost_relax = cost_relax
+        self.w_tr_adapt = w_tr_adapt
+        self.w_tr_max = w_tr_max
+        self.w_tr_max_scaling_factor = w_tr_max_scaling_factor
+
+    def __post_init__(self):
+        keys_to_scale = ["w_tr", "lam_vc", "lam_cost", "lam_vb"]
+        # Handle lam_vc which might be scalar or array
+        scale_values = []
+        for key in keys_to_scale:
+            val = getattr(self, key)
+            if isinstance(val, np.ndarray):
+                scale_values.append(np.max(val))
+            else:
+                scale_values.append(val)
+        scale = max(scale_values)
+        for key in keys_to_scale:
+            val = getattr(self, key)
+            if isinstance(val, np.ndarray):
+                setattr(self, key, val / scale)
+            else:
+                setattr(self, key, val / scale)
+
+        if self.w_tr_max_scaling_factor is not None and self.w_tr_max is None:
+            self.w_tr_max = self.w_tr_max_scaling_factor * self.w_tr
+
+
+@dataclass
+class Config:
+    sim: SimConfig
+    scp: ScpConfig
+    cvx: ConvexSolverConfig
+    dis: DiscretizationConfig
+    prp: PropagationConfig
+    dev: DevConfig
+
+    def __post_init__(self):
+        pass

openscvx/discretization/__init__.py

@@ -0,0 +1,47 @@
+"""Discretization methods for trajectory optimization.
+
+This module provides implementations of discretization schemes that convert
+continuous-time optimal control problems into discrete-time approximations
+suitable for numerical optimization. Discretization is a critical step in
+trajectory optimization that linearizes the nonlinear dynamics around a
+reference trajectory.
+
+Planned Architecture (ABC-based):
+
+A base class will be introduced to enable pluggable discretization methods.
+This will enable users to implement custom discretization methods.
+Future discretizers will implement the Discretizer interface:
+
+```python
+# discretization/base.py (planned):
+class Discretizer(ABC):
+    def __init__(self, integrator: Integrator):
+        '''Initialize with a numerical integrator.'''
+        self.integrator = integrator
+
+    @abstractmethod
+    def discretize(self, dynamics, x, u, dt) -> tuple[A_d, B_d, C_d]:
+        '''Discretize continuous dynamics around trajectory (x, u).
+
+        Args:
+            dynamics: Continuous-time dynamics object
+            x: State trajectory
+            u: Control trajectory
+            dt: Time step
+
+        Returns:
+            A_d: Discretized state transition matrix
+            B_d: Discretized control influence matrix (current node)
+            C_d: Discretized control influence matrix (next node)
+        '''
+        ...
+```
+"""
+
+from .discretization import calculate_discretization, dVdt, get_discretization_solver
+
+__all__ = [
+    "calculate_discretization",
+    "get_discretization_solver",
+    "dVdt",
+]