Add non-negative least squares (NNLS) solver.

[carlosgmartin]

Fixes #1152.

[carlosgmartin]

@rdyro How does this look?

[rdyro]

@carlosgmartin This looks great, I left one comment!

[carlosgmartin]

@rdyro Where is this comment?

[rdyro]

@rdyro Where is this comment?

Oops, should be up now!

[rdyro]

# We use lstsq with a pre-computed AtA to reduce computation time.
s = jnp.linalg.lstsq(AtA * p[:, None] * p[None, :], Atb * p)[0]

If speed is a consideration, perhaps we should use the jnp.linlag.lsqt(A * p, b) directly, letting XLA optimize as necessary?

I'd be curious if we could introduce a once-factorized cholesky version of this algorithm where we repeatedly apply cho_solve to a masked L or U factor of A^T A?

@carlosgmartin @fabianp ?

[carlosgmartin]

@rdyro Edited.

[rdyro]

Thank you for this PR!

[rdyro]

I ran the tests internally and I can't pass the tolerance checks for this algorithm for larger dimensions (e.g., n=24 d=64). Could you take a look?

[carlosgmartin]

Possibly related: jax.numpy.linalg.lstsq seems to be buggy.

[rdyro]

NaNs propagate so we'd see NaNs in the output, but the algorithm seems to not be converging instead.

[carlosgmartin]

True. I also tried the JAX_DEBUG_NANS=True flag.

The issue might be that the arithmetic operations of the algorithm accumulate numerical error, and thus the inequality comparisons (which are discontinuous) yield incorrect results. I'm taking a closer look.

[carlosgmartin]

Due to the above difficulties, I switched to the fast projected gradient (FPG) algorithm of Polyak 2015. A clear presentation of the algorithm can be found here.

It has some advantages described in the paper's introduction, including not solving a linear system on each iteration (only doing matrix-vector multiplications instead).

[rdyro]

Nice! Maybe we can explore also adding an ADMM version in a future PR https://stanford.edu/class/ee364b/lectures/admm_slides.pdf since it's also promising a $\mathcal{O}(1/k)$ convergence (requires resolves, but with only a single factorization at the beginning).

[carlosgmartin]

@rdyro Done.

[rdyro]

@rdyro Done.

Nice, this PR looks really good! Thank you!

[carlosgmartin]

@rdyro Made some minor changes to handle zero matrices without producing nans.