MPax is very very slow at solving some linear programming problems.

[fromseto]

Thanks for making this package available. As someone using LP in another research domain (Metabolic Engineering), I was very excited to try it out.

The LP problem I am working on can be formulated using the code below (you would need to pip install cobra for easy access to a 'metabolic model', but it is just a wrapper that connects biology to LP. However, I cannot solve this problem in MPAX for a long time, even though GLPK can solve it within seconds. I wonder if you have any thoughts on this. Appreciate your help!

import jax
import cobra # need to pip install
from cobra.io import load_model
import numpy as np
from mpax import create_lp, create_qp, raPDHG, r2HPDHG

jax.config.update("jax_enable_x64", True)

# Load the LP information from a metabolic model
model = load_model("ijo1366")
S = cobra.util.create_stoichiometric_matrix(model)
c = np.array([r.objective_coefficient for r in model.reactions])  # Objective vector
lb = np.array([r.lower_bound for r in model.reactions])
ub = np.array([r.upper_bound for r in model.reactions])
model.optimize() # this is calling glpk, which is very fast to solve.

# convert to mpax format
c = -c
A = S
b = np.zeros(A.shape[0])
G = np.zeros((0, A.shape[1])) # no inequality constrants
h = np.zeros((0,)) # no inequality constrants
l = lb
u = ub

lp = create_lp(c, A, b, G, h, l, u, use_sparse_matrix=True)
solver = r2HPDHG(eps_abs=1e-4, eps_rel=1e-4, verbose=True)
result = solver.optimize(lp)

[ZedongPeng]

Thank you for providing the example. I tried it out and noticed that MPAX is slow on this instance in terms of time per iteration, which is unusual given the problem size. I’ll look into it more closely and get back to you.

The instance has 2,583 variables, 1,805 constraints (all equalities), and 1,805 nonzero coefficients.

[ZedongPeng]

Hi @fromseto, I found the reason for the slow iteration. In MPAX, we use BCOO.fromdense to convert a numpy.array to BCOO. To make the function safely jittable, we had set nse=matrix.shape[0] * matrix.shape[1], which stores the matrix as fully dense in a sparse wrapper. Honestly, this does not make sense. I'm considering removing support for numpy.array when use_sparse_matrix=True, to avoid silent performance issues like this in the future.

For now, you can fix this quickly by converting the matrix with from scipy.sparse import csr_matrix; A = csr_matrix(S), which will not be affected by the transformation.

That said, PDHG still shows slow convergence on this example, which I don't know why yet. I’ll keep you posted.

[fromseto]

Thank you so much for the update! Changing to csr_matrix makes the solver run blazingly fast per iteration!

I played around with the tolerance for the settings after using A = csr_matrix(S) as you suggsted. It seems that lowering the tolerance to eps_abs=1e-2, eps_rel=1e-2 can terminate faster. Is the problem sensitive to tolerances? For your debugging purposes, the objective value from glpk is 0.9823718127269739.

Terminated after 72010 iterations. Termination Status: 2
 - Primal objective: -9.8497770943e-01
 - Dual objective: -1.0015428749e+00
 - Corrected dual objective: -1.0015428749e+00
 - Abs gap: 1.65652e-02
 - Rel gap: 8.27620e-03
 - Abs Primal residual: 6.28937e-04
 - Abs Dual residual: 0.00000e+00
 - Rel primal residual: 6.28542e-04
 - Rel dual residual: 0.00000e+00
Timing Information:
 - Preconditioning: 1.9168 seconds
 - Iteration loop: 10.4483 seconds
 - Feasibility polishing: 0.0000 seconds

[ZedongPeng]

To accelerate MPAX, you can try using jax.jit. We also found parameter tuning to be helpful.

One issue is that your constraint matrix contains very small values — the smallest being 0.000002 — which can cause numerical difficulties for MPAX. Additionally, since your problem only has equality constraints, presolve techniques could be particularly effective. Unfortunately, MPAX doesn’t currently support presolve. You might want to check the solver logs from a commercial solver to see if presolve is playing a role.

There are 2583 variables, 1805 constraints (including 1805 equalities) and 10183 nonzero coefficients.
Absolute value of nonzero constraint matrix elements: largest=54.124831, smallest=0.000002, avg=1.097226