openscvx 0.3.2.dev170__py3-none-any.whl

openscvx/symbolic/problem.py

@@ -0,0 +1,110 @@
+"""SymbolicProblem dataclass - container for symbolic problem specification.
+
+This module provides the SymbolicProblem dataclass that represents a trajectory
+optimization problem in symbolic form, before lowering to executable code.
+
+The SymbolicProblem can represent two lifecycle stages:
+
+1. **Before preprocessing**: Raw user input with unsorted constraints
+2. **After preprocessing**: Augmented and validated, ready for lowering
+
+Use `is_preprocessed` to check which stage the problem is in.
+"""
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+
+from openscvx.symbolic.constraint_set import ConstraintSet
+
+if TYPE_CHECKING:
+    from openscvx.symbolic.expr import Expr
+    from openscvx.symbolic.expr.control import Control
+    from openscvx.symbolic.expr.state import State
+
+
+@dataclass
+class SymbolicProblem:
+    """Container for symbolic problem specification.
+
+    This dataclass holds a trajectory optimization problem in symbolic form,
+    either as raw user input or after preprocessing/augmentation. It provides
+    a typed interface for the preprocessing and lowering pipeline.
+
+    Lifecycle Stages:
+        1. **Before preprocessing**: User creates with raw dynamics, states,
+           controls, and unsorted constraints. Propagation fields are None.
+        2. **After preprocessing**: Dynamics and states are augmented (CTCS,
+           time dilation), constraints are categorized, propagation fields
+           are populated.
+
+    Use `is_preprocessed` to check whether preprocessing has completed.
+
+    Attributes:
+        dynamics: Symbolic dynamics expression (dx/dt = f(x, u)).
+            After preprocessing, includes CTCS augmented state dynamics.
+        states: List of State objects. After preprocessing, includes
+            time state and CTCS augmented states.
+        controls: List of Control objects. After preprocessing, includes
+            time dilation control.
+        constraints: ConstraintSet holding all constraints. Before preprocessing,
+            raw constraints live in `constraints.unsorted`. After preprocessing,
+            constraints are categorized into ctcs, nodal, nodal_convex, etc.
+        parameters: Dictionary mapping parameter names to numpy arrays.
+        N: Number of discretization nodes.
+        node_intervals: List of (start, end) tuples for CTCS constraint intervals.
+            Populated during preprocessing when CTCS constraints are sorted.
+
+        dynamics_prop: Propagation dynamics (may include extra states).
+            None before preprocessing, populated after.
+        states_prop: Propagation states (may include extra states).
+            None before preprocessing, populated after.
+        controls_prop: Propagation controls (typically same as controls).
+            None before preprocessing, populated after.
+
+    Example:
+        Before preprocessing::
+
+            problem = SymbolicProblem(
+                dynamics=dynamics_expr,
+                states=[x, v],
+                controls=[u],
+                constraints=ConstraintSet(unsorted=[c1, c2, c3]),
+                parameters={"mass": 1.0},
+                N=50,
+            )
+            assert not problem.is_preprocessed
+
+        After preprocessing::
+
+            processed = preprocess_symbolic_problem(problem, time=time_config)
+            assert processed.is_preprocessed
+            assert processed.constraints.is_categorized
+            # Now ready for lowering
+            lowered = lower_symbolic_problem(processed)
+    """
+
+    # Core problem specification
+    dynamics: "Expr"
+    states: List["State"]
+    controls: List["Control"]
+    constraints: ConstraintSet
+    parameters: Dict[str, any]
+    N: int
+
+    # CTCS node intervals (populated during preprocessing)
+    node_intervals: List[Tuple[int, int]] = field(default_factory=list)
+
+    # Propagation (None before preprocessing, populated after)
+    dynamics_prop: Optional["Expr"] = None
+    states_prop: Optional[List["State"]] = None
+    controls_prop: Optional[List["Control"]] = None
+
+    @property
+    def is_preprocessed(self) -> bool:
+        """True if the problem has been preprocessed and is ready for lowering.
+
+        A problem is considered preprocessed when:
+        1. All constraints have been categorized (unsorted is empty)
+        2. Propagation dynamics have been set up
+        """
+        return self.constraints.is_categorized and self.dynamics_prop is not None

openscvx/symbolic/time.py

@@ -0,0 +1,116 @@
+from typing import Union
+
+
+class Time:
+    """Time configuration for trajectory optimization problems.
+
+    This class encapsulates time-related parameters for trajectory optimization.
+    The time derivative is internally assumed to be 1.0.
+
+    Attributes:
+        initial (float or tuple): Initial time boundary condition.
+            Can be a float (fixed) or tuple like ("free", value), ("minimize", value),
+            or ("maximize", value).
+        final (float or tuple): Final time boundary condition.
+            Can be a float (fixed) or tuple like ("free", value), ("minimize", value),
+            or ("maximize", value).
+        min (float): Minimum bound for time variable (required).
+        max (float): Maximum bound for time variable (required).
+
+    Example:
+        ```python
+        # Fixed initial and final time
+        time = Time(initial=0.0, final=10.0, min=0.0, max=20.0)
+
+        # Free final time
+        time = Time(initial=0.0, final=("free", 10.0), min=0.0, max=20.0)
+
+        # Minimize final time
+        time = Time(initial=0.0, final=("minimize", 10.0), min=0.0, max=20.0)
+
+        # Maximize initial time
+        time = Time(initial=("maximize", 0.0), final=10.0, min=0.0, max=20.0)
+        ```
+    """
+
+    def __init__(
+        self,
+        initial: Union[float, tuple],
+        final: Union[float, tuple],
+        min: float,
+        max: float,
+    ):
+        """Initialize a Time object.
+
+        Args:
+            initial: Initial time boundary condition (float or tuple).
+                Tuple format: ("free", value), ("minimize", value), or ("maximize", value).
+            final: Final time boundary condition (float or tuple).
+                Tuple format: ("free", value), ("minimize", value), or ("maximize", value).
+            min: Minimum bound for time variable (required).
+            max: Maximum bound for time variable (required).
+
+        Raises:
+            ValueError: If tuple format is invalid.
+        """
+        # Validate tuple format if provided
+        for name, value in [("initial", initial), ("final", final)]:
+            if isinstance(value, tuple):
+                if len(value) != 2:
+                    raise ValueError(f"{name} tuple must have exactly 2 elements: (type, value)")
+                bc_type, bc_value = value
+                if bc_type not in ["free", "minimize", "maximize"]:
+                    raise ValueError(
+                        f"{name} boundary condition type must be 'free', "
+                        f"'minimize', or 'maximize', got '{bc_type}'"
+                    )
+                if not isinstance(bc_value, (int, float)):
+                    raise ValueError(
+                        f"{name} boundary condition value must be a number, "
+                        f"got {type(bc_value).__name__}"
+                    )
+
+        self.initial = initial
+        self.final = final
+        self.min = min
+        self.max = max
+        # Time derivative is always 1.0 internally
+        self.derivative = 1.0
+        self._scaling_min = None
+        self._scaling_max = None
+
+    @property
+    def scaling_min(self):
+        """Get the scaling minimum bound for the time variable.
+
+        Returns:
+            Scaling minimum value, or None if not set.
+        """
+        return self._scaling_min
+
+    @scaling_min.setter
+    def scaling_min(self, val):
+        """Set the scaling minimum bound for the time variable.
+
+        Args:
+            val: Scaling minimum value (float or None)
+        """
+        self._scaling_min = float(val) if val is not None else None
+
+    @property
+    def scaling_max(self):
+        """Get the scaling maximum bound for the time variable.
+
+        Returns:
+            Scaling maximum value, or None if not set.
+        """
+        return self._scaling_max
+
+    @scaling_max.setter
+    def scaling_max(self, val):
+        """Set the scaling maximum bound for the time variable.
+
+        Args:
+            val: Scaling maximum value (float or None)
+        """
+        self._scaling_max = float(val) if val is not None else None

openscvx/symbolic/unified.py

@@ -0,0 +1,420 @@
+"""Unification functions for aggregating symbolic State and Control objects.
+
+This module provides the unification layer that transforms multiple symbolic State
+and Control objects into unified representations for numerical optimization.
+
+The unification process:
+    1. **Collection**: Gathers all State and Control objects from expression trees
+    2. **Sorting**: Organizes variables (user-defined first, then augmented)
+    3. **Aggregation**: Concatenates bounds, guesses, and boundary conditions
+    4. **Slice Assignment**: Assigns each State/Control a slice for indexing
+    5. **Unified Representation**: Creates UnifiedState/UnifiedControl objects
+
+This separation allows users to define problems with natural variable names
+while maintaining efficient vectorized operations during optimization.
+
+Example:
+    Creating and unifying multiple states::
+
+        import openscvx as ox
+        from openscvx.symbolic.unified import unify_states
+
+        # Define separate symbolic states
+        position = ox.State("position", shape=(3,), min=-10, max=10)
+        velocity = ox.State("velocity", shape=(3,), min=-5, max=5)
+        mass = ox.State("mass", shape=(1,), min=0.1, max=10.0)
+
+        # Unify into single state vector
+        unified_x = unify_states([position, velocity, mass], name="x")
+
+        # Access unified properties
+        print(unified_x.shape)    # (7,) - combined shape
+        print(unified_x.min)      # Combined bounds: [-10, -10, -10, -5, -5, -5, 0.1]
+        print(unified_x.true)     # Access only user-defined states
+
+    Accessing slices after unification::
+
+        # After unification, each State has a slice assigned
+        print(position._slice)    # slice(0, 3)
+        print(velocity._slice)    # slice(3, 6)
+        print(mass._slice)        # slice(6, 7)
+
+        # During lowering, these slices extract values from unified vector
+        x_unified = jnp.array([1, 2, 3, 4, 5, 6, 7])
+        position_val = x_unified[position._slice]  # [1, 2, 3]
+
+See Also:
+    - UnifiedState: Dataclass for unified state representation (in openscvx.lowered.unified)
+    - UnifiedControl: Dataclass for unified control representation (in openscvx.lowered.unified)
+    - State: Individual symbolic state variable (symbolic/expr/state.py)
+    - Control: Individual symbolic control variable (symbolic/expr/control.py)
+"""
+
+from typing import List
+
+import numpy as np
+
+from openscvx.lowered.unified import UnifiedControl, UnifiedState
+from openscvx.symbolic.expr.control import Control
+from openscvx.symbolic.expr.state import State
+
+# Re-export for backwards compatibility
+__all__ = ["unify_states", "unify_controls", "UnifiedState", "UnifiedControl"]
+
+
+def unify_states(states: List[State], name: str = "unified_state") -> UnifiedState:
+    """Create a UnifiedState from a list of State objects.
+
+    This function is the primary way to aggregate multiple symbolic State objects into
+    a single unified state vector for numerical optimization. It:
+
+    1. Sorts states (user-defined first, augmented states second)
+    2. Concatenates all state properties (bounds, guesses, boundary conditions)
+    3. Assigns slices to each State for extracting values from unified vector
+    4. Identifies special states (time, CTCS augmented states)
+    5. Returns a UnifiedState with all aggregated data
+
+    Args:
+        states (List[State]): List of State objects to unify. Can include both
+            user-defined states and augmented states (names starting with '_').
+        name (str): Name identifier for the unified state vector (default: "unified_state")
+
+    Returns:
+        UnifiedState: Unified state object containing:
+            - Aggregated bounds, guesses, and boundary conditions
+            - Shape equal to sum of all state shapes
+            - Slices for extracting individual state components
+            - Properties for accessing true vs augmented states
+
+    Example:
+        Basic unification::
+
+            import openscvx as ox
+            from openscvx.symbolic.unified import unify_states
+
+            position = ox.State("pos", shape=(3,), min=-10, max=10)
+            velocity = ox.State("vel", shape=(3,), min=-5, max=5)
+
+            unified = unify_states([position, velocity], name="x")
+            print(unified.shape)       # (6,)
+            print(unified._true_dim)   # 6 (all are user states)
+            print(position._slice)     # slice(0, 3) - assigned during unification
+            print(velocity._slice)     # slice(3, 6)
+
+        With augmented states::
+
+            # CTCS or other features may add augmented states
+            time_state = ox.State("time", shape=(1,))
+            ctcs_aug = ox.State("_ctcs_aug_0", shape=(2,))  # Augmented state
+
+            unified = unify_states([position, velocity, time_state, ctcs_aug])
+            print(unified._true_dim)         # 7 (pos + vel + time)
+            print(unified.true.shape)        # (7,)
+            print(unified.augmented.shape)   # (2,) - only CTCS augmented
+
+    Note:
+        After unification, each State object has its `_slice` attribute set,
+        which is used during JAX lowering to extract the correct values from
+        the unified state vector.
+
+    See Also:
+        - UnifiedState: Return type with detailed documentation
+        - unify_controls(): Analogous function for Control objects
+        - State: Individual symbolic state variable
+    """
+    if not states:
+        return UnifiedState(name=name, shape=(0,))
+
+    # Sort states: true states (not starting with '_') first, then augmented states
+    # (starting with '_')
+    true_states = [state for state in states if not state.name.startswith("_")]
+    augmented_states = [state for state in states if state.name.startswith("_")]
+    sorted_states = true_states + augmented_states
+
+    # Calculate total shape
+    total_shape = sum(state.shape[0] for state in sorted_states)
+
+    # Concatenate all arrays, handling None values properly
+    min_arrays = []
+    max_arrays = []
+    guess_arrays = []
+    initial_arrays = []
+    final_arrays = []
+    _initial_arrays = []
+    _final_arrays = []
+    initial_type_arrays = []
+    final_type_arrays = []
+
+    for state in sorted_states:
+        if state.min is not None:
+            min_arrays.append(state.min)
+        else:
+            # If min is None, fill with -inf for this state's dimensions
+            min_arrays.append(np.full(state.shape[0], -np.inf))
+
+        if state.max is not None:
+            max_arrays.append(state.max)
+        else:
+            # If max is None, fill with +inf for this state's dimensions
+            max_arrays.append(np.full(state.shape[0], np.inf))
+
+        if state.guess is not None:
+            guess_arrays.append(state.guess)
+        if state.initial is not None:
+            initial_arrays.append(state.initial)
+        if state.final is not None:
+            final_arrays.append(state.final)
+        if state._initial is not None:
+            _initial_arrays.append(state._initial)
+        if state._final is not None:
+            _final_arrays.append(state._final)
+        if state.initial_type is not None:
+            initial_type_arrays.append(state.initial_type)
+        else:
+            # If initial_type is None, fill with "Free" for this state's dimensions
+            initial_type_arrays.append(np.full(state.shape[0], "Free", dtype=object))
+
+        if state.final_type is not None:
+            final_type_arrays.append(state.final_type)
+        else:
+            # If final_type is None, fill with "Free" for this state's dimensions
+            final_type_arrays.append(np.full(state.shape[0], "Free", dtype=object))
+
+    # Concatenate arrays if they exist
+    unified_min = np.concatenate(min_arrays) if min_arrays else None
+    unified_max = np.concatenate(max_arrays) if max_arrays else None
+    unified_guess = np.concatenate(guess_arrays, axis=1) if guess_arrays else None
+    unified_initial = np.concatenate(initial_arrays) if initial_arrays else None
+    unified_final = np.concatenate(final_arrays) if final_arrays else None
+    unified__initial = np.concatenate(_initial_arrays) if _initial_arrays else None
+    unified__final = np.concatenate(_final_arrays) if _final_arrays else None
+    unified_initial_type = np.concatenate(initial_type_arrays) if initial_type_arrays else None
+    unified_final_type = np.concatenate(final_type_arrays) if final_type_arrays else None
+
+    # Calculate true dimension (only from user-defined states, not augmented ones)
+    # Since we simplified State/Control classes, all user states are "true" dimensions
+    true_dim = sum(state.shape[0] for state in true_states)
+
+    # Find time state slice
+    time_state = next((s for s in sorted_states if s.name == "time"), None)
+    time_slice = time_state._slice if time_state else None
+
+    # Find CTCS augmented states slice
+    ctcs_states = [s for s in sorted_states if s.name.startswith("_ctcs_aug_")]
+    ctcs_slice = (
+        slice(ctcs_states[0]._slice.start, ctcs_states[-1]._slice.stop) if ctcs_states else None
+    )
+
+    # Aggregate scaling_min and scaling_max from individual states
+    # Build full arrays using scaling where available, min/max otherwise
+    unified_scaling_min = None
+    unified_scaling_max = None
+
+    # Check if any state has scaling
+    has_any_scaling = any(
+        state.scaling_min is not None or state.scaling_max is not None for state in sorted_states
+    )
+
+    if has_any_scaling:
+        # Build full scaling arrays
+        scaling_min_list = []
+        scaling_max_list = []
+        for state in sorted_states:
+            if state.scaling_min is not None:
+                scaling_min_list.append(state.scaling_min)
+            else:
+                # Use min as fallback
+                if state.min is not None:
+                    scaling_min_list.append(state.min)
+                else:
+                    scaling_min_list.append(np.full(state.shape[0], -np.inf))
+
+            if state.scaling_max is not None:
+                scaling_max_list.append(state.scaling_max)
+            else:
+                # Use max as fallback
+                if state.max is not None:
+                    scaling_max_list.append(state.max)
+                else:
+                    scaling_max_list.append(np.full(state.shape[0], np.inf))
+
+        unified_scaling_min = np.concatenate(scaling_min_list)
+        unified_scaling_max = np.concatenate(scaling_max_list)
+
+    return UnifiedState(
+        name=name,
+        shape=(total_shape,),
+        min=unified_min,
+        max=unified_max,
+        guess=unified_guess,
+        initial=unified_initial,
+        final=unified_final,
+        _initial=unified__initial,
+        _final=unified__final,
+        initial_type=unified_initial_type,
+        final_type=unified_final_type,
+        _true_dim=true_dim,
+        _true_slice=slice(0, true_dim),
+        _augmented_slice=slice(true_dim, total_shape),
+        time_slice=time_slice,
+        ctcs_slice=ctcs_slice,
+        scaling_min=unified_scaling_min,
+        scaling_max=unified_scaling_max,
+    )
+
+
+def unify_controls(controls: List[Control], name: str = "unified_control") -> UnifiedControl:
+    """Create a UnifiedControl from a list of Control objects.
+
+    This function is the primary way to aggregate multiple symbolic Control objects into
+    a single unified control vector for numerical optimization. It:
+
+    1. Sorts controls (user-defined first, augmented controls second)
+    2. Concatenates all control properties (bounds, guesses)
+    3. Assigns slices to each Control for extracting values from unified vector
+    4. Identifies special controls (time dilation)
+    5. Returns a UnifiedControl with all aggregated data
+
+    Args:
+        controls (List[Control]): List of Control objects to unify. Can include both
+            user-defined controls and augmented controls (names starting with '_').
+        name (str): Name identifier for the unified control vector (default: "unified_control")
+
+    Returns:
+        UnifiedControl: Unified control object containing:
+            - Aggregated bounds and guesses
+            - Shape equal to sum of all control shapes
+            - Slices for extracting individual control components
+            - Properties for accessing true vs augmented controls
+
+    Example:
+        Basic unification::
+
+            import openscvx as ox
+            from openscvx.symbolic.unified import unify_controls
+
+            thrust = ox.Control("thrust", shape=(3,), min=0, max=10)
+            torque = ox.Control("torque", shape=(3,), min=-1, max=1)
+
+            unified = unify_controls([thrust, torque], name="u")
+            print(unified.shape)       # (6,)
+            print(unified._true_dim)   # 6 (all are user controls)
+            print(thrust._slice)       # slice(0, 3) - assigned during unification
+            print(torque._slice)       # slice(3, 6)
+
+        With augmented controls::
+
+            # Time-optimal problems may add time dilation control
+            time_dilation = ox.Control("_time_dilation", shape=(1,))
+
+            unified = unify_controls([thrust, torque, time_dilation])
+            print(unified._true_dim)         # 6 (thrust + torque)
+            print(unified.true.shape)        # (6,)
+            print(unified.augmented.shape)   # (1,) - time dilation
+
+    Note:
+        After unification, each Control object has its `_slice` attribute set,
+        which is used during JAX lowering to extract the correct values from
+        the unified control vector.
+
+    See Also:
+        - UnifiedControl: Return type with detailed documentation
+        - unify_states(): Analogous function for State objects
+        - Control: Individual symbolic control variable
+    """
+    if not controls:
+        return UnifiedControl(name=name, shape=(0,))
+
+    # Sort controls: true controls (not starting with '_') first, then augmented controls
+    # (starting with '_')
+    true_controls = [control for control in controls if not control.name.startswith("_")]
+    augmented_controls = [control for control in controls if control.name.startswith("_")]
+    sorted_controls = true_controls + augmented_controls
+
+    # Calculate total shape
+    total_shape = sum(control.shape[0] for control in sorted_controls)
+
+    # Concatenate all arrays, handling None values properly
+    min_arrays = []
+    max_arrays = []
+    guess_arrays = []
+
+    for control in sorted_controls:
+        if control.min is not None:
+            min_arrays.append(control.min)
+        else:
+            # If min is None, fill with -inf for this control's dimensions
+            min_arrays.append(np.full(control.shape[0], -np.inf))
+
+        if control.max is not None:
+            max_arrays.append(control.max)
+        else:
+            # If max is None, fill with +inf for this control's dimensions
+            max_arrays.append(np.full(control.shape[0], np.inf))
+
+        if control.guess is not None:
+            guess_arrays.append(control.guess)
+
+    # Concatenate arrays if they exist
+    unified_min = np.concatenate(min_arrays) if min_arrays else None
+    unified_max = np.concatenate(max_arrays) if max_arrays else None
+    unified_guess = np.concatenate(guess_arrays, axis=1) if guess_arrays else None
+
+    # Calculate true dimension (only from user-defined controls, not augmented ones)
+    # Since we simplified State/Control classes, all user controls are "true" dimensions
+    true_dim = sum(control.shape[0] for control in true_controls)
+
+    # Find time dilation control slice
+    time_dilation_control = next((c for c in sorted_controls if c.name == "_time_dilation"), None)
+    time_dilation_slice = time_dilation_control._slice if time_dilation_control else None
+
+    # Aggregate scaling_min and scaling_max from individual controls
+    # Build full arrays using scaling where available, min/max otherwise
+    unified_scaling_min = None
+    unified_scaling_max = None
+
+    # Check if any control has scaling
+    has_any_scaling = any(
+        control.scaling_min is not None or control.scaling_max is not None
+        for control in sorted_controls
+    )
+
+    if has_any_scaling:
+        # Build full scaling arrays
+        scaling_min_list = []
+        scaling_max_list = []
+        for control in sorted_controls:
+            if control.scaling_min is not None:
+                scaling_min_list.append(control.scaling_min)
+            else:
+                # Use min as fallback
+                if control.min is not None:
+                    scaling_min_list.append(control.min)
+                else:
+                    scaling_min_list.append(np.full(control.shape[0], -np.inf))
+
+            if control.scaling_max is not None:
+                scaling_max_list.append(control.scaling_max)
+            else:
+                # Use max as fallback
+                if control.max is not None:
+                    scaling_max_list.append(control.max)
+                else:
+                    scaling_max_list.append(np.full(control.shape[0], np.inf))
+
+        unified_scaling_min = np.concatenate(scaling_min_list)
+        unified_scaling_max = np.concatenate(scaling_max_list)
+
+    return UnifiedControl(
+        name=name,
+        shape=(total_shape,),
+        min=unified_min,
+        max=unified_max,
+        guess=unified_guess,
+        _true_dim=true_dim,
+        _true_slice=slice(0, true_dim),
+        _augmented_slice=slice(true_dim, total_shape),
+        time_dilation_slice=time_dilation_slice,
+        scaling_min=unified_scaling_min,
+        scaling_max=unified_scaling_max,
+    )

openscvx/utils/__init__.py

@@ -0,0 +1,20 @@
+"""Utility functions for caching, printing, and output formatting.
+
+This module provides utilities for OpenSCvx.
+"""
+
+from .utils import (
+    calculate_cost_from_boundaries,
+    gen_vertices,
+    generate_orthogonal_unit_vectors,
+    get_kp_pose,
+    rot,
+)
+
+__all__ = [
+    "generate_orthogonal_unit_vectors",
+    "rot",
+    "gen_vertices",
+    "get_kp_pose",
+    "calculate_cost_from_boundaries",
+]