add nnls algorithm implementation

glass/core/algorithm.py

@@ -0,0 +1,70 @@
+# author: Nicolas Tessore <[email protected]>
+# license: MIT
+'''core module for algorithms'''
+
+from __future__ import annotations
+
+import numpy as np
+from numpy.typing import ArrayLike
+
+
+def nnls(
+    a: ArrayLike,
+    b: ArrayLike,
+    *,
+    tol: float = 1e-10,
+    maxiter: int | None = None,
+) -> ArrayLike:
+    """Compute a non-negative least squares solution.
+
+    Implementation of the algorithm due to [1]_ as described in [2]_.
+
+    References
+    ----------
+    .. [1] Lawson, C. L. and Hanson, R. J. (1995), Solving Least Squares
+        Problems. doi: 10.1137/1.9781611971217
+    .. [2] Bro, R. and De Jong, S. (1997), A fast
+        non-negativity-constrained least squares algorithm. J.
+        Chemometrics, 11, 393-401.
+
+    """
+
+    a = np.asanyarray(a)
+    b = np.asanyarray(b)
+
+    if a.ndim != 2:
+        raise ValueError("input `a` is not a matrix")
+    if b.ndim != 1:
+        raise ValueError("input `b` is not a vector")
+    if a.shape[0] != b.shape[0]:
+        raise ValueError("the shapes of `a` and `b` do not match")
+
+    m, n = a.shape
+
+    if maxiter is None:
+        maxiter = 3 * n
+
+    index = np.arange(n)
+    p = np.full(n, False)
+    x = np.zeros(n)
+    for i in range(maxiter):
+        if np.all(p):
+            break
+        w = np.dot(b - a @ x, a)
+        m = index[~p][np.argmax(w[~p])]
+        if w[m] <= tol:
+            break
+        p[m] = True
+        while True:
+            ap = a[:, p]
+            xp = x[p]
+            sp = np.linalg.solve(ap.T @ ap, b @ ap)
+            t = (sp <= 0)
+            if not np.any(t):
+                break
+            alpha = -np.min(xp[t]/(xp[t] - sp[t]))
+            x[p] += alpha * (sp - xp)
+            p[x <= 0] = False
+        x[p] = sp
+        x[~p] = 0
+    return x

glass/core/test/test_algorithm.py

@@ -0,0 +1,25 @@
+import pytest
+
+# check if scipy is available for testing
+try:
+    import scipy
+except ImportError:
+    HAVE_SCIPY = False
+else:
+    del scipy
+    HAVE_SCIPY = True
+
+
+@pytest.mark.skipif(not HAVE_SCIPY, reason="test requires SciPy")
+def test_nnls():
+    import numpy as np
+    from scipy.optimize import nnls as nnls_scipy
+    from glass.core.algorithm import nnls as nnls_glass
+
+    a = np.random.randn(100, 20)
+    b = np.random.randn(100)
+
+    x_glass = nnls_glass(a, b)
+    x_scipy, _ = nnls_scipy(a, b)
+
+    np.testing.assert_allclose(x_glass, x_scipy)

glass/shells.py

@@ -335,7 +335,7 @@ def partition(z: ArrayLike,
               fz: ArrayLike,
               shells: Sequence[RadialWindow],
               *,
-              method: str = "lstsq",
+              method: str = "nnls",
               ) -> ArrayLike:
     """Partition a function by a sequence of windows.
 
@@ -389,13 +389,16 @@ def partition(z: ArrayLike,
     redshift arrays of all window functions.  Intermediate function
     values are found by linear interpolation.
 
-    If ``method="lstsq"``, obtain a partition from a least-squares
-    solution.  This will more closely match the shape of the input
-    function, but the normalisation might differ.
+    If ``method="nnls"`` (the default), obtain a partition from a
+    non-negative least-squares solution.  This will match the shape of
+    the input function well, but the overall normalisation might be
+    differerent.  The contribution from each shell is a positive number,
+    which is required to partition e.g. density distributions.
 
-    If ``method="nnls"``, obtain a partition from a non-negative
-    least-squares solution.  This method should be used to partition
-    e.g. density distributions.
+    If ``method="lstsq"``, obtain a partition from an unconstrained
+    least-squares solution.  This will more closely match the shape of
+    the input function, but might lead to shells with negative
+    contributions.
 
     If ``method="restrict"``, obtain a partition by integrating the
     restriction (using :func:`restrict`) of the function :math:`f` to
@@ -461,16 +464,7 @@ def partition_nnls(
 
     """
 
-    try:
-        from scipy.optimize import nnls
-    except ImportError as e:
-        try:
-            e.addNote("Calling partition() with the 'nnls' method requires "
-                      "SciPy.  You can try another method using the `method=` "
-                      "keyword.  Example: partition(..., method='lstsq')")
-        except AttributeError:
-            pass
-        raise
+    from .core.algorithm import nnls
 
     # compute the union of all given redshift grids
     zp = z
@@ -496,7 +490,7 @@ def partition_nnls(
     # and b is a matrix of shape (*dims, len(zp))
     # for each dim, find weights x such that b == a.T @ x
     # the output is more conveniently shapes with len(shells) first
-    x = np.empty((len(shells), *dims))
+    x = np.empty([len(shells)] + dims)
     for i in np.ndindex(*dims):
         x[:, *i] = nnls(a.T, b[i])[0]
 

glass/test/test_shells.py

@@ -49,7 +49,7 @@ def test_restrict():
             assert fr[i] == fi*np.interp(zi, w.za, w.wa)
 
 
-@pytest.mark.parametrize("method", ["lstsq", "restrict"])
+@pytest.mark.parametrize("method", ["lstsq", "nnls", "restrict"])
 def test_partition(method):
     import numpy as np
     from glass.shells import RadialWindow, partition

pyproject.toml

@@ -28,6 +28,7 @@ dynamic = ["version"]
 [project.optional-dependencies]
 test = [
     "pytest",
+    "scipy",
 ]
 docs = [
     "sphinx",