diffapqp 0.1.0__py3-none-any.whl

lao/utils.py

@@ -0,0 +1,459 @@
+"""
+Helper functions
+"""
+
+import cvxpy as cp
+import numpy as np
+import scipy.sparse as ssparse
+
+
+def _to_dense_array(matrix):
+    return matrix.toarray() if ssparse.issparse(matrix) else np.asarray(matrix)
+
+
+def _to_dense_vector(matrix):
+    return np.asarray(_to_dense_array(matrix)).reshape(-1)
+
+
+def _matrix_column(matrix, col_idx):
+    if ssparse.issparse(matrix):
+        return matrix[:, [col_idx]]
+    return matrix[:, col_idx]
+
+
+def _has_nonzero_entries(matrix):
+    if ssparse.issparse(matrix):
+        return bool(np.count_nonzero(matrix.data))
+    return bool(np.count_nonzero(np.asarray(matrix)))
+
+
+def _expr_has_parameters(expr):
+    return len(expr.parameters()) > 0
+
+
+def _iter_expression_tree(expr):
+    # Walk CVXPY expression nodes depth-first through their child args.
+    # including objective and constraints
+    yield expr
+    for child in getattr(expr, "args", []):
+        yield from _iter_expression_tree(child)
+
+
+def validate_apqp_problem(
+    prob,
+    *,
+    canonicalize_solver="OSQP",
+    require_dpp=True,
+    require_continuous=True,
+):
+    """Validate that a CVXPY problem matches the APQP/APLP form expected by DiffAPQP.
+
+    Accepted form:
+
+        min 1/2 x^T P x + q(theta)^T x
+        s.t. A x = b(theta)
+             G x <= h(theta)
+
+    where parameters may appear affinely in the linear objective and constraint
+    offsets. Equivalent syntactic placements such as ``x + p <= c`` are
+    accepted if CVXPY canonicalization leaves the variable coefficient matrices
+    parameter-free. Parameterized P, A, or G are rejected.
+    """
+    
+    # minimum test
+    if prob is None:
+        raise ValueError("A CVXPY problem instance is required for validation.")
+
+    if not isinstance(prob.objective, cp.Minimize):
+        raise ValueError("DiffAPQP only supports minimization problems.")
+
+    if not prob.is_dcp():
+        raise ValueError("DiffAPQP requires a DCP-compliant CVXPY problem.")
+
+    if not prob.is_qp():
+        raise ValueError("DiffAPQP only supports QP/LP problems in CVXPY's QP class.")
+
+    if prob.parameters() and require_dpp and not prob.is_dpp():
+        raise ValueError(
+            "DiffAPQP requires a DPP-compliant parametric problem. "
+            "Parameters must enter affinely in the supported APQP/APLP form."
+        )
+
+    # Parameters may not appear inside the quadratic matrix P(theta).
+    for node in _iter_expression_tree(prob.objective.expr):
+        if type(node).__name__ == "QuadForm" and _expr_has_parameters(node):
+            raise ValueError(
+                "DiffAPQP does not support parameterized quadratic matrices. "
+                "Parameters may appear only in the linear objective term and "
+                "constraint right-hand sides."
+            )
+
+    # Constraints must be affine in variables/parameters. We do not require the
+    # original CVXPY left-hand side to be parameter-free, because expressions
+    # like x + p <= c are equivalent to x <= c - p and still have constant
+    # variable coefficients after canonicalization.
+    for idx, constr in enumerate(prob.constraints):
+        if not constr.expr.is_affine():
+            raise ValueError(
+                f"Constraint {idx} is not affine. DiffAPQP supports only "
+                "affine equality and inequality constraints."
+            )
+
+    # Currently we only support continuous problems: all variables must be continuous.
+    if require_continuous:
+        has_discrete_vars = any(
+            bool(getattr(var, "boolean_idx", [])) or bool(getattr(var, "integer_idx", []))
+            for var in prob.variables()
+        )
+        if has_discrete_vars:
+            raise ValueError(
+                "DiffAPQP layers currently support only continuous problems. "
+                "Relax integer/boolean variables before using the differentiable layers."
+            )
+
+    param_qp_prog, _, zero_dim, int_dim, bool_dim = return_compiler(
+        prob, canonicalize_solver=canonicalize_solver
+    )
+
+    if _has_nonzero_entries(param_qp_prog.P[:, :-1]):
+        raise ValueError(
+            "DiffAPQP does not support parameterized quadratic matrices. "
+            "Parameters may appear only in the linear objective term and "
+            "constraint offsets."
+        )
+
+    no_cons = param_qp_prog.constr_size
+    if no_cons > 0 and _has_nonzero_entries(param_qp_prog.A[:-no_cons, :-1]):
+        raise ValueError(
+            "DiffAPQP does not support parameterized constraint matrices. "
+            "Parameters may shift constraint offsets, but A and G must remain constant."
+        )
+    
+    if require_continuous:
+        if len(int_dim) > 0 or len(bool_dim) > 0:
+            raise ValueError(
+                "DiffAPQP layers currently support only continuous problems. "
+                "Relax integer/boolean variables before using the differentiable layers."
+            )
+
+
+def is_apqp_problem(prob, **kwargs):
+    """Return True iff ``validate_apqp_problem`` succeeds."""
+    try:
+        validate_apqp_problem(prob, **kwargs)
+    except Exception:
+        return False
+    return True
+
+
+def _flattened_size(shape):
+    """Return flattened size for scalar/vector/matrix CVXPY objects. This is also the number of elements"""
+    if len(shape) == 0:
+        return 1
+    return int(np.prod(shape))
+
+
+def flatten_fortran_numpy_sample(value):
+    """Flatten a single CVXPY-shaped value using default CVXPY/Fortran order."""
+    arr = np.asarray(value)
+    if arr.ndim == 0:
+        # scalar
+        return np.array([arr.item()])
+    if arr.ndim == 1:
+        # vector
+        return arr
+    # matrix
+    return arr.reshape(-1, order="F")
+
+
+def unflatten_fortran_torch_batch(x_flat, shape):
+    """Recover batched flat CVXPY data into its original shape."""
+    # x_flat.shape == (batch_size, flattened_size)
+    # shape is the original shape of the variable
+    if len(shape) == 0:
+        return x_flat[:, 0]
+    if len(shape) == 1:
+        return x_flat
+
+    view_shape = (x_flat.shape[0], *reversed(shape))
+    permute_dims = (0, *range(len(shape), 0, -1))
+    return x_flat.reshape(view_shape).permute(permute_dims)
+
+
+def flatten_fortran_torch_batch(x, shape):
+    """Flatten batched original-shaped data using CVXPY/Fortran order."""
+    if len(shape) == 0:
+        return x.reshape(x.shape[0], 1)
+    if len(shape) == 1:
+        return x.reshape(x.shape[0], -1)
+
+    permute_dims = (0, *range(x.ndim - 1, 0, -1))
+    return x.permute(permute_dims).reshape(x.shape[0], -1)
+
+def return_compiler(prob, canonicalize_solver="OSQP"):
+    """
+    return the compiler of the problem given by cvxpy
+    args:
+        - prob: the cvxpy problem
+        - canonicalize_solver: cvxpy solver name used for problem stuffing
+            default is open-source solver OSQP
+    return:
+        - param_qp_prog: the qp problem in standard form
+        - params_idx: {param_id: param_name}, link the id to the parameter name
+        - zero_dim: the number of equality constraint in the standard form 
+                                        (always be the first rows in the matrix)
+        - int_vars_idx: the index of integer variables
+        - bool_vars_idx: the index of boolean variables
+    """
+
+    # canonicalize the problem
+    solver_name = canonicalize_solver.upper()
+    if not hasattr(cp, solver_name):
+        raise ValueError(f"Unknown CVXPY solver for canonicalization: {canonicalize_solver}")
+    solver = getattr(cp, solver_name)
+    solver_opts = {"use_quad_obj": True} if solver_name == "GUROBI" else {}
+    data, _, _ = prob.get_problem_data(
+                solver=solver, solver_opts=solver_opts)
+
+    assert prob.is_qp(), 'only support QP (and LP) for now'
+    assert data['dims'].exp == 0, 'does not support exponential cone'
+    assert len(data['dims'].psd) == 0, 'does not support positive semidefinite cone'
+    # assert len(data['dims'].soc) == 0, 'does not support second-order cone'
+
+    # parametric QP problem
+    param_qp_prog = data[cp.settings.PARAM_PROB]
+    
+    # ! the order of parameter idx is changed internally in cvxpy so we link the parameter id to its name
+    params_idx = {p.id: p.name() for p in prob.parameters()}  
+
+    return param_qp_prog, params_idx, data['dims'].zero, data['int_vars_idx'], data['bool_vars_idx']
+
+def return_standard_form(prob, sparse=False, canonicalize_solver="OSQP"): 
+    """
+    This function will identify all the 
+    parameters (data) from the param_qp_prog and return the standard form.
+    
+    The standard form we consider is:
+    min 1/2 x^T P x + (q + sum Q_i z_i)^T x + c + sum d_i^T z_i
+    s.t. A x = b + sum B_i z_i
+         G x <= h + sum H_i z_i
+    in which x is the decision variable, z_i is the i-th parameter.
+    args:
+        - prob: the cvxpy problem
+        - sparse: whether return sparse matrices
+        - canonicalize_solver: cvxpy solver name used for problem stuffing
+    """
+    if prob is not None:
+        validate_apqp_problem(prob, canonicalize_solver=canonicalize_solver)
+    
+    # canonicalize the problem
+    param_qp_prog, param_id_to_name, zero_dim, int_dim, bool_dim = return_compiler(
+        prob, canonicalize_solver=canonicalize_solver
+    )
+
+    # Get the total number of constraints and variables
+    no_cons = param_qp_prog.constr_size
+    no_var = param_qp_prog.reduced_A.var_len
+
+    # Extract the constant quadratic cost matrix P without densifying the full
+    # parameter-to-P map, which has no_var**2 rows and can be enormous.
+    P = _to_dense_vector(_matrix_column(param_qp_prog.P, -1)).reshape(no_var, no_var)
+    
+    # Extract the linear cost vector q
+    q_raw = _to_dense_array(param_qp_prog.q)
+    q = q_raw[:-1,-1]         # q vector
+    Q_tilde = q_raw[:-1,:-1]  # parameter matrices
+    c = float(q_raw[-1, -1])   # scalar objective offset
+    d_tilde = q_raw[-1, :-1]   # parameter-to-offset vectors
+
+    # Extract and reshape the constraint matrix A_tilde
+    # This contains both equality (Ax=b) and inequality constraints (Gx<=h)
+    if no_cons > 0:
+        A_const_col = _to_dense_vector(_matrix_column(param_qp_prog.A, -1))
+        # A_tilde = A_raw[:int(no_cons * no_var), -1]
+        A_tilde = A_const_col[:-no_cons].reshape(no_var, no_cons).T
+        # b_tilde = A_raw[int(no_cons * no_var):, -1]
+        b_tilde = A_const_col[-no_cons:]
+    else:
+        # Keep explicit empty objects for unconstrained problems.
+        A_tilde = np.zeros((0, no_var))
+        b_tilde = np.zeros((0,))
+
+    # Split A_tilde into:
+    # A: equality constraints (first zero_dim rows)
+    # G: inequality constraints (remaining rows, with negative sign)
+    A = A_tilde[:zero_dim,:]
+    G = -A_tilde[zero_dim:,:]
+
+    # Split b_tilde into:
+    # b: equality constraints right-hand side (with negative sign)
+    # h: inequality constraints right-hand side
+    b = -b_tilde[:zero_dim]
+    h = b_tilde[zero_dim:]
+
+    # Extract the parameter coefficient matrix B_tilde
+    # B_tilde = A_raw[int(no_cons * no_var):, :-1]
+    if no_cons > 0:
+        B_tilde = _to_dense_array(param_qp_prog.A[-no_cons:, :-1])
+    else:
+        B_tilde = np.zeros((0, param_qp_prog.A.shape[1] - 1))
+
+    # Initialize dictionaries for parameter matrices
+    B = {}  # For equality constraints (B_i)
+    H = {}  # For inequality constraints (H_i)
+    Q = {}  # For linear cost parameters (Q_i)
+    d = {}  # For scalar objective-offset parameters (d_i)
+
+    # Get mappings for parameters
+    param_id_to_col = param_qp_prog.param_id_to_col    # Maps parameter ID to starting column
+    param_id_to_size = param_qp_prog.param_id_to_size  # Maps parameter ID to its size
+    var_name_to_size = {
+        v.name(): _flattened_size(v.shape) for v in param_qp_prog.variables
+    }
+
+    # For each parameter:
+    for key, start_idx in param_id_to_col.items():
+        if key == -1:  # Skip if key is -1 (special value)
+            continue
+        
+        size = param_id_to_size[key]        # Get parameter size
+        name = param_id_to_name[key]        # Get parameter name
+        
+        # Extract parameter matrices:
+        # B[name]: how parameter affects equality constraints
+        # H[name]: how parameter affects inequality constraints
+        B[name] = -B_tilde[:zero_dim, start_idx:start_idx+size]
+        H[name] = B_tilde[zero_dim:, start_idx:start_idx+size]
+        Q[name] = Q_tilde[:, start_idx:start_idx+size]
+        d[name] = d_tilde[start_idx:start_idx+size]
+
+    if sparse:
+        P = ssparse.csc_matrix(P)
+        A = ssparse.csc_matrix(A)
+        G = ssparse.csc_matrix(G)
+        for key in B.keys():
+            B[key] = ssparse.csc_matrix(B[key])
+        for key in H.keys():
+            H[key] = ssparse.csc_matrix(H[key])
+
+    return P, q, Q, c, d, A, G, b, h, B, H, var_name_to_size, int_dim, bool_dim # Q, d, B, H are dictionaries
+
+def return_standard_form_in_cvxpy(prob, sparse=False, canonicalize_solver="OSQP"):
+    """
+    Return the standard form of the problem formulated as cvxpy problem
+    args:
+        - prob: the cvxpy problem, the parameters must be assigned a name
+        - sparse: whether to return the sparse matrix, default is False
+        - canonicalize_solver: cvxpy solver name used for problem stuffing
+    return:
+        - prob: the cvxpy problem in standard form
+        - prob_data: the data of the problem, including P, q, c, A, G, b, h, Q, d, B, H
+    Standard form:
+    min 1/2 x^T P x + q^T x + c
+    s.t. A x = b + sum B_i z_i
+         G x <= h + sum H_i z_i
+    in which x is the decision variable, z_i is the i-th parameter
+    NOTE: the formulation accepts multiple parameters for both inequality and equality constraints
+    NOTE: you still need to assign the parameters by param.value = your_value to the problem after this conversion
+    NOTE: Only support continuous variables, So for discrete QP, the discrete variables must be relaxed first.
+    """
+    
+    P, q, Q, c, d, A, G, b, h, B, H, var_name_to_size, int_dim, bool_dim = return_standard_form(
+        prob, sparse=sparse, canonicalize_solver=canonicalize_solver
+    )
+    
+    x = cp.Variable(P.shape[1])
+    # Build parameters from the union of objective/constraint parameter sets.
+    # Some parameters may appear only in Q (objective), or only in B/H (constraints).
+    param_keys = list(dict.fromkeys(list(Q.keys()) + list(d.keys()) + list(B.keys()) + list(H.keys())))
+    parameters = {}
+    for key in param_keys:
+        if key in Q:
+            size = Q[key].shape[1]
+        elif key in d:
+            size = d[key].shape[0]
+        elif key in B:
+            size = B[key].shape[1]
+        else:
+            size = H[key].shape[1]
+        parameters[key] = cp.Parameter(size, name=key)
+    
+    # formulate the cvxpy problem
+    constraints = []
+    
+    q_ = cp.Constant(np.zeros(P.shape[0]))
+    c_ = cp.Constant(c)
+    b_ = cp.Constant(np.zeros(A.shape[0])) if A.shape[0] > 0 else None
+    h_ = cp.Constant(np.zeros(G.shape[0])) if G.shape[0] > 0 else None
+    
+    for key in parameters.keys():
+        if key in B and b_ is not None:
+            b_ += B[key] @ parameters[key]
+        if key in H and h_ is not None:
+            h_ += H[key] @ parameters[key]
+        if key in Q:
+            q_ += Q[key] @ parameters[key]
+        if key in d:
+            c_ += d[key] @ parameters[key]
+    
+    if A.shape[0] > 0:
+        constraints.append(A @ x == b_ + b)
+    if G.shape[0] > 0:
+        constraints.append(G @ x <= h_ + h)
+    
+    objective = cp.Minimize(0.5 * cp.quad_form(x, cp.psd_wrap(P)) + (q_ + q) @ x + c_)
+    prob = cp.Problem(objective, constraints)
+    
+    return prob, var_name_to_size, {"P": P, "Q": Q, "q": q, "c": c, "d": d, "A": A, "G": G, "b": b, "h": h, "B": B, "H": H}
+
+def get_default_solver_args(solver, eps = None):
+    
+    if solver.upper() == "OSQP":
+        # see https://osqp.org/docs/interfaces/solver_settings.html
+        solver_args = {
+            "eps_abs": 1e-5 if eps is None else eps, # in cvxpy, the default is 1e-5
+            "eps_rel": 1e-5 if eps is None else eps, # in cvxpy, the default is 1e-5
+            "max_iter": 100000,
+        }
+    elif solver.upper() == "GUROBI":
+        solver_args = {
+            "Threads": 1, # this will limit the number of threads to 1, each problem is solved in a single cpu core
+            "OptimalityTol": 1e-8 if eps is None else eps,  # set as default # ! for simplex method
+            "FeasibilityTol": 1e-8 if eps is None else eps, # set as default
+            "BarConvTol": 1e-8 if eps is None else eps, # ! this is for the barrier method
+        }
+    elif solver.upper() == "SCS":
+        solver_args = {
+            "eps_abs": 1e-6 if eps is None else eps,  # default 1e-4
+            "eps_rel": 1e-6 if eps is None else eps,  # default 1e-4
+            "max_iters": 10000,
+        }
+    elif solver.upper() == "ECOS":
+        solver_args = {
+            "abstol": 1e-8 if eps is None else eps,
+            "reltol": 1e-8 if eps is None else eps,
+            "max_iters": 200,
+        }
+    elif solver.upper() == "CLARABEL":
+        solver_args = {
+            "tol_gap_abs": 1e-8 if eps is None else eps,
+            "tol_gap_rel": 1e-8 if eps is None else eps,
+            "max_iter": 1000
+        }
+    elif solver.upper() == "QPALM":
+        solver_args = {
+            "eps_abs": 1e-6 if eps is None else eps,
+            "eps_rel": 1e-6 if eps is None else eps,
+            "max_iter": 10000,
+        }
+    elif solver.upper() == "PROXQP":
+        solver_args = {
+            "eps_abs": 1e-6 if eps is None else eps,
+            "eps_rel": 1e-6 if eps is None else eps,
+            "max_iter": 10000,
+            "backend": "sparse"  # solve the memory issue
+        }
+    else:
+        solver_args = {}
+    
+    return solver_args