Benchmarks, Tests, and Style

[alberthli]

Action Items

This PR adds benchmarks and tests for this implementation vs. OSQP and cvxpy as a front end to ECOS, Mosek, and Gurobi (if we really wanted to split hairs, we could forgo the cvxpy interface to get a fully optimized benchmark). Also adds linters for style checking.

Style

• flake8, black, and isort + pre-commit hooks

Timing benchmarks

• Random QPs for OSQP, ECOS, Mosek, and Gurobi

Correctness tests (compare vs. any cvxpy implementation when applicable)

• Convergence checks
  • Change the hypothesis strategy to search over input matrices instead of the batch size/dimensions
• Primal infeasibility and dual infeasibility checks
• ADMM steps (Q: what do I use as ground truth?)

Convergence Bugs: Our Random Sampling

1. For the case

   P = jnp.array([[[ 1.92373456e-09,  4.96351864e-11, -8.26731886e-11,
         8.77930252e-12],
       [ 4.96351864e-11,  2.18795443e-10,  4.03629116e-10,
        -4.28455352e-11],
       [-8.26731886e-11,  4.03629116e-10,  7.60485468e-10,
        -8.07259096e-11],
       [ 8.77930252e-12, -4.28455352e-11, -8.07259096e-11,
         8.56914989e-12]]], dtype=jnp.float64)
   q = jnp.array([[1.00000000e+00, 1.90283739e-17, 5.17593600e-10, 1.90283739e-17]], dtype=jnp.float64)
   A = jnp.array([[[ 3.49829621e-048, -7.26732123e-011, -9.99710513e-001,
         1.59881067e-010],
       [ 9.99999990e-001,  4.55592754e-319,  3.31970124e-010,
        -2.25229027e-010],
       [-1.45230074e-038, -1.33981414e-016, -9.78299107e-001,
        -1.34260656e-309],
       [-1.42767930e-016,  1.30147266e-318, -9.98350799e-001,
        -3.48633396e-017],
       [ 1.97980986e-038,  0.00000000e+000, -9.78002597e-001,
         1.83027221e-311],
       [-1.02376831e-010, -9.99999990e-001,  1.62131470e-010,
         1.11253702e-318],
       [-3.37974879e-147,  1.36860783e-017, -9.99941036e-001,
        -2.55454172e-318],
       [ 4.55592754e-319, -2.22507404e-319,  2.40501275e-319,
        -9.99999990e-001],
       [-3.07130492e-010, -1.11253702e-318,  1.27056418e-048,
        -9.99999990e-001],
       [ 0.00000000e+000, -1.27056418e-048,  9.99999990e-001,
         2.16176374e-010]]], dtype=jnp.float64)
   l = jnp.array([[-7.26732127e-006, -2.86920992e-050, -3.65475180e-319,
       -1.94969927e-005, -5.34305159e+014,  9.99999995e+001,
       -2.29615441e-005, -1.09999999e-005, -9.99999995e+004,
       -4.84254509e+004]], dtype=jnp.float64)
   u = jnp.array([[ 2.18019638e-005, -9.91416417e-286,  3.22199312e+014,
        0.00000000e+000,  2.74188986e+004,  6.23567316e+004,
        0.00000000e+000,  2.22507385e-313,  4.18553362e+004,
        1.08087647e+005]], dtype=jnp.float64)
   

   jaxosqp does not converge while cvxpy using OSQP does. The optimal value returned is 6.610848939084965 and the optimal solution is array([ 4.26391780e-09, -1.60000255e-03, -2.47453266e-13, -5.24629431e-14]).

2. Similar failure case:

   P = jnp.array([[[ 5.11510887e-005, -1.44462998e-313,  4.68912166e-005],
       [-1.44462998e-313,  1.26933527e+000,  9.55236872e-315],
       [ 4.68912166e-005,  9.55236872e-315,  1.26933527e+000]]],  dtype=jnp.float64)
   q = jnp.array([[-6.51595591e-45,  1.00000000e+00, -3.84133271e-01]], dtype=jnp.float64)
   A = jnp.array([[[ 5.64539728e-10,  6.67376968e-01, -1.86550084e-10],
       [ 9.95239602e-01,  2.94514822e-01, -1.79768562e-01],
       [ 9.99999990e-01,  0.00000000e+00, -1.94215985e-49]]],  dtype=jnp.float64)
   l = jnp.array([[-1.25354255e-020, -2.86220631e-005, -2.22507384e-316]], dtype=jnp.float64)
   u = jnp.array([[1.25615516e-313, 1.50642438e-010, 2.22044604e-021]], dtype=jnp.float64)
   

   Optimal value: -1.3149429255238995e-10
   Optimal solution: array([4.94630958e-22, 8.77618090e-20, 2.63057785e-10])

3. Similar failure case:

   P = jnp.array([[[ 3.15762871e-08, -1.08627949e-04,  1.46096878e-09,
        -6.07272241e-09, -6.96343076e-09,  6.58799401e-10,
        -3.87654539e-05],
       [-1.08627949e-04,  7.88703719e-01, -1.73691110e-05,
         6.92811797e-05, -3.62857075e-05,  1.00194573e-06,
        -4.43349031e-01],
       [ 1.46096878e-09, -1.73691110e-05,  4.83564964e-10,
        -1.90220820e-09,  1.50426046e-10, -1.08516540e-10,
         1.80749277e-05],
       [-6.07272241e-09,  6.92811797e-05, -1.90220820e-09,
         7.48829550e-09, -7.70854759e-10,  4.10077351e-10,
        -6.99075379e-05],
       [-6.96343076e-09, -3.62857075e-05,  1.50426046e-10,
        -7.70854759e-10,  4.58938962e-08,  5.08585555e-10,
         1.89335957e-05],
       [ 6.58799401e-10,  1.00194573e-06, -1.08516540e-10,
         4.10077351e-10,  5.08585555e-10,  7.52302780e-11,
        -7.67376140e-06],
       [-3.87654539e-05, -4.43349031e-01,  1.80749277e-05,
        -6.99075379e-05,  1.89335957e-05, -7.67376140e-06,
         9.99999990e-01]]], dtype=jnp.float64)
   q = jnp.array([[-1.00000000e+00,  0.00000000e+00, -1.65132587e-09,
        1.18552977e-46,  0.00000000e+00,  5.82337270e-47,
        1.68853736e-05]], dtype=jnp.float64)
   A = jnp.array([[[ 2.22329541e-321,  8.60497183e-265, -2.00000998e-010,
         6.17904100e-001, -1.17549434e-048,  9.99999990e-001,
        -1.51680959e-315],
       [-1.21520939e-058, -9.99999990e-001,  2.50382818e-039,
        -5.16930259e-008,  0.00000000e+000, -4.73945500e-314,
        -2.03470406e-042],
       [ 2.91411018e-010,  7.90466970e-043, -1.66104280e-009,
         9.98800506e-001, -9.61656358e-010,  1.94118733e-025,
         3.55217174e-013],
       [ 0.00000000e+000,  1.62938082e-003,  0.00000000e+000,
         9.99999990e-001,  5.12187532e-010,  5.12187532e-010,
        -1.92033125e-006],
       [-9.99999990e-001, -3.41111378e-254, -1.88635361e-001,
         4.52984087e-319, -2.22507404e-319,  4.99999995e-011,
         1.51680959e-315],
       [ 7.50003212e-001,  0.00000000e+000, -9.99999990e-001,
        -2.26492044e-318,  4.09759163e-001, -9.99999990e-001,
         6.81689532e-012],
       [-1.17549434e-048, -1.70795404e-021,  1.11253702e-318,
         2.26492044e-318,  9.99999990e-001, -9.99999990e-004,
         1.51608984e-319]]], dtype=jnp.float64)
   l = jnp.array([[-1.19209289e-012, -2.13002146e-001, -3.59549592e+005,
       -5.69824698e-313,  6.10351559e-010, -1.19209289e-012,
       -9.99999995e+004]], dtype=jnp.float64)
   u = jnp.array([[ 9.99999995e+001, -2.98479579e-041,  5.33589512e-009,
        5.69824699e-316,  3.33333332e-006,  1.49999999e-005,
       -1.89999999e-005]], dtype=jnp.float64)
   

   Optimal value: -13243690.119393358
   Optimal solution: array([ 6.78801264e-02, -2.95045580e+00, -3.59848352e-01, 9.69489208e-08, -9.99999150e-01, 9.99849017e-04, 7.54688174e-02])

4. Similar failure case:

   P = jnp.array([[[ 9.54672887e-17,  1.05715714e-32, -5.76389386e-19,
        -4.98635559e-14,  9.95718360e-19, -9.52206754e-17],
       [ 1.05715714e-32,  6.52688679e-13,  1.66886960e-34,
        -5.53593649e-24, -9.70481311e-41, -1.04991087e-26],
       [-5.76389386e-19,  1.66886960e-34,  6.39212927e-16,
        -7.87161384e-16,  5.27434780e-17, -1.50319149e-18],
       [-4.98635559e-14, -5.53593649e-24, -7.87161384e-16,
         2.63636716e-05, -1.82756717e-19,  4.98635560e-08],
       [ 9.95718360e-19, -9.70481311e-41,  5.27434780e-17,
        -1.82756717e-19,  4.36343314e-18,  1.98688881e-24],
       [-9.52206754e-17, -1.04991087e-26, -1.50319149e-18,
         4.98635560e-08,  1.98688881e-24,  9.52207709e-11]]],  dtype=jnp.float64)
   q = jnp.array([[-1.61187347e-002,  1.13819874e-018,  3.79316743e-011,
       -1.00000000e+000,  1.12235061e-049,  2.08105391e-319]], dtype=jnp.float64)
   A = jnp.array([[[ 1.49999999e-010, -9.99999990e-001, -9.99048342e-001,
         2.33043199e-315,  9.08519450e-001,  2.22329541e-321],
       [ 2.12447892e-312,  9.54790292e-005, -2.52877826e-004,
         9.99999990e-001,  1.16551549e-008,  2.86437088e-004],
       [-5.96046442e-018, -2.47032823e-323,  2.64851691e-003,
        -1.46765049e-051,  5.93564137e-001,  9.99999990e-001],
       [ 7.29573026e-001,  5.15185572e-010,  8.44176761e-048,
        -9.46631250e-007, -9.99025389e-001, -3.01665614e-318],
       [ 7.93501629e-287, -4.99999995e-011, -2.48211644e-001,
         2.33060646e-319,  9.99999990e-001, -9.99999990e-001],
       [-9.99999990e-001, -9.99999990e-001,  2.64851691e-015,
         1.98996573e-006,  2.22507404e-318,  9.99999990e-001]]], dtype=jnp.float64)
   l = jnp.array([[-1.09999999e-005,  1.07958898e-185, -2.65201840e+004,
       -2.71150302e+002, -1.17549435e-043, -9.99999995e+004]], dtype=jnp.float64)
   u = jnp.array([[ 9.99999995e+04,  1.30329561e+11,  5.96046445e-13,
        1.61618175e-12,  1.40129846e-50, -9.99989995e-06]], dtype=jnp.float64)
   

   Optimal value: -19863515953.974216
   Optimal solution: array([ 3.92104260e-07, 4.07435525e-07, -1.25582277e-06, 3.97270312e+03, -7.30674681e-08, 4.46252001e-08])

5. Similar failure case:

   P = jnp.array([[[1.02399994e-07, 0.00000000e+00],
       [0.00000000e+00, 1.99999990e+00]]], dtype=jnp.float64)
   q = jnp.array([[1.02399994e-07, 3.25900366e-09]], dtype=jnp.float64)
   A = jnp.array([[[-9.99999990e-001,  3.25900383e-014],
       [ 3.29036219e-004, -3.32561833e-278],
       [-2.22507385e-278,  2.28541858e-073],
       [-4.16631871e-294,  4.44000126e-001],
       [ 8.33263742e-297, -4.44000126e-001]]], dtype=jnp.float64)
   l = jnp.array([[-9.99999995e+001, -6.22827232e-272,  1.19209290e+033,
        8.33263747e-289, -3.74484384e+019]], dtype=jnp.float64)
   u = jnp.array([[0.00000000e+00, 0.00000000e+00, 9.23504030e+48, 7.11527445e+19,
       5.61732194e+19]], dtype=jnp.float64)
   

   Optimal value: -2.616281296751756e-21
   Optimal solution: array([-2.32378377e-18, -8.70083478e-14])