openscvx 0.3.2.dev170__py3-none-any.whl

openscvx/symbolic/expr/state.py

@@ -0,0 +1,607 @@
+import hashlib
+from enum import Enum
+
+import numpy as np
+
+from .variable import Variable
+
+
+class BoundaryType(str, Enum):
+    """Enumeration of boundary condition types for state variables.
+
+    This enum allows users to specify boundary conditions using plain strings
+    while maintaining type safety internally. Boundary conditions control how
+    the optimizer handles initial and final state values.
+
+    Attributes:
+        FIXED (str): State value is fixed to a specific value
+        FREE (str): State value is free to be optimized within bounds
+        MINIMIZE (str): Objective term to minimize the state value
+        MAXIMIZE (str): Objective term to maximize the state value
+
+    Example:
+        Can use either enum or string:
+
+            BoundaryType.FIXED
+            "fixed"  # Equivalent
+    """
+
+    FIXED = "fixed"
+    FREE = "free"
+    MINIMIZE = "minimize"
+    MAXIMIZE = "maximize"
+
+
+def Free(guess):
+    """Create a free boundary condition tuple.
+
+    This is a convenience function that returns a tuple ("free", guess) which
+    can be used to specify free boundary conditions for State or Time objects.
+
+    Args:
+        guess: Initial guess value for the free variable.
+
+    Returns:
+        tuple: ("free", guess) tuple suitable for use in State.initial, State.final,
+            or Time.initial, Time.final.
+
+    Example:
+        ```python
+        pos = ox.State("pos", (3,))
+        pos.final = [ox.Free(5.0), ox.Free(3.0), 10]  # First two free, third fixed
+
+        time = ox.Time(
+            initial=0.0,
+            final=ox.Free(10.0),
+            min=0.0,
+            max=20.0
+        )
+        ```
+    """
+    return ("free", guess)
+
+
+def Fixed(value):
+    """Create a fixed boundary condition tuple.
+
+    This is a convenience function that returns a tuple ("fixed", value) which
+    can be used to explicitly specify fixed boundary conditions for State or Time objects.
+    Note that plain numbers default to fixed, so this is mainly for clarity.
+
+    Args:
+        value: Fixed value for the boundary condition.
+
+    Returns:
+        tuple: ("fixed", value) tuple suitable for use in State.initial, State.final,
+            or Time.initial, Time.final.
+
+    Example:
+        ```python
+        pos = ox.State("pos", (3,))
+        pos.final = [ox.Fixed(10.0), ox.Free(5.0), ox.Fixed(2.0)]
+
+        # Equivalent to:
+        pos.final = [10.0, ox.Free(5.0), 2.0]  # Plain numbers default to fixed
+        ```
+    """
+    return ("fixed", value)
+
+
+def Minimize(guess):
+    """Create a minimize boundary condition tuple.
+
+    This is a convenience function that returns a tuple ("minimize", guess) which
+    can be used to specify that a boundary value should be minimized in the objective
+    function for State or Time objects.
+
+    Args:
+        guess: Initial guess value for the variable to be minimized.
+
+    Returns:
+        tuple: ("minimize", guess) tuple suitable for use in State.initial, State.final,
+            or Time.initial, Time.final.
+
+    Example:
+        ```python
+        time = ox.Time(
+            initial=0.0,
+            final=ox.Minimize(10.0),  # Minimize final time
+            min=0.0,
+            max=20.0
+        )
+
+        fuel = ox.State("fuel", (1,))
+        fuel.final = [ox.Minimize(0)]  # Minimize final fuel consumption
+        ```
+    """
+    return ("minimize", guess)
+
+
+def Maximize(guess):
+    """Create a maximize boundary condition tuple.
+
+    This is a convenience function that returns a tuple ("maximize", guess) which
+    can be used to specify that a boundary value should be maximized in the objective
+    function for State or Time objects.
+
+    Args:
+        guess: Initial guess value for the variable to be maximized.
+
+    Returns:
+        tuple: ("maximize", guess) tuple suitable for use in State.initial, State.final,
+            or Time.initial, Time.final.
+
+    Example:
+        ```python
+        altitude = ox.State("altitude", (1,))
+        altitude.final = [ox.Maximize(100.0)]  # Maximize final altitude
+
+        time = ox.Time(
+            initial=ox.Maximize(0.0),  # Maximize initial time
+            final=10.0,
+            min=0.0,
+            max=20.0
+        )
+        ```
+    """
+    return ("maximize", guess)
+
+
+class State(Variable):
+    """State variable with boundary conditions for trajectory optimization.
+
+    State represents a dynamic state variable in a trajectory optimization problem.
+    Unlike control inputs, states evolve according to dynamics constraints and can
+    have boundary conditions specified at the initial and final time points.
+    Like all Variables, States also support min/max bounds and initial trajectory
+    guesses to help guide the optimization solver toward good solutions.
+
+    States support four types of boundary conditions:
+
+    - **fixed**: State value is constrained to a specific value
+    - **free**: State value is optimized within the specified bounds
+    - **minimize**: Adds a term to the objective function to minimize the state value
+    - **maximize**: Adds a term to the objective function to maximize the state value
+
+    Each element of a multi-dimensional state can have different boundary condition
+    types, allowing for fine-grained control over the optimization.
+
+    Attributes:
+        name (str): Unique name identifier for this state variable
+        _shape (tuple[int, ...]): Shape of the state vector (typically 1D like (3,) for 3D position)
+        _slice (slice | None): Internal slice information for variable indexing
+        _min (np.ndarray | None): Minimum bounds for state variables
+        _max (np.ndarray | None): Maximum bounds for state variables
+        _guess (np.ndarray | None): Initial trajectory guess
+        _initial (np.ndarray | None): Initial state values with boundary condition types
+        initial_type (np.ndarray | None): Array of boundary condition types for initial state
+        _final (np.ndarray | None): Final state values with boundary condition types
+        final_type (np.ndarray | None): Array of boundary condition types for final state
+
+    Example:
+        Scalar time state with fixed initial time, minimize final time:
+
+            time = State("time", (1,))
+            time.min = [0.0]
+            time.max = [10.0]
+            time.initial = [("fixed", 0.0)]
+            time.final = [("minimize", 5.0)]
+
+        3D position state with mixed boundary conditions:
+
+            pos = State("pos", (3,))
+            pos.min = [0, 0, 10]
+            pos.max = [10, 10, 200]
+            pos.initial = [0, ("free", 1), 50]  # x fixed, y free, z fixed
+            pos.final = [10, ("free", 5), ("maximize", 150)]  # Maximize final altitude
+    """
+
+    def __init__(self, name, shape):
+        """Initialize a State object.
+
+        Args:
+            name: Name identifier for the state variable
+            shape: Shape of the state vector (typically 1D tuple)
+        """
+        super().__init__(name, shape)
+        self._initial = None
+        self.initial_type = None
+        self._final = None
+        self.final_type = None
+        self._scaling_min = None
+        self._scaling_max = None
+
+    def _hash_into(self, hasher: "hashlib._Hash") -> None:
+        """Hash State including boundary condition types.
+
+        Extends Variable._hash_into to include the structural metadata that
+        affects the compiled problem: boundary condition types (fixed, free,
+        minimize, maximize). Values are not hashed as they are runtime parameters.
+
+        Args:
+            hasher: A hashlib hash object to update
+        """
+        # Hash the base Variable attributes (class name, shape, slice)
+        super()._hash_into(hasher)
+        # Hash boundary condition types (these affect constraint structure)
+        if self.initial_type is not None:
+            hasher.update(b"initial_type:")
+            hasher.update(str(self.initial_type.tolist()).encode())
+        if self.final_type is not None:
+            hasher.update(b"final_type:")
+            hasher.update(str(self.final_type.tolist()).encode())
+
+    @property
+    def min(self):
+        """Get the minimum bounds for the state variables.
+
+        Returns:
+            Array of minimum values for each state variable element.
+
+        Example:
+            Get lower bounds:
+
+                pos = State("pos", (3,))
+                pos.min = [0, 0, 10]
+                print(pos.min)  # [0. 0. 10.]
+        """
+        return self._min
+
+    @min.setter
+    def min(self, val):
+        """Set the minimum bounds for the state variables.
+
+        Bounds are validated against any fixed initial/final conditions to ensure
+        consistency.
+
+        Args:
+            val: Array of minimum values, must match the state shape exactly
+
+        Raises:
+            ValueError: If the shape doesn't match the state shape, or if fixed
+                boundary conditions violate the bounds
+
+        Example:
+            Set lower bounds:
+
+                pos = State("pos", (3,))
+                pos.min = [0, 0, 10]
+                pos.initial = [0, 5, 15]  # Must satisfy: 0>=0, 5>=0, 15>=10
+        """
+        val = np.asarray(val, dtype=float)
+        if val.shape != self.shape:
+            raise ValueError(f"Min shape {val.shape} does not match State shape {self.shape}")
+        self._min = val
+        self._check_bounds_against_initial_final()
+
+    @property
+    def max(self):
+        """Get the maximum bounds for the state variables.
+
+        Returns:
+            Array of maximum values for each state variable element.
+
+        Example:
+            Get upper bounds:
+
+                vel = State("vel", (3,))
+                vel.max = [10, 10, 5]
+                print(vel.max)  # [10. 10. 5.]
+        """
+        return self._max
+
+    @max.setter
+    def max(self, val):
+        """Set the maximum bounds for the state variables.
+
+        Bounds are validated against any fixed initial/final conditions to ensure
+        consistency.
+
+        Args:
+            val: Array of maximum values, must match the state shape exactly
+
+        Raises:
+            ValueError: If the shape doesn't match the state shape, or if fixed
+                boundary conditions violate the bounds
+
+        Example:
+            Set upper bounds:
+
+                vel = State("vel", (3,))
+                vel.max = [10, 10, 5]
+                vel.final = [8, 9, 4]  # Must satisfy: 8<=10, 9<=10, 4<=5
+        """
+        val = np.asarray(val, dtype=float)
+        if val.shape != self.shape:
+            raise ValueError(f"Max shape {val.shape} does not match State shape {self.shape}")
+        self._max = val
+        self._check_bounds_against_initial_final()
+
+    def _check_bounds_against_initial_final(self):
+        """Validate that fixed boundary conditions respect min/max bounds.
+
+        This internal method is automatically called when bounds or boundary
+        conditions are set to ensure consistency.
+
+        Raises:
+            ValueError: If any fixed initial or final value violates the min/max bounds
+        """
+        for field_name, data, types in [
+            ("initial", self._initial, self.initial_type),
+            ("final", self._final, self.final_type),
+        ]:
+            if data is None or types is None:
+                continue
+            for i, val in np.ndenumerate(data):
+                if types[i] != "Fix":
+                    continue
+                min_i = self._min[i] if self._min is not None else -np.inf
+                max_i = self._max[i] if self._max is not None else np.inf
+                if val < min_i:
+                    raise ValueError(
+                        f"{field_name.capitalize()} Fixed value at index {i[0]} is lower then the "
+                        f"min: {val} < {min_i}"
+                    )
+                if val > max_i:
+                    raise ValueError(
+                        f"{field_name.capitalize()} Fixed value at index {i[0]} is greater then "
+                        f"the max: {val} > {max_i}"
+                    )
+
+    @property
+    def initial(self):
+        """Get the initial state boundary condition values.
+
+        Returns:
+            Array of initial state values (regardless of boundary condition type),
+            or None if not set.
+
+        Note:
+            Use `initial_type` to see the boundary condition types for each element.
+
+        Example:
+            Get initial state boundary conditions:
+
+                x = State("x", (2,))
+                x.initial = [0, ("free", 1)]
+                print(x.initial)  # [0. 1.]
+                print(x.initial_type)  # ['Fix' 'Free']
+        """
+        return self._initial
+
+    @initial.setter
+    def initial(self, arr):
+        """Set the initial state boundary conditions.
+
+        Each element can be specified as either a simple number (defaults to "fixed")
+        or a tuple of (type, value) where type specifies the boundary condition.
+
+        Args:
+            arr: Array-like of initial conditions. Each element can be:
+                - A number: Defaults to fixed boundary condition at that value
+                - A tuple (type, value): Where type is one of:
+                    - "fixed": Constrain state to this exact value
+                    - "free": Let optimizer choose within bounds, initialize at value
+                    - "minimize": Add objective term to minimize, initialize at value
+                    - "maximize": Add objective term to maximize, initialize at value
+
+        Raises:
+            ValueError: If the shape doesn't match the state shape, if boundary
+                condition type is invalid, or if fixed values violate bounds
+
+        Example:
+            Set initial state boundary conditions:
+
+                pos = State("pos", (3,))
+                pos.min = [0, 0, 0]
+                pos.max = [10, 10, 10]
+                # x fixed at 0, y free (starts at 5), z fixed at 2
+                pos.initial = [0, ("free", 5), 2]
+
+            Can also minimize/maximize boundary values:
+
+                time = State("t", (1,))
+                time.initial = [("minimize", 0)]  # Minimize initial time
+        """
+        # Convert to list first to handle mixed types properly
+        if not isinstance(arr, (list, tuple)):
+            arr = np.asarray(arr)
+            if arr.shape != self.shape:
+                raise ValueError(f"Shape mismatch: {arr.shape} != {self.shape}")
+            arr = arr.tolist()
+
+        # Ensure we have the right number of elements
+        if len(arr) != self.shape[0]:
+            raise ValueError(f"Length mismatch: got {len(arr)} elements, expected {self.shape[0]}")
+
+        self._initial = np.zeros(self.shape, dtype=float)
+        self.initial_type = np.full(self.shape, "Fix", dtype=object)
+
+        for i, v in enumerate(arr):
+            if isinstance(v, tuple) and len(v) == 2:
+                # Tuple API: (type, value)
+                bc_type_str, bc_value = v
+                try:
+                    bc_type = BoundaryType(bc_type_str)  # Validates the string
+                except ValueError:
+                    valid_types = [t.value for t in BoundaryType]
+                    raise ValueError(
+                        f"Invalid boundary condition type: {bc_type_str}. "
+                        f"Valid types are: {valid_types}"
+                    )
+                self._initial[i] = float(bc_value)
+                self.initial_type[i] = bc_type.value.capitalize()
+            elif isinstance(v, (int, float, np.number)):
+                # Simple number defaults to fixed
+                self._initial[i] = float(v)
+                self.initial_type[i] = "Fix"
+            else:
+                raise ValueError(
+                    f"Invalid boundary condition format: {v}. "
+                    f"Use a number (defaults to fixed) or tuple ('type', value) "
+                    f"where type is 'fixed', 'free', 'minimize', or 'maximize'."
+                )
+
+        self._check_bounds_against_initial_final()
+
+    @property
+    def final(self):
+        """Get the final state boundary condition values.
+
+        Returns:
+            Array of final state values (regardless of boundary condition type),
+            or None if not set.
+
+        Note:
+            Use `final_type` to see the boundary condition types for each element.
+
+        Example:
+            Get final state boundary conditions:
+
+                x = State("x", (2,))
+                x.final = [10, ("minimize", 0)]
+                print(x.final)  # [10. 0.]
+                print(x.final_type)  # ['Fix' 'Minimize']
+        """
+        return self._final
+
+    @final.setter
+    def final(self, arr):
+        """Set the final state boundary conditions.
+
+        Each element can be specified as either a simple number (defaults to "fixed")
+        or a tuple of (type, value) where type specifies the boundary condition.
+
+        Args:
+            arr: Array-like of final conditions. Each element can be:
+                - A number: Defaults to fixed boundary condition at that value
+                - A tuple (type, value): Where type is one of:
+                    - "fixed": Constrain state to this exact value
+                    - "free": Let optimizer choose within bounds, initialize at value
+                    - "minimize": Add objective term to minimize, initialize at value
+                    - "maximize": Add objective term to maximize, initialize at value
+
+        Raises:
+            ValueError: If the shape doesn't match the state shape, if boundary
+                condition type is invalid, or if fixed values violate bounds
+
+        Example:
+            Set final state boundary conditionis:
+
+                pos = State("pos", (3,))
+                pos.min = [0, 0, 0]
+                pos.max = [10, 10, 10]
+                # x fixed at 10, y free (starts at 5), z maximize altitude
+                pos.final = [10, ("free", 5), ("maximize", 8)]
+
+            Minimize final time in time-optimal problem:
+
+                time = State("t", (1,))
+                time.final = [("minimize", 10)]
+        """
+        # Convert to list first to handle mixed types properly
+        if not isinstance(arr, (list, tuple)):
+            arr = np.asarray(arr)
+            if arr.shape != self.shape:
+                raise ValueError(f"Shape mismatch: {arr.shape} != {self.shape}")
+            arr = arr.tolist()
+
+        # Ensure we have the right number of elements
+        if len(arr) != self.shape[0]:
+            raise ValueError(f"Length mismatch: got {len(arr)} elements, expected {self.shape[0]}")
+
+        self._final = np.zeros(self.shape, dtype=float)
+        self.final_type = np.full(self.shape, "Fix", dtype=object)
+
+        for i, v in enumerate(arr):
+            if isinstance(v, tuple) and len(v) == 2:
+                # Tuple API: (type, value)
+                bc_type_str, bc_value = v
+                try:
+                    bc_type = BoundaryType(bc_type_str)  # Validates the string
+                except ValueError:
+                    valid_types = [t.value for t in BoundaryType]
+                    raise ValueError(
+                        f"Invalid boundary condition type: {bc_type_str}. "
+                        f"Valid types are: {valid_types}"
+                    )
+                self._final[i] = float(bc_value)
+                self.final_type[i] = bc_type.value.capitalize()
+            elif isinstance(v, (int, float, np.number)):
+                # Simple number defaults to fixed
+                self._final[i] = float(v)
+                self.final_type[i] = "Fix"
+            else:
+                raise ValueError(
+                    f"Invalid boundary condition format: {v}. "
+                    f"Use a number (defaults to fixed) or tuple ('type', value) "
+                    f"where type is 'fixed', 'free', 'minimize', or 'maximize'."
+                )
+
+        self._check_bounds_against_initial_final()
+
+    @property
+    def scaling_min(self):
+        """Get the scaling minimum bounds for the state variables.
+
+        Returns:
+            Array of scaling minimum values for each state variable element, or None if not set.
+        """
+        return self._scaling_min
+
+    @scaling_min.setter
+    def scaling_min(self, val):
+        """Set the scaling minimum bounds for the state variables.
+
+        Args:
+            val: Array of scaling minimum values, must match the state shape exactly
+
+        Raises:
+            ValueError: If the shape doesn't match the state shape
+        """
+        if val is None:
+            self._scaling_min = None
+            return
+        val = np.asarray(val, dtype=float)
+        if val.shape != self.shape:
+            raise ValueError(
+                f"Scaling min shape {val.shape} does not match State shape {self.shape}"
+            )
+        self._scaling_min = val
+
+    @property
+    def scaling_max(self):
+        """Get the scaling maximum bounds for the state variables.
+
+        Returns:
+            Array of scaling maximum values for each state variable element, or None if not set.
+        """
+        return self._scaling_max
+
+    @scaling_max.setter
+    def scaling_max(self, val):
+        """Set the scaling maximum bounds for the state variables.
+
+        Args:
+            val: Array of scaling maximum values, must match the state shape exactly
+
+        Raises:
+            ValueError: If the shape doesn't match the state shape
+        """
+        if val is None:
+            self._scaling_max = None
+            return
+        val = np.asarray(val, dtype=float)
+        if val.shape != self.shape:
+            raise ValueError(
+                f"Scaling max shape {val.shape} does not match State shape {self.shape}"
+            )
+        self._scaling_max = val
+
+    def __repr__(self):
+        """String representation of the State object.
+
+        Returns:
+            Concise string showing the state name and shape.
+        """
+        return f"State('{self.name}', shape={self.shape})"