New: non-negative least squares

Author: heptaflar

test/b/test_b_jax.py

@@ -11,6 +11,7 @@
 from uncertaintyx.b.jax import BernsteinPoly
 from uncertaintyx.b.jax import b_basis
 from uncertaintyx.b.jax import b_poly
+from uncertaintyx.b.jax import solve
 
 
 class BBasisTest(unittest.TestCase):
@@ -277,6 +278,72 @@ def test_bernstein_poly(self):
         self.assertEqual((3,) + d, g.shape)
         self.assertTrue(np.all(g > 0.0))
 
+    def test_from_lookup_table(self):
+        k = 5
+        x = np.array([0.00, 0.20, 0.40, 0.60, 0.80, 1.00])
+        y = np.array(  # y = x ** 2 + 2 x + 3
+            [3.00, 3.44, 3.96, 4.56, 5.24, 6.00]
+        )
+
+        f = BernsteinPoly.from_lookup_table((k,), (x,), y, non_negative=True)
+        c = f.prior()
+        self.assertEqual((k + 1,), c.shape)
+        self.assertTrue(jnp.allclose(f.eval(c, x), y))
+        self.assertAlmostEqual(3.0, c[0])
+        self.assertAlmostEqual(3.4, c[1])
+        self.assertAlmostEqual(3.9, c[2])
+        self.assertAlmostEqual(4.5, c[3])
+        self.assertAlmostEqual(5.2, c[4])
+        self.assertAlmostEqual(6.0, c[5])
+
+
+class SolveTest(unittest.TestCase):
+    """Tests the solving function."""
+
+    def test_solve_degree_2(self):
+        k = 2
+        x = jnp.array([0.00, 0.20, 0.40, 0.60, 0.80, 1.00])
+        y = jnp.array(  # y = x ** 2 + 2 x + 3
+            [3.00, 3.44, 3.96, 4.56, 5.24, 6.00]
+        )
+
+        c = solve((k,), (x,), y, non_negative=True)
+        self.assertEqual((k + 1,), c.shape)
+        self.assertTrue(jnp.allclose(b_poly(c, x), y))
+        self.assertAlmostEqual(3.0, c[0].item())
+        self.assertAlmostEqual(4.0, c[1].item())
+        self.assertAlmostEqual(6.0, c[2].item())
+
+    def test_solve_degree_5(self):
+        k = 5
+        x = jnp.array([0.00, 0.20, 0.40, 0.60, 0.80, 1.00])
+        y = (
+            jnp.array(  # y = x ** 2 + 2 x - 1 / 100
+                [0.00, 0.44, 0.96, 1.56, 2.24, 3.00]
+            )
+            - 0.01
+        )
+
+        c = solve((k,), (x,), y)
+        self.assertEqual((k + 1,), c.shape)
+        self.assertTrue(jnp.allclose(b_poly(c, x), y))
+        self.assertAlmostEqual(-0.01, c[0].item())
+        self.assertAlmostEqual(0.39, c[1].item())
+        self.assertAlmostEqual(0.89, c[2].item())
+        self.assertAlmostEqual(1.49, c[3].item())
+        self.assertAlmostEqual(2.19, c[4].item())
+        self.assertAlmostEqual(2.99, c[5].item())
+
+        c = solve((k,), (x,), y, non_negative=True)
+        self.assertEqual((k + 1,), c.shape)
+        self.assertTrue(jnp.allclose(b_poly(c, x), y, atol=0.1))
+        self.assertAlmostEqual(0.00, c[0].item())
+        self.assertAlmostEqual(0.38, c[1].item(), places=2)
+        self.assertAlmostEqual(0.90, c[2].item(), places=2)
+        self.assertAlmostEqual(1.48, c[3].item(), places=2)
+        self.assertAlmostEqual(2.19, c[4].item(), places=2)
+        self.assertAlmostEqual(2.99, c[5].item(), places=2)
+
 
 def to_var(u: np.ndarray) -> np.ndarray:
     """

uncertaintyx/b/jax.py

@@ -8,6 +8,8 @@
 import jax.numpy as jnp
 import jax.numpy.linalg as jli
 import numpy as np
+import optax
+import optimistix
 from jax import Array
 from jax.scipy.special import gammaln
 
@@ -97,7 +99,7 @@ def b_basis(k: int, x: Array) -> Array:
 
 
 @jax.jit(static_argnums=(0,), static_argnames=("non_negative", "max_steps"))
-def linear_solve(
+def solve(
     k: tuple[int, ...],
     x: tuple[Array, ...],
     y: Array,
@@ -123,37 +125,60 @@ def linear_solve(
 
     def nnls(c: Array):
         """
-        Non-negative linear least squares solver (yet incomplete).
+        Non-negative least-squares solver.
+
+        Uses a quadratic transformation and an L-BFGS minimizer.
         """
 
         def hvp(c: Array):
             """The Hessian-vector product."""
-            hvp = c
+            res = c
             for i in range(N):
                 G = grams[i]  # noqa: N806
-                hvp = jnp.tensordot(hvp, G, axes=(0, 1))
-            return hvp
-
-        return c
+                res = jnp.tensordot(res, G, axes=(0, 1))
+            return res
+
+        def misfit(u: Array, _: None = None) -> Array:
+            """The misfit function with quadratic transformation."""
+            c_ = jnp.square(u)
+            return 0.5 * jnp.sum(c_ * hvp(c_)) - jnp.sum(c_ * rhs)
+
+        def make_minimizer():
+            """Returns the minimizer."""
+            return optimistix.OptaxMinimiser(
+                optax.lbfgs(), atol=atol, rtol=rtol, norm=optimistix.max_norm
+            )
+
+        u = jnp.sqrt(jnp.maximum(0.0, c))
+        minimum = optimistix.minimise(
+            misfit, make_minimizer(), u, max_steps=max_steps, throw=False
+        )
+        return jnp.square(minimum.value)
 
     N = len(k)  # noqa: N806
     bases = [b_basis(k[i], x[i]) for i in range(N)]
     grams = [jnp.dot(B, B.T) for B in bases]  # noqa: N806
 
     # compute the right hand side of the normal equation
+    rhs = y
     for i in range(N):
         B = bases[i]  # noqa: N806
-        y = jnp.tensordot(y, B, axes=(0, 1))
+        rhs = jnp.tensordot(rhs, B, axes=(0, 1))
     # solve the normal equation
-    c = y
+    c_unconstrained = rhs
     for i in range(N):
         G = grams[i]  # noqa: N806
-        c = jnp.tensordot(c, jli.inv(G), axes=(0, 1))
-    # solve with non-negativity constraint, if requested
-    if non_negative:
-        c = nnls(jnp.maximum(0.0, c))
-
-    return c
+        c_unconstrained = jnp.tensordot(
+            c_unconstrained, jli.pinv(G), axes=(0, 1)
+        )
+    # solve iteratively with non-negativity constraint, if needed
+    nnls_needed = non_negative and jnp.any(c_unconstrained < 0.0)
+    return jax.lax.cond(
+        nnls_needed,
+        nnls,
+        lambda _: _,  # forwards the unconstrained solution
+        c_unconstrained,
+    )
 
 
 @jax.jit
@@ -380,7 +405,7 @@ def from_lookup_table(
         b = _upper_bounds(b, x)
         x_ = tuple(jnp.asarray((x_ - a) / (b - a)) for x_ in x)
         y_ = jnp.asarray(y)
-        c_ = linear_solve(
+        c_ = solve(
             k,
             x_,
             y_,