ENH(shells): non-negative least squares partition (#141)

Adds the `"nnls"` method to the `partition()` function, which returns a
non-linear least-squares solution. This should be used when each shell's
contribution is required to be positive, e.g. when partitioning
densities.

Closes: #140
Added: The glass.core.algorithm module.
Added: The new partition(method="nnls") function computes a partition
  with non-negative contributions for each shell.
Changed: The default method for partition() is now "nnls".

.readthedocs.yml

@@ -1,5 +1,10 @@
 version: 2
 
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "3.11"
+
 python:
   install:
     - method: pip

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.
 
@@ -358,7 +358,7 @@ def partition(z: ArrayLike,
         its last axis must agree with *z*.
     shells : sequence of :class:`RadialWindow`
         Ordered sequence of window functions for the partition.
-    method : {"lstsq", "restrict"}
+    method : {"lstsq", "nnls", "restrict"}
         Method for the partition.  See notes for description.
 
     Returns
@@ -389,9 +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="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
@@ -445,6 +452,52 @@ def partition_lstsq(
     return x
 
 
+def partition_nnls(
+    z: ArrayLike,
+    fz: ArrayLike,
+    shells: Sequence[RadialWindow],
+) -> ArrayLike:
+    """Non-negative least-squares partition.
+
+    Uses the ``nnls()`` algorithm from ``scipy.optimize`` and thus
+    requires SciPy.
+
+    """
+
+    from .core.algorithm import nnls
+
+    # compute the union of all given redshift grids
+    zp = z
+    for w in shells:
+        zp = np.union1d(zp, w.za)
+
+    # get extra leading axes of fz
+    *dims, _ = np.shape(fz)
+
+    # compute grid spacing
+    dz = np.gradient(zp)
+
+    # create the window function matrix
+    a = [np.interp(zp, za, wa, left=0., right=0.) for za, wa, _ in shells]
+    a = a/np.trapz(a, zp, axis=-1)[..., None]
+    a = a*dz
+
+    # create the target vector of distribution values
+    b = ndinterp(zp, z, fz, left=0., right=0.)
+    b = b*dz
+
+    # now a is a matrix of shape (len(shells), len(zp))
+    # 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)
+    for i in np.ndindex(*dims):
+        x[(slice(None),) + i] = nnls(a.T, b[i])[0]
+
+    # all done
+    return x
+
+
 def partition_restrict(
     z: ArrayLike,
     fz: ArrayLike,

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",