jaxopt.OSQP

class jaxopt.OSQP(*, matvec_Q=None, matvec_A=None, matvec_G=None, fun=None, **kwargs)[source]

OSQP solver for general quadratic programming.

Meant as drop-in replacement for CvxpyQP. Depending on the format of your problem, BoxOSQP API may be more appropriate. Support for matvec and pytrees. Supports jit and vmap.

CvxpyQP is more precise and should be preferred on CPU. OSQP can be quicker than CvxpyQP when GPU/TPU are available.

The objective function is:

0.5 * x^T Q x + c^T x subject to Gx <= h, Ax = b.

The attributes must be given as keyword arguments. Hyper-parameters defaults to the same values as in BoxQP.

Parameters

  • matvec_Q (Optional[Callable]) –

  • matvec_A (Optional[Callable]) –

  • matvec_G (Optional[Callable]) –

  • fun (Optional[Callable]) –

matvec_Q

(optional) a Callable matvec_Q(params_Q, x). By default, matvec_Q(P, x) = tree_dot(P, x), where the pytree Q = params_Q matches x structure. matvec_Q incompatible with the specification of fun. The shape of primal variables may be inferred from params_obj = (matvec_Q, c).

matvec_A

(optional) a Callable matvec_A(params_A, x). By default, matvec_A(A, x) = tree_dot(A, x), where tree pytree A = params_A matches x structure.

matvec_G

(optional) a Callable matvec_G(params_G, x). By default, matvec_G(G, x) = tree_dot(G, x), where tree pytree G = params_G matches x structure.

fun

(optional) a function with signature fun(params, params_obj) that is promised to be quadratic polynomial convex with respect to params, i.e fun can be written

fun(x, params_obj) = 0.5*jnp.dot(x, jnp.dot(Q, x)) + jnp.dot(c, x) + cste

with params_obj a pytree that contains the parameters of the objective function. (Q, c) do not need to be explicited in params_obj by the user: c will be inferred by Jaxopt,

and the operator x -> Qx will be computed upon request.

fun incompatible with the specification of matvec_Q. Note that the shape of primal cannot be inferred from params_obj anymore, so the user should provide it in init_params.

check_primal_dual_infeasability

if True populates the status field of state with one of BoxOSQP.PRIMAL_INFEASIBLE, BoxOSQP.DUAL_INFEASIBLE. (default: True). If False it improves speed but does not check feasability. If jit=False, and if the problem is primal or dual infeasible, then a ValueError exception is raised.

sigma

ridge regularization parameter in linear system.

momentum

relaxation parameter (default: 1.6), must belong to the open interval (0,2). momentum=1 => no relaxation. momentum<1 => under-relaxation. momentum>1 => over-relaxation. Boyd [2, p21] suggests chosing momentum in [1.5, 1.8].

eq_qp_solve

‘cg’, ‘cg+jacobi’ or ‘lu’ (default: ‘cg’). ‘cg’ is conjugate gradient: an indirect solver that works with matvecs or pytree of matrices. ‘cg+jacobi’ is conjugate gradient with Jacobi preconditioning: only works on pytree of matrices

but can provide speedup.

‘lu’ is LU factorization: a direct solver that only work on pytree of matrices.

rho_start

initial learning rate. (default: 1e-1)

rho_min

minimum learning rate. (default: 1e-6)

rho_max

maximum learning rate. (default: 1e6)

stepsize_updates_frequency

frequency of stepsize updates. (default: 10). One every stepsize_updates_frequency updates computes a new stepsize.

primal_infeasible_tol

relative tolerance for primal infeasability detection. (default: 1e-4)

dual_infeasible_tol

relative tolerance for dual infeasability detection. (default: 1e-4)

maxiter

maximum number of iterations. (default: 4000)

tol

absolute tolerance for stoping criterion (default: 1e-3).

termination_check_frequency

frequency of termination check. (default: 5). One every termination_check_frequency the error is computed.

implicit_diff_solve

the linear system solver to use.

__init__(*, matvec_Q=None, matvec_A=None, matvec_G=None, fun=None, **kwargs)[source]
Parameters

  • matvec_Q (Optional[Callable]) –

  • matvec_A (Optional[Callable]) –

  • matvec_G (Optional[Callable]) –

  • fun (Optional[Callable]) –

Methods

__init__(*[, matvec_Q, matvec_A, matvec_G, fun])

attribute_names()

attribute_values()

init_params(init_x, params_obj, params_eq, ...)

Return default params for initialization.

l2_optimality_error(params, params_obj, ...)

Computes the L2 norm of the KKT residuals.

run([init_params, params_obj, params_eq, ...])

Runs the quadratic programming solver in BoxOSQP.

Attributes

matvec_A_box

init_params(init_x, params_obj, params_eq, params_ineq)[source]

Return default params for initialization.

Parameters

  • init_x – initial primal solution.

  • params_obj – see the doc of run method.

  • params_eq – see the doc of run method.

  • params_ineq – see the doc of run method.

Returns

a pytree KKTSolution of parameters for BoxOSQP.

Return type

init_params

l2_optimality_error(params, params_obj, params_eq, params_ineq)[source]

Computes the L2 norm of the KKT residuals.

Note that this function is exposed for consistency of the API, but the differentiation is actually performed in the BoxOSQP class.

Parameters

  • params (Array) –

  • params_obj (Union[Tuple[Array, Array], Any]) –

  • params_eq (Optional[Tuple[Array, Array]]) –

  • params_ineq (Optional[Tuple[Array, Array]]) –

run(init_params=None, params_obj=None, params_eq=None, params_ineq=None)[source]

Runs the quadratic programming solver in BoxOSQP.

The returned params contains both the primal and dual solutions.

Parameters

  • init_params (Optional[Any]) – init_params: (optional) initial KKTSolution for warm-start. Must be provided if fun is not None.

  • params_obj (Union[Tuple[Any, Any], Any, None]) – parameters of the quadratic objective, can be: a tuple (Q, c) with Q a pytree of matrices, or a tuple (params_Q, c) if matvec_Q is provided, or an arbitrary pytree if fun is provided.

  • params_eq (Optional[Tuple[Any, Any]]) – (A, b) or None if no equality constraints.

  • params_ineq (Optional[Tuple[Any, Any]]) – (G, h) or None if no inequality constraints.

Return type

OptStep

Returns

(params, state), params = (primal_var, dual_var_eq, dual_var_ineq)