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

openscvx/_version.py

@@ -0,0 +1,21 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.1.2'
+__version_tuple__ = version_tuple = (0, 1, 2)

openscvx/augmentation/ctcs.py

@@ -0,0 +1,44 @@
+from typing import List
+
+from openscvx.constraints.ctcs import CTCSConstraint
+
+def sort_ctcs_constraints(constraints_ctcs: List[CTCSConstraint], N: int):        
+    idx_to_nodes: dict[int, tuple] = {}
+    next_idx = 0
+    for c in constraints_ctcs:
+        # normalize None to full horizon
+        c.nodes = c.nodes or (0, N)
+        key = c.nodes
+
+        if c.idx is not None:
+            # user supplied an identifier: ensure it always points to the same interval
+            if c.idx in idx_to_nodes:
+                if idx_to_nodes[c.idx] != key:
+                    raise ValueError(
+                        f"idx={c.idx} was first used with interval={idx_to_nodes[c.idx]}, "
+                        f"but now you gave it interval={key}"
+                    )
+            else:
+                idx_to_nodes[c.idx] = key
+
+        else:
+            # no identifier: see if this interval already has one
+            for existing_id, nodes in idx_to_nodes.items():
+                if nodes == key:
+                    c.idx = existing_id
+                    break
+            else:
+                # brand-new interval: pick the next free auto-id
+                while next_idx in idx_to_nodes:
+                    next_idx += 1
+                c.idx = next_idx
+                idx_to_nodes[next_idx] = key
+                next_idx += 1
+    
+    # Extract your intervals in ascending‐idx order
+    ordered_ids       = sorted(idx_to_nodes.keys())
+    node_intervals    = [ idx_to_nodes[i] for i in ordered_ids ]
+    id_to_position    = { ident: pos for pos, ident in enumerate(ordered_ids) }
+    num_augmented_states = len(ordered_ids)
+
+    return constraints_ctcs, node_intervals, num_augmented_states,

openscvx/augmentation/dynamics_augmentation.py

@@ -0,0 +1,122 @@
+from typing import Callable, List, Tuple
+
+import jax
+import jax.numpy as jnp
+
+from openscvx.constraints.violation import CTCSViolation
+from openscvx.dynamics import Dynamics
+
+def build_augmented_dynamics(
+    dynamics_non_augmented: Dynamics,
+    violations: List[CTCSViolation],
+    idx_x_true: slice,
+    idx_u_true: slice,
+) -> Dynamics:
+    dynamics_augmented = Dynamics(
+        f=get_augmented_dynamics(
+            dynamics_non_augmented.f, violations, idx_x_true, idx_u_true
+        ),
+    )
+    A, B = get_jacobians(
+        dynamics_augmented.f, dynamics_non_augmented, violations, idx_x_true, idx_u_true
+    )
+    dynamics_augmented.A = A
+    dynamics_augmented.B = B
+    return dynamics_augmented
+
+
+def get_augmented_dynamics(
+    dynamics: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray],
+    violations: List[CTCSViolation],
+    idx_x_true: slice,
+    idx_u_true: slice,
+) -> Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]:
+    def dynamics_augmented(x: jnp.array, u: jnp.array, node: int) -> jnp.array:
+        x_dot = dynamics(x[idx_x_true], u[idx_u_true])
+
+        # Iterate through the g_func dictionary and stack the output each function
+        # to x_dot
+        for v in violations:
+            x_dot = jnp.hstack([x_dot, v.g(x[idx_x_true], u[idx_u_true], node)])
+
+        return x_dot
+
+    return dynamics_augmented
+
+
+def get_jacobians(
+    dyn_augmented: Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray],
+    dynamics_non_augmented: Dynamics,
+    violations: List[CTCSViolation],
+    idx_x_true: slice,
+    idx_u_true: slice,
+) -> Tuple[
+    Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray],
+    Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray],
+]:
+    # 1) Early return if absolutely no custom grads anywhere
+    no_dyn_grads = dynamics_non_augmented.A is None and dynamics_non_augmented.B is None
+    no_vio_grads = all(v.g_grad_x is None and v.g_grad_u is None for v in violations)
+
+    if no_dyn_grads and no_vio_grads:
+        return (
+            jax.jacfwd(dyn_augmented, argnums=0),
+            jax.jacfwd(dyn_augmented, argnums=1),
+        )
+
+    # 2) Build the *true‐state* Jacobians of f(x_true,u_true)
+    f_fn = dynamics_non_augmented.f
+    if dynamics_non_augmented.A is None:
+        A_f = lambda x_true, u_true: jax.jacfwd(f_fn, argnums=0)(x_true, u_true)
+    else:
+        A_f = dynamics_non_augmented.A
+
+    if dynamics_non_augmented.B is None:
+        B_f = lambda x_true, u_true: jax.jacfwd(f_fn, argnums=1)(x_true, u_true)
+    else:
+        B_f = dynamics_non_augmented.B
+
+    # 3) Build per-violation gradients
+    def make_violation_grad_x(i: int) -> Callable:
+        viol = violations[i]
+        # use user‐provided if present, otherwise autodiff viol.g in argnum=0
+        return viol.g_grad_x or jax.jacfwd(viol.g, argnums=0)
+
+    def make_violation_grad_u(i: int) -> Callable:
+        viol = violations[i]
+        # use user‐provided if present, otherwise autodiff viol.g in argnum=0
+        return viol.g_grad_u or jax.jacfwd(viol.g, argnums=1)
+
+    # 4) Assemble A_aug, B_aug
+    def A(x_aug, u_aug, node):
+        # dynamics block + zero‐pad
+        Af = A_f(x_aug[idx_x_true], u_aug[idx_u_true])  # (n_f, n_x_true)
+        nv = sum(
+            v.g(x_aug[idx_x_true], u_aug[idx_u_true], node).shape[0] for v in violations
+        )  # total # rows of violations
+        zero_pad = jnp.zeros((Af.shape[0], nv))  # (n_f, n_v)
+        top = jnp.hstack([Af, zero_pad])  # (n_f, n_x_true + n_v)
+
+        # violation blocks
+        rows = [top]
+        for i in range(len(violations)):
+            dx_i = make_violation_grad_x(i)(
+                x_aug[idx_x_true], u_aug[idx_u_true], node
+            )  # (n_gi, n_x_true)
+            pad_i = jnp.zeros((dx_i.shape[0], nv))  # (n_gi, n_v)
+            rows.append(jnp.hstack([dx_i, pad_i]))  # (n_gi, n_x_true + n_v)
+
+        return jnp.vstack(rows)
+
+    def B(x_aug, u_aug, node):
+        Bf = B_f(x_aug[idx_x_true], u_aug[idx_u_true])  # (n_f, n_u_true)
+        rows = [Bf]
+        for i in range(len(violations)):
+            du_i = make_violation_grad_u(i)(
+                x_aug[idx_x_true], u_aug[idx_u_true], node
+            )  # (n_gi, n_u_true)
+            rows.append(du_i)
+
+        return jnp.vstack(rows)
+
+    return A, B

openscvx/config.py

@@ -0,0 +1,247 @@
+import numpy as np
+from dataclasses import dataclass, field
+from typing import Dict, List
+
+
+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:
+    dis_type: str = "FOH"
+    custom_integrator: bool = True
+    solver: str = "Tsit5"
+    args: Dict = field(default_factory=dict)
+    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 algorthim. This tends to be faster then Diffrax but unless your going for speed, its reccomended 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.
+    """
+
+
+@dataclass
+class DevConfig:
+    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.
+    """
+
+
+@dataclass
+class ConvexSolverConfig:
+    solver: str = "QOCO"
+    solver_args: dict = field(default_factory=lambda: {"abstol": 1e-6, "reltol": 1e-9})
+    cvxpygen: 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 and allows. 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): 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.
+    """
+
+
+@dataclass
+class PropagationConfig:
+    inter_sample: int = 30
+    dt: float = 0.1
+    solver: str = "Dopri8"
+    args: Dict = field(default_factory=dict)
+    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.
+    
+    Args:
+        dt (float): The time step for propagation. Defaults to 0.1.
+        inter_sample (int): How dense the propagation within multishot discretization should be.
+
+    Other arguments:
+    The solver should likley not to be changed as it is a high accuracy 8th order runga kutta method.
+    
+    Args:
+        solver (str): The numerical solver to use for propagation (e.g., "Dopri8"). Defaults to "Dopri8".
+        args (Dict): 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.
+    """
+
+
+@dataclass
+class SimConfig:
+    x_bar: np.ndarray
+    u_bar: np.ndarray
+    initial_state: np.ndarray
+    final_state: np.ndarray
+    max_state: np.ndarray
+    min_state: np.ndarray
+    max_control: np.ndarray
+    min_control: np.ndarray
+    total_time: float
+    idx_x_true: slice
+    idx_u_true: slice
+    idx_t: slice
+    idx_y: slice
+    idx_s: slice
+    ctcs_node_intervals: list = None
+    constraints_ctcs: List[callable] = field(
+        default_factory=list
+    )  # TODO (norrisg): clean this up, consider moving to dedicated `constraints` dataclass
+    constraints_nodal: List[callable] = field(default_factory=list)
+    n_states: int = None
+    n_controls: int = None
+    S_x: np.ndarray = None
+    inv_S_x: np.ndarray = None
+    c_x: np.ndarray = None
+    S_u: np.ndarray = None
+    inv_S_u: np.ndarray = None
+    c_u: np.ndarray = None
+
+    def __post_init__(self):
+        self.n_states = len(self.max_state)
+        self.n_controls = len(self.max_control)
+
+        assert (
+            len(self.initial_state.value) == self.n_states - (self.idx_y.stop - self.idx_y.start)
+        ), f"Initial state must have {self.n_states - (self.idx_y.stop - self.idx_y.start)} elements"
+        assert (
+            len(self.final_state.value) == self.n_states - (self.idx_y.stop - self.idx_y.start)
+        ), f"Final state must have {self.n_states - (self.idx_y.stop - self.idx_y.start)} elements"
+        assert (
+            self.max_state.shape[0] == self.n_states
+        ), f"Max state must have {self.n_states} elements"
+        assert (
+            self.min_state.shape[0] == self.n_states
+        ), f"Min state must have {self.n_states} elements"
+        assert (
+            self.max_control.shape[0] == self.n_controls
+        ), f"Max control must have {self.n_controls} elements"
+        assert (
+            self.min_control.shape[0] == self.n_controls
+        ), f"Min control must have {self.n_controls} elements"
+
+        if self.S_x is None or self.c_x is None:
+            self.S_x, self.c_x = get_affine_scaling_matrices(
+                self.n_states, self.min_state, self.max_state
+            )
+            # Use the fact that S_x is diagonal to compute the inverse
+            self.inv_S_x = np.diag(1 / np.diag(self.S_x))
+        if self.S_u is None or self.c_u is None:
+            self.S_u, self.c_u = get_affine_scaling_matrices(
+                self.n_controls, self.min_control, self.max_control
+            )
+            self.inv_S_u = np.diag(1 / np.diag(self.S_u))
+
+
+@dataclass
+class ScpConfig:
+    n: int = None
+    k_max: int = 200
+    w_tr: float = 1e0
+    lam_vc: float = 1e0
+    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: float = None
+    w_tr_max_scaling_factor: 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. TODO haynec add a link to the time dilation page. 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`.
+    """
+
+    def __post_init__(self):
+        keys_to_scale = ["w_tr", "lam_vc", "lam_cost", "lam_vb"]
+        scale = max(getattr(self, key) for key in keys_to_scale)
+        for key in keys_to_scale:
+            setattr(self, key, getattr(self, key) / 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

{constraints → openscvx/constraints}/ctcs.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 from typing import Callable, Sequence, Tuple, Optional
 import functools
+import types
 
 from jax.lax import cond
 import jax.numpy as jnp
@@ -10,8 +11,18 @@ import jax.numpy as jnp
 class CTCSConstraint:
     func: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]
     penalty: Callable[[jnp.ndarray], jnp.ndarray]
-    nodes: Optional[Sequence[Tuple[int, int]]] = None
+    nodes: Optional[Tuple[int, int]] = None
     idx: Optional[int] = None
+    grad_f_x: Optional[Callable] = None
+    grad_f_u: Optional[Callable] = None
+
+    def __post_init__(self):
+        if self.grad_f_x is not None:
+            _grad_f_x = self.grad_f_x
+            self.grad_f_x = lambda x, u, nodes: _grad_f_x(x, u)
+        if self.grad_f_u is not None:
+            _grad_f_u = self.grad_f_u
+            self.grad_f_u = lambda x, u, nodes: _grad_f_u(x, u)
 
     def __call__(self, x, u, node):
         return cond(
@@ -28,6 +39,8 @@ def ctcs(
     penalty: str = "squared_relu",
     nodes: Optional[Sequence[Tuple[int, int]]] = None,
     idx: Optional[int] = None,
+    grad_f_x: Optional[Callable] = None,
+    grad_f_u: Optional[Callable] = None,
 ):
     """Decorator to mark a function as a 'ctcs' constraint.
 
@@ -44,14 +57,25 @@ def ctcs(
         def pen(x):
             r = jnp.maximum(0, x)
             return jnp.where(r < delta, 0.5 * r**2, r - 0.5 * delta)
-
+    elif penalty == "smooth_relu":
+        c = 1e-8
+        pen = lambda x: (jnp.maximum(0, x) ** 2 + c**2) ** 0.5 - c
+    elif isinstance(penalty, types.LambdaType):
+        pen = penalty
     else:
         raise ValueError(f"Unknown penalty {penalty}")
 
     def decorator(f: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]):
         # wrap so name, doc, signature stay on f
         wrapped = functools.wraps(f)(f)
-        return CTCSConstraint(func=wrapped, penalty=pen, nodes=nodes, idx=idx)
+        return CTCSConstraint(
+            func=wrapped,
+            penalty=pen,
+            nodes=nodes,
+            idx=idx,
+            grad_f_x=grad_f_x,
+            grad_f_u=grad_f_u,
+        )
 
     # if called as @ctcs or @ctcs(...), _func will be None and we return decorator
     if _func is None:

{constraints → openscvx/constraints}/nodal.py

@@ -11,20 +11,22 @@ class NodalConstraint:
     nodes: Optional[List[int]] = None
     convex: bool = False
     vectorized: bool = False
+    grad_g_x: Optional[Callable] = None
+    grad_g_u: Optional[Callable] = None
 
     def __post_init__(self):
         if not self.convex:
-            # TODO: (haynec) switch to AOT instead of JIT
-            if self.vectorized:
-                # single-node but still using JAX
-                self.g = jit(self.func)
-                self.grad_g_x = jit(jacfwd(self.func, argnums=0))
-                self.grad_g_u = jit(jacfwd(self.func, argnums=1))
-            else:
-                self.g = vmap(jit(self.func), in_axes=(0, 0))
-                self.grad_g_x = jit(vmap(jacfwd(self.func, argnums=0), in_axes=(0, 0)))
-                self.grad_g_u = jit(vmap(jacfwd(self.func, argnums=1), in_axes=(0, 0)))
-        # if convex=True and inter_nodal=False, assume an external solver (e.g. CVX) will handle it
+            # single-node but still using JAX
+            self.g = self.func
+            if self.grad_g_x is None:
+                self.grad_g_x = jacfwd(self.func, argnums=0)
+            if self.grad_g_u is None:
+                self.grad_g_u = jacfwd(self.func, argnums=1)
+            if not self.vectorized:
+                self.g = vmap(self.g, in_axes=(0, 0))
+                self.grad_g_x = vmap(self.grad_g_x, in_axes=(0, 0))
+                self.grad_g_u = vmap(self.grad_g_u, in_axes=(0, 0))
+        # if convex=True assume an external solver (e.g. CVX) will handle it
 
     def __call__(self, x: jnp.ndarray, u: jnp.ndarray):
         return self.func(x, u)
@@ -36,6 +38,8 @@ def nodal(
     nodes: Optional[List[int]] = None,
     convex: bool = False,
     vectorized: bool = False,
+    grad_g_x: Optional[Callable] = None,
+    grad_g_u: Optional[Callable] = None,
 ):
     def decorator(f: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]):
         return NodalConstraint(
@@ -43,6 +47,8 @@ def nodal(
             nodes=nodes,
             convex=convex,
             vectorized=vectorized,
+            grad_g_x=grad_g_x,
+            grad_g_u=grad_g_u,
         )
 
     return decorator if _func is None else decorator(_func)

openscvx/constraints/violation.py

@@ -0,0 +1,67 @@
+from collections import defaultdict
+from typing import List, Optional, Tuple, Callable
+from dataclasses import dataclass
+
+import jax.numpy as jnp
+
+from openscvx.constraints.ctcs import CTCSConstraint
+
+
+@dataclass
+class CTCSViolation:
+    g: Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]
+    g_grad_x: Optional[Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]] = None
+    g_grad_u: Optional[Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]] = None
+
+
+def get_g_grad_x(constraints_ctcs: List[CTCSConstraint]) -> Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]:
+    def g_grad_x(x: jnp.ndarray, u: jnp.ndarray, node: int) -> jnp.ndarray:
+        grads = [
+            c.grad_f_x(x, u, node) for c in constraints_ctcs if c.grad_f_x is not None
+        ]
+        return sum(grads) if grads else None
+
+    return g_grad_x
+
+
+def get_g_grad_u(constraints_ctcs: List[CTCSConstraint]) -> Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]:
+    def g_grad_u(x: jnp.ndarray, u: jnp.ndarray, node: int) -> jnp.ndarray:
+        grads = [
+            c.grad_f_u(x, u, node) for c in constraints_ctcs if c.grad_f_u is not None
+        ]
+        return sum(grads) if grads else None
+
+    return g_grad_u
+
+
+def get_g_func(constraints_ctcs: List[CTCSConstraint]) -> Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]:
+    def g_func(x: jnp.array, u: jnp.array, node: int) -> jnp.array:
+        return sum(c(x, u, node) for c in constraints_ctcs)
+
+    return g_func
+
+
+def get_g_funcs(constraints_ctcs: List[CTCSConstraint]) -> List[CTCSViolation]:
+    # Bucket by idx
+    groups: dict[int, List[CTCSConstraint]] = defaultdict(list)
+    for c in constraints_ctcs:
+        if c.idx is None:
+            raise ValueError(f"CTCS constraint {c} has no .idx assigned")
+        groups[c.idx].append(c)
+
+    # For each bucket, build one CTCSViolation
+    violations: List[CTCSViolation] = []
+    for idx, bucket in sorted(groups.items(), key=lambda kv: kv[0]):
+        g = get_g_func(bucket)
+        g_grad_x = get_g_grad_u(bucket) if all(c.grad_f_x for c in bucket) else None
+        g_grad_u = get_g_grad_x(bucket) if all(c.grad_f_u for c in bucket) else None
+
+        violations.append(
+            CTCSViolation(
+                g=g,
+                g_grad_x=g_grad_x,
+                g_grad_u=g_grad_u,
+            )
+        )
+
+    return violations