openscvx 0.3.2.dev333__tar.gz → 0.3.2.dev342__tar.gz

{openscvx-0.3.2.dev333/openscvx.egg-info → openscvx-0.3.2.dev342}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.3.2.dev333
+Version: 0.3.2.dev342
 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-0.3.2.dev333 → openscvx-0.3.2.dev342}/openscvx/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.3.2.dev333'
-__version_tuple__ = version_tuple = (0, 3, 2, 'dev333')
+__version__ = version = '0.3.2.dev342'
+__version_tuple__ = version_tuple = (0, 3, 2, 'dev342')
 
-__commit_id__ = commit_id = 'g298f4bc4f'
+__commit_id__ = commit_id = 'gd4121da9a'

{openscvx-0.3.2.dev333 → openscvx-0.3.2.dev342}/openscvx/problem.py

@@ -50,7 +50,6 @@ from openscvx.lowered.jax_constraints import (
 from openscvx.propagation import get_propagation_solver, propagate_trajectory_results
 from openscvx.solvers import PTRSolver
 from openscvx.symbolic.builder import preprocess_symbolic_problem
-from openscvx.symbolic.constraint_set import ConstraintSet
 from openscvx.symbolic.expr import CTCS, Constraint
 from openscvx.symbolic.expr.control import Control
 from openscvx.symbolic.expr.state import State
@@ -128,7 +127,7 @@ class Problem:
         # Symbolic Preprocessing & Augmentation
         self.symbolic: SymbolicProblem = preprocess_symbolic_problem(
             dynamics=dynamics,
-            constraints=ConstraintSet(unsorted=list(constraints)),
+            constraints=constraints,
             states=states,
             controls=controls,
             N=N,

{openscvx-0.3.2.dev333 → openscvx-0.3.2.dev342}/openscvx/symbolic/builder.py

@@ -55,6 +55,8 @@ from openscvx.symbolic.preprocessing import (
     validate_dynamics_dict_dimensions,
     validate_dynamics_dimension,
     validate_guesses,
+    validate_input_types,
+    validate_propagation_input_types,
     validate_shapes,
     validate_variable_names,
 )
@@ -64,7 +66,7 @@ from openscvx.symbolic.time import Time
 
 def preprocess_symbolic_problem(
     dynamics: dict,
-    constraints: ConstraintSet,
+    constraints: list,
     states: List[State],
     controls: List[Control],
     N: int,
@@ -99,8 +101,8 @@ def preprocess_symbolic_problem(
     Args:
         dynamics: Dictionary mapping state names to dynamics expressions.
             Example: {"x": v, "v": u}
-        constraints: ConstraintSet with raw constraints in `unsorted` field.
-            Create with: ConstraintSet(unsorted=[c1, c2, c3])
+        constraints: List of constraint objects (Constraint, NodalConstraint,
+            CrossNodeConstraint, or CTCS).
         states: List of user-defined State objects (should NOT include time or CTCS states)
         controls: List of user-defined Control objects (should NOT include time dilation)
         N: Number of discretization nodes in the trajectory
@@ -140,16 +142,13 @@ def preprocess_symbolic_problem(
         Basic usage with CTCS constraint::
 
             import openscvx as ox
-            from openscvx.symbolic.constraint_set import ConstraintSet
 
             x = ox.State("x", shape=(2,))
             v = ox.State("v", shape=(2,))
             u = ox.Control("u", shape=(2,))
 
             dynamics = {"x": v, "v": u}
-            constraints = ConstraintSet(unsorted=[
-                (ox.Norm(x) <= 5.0).over((0, 50))
-            ])
+            constraints = [(ox.Norm(x) <= 5.0).over((0, 50))]
 
             problem = preprocess_symbolic_problem(
                 dynamics=dynamics,
@@ -187,6 +186,13 @@ def preprocess_symbolic_problem(
             print([s.name for s in problem.states_prop])
     """
 
+    # Validate input types before anything else
+    validate_input_types(dynamics, states, controls, constraints, N, time)
+    validate_propagation_input_types(dynamics_prop_extra, states_prop_extra)
+
+    # Wrap validated constraints into a ConstraintSet
+    constraints = ConstraintSet(unsorted=list(constraints))
+
     # Validate user-provided variables have required attributes
     validate_boundary_conditions(states)
     validate_bounds(states + controls)

{openscvx-0.3.2.dev333 → openscvx-0.3.2.dev342}/openscvx/symbolic/lowerers/jax/constraint.py

@@ -3,9 +3,6 @@
 Visitors: Equality, Inequality, NodalConstraint, CrossNodeConstraint, CTCS
 """
 
-import jax.numpy as jnp
-from jax.lax import cond
-
 # Expression types to handle — uncomment as you paste visitors:
 from openscvx.symbolic.expr.constraint import (
     CTCS,
@@ -15,6 +12,7 @@ from openscvx.symbolic.expr.constraint import (
     Inequality,
     NodalConstraint,
 )
+from openscvx.symbolic.expr.logic import Cond
 from openscvx.symbolic.lowerers.jax._registry import visitor  # noqa: F401
 
 
@@ -106,17 +104,13 @@ def _visit_cross_node_constraint(lowerer, node: CrossNodeConstraint):
     return trajectory_constraint
 
 
-# TODO: (norrisg) CTCS is playing 2 roles here: both as a constraint wrapper and as the penalty
-# expression w/ conditional logic. Consider adding conditional logic as separate AST nodes.
-# Then, CTCS remains a wrapper and we just wrap the penalty expression with the conditional
-# logic when we lower it.
 @visitor(CTCS)
 def _visit_ctcs(lowerer, node: CTCS):
     """Lower CTCS (Continuous-Time Constraint Satisfaction) to JAX function.
 
     CTCS constraints use penalty methods to enforce constraints over continuous
-    time intervals. The lowered function includes conditional logic to activate
-    the penalty only within the specified node interval.
+    time intervals. When a node range is specified, the penalty expression is
+    wrapped in a Cond node to activate it only within that interval.
 
     Args:
         node: CTCS constraint node with penalty expression and optional node range
@@ -125,35 +119,20 @@ def _visit_ctcs(lowerer, node: CTCS):
         Function (x, u, current_node, params) -> penalty value or 0
 
     Note:
-        Uses jax.lax.cond for JAX-traceable conditional evaluation. The penalty
-        is active only when current_node is in [start_node, end_node).
+        The penalty is active only when current_node is in [start_node, end_node).
         If no node range is specified, the penalty is always active.
+        Conditional node-range logic is delegated to the Cond AST node.
 
     See Also:
         - CTCS: The symbolic CTCS constraint class
+        - Cond: Conditional AST node used for node-range gating
         - penalty functions: PositivePart, Huber, SmoothReLU
     """
-    # Lower the penalty expression (which includes the constraint residual)
-    penalty_expr_fn = lowerer.lower(node.penalty_expr())
-
-    def ctcs_fn(x, u, current_node, params):
-        # Check if constraint is active at this node
-        if node.nodes is not None:
-            start_node, end_node = node.nodes
-            # Extract scalar value from current_node (which may be array or scalar)
-            # Keep as JAX array for tracing compatibility
-            node_scalar = jnp.atleast_1d(current_node)[0]
-            is_active = (start_node <= node_scalar) & (node_scalar < end_node)
-
-            # Use jax.lax.cond for conditional evaluation
-            return cond(
-                is_active,
-                lambda _: penalty_expr_fn(x, u, current_node, params),
-                lambda _: 0.0,
-                operand=None,
-            )
-        else:
-            # Always active if no node range specified
-            return penalty_expr_fn(x, u, current_node, params)
-
-    return ctcs_fn
+    penalty_expr = node.penalty_expr()
+
+    if node.nodes is not None:
+        # Wrap penalty in a Cond node that gates on the node range
+        gated_expr = Cond(None, penalty_expr, 0.0, node_ranges=[node.nodes])
+        return lowerer.lower(gated_expr)
+    else:
+        return lowerer.lower(penalty_expr)

{openscvx-0.3.2.dev333 → openscvx-0.3.2.dev342}/openscvx/symbolic/lowerers/jax/logic.py

@@ -143,10 +143,12 @@ def _visit_cond(lowerer, node: Cond):
 
         # If node_ranges is specified, check if current node is in range
         if node_ranges is not None:
-            # Check if node_arg is within any of the specified ranges [start, end)
+            # Extract scalar from node_arg (which may be array or scalar from vmap)
+            node_scalar = jnp.atleast_1d(node_arg)[0]
+            # Check if node_scalar is within any of the specified ranges [start, end)
             in_range = jnp.array(False)
             for start, end in node_ranges:
-                in_range = in_range | ((node_arg >= start) & (node_arg < end))
+                in_range = in_range | ((node_scalar >= start) & (node_scalar < end))
             # Combined predicate: must be in range AND predicate satisfied
             pred_bool = in_range & pred_bool
 

{openscvx-0.3.2.dev333 → openscvx-0.3.2.dev342}/openscvx/symbolic/preprocessing.py

@@ -748,6 +748,142 @@ def validate_bounds(variables: List[Variable]) -> None:
             )
 
 
+def validate_input_types(
+    dynamics: any,
+    states: any,
+    controls: any,
+    constraints: any,
+    N: any,
+    time: any,
+) -> None:
+    """Validate that all user-facing inputs have correct types.
+
+    This catches common user errors like passing a single State or Control
+    instead of a list, or passing wrong types for dynamics, N, or time.
+    Should be called before any other validation in the preprocessing pipeline.
+
+    Raises:
+        TypeError: If any input has the wrong type
+        ValueError: If N is not positive
+    """
+    from openscvx.symbolic.expr import CTCS, Constraint, CrossNodeConstraint, NodalConstraint
+    from openscvx.symbolic.time import Time
+
+    if not isinstance(dynamics, dict):
+        raise TypeError(
+            f"'dynamics' must be a dict mapping state names to expressions, "
+            f"got {type(dynamics).__name__}"
+        )
+
+    if not isinstance(states, list):
+        hint = ""
+        if isinstance(states, State):
+            hint = f" Hint: use states=[{states.name}] instead of states={states.name}"
+        raise TypeError(
+            f"'states' must be a list of State objects, got {type(states).__name__}.{hint}"
+        )
+
+    for i, s in enumerate(states):
+        if not isinstance(s, State):
+            raise TypeError(f"states[{i}] must be a State, got {type(s).__name__}")
+
+    if not isinstance(controls, list):
+        hint = ""
+        if isinstance(controls, Control):
+            hint = f" Hint: use controls=[{controls.name}] instead of controls={controls.name}"
+        raise TypeError(
+            f"'controls' must be a list of Control objects, got {type(controls).__name__}.{hint}"
+        )
+
+    for i, c in enumerate(controls):
+        if not isinstance(c, Control):
+            raise TypeError(f"controls[{i}] must be a Control, got {type(c).__name__}")
+
+    if not isinstance(constraints, list):
+        raise TypeError(
+            f"'constraints' must be a list of Constraint objects, got {type(constraints).__name__}"
+        )
+
+    valid_constraint_types = (Constraint, NodalConstraint, CrossNodeConstraint, CTCS)
+    for i, c in enumerate(constraints):
+        if not isinstance(c, valid_constraint_types):
+            raise TypeError(
+                f"constraints[{i}] must be a Constraint, NodalConstraint, "
+                f"CrossNodeConstraint, or CTCS, got {type(c).__name__}"
+            )
+
+    if not isinstance(N, int):
+        raise TypeError(f"'N' must be an integer, got {type(N).__name__}")
+
+    if N < 1:
+        raise ValueError(f"'N' must be positive, got {N}")
+
+    if not isinstance(time, Time):
+        raise TypeError(f"'time' must be a Time object, got {type(time).__name__}")
+
+
+def validate_propagation_input_types(
+    dynamics_prop_extra: any,
+    states_prop_extra: any,
+) -> None:
+    """Validate types for optional propagation inputs.
+
+    These parameters must either both be None or both be provided.
+    When provided, dynamics_prop_extra must be a dict and states_prop_extra
+    must be a list of State objects.
+
+    Args:
+        dynamics_prop_extra: Should be None or a dict mapping state names to expressions
+        states_prop_extra: Should be None or a list of State objects
+
+    Raises:
+        TypeError: If either input has the wrong type
+        ValueError: If only one of the two is provided
+
+    Example:
+            distance = ox.State("distance", shape=(1,))
+
+            # Wrong: passing bare State instead of list
+            validate_propagation_input_types({"distance": expr}, distance)
+            # Raises TypeError: 'states_prop_extra' must be a list ...
+    """
+    both_none = dynamics_prop_extra is None and states_prop_extra is None
+    both_set = dynamics_prop_extra is not None and states_prop_extra is not None
+
+    if not both_none and not both_set:
+        provided = "dynamics_prop" if dynamics_prop_extra is not None else "states_prop"
+        missing = "states_prop" if dynamics_prop_extra is not None else "dynamics_prop"
+        raise ValueError(
+            f"'{provided}' was provided but '{missing}' was not. "
+            f"Both must be provided together, or both omitted."
+        )
+
+    if both_none:
+        return
+
+    if not isinstance(dynamics_prop_extra, dict):
+        raise TypeError(
+            f"'dynamics_prop' must be a dict mapping state names to expressions, "
+            f"got {type(dynamics_prop_extra).__name__}"
+        )
+
+    if not isinstance(states_prop_extra, list):
+        hint = ""
+        if isinstance(states_prop_extra, State):
+            hint = (
+                f" Hint: use states_prop=[{states_prop_extra.name}]"
+                f" instead of states_prop={states_prop_extra.name}"
+            )
+        raise TypeError(
+            f"'states_prop' must be a list of State objects, "
+            f"got {type(states_prop_extra).__name__}.{hint}"
+        )
+
+    for i, s in enumerate(states_prop_extra):
+        if not isinstance(s, State):
+            raise TypeError(f"states_prop[{i}] must be a State, got {type(s).__name__}")
+
+
 def validate_guesses(variables: List[Variable]) -> None:
     """Validate that all variables have initial guesses set.
 

{openscvx-0.3.2.dev333 → openscvx-0.3.2.dev342/openscvx.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.3.2.dev333
+Version: 0.3.2.dev342
 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-0.3.2.dev333 → openscvx-0.3.2.dev342}/tests/symbolic/test_preprocessing.py

@@ -15,6 +15,8 @@ from openscvx.symbolic.preprocessing import (
     validate_dynamics_dict_dimensions,
     validate_dynamics_dimension,
     validate_guesses,
+    validate_input_types,
+    validate_propagation_input_types,
     validate_variable_names,
 )
 
@@ -976,3 +978,176 @@ def test_validate_guesses_raises_missing():
         match="Control 'thrust' is missing initial guess.*controls require explicit guesses",
     ):
         validate_guesses([u])
+
+
+# =============================================================================
+# validate_input_types Tests
+# =============================================================================
+
+
+@pytest.fixture
+def valid_inputs():
+    """Provide a minimal set of valid inputs for validate_input_types."""
+    from openscvx.symbolic.time import Time
+
+    x = State("x", shape=(3,))
+    u = Control("u", shape=(2,))
+    dynamics = {"x": x}
+    constraints = [x <= 5]
+    time = Time(initial=0.0, final=10.0, min=0.0, max=20.0)
+    return dynamics, [x], [u], constraints, 50, time
+
+
+def test_validate_input_types_passes(valid_inputs):
+    """Test that valid inputs pass validation."""
+    dynamics, states, controls, constraints, N, time = valid_inputs
+    validate_input_types(dynamics, states, controls, constraints, N, time)
+
+
+def test_validate_input_types_bare_control(valid_inputs):
+    """Test that passing a bare Control instead of a list raises TypeError with hint."""
+    dynamics, states, _, constraints, N, time = valid_inputs
+    u = Control("thrust", shape=(2,))
+
+    with pytest.raises(
+        TypeError, match=r"'controls' must be a list.*Control.*Hint.*controls=\[thrust\]"
+    ):
+        validate_input_types(dynamics, states, u, constraints, N, time)
+
+
+def test_validate_input_types_bare_state(valid_inputs):
+    """Test that passing a bare State instead of a list raises TypeError with hint."""
+    dynamics, _, controls, constraints, N, time = valid_inputs
+    x = State("position", shape=(3,))
+
+    with pytest.raises(
+        TypeError, match=r"'states' must be a list.*State.*Hint.*states=\[position\]"
+    ):
+        validate_input_types(dynamics, x, controls, constraints, N, time)
+
+
+def test_validate_input_types_wrong_element_in_states(valid_inputs):
+    """Test that a non-State element in the states list raises TypeError."""
+    dynamics, _, controls, constraints, N, time = valid_inputs
+
+    with pytest.raises(TypeError, match=r"states\[0\] must be a State, got str"):
+        validate_input_types(dynamics, ["not_a_state"], controls, constraints, N, time)
+
+
+def test_validate_input_types_wrong_element_in_controls(valid_inputs):
+    """Test that a non-Control element in the controls list raises TypeError."""
+    dynamics, states, _, constraints, N, time = valid_inputs
+
+    with pytest.raises(TypeError, match=r"controls\[0\] must be a Control, got int"):
+        validate_input_types(dynamics, states, [42], constraints, N, time)
+
+
+def test_validate_input_types_dynamics_not_dict(valid_inputs):
+    """Test that passing non-dict dynamics raises TypeError."""
+    _, states, controls, constraints, N, time = valid_inputs
+
+    with pytest.raises(TypeError, match="'dynamics' must be a dict"):
+        validate_input_types([1, 2, 3], states, controls, constraints, N, time)
+
+
+def test_validate_input_types_constraints_not_list(valid_inputs):
+    """Test that passing a non-list constraints raises TypeError."""
+    dynamics, states, controls, _, N, time = valid_inputs
+
+    with pytest.raises(TypeError, match="'constraints' must be a list"):
+        validate_input_types(dynamics, states, controls, "not_a_list", N, time)
+
+
+def test_validate_input_types_constraints_invalid_element(valid_inputs):
+    """Test that non-constraint elements raise TypeError."""
+    dynamics, states, controls, _, N, time = valid_inputs
+
+    with pytest.raises(TypeError, match=r"constraints\[0\] must be a Constraint.*got int"):
+        validate_input_types(dynamics, states, controls, [42], N, time)
+
+
+def test_validate_input_types_constraints_mixed_valid_and_invalid(valid_inputs):
+    """Test that invalid element is caught even after valid ones."""
+    dynamics, states, controls, _, N, time = valid_inputs
+    x = State("x", shape=(3,))
+
+    with pytest.raises(TypeError, match=r"constraints\[1\] must be a Constraint.*got str"):
+        validate_input_types(dynamics, states, controls, [x <= 5, "bad"], N, time)
+
+
+def test_validate_input_types_N_not_int(valid_inputs):
+    """Test that passing non-int N raises TypeError."""
+    dynamics, states, controls, constraints, _, time = valid_inputs
+
+    with pytest.raises(TypeError, match="'N' must be an integer, got float"):
+        validate_input_types(dynamics, states, controls, constraints, 50.0, time)
+
+
+def test_validate_input_types_N_not_positive(valid_inputs):
+    """Test that passing non-positive N raises ValueError."""
+    dynamics, states, controls, constraints, _, time = valid_inputs
+
+    with pytest.raises(ValueError, match="'N' must be positive, got 0"):
+        validate_input_types(dynamics, states, controls, constraints, 0, time)
+
+    with pytest.raises(ValueError, match="'N' must be positive, got -5"):
+        validate_input_types(dynamics, states, controls, constraints, -5, time)
+
+
+def test_validate_input_types_time_not_time(valid_inputs):
+    """Test that passing non-Time object raises TypeError."""
+    dynamics, states, controls, constraints, N, _ = valid_inputs
+
+    with pytest.raises(TypeError, match="'time' must be a Time object, got float"):
+        validate_input_types(dynamics, states, controls, constraints, N, 10.0)
+
+
+# =============================================================================
+# validate_propagation_input_types Tests
+# =============================================================================
+
+
+def test_validate_propagation_input_types_both_none():
+    """Test that both None passes validation."""
+    validate_propagation_input_types(None, None)
+
+
+def test_validate_propagation_input_types_both_valid():
+    """Test that valid dict + list passes validation."""
+    distance = State("distance", shape=(1,))
+    validate_propagation_input_types({"distance": distance}, [distance])
+
+
+def test_validate_propagation_input_types_only_dynamics():
+    """Test that providing dynamics_prop without states_prop raises ValueError."""
+    with pytest.raises(ValueError, match="'dynamics_prop' was provided but 'states_prop' was not"):
+        validate_propagation_input_types({"distance": 1.0}, None)
+
+
+def test_validate_propagation_input_types_only_states():
+    """Test that providing states_prop without dynamics_prop raises ValueError."""
+    distance = State("distance", shape=(1,))
+    with pytest.raises(ValueError, match="'states_prop' was provided but 'dynamics_prop' was not"):
+        validate_propagation_input_types(None, [distance])
+
+
+def test_validate_propagation_input_types_bare_state():
+    """Test that passing a bare State instead of a list raises TypeError with hint."""
+    distance = State("distance", shape=(1,))
+    with pytest.raises(
+        TypeError, match=r"'states_prop' must be a list.*Hint.*states_prop=\[distance\]"
+    ):
+        validate_propagation_input_types({"distance": distance}, distance)
+
+
+def test_validate_propagation_input_types_dynamics_not_dict():
+    """Test that passing non-dict dynamics_prop raises TypeError."""
+    distance = State("distance", shape=(1,))
+    with pytest.raises(TypeError, match="'dynamics_prop' must be a dict"):
+        validate_propagation_input_types([distance], [distance])
+
+
+def test_validate_propagation_input_types_wrong_element_in_states():
+    """Test that a non-State element in states_prop raises TypeError."""
+    with pytest.raises(TypeError, match=r"states_prop\[0\] must be a State, got str"):
+        validate_propagation_input_types({"x": 1.0}, ["not_a_state"])