Update: grid and Bernstein functions

Author: heptaflar

test/b/test_b_jax.py

@@ -5,12 +5,12 @@
 
 import jax
 import jax.numpy as jnp
+import numpy as np
 
+from uncertaintyx.b.jax import BernsteinGrid
+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 b_poly_grid
-from uncertaintyx.b.jax import b_poly_point
-from uncertaintyx.b.jax import b_poly_points
 
 
 class BBasisTest(unittest.TestCase):
@@ -156,23 +156,24 @@ def b_poly_grad(b, x):
         self.assertTrue(jnp.allclose(g, 1.0))
 
 
-class BPolyGridTest(unittest.TestCase):
+class BernsteinGridTest(unittest.TestCase):
     """
     Tests the evaluation of multivariate Bernstein polynomials
     against values precalculated with Mathematica.
     """
 
-    def test_b_poly_grid(self):
+    def test_bernstein_grid(self):
         k = (4, 3, 2)
         d = tuple([k_ + 1 for k_ in k])
-        b = jnp.arange(jnp.prod(jnp.asarray(d))).reshape(d) + 1.0
+        b = np.arange(np.prod(np.asarray(d))).reshape(d) + 1.0
         x = (
-            jnp.asarray([0.2718, 0.5772, 0.3141]),
-            jnp.asarray([0.5772, 0.3141, 0.2718]),
-            jnp.asarray([0.3141, 0.2718, 0.5772]),
+            np.asarray([0.2718, 0.5772, 0.3141]),
+            np.asarray([0.5772, 0.3141, 0.2718]),
+            np.asarray([0.3141, 0.2718, 0.5772]),
         )
-        y = b_poly_grid(b, x)
-        precalculated = jnp.asarray(
+        f = BernsteinGrid(k, x)
+        y = f.eval(b)
+        precalculated = np.asarray(
             [
                 [
                     [19.8694, 19.7848, 20.3956],
@@ -192,41 +193,57 @@ def test_b_poly_grid(self):
             ]
         )
         self.assertEqual((3, 3, 3), y.shape)
-        self.assertTrue(jnp.allclose(y, precalculated))
+        self.assertTrue(np.allclose(y, precalculated))
+
+        g = f.jac(b)
+        self.assertEqual(y.shape + b.shape, g.shape)
+        self.assertTrue(np.all(g > 0.0))
+
+        u = to_var(0.1 * b)
+        u = f.lpu(b, u, diag=True)
+        self.assertEqual(y.shape, u.shape)
+        self.assertTrue(np.all(u > 0.0))
+
+        u = to_var(0.1 * b)
+        u = f.lpu(b, u)
+        self.assertEqual(y.shape + y.shape, u.shape)
+        self.assertTrue(np.all(u > 0.0))
 
 
-class BPolyPointsTest(unittest.TestCase):
+class BernsteinPolyTest(unittest.TestCase):
     """
     Tests the evaluation of multivariate Bernstein polynomials
     against values precalculated with Mathematica.
     """
 
-    def test_b_poly_point(self):
+    def test_bernstein_poly(self):
         k = (4, 3, 2)
         d = tuple([k_ + 1 for k_ in k])
-        b = jnp.arange(jnp.prod(jnp.asarray(d))).reshape(d) + 1.0
-        x = jnp.asarray([0.2718, 0.5772, 0.3141])
-        y = b_poly_point(b, x)
-        precalculated = 19.8694
-        self.assertEqual((), y.shape)
-        self.assertTrue(jnp.allclose(y, precalculated))
-
-    def test_b_poly_points(self):
-        k = (4, 3, 2)
-        d = tuple([k_ + 1 for k_ in k])
-        b = jnp.arange(jnp.prod(jnp.asarray(d))).reshape(d) + 1.0
-        x = jnp.asarray(
+        b = np.arange(np.prod(np.asarray(d))).reshape(d) + 1.0
+        x = np.asarray(
             [
                 [0.2718, 0.5772, 0.3141],
                 [0.5772, 0.3141, 0.2718],
                 [0.3141, 0.2718, 0.5772],
             ]
         )
-        y = b_poly_points(b, x)
-        precalculated = jnp.asarray([19.8694, 32.0761, 19.6774])
+        f = BernsteinPoly(b)
+        y = f.eval(b, x)
+        precalculated = np.asarray([19.8694, 32.0761, 19.6774])
         self.assertEqual((3,), y.shape)
         self.assertTrue(jnp.allclose(y, precalculated))
 
+        g = f.jac_p(b, x)
+        self.assertEqual((3,) + d, g.shape)
+        self.assertTrue(np.all(g > 0.0))
+
+
+def to_var(u: np.ndarray) -> np.ndarray:
+    """
+    Converts standard uncertainty to a diagonal uncertainty tensor.
+    """
+    return np.square(u)
+
 
 if __name__ == "__main__":
     unittest.main()

uncertaintyx/b/__init__.py

@@ -1,5 +1,5 @@
 #  Copyright (c) Brockmann Consult GmbH, 2026.
 #  License: MIT
 """
-A package to provide univariate and multi-variate Bernstein polynomials.
+A package to provide univariate and N-variate Bernstein polynomials.
 """

uncertaintyx/b/jax.py

@@ -2,9 +2,13 @@
 #  License: MIT
 import jax
 import jax.numpy as jnp
+import numpy as np
 from jax import Array
 from jax.scipy.special import gammaln
 
+from uncertaintyx.g.jax import ToG
+from uncertaintyx.m.jax import ToM
+
 
 def _binom(i: Array, k: int) -> Array:
     """
@@ -154,19 +158,78 @@ def b_poly_point(b: Array, x: Array) -> Array:
     return b
 
 
-@jax.jit
-def b_poly_points(b: Array, x: Array) -> Array:
+class BernsteinGrid(ToG):
     r"""
-    Evaluates a multivariate Bernstein polynomial on a batch of
-    irregularly distributed N-dimensional points.
+    Evaluates an N-variate Bernstein polynomial on a regular
+    grid of points.
 
-    Under the same notation as :meth:`b_poly_point` let
-    :math:`X \in \mathbb{R}^{M \times N}` be a batch of
-    points over the outer batch dimension :math:`M `.
-    Then:
+    Encapsulates the multivariate Bernstein polynomial
 
-    :param b: The Bernstein coefficients :math:`b \in \mathbb{R}^{k + 1}`.
-    :param x: The points :math:`X \in \mathbb{R}^{M \times N}`.
-    :returns: The polynomial values :math:`B(x) \in \mathbb{R}^{M}`.
+    .. math::
+        B: \mathbb{R}^{k + 1} \to \mathbb{R}^{m},
+        b \mapsto B_{k}(b, x)
+
+    where :math:`N \in \mathbb{N}` is the arity of the Bernstein
+    polynomial, :math:`k = (k_{1}, \dots, k_{N}) \in \mathbb{N}^{N}`
+    are its degrees, :math:`\mathbb{R}^{k + 1}` is the tensor space
+    with dimensions :math:`(k_{1} + 1, \dots, k_{N} + 1)`,
+    :math:`m = (m_{1}, \dots , m_{N}) \in \mathbb{N}^{N}` are the
+    dimensions of the grid coordinates :math:`x = (x_{1}, \dots, x_{N})`,
+    and :math:`\mathbb{R}^{m}` is the tensor space with dimensions
+    :math:`m`.
+    """
+
+    def __init__(self, k: tuple[int, ...], x: tuple[np.ndarray, ...]):
+        """
+        Creates a new instance of this class.
+
+        :param k: The degrees :math:`k`.
+        :param x: The grid coordinates :math:`x`.
+        """
+        self._d = tuple([k_ + 1 for k_ in k])
+        self._x = tuple([jnp.asarray(x_) for x_ in x])
+
+        def f(b: Array) -> Array:
+            return b_poly_grid(b, self._x)
+
+        super().__init__(f, rev=False)
+
+    def prior(self, preset: str | None = None) -> np.ndarray:
+        return np.ones(self._d)
+
+
+class BernsteinPoly(ToM):
+    r"""
+    Evaluates an N-variate Bernstein polynomial on a batch of
+    irregularly distributed points.
+
+    Encapsulates the multivariate Bernstein polynomial
+
+    .. math::
+        B: \mathbb{R}^{k + 1} \times \mathbb{R}^{N} \to
+        \mathbb{R},
+        (b, x) \mapsto B(b, x)
+
+    where :math:`N \in \mathbb{N}` is the arity of the Bernstein
+    polynomial, :math:`k = (k_{1}, \dots, k_{N}) \in \mathbb{N}^{N}`
+    denotes its degrees, and :math:`\mathbb{R}^{k + 1}` denotes the
+    tensor space with dimensions :math:`(k_{1} + 1, \dots, k_{N} + 1)`.
     """
-    return jax.vmap(b_poly_point, in_axes=(None, 0))(b, x)
+
+    def __init__(self, b: np.ndarray):
+        """
+        Creates a new instance of this class.
+
+        :param b: The Bernstein coefficients :math:`b \in \mathbb{R}^{k + 1}`.
+        """
+        self._b = b
+
+        super().__init__(b_poly_point)
+
+    def prior(
+        self,
+        x: np.ndarray | None = None,
+        y: np.ndarray | None = None,
+        preset: str | None = None,
+    ) -> np.ndarray:
+        return self._b

uncertaintyx/g/__init__.py

@@ -0,0 +1,2 @@
+#  Copyright (c) Brockmann Consult GmbH, 2026.
+#  License: MIT

uncertaintyx/g/jax.py

@@ -0,0 +1,144 @@
+#  Copyright (c) Brockmann Consult GmbH, 2026.
+#  License: MIT
+"""
+Interface adapters for pure JAX functions.
+
+Adapters employ JAX algorithmic differentiation to compute
+derivatives of generic model functions.
+"""
+
+from abc import ABC
+from typing import Callable
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from jax import Array
+
+from ..tyx import G
+
+
+@jax.jit(static_argnums=(0, 2))
+def jac(f: Callable[[Array], Array], p: Array, rev: bool = True) -> Array:
+    """Returns the Jacobian (does not belong to public API)."""
+    return jax.jacrev(f)(p) if rev else jax.jacfwd(f)(p)
+
+
+@jax.jit(static_argnums=(0, 3))
+def lpu(d: int, g: Array, u: Array, diag: bool = True) -> Array:
+    r"""
+    Implementation of the law of propagation of uncertainty in
+    general tensor form.
+
+    Using Einstein's summation convention and the symmetry of the
+    parameter uncertainty tensor :math:`U`:, the output uncertainty
+    tensor reads:
+
+    .. math::
+        V_{\dots ij} = G_{\dots ik}U_{\dots lk}G_{\dots jl}
+
+    with multi-indices :math:`k, l \in D \subset \mathbb{N}^d`
+    for some :math:`d \in \mathbb{N}`. The summation is taken over
+    all :math:`k, l \in D`.
+
+    Here, :math:`D` denotes the set of inner tensor indices
+    (multi-indices of length :math:`d`), and the trailing tensor
+    dimensions of :math:`G` and :math:`U` correspond to these
+    indices.
+
+    In what follows, we write :math:`\mathbb{R}^{\cdots \times D}`
+    for a tensor space whose trailing indices are labelled by the
+    index set :math:`D`.
+
+    :param d: The number of inner tensor dimensions.
+    :param g: Jacobian :math:`G \in \mathbb{R}^{M \times \cdots \times D}`.
+    :param u: Tensor :math:`U \in \mathbb{R}^{\cdots \times D}`.
+    :param diag: To return only variance elements of :math:`V`.
+    :returns: Tensor :math:`V \in \mathbb{R}^{M \times \cdots}`.
+    """
+    return make_lpu(d, diag)(g, u)
+
+
+def make_lpu(d: int, diag: bool = False) -> Callable[[Array, Array], Array]:
+    """
+    Returns the law of propagation of uncertainty.
+
+    :param d: The number of inner tensor dimensions.
+    :param diag: To return only the diagonal elements .
+    :returns: The law of propagation of uncertainty.
+    """
+
+    def lpu(g: Array, u: Array) -> Array:
+        """The law of propagation of uncertainty."""
+        dims = tuple(range(-d, 0))
+        gu = jnp.tensordot(g, u, (dims, dims)) if u.ndim != d else g * u
+        return (
+            jnp.tensordot(gu, g, (dims, dims))
+            if not diag
+            else jnp.sum(gu * g, dims)
+        )
+
+    return lpu
+
+
+def jac_no_jit(  # pragma: no cover
+    f: Callable[[Array], Array], p: Array, rev: bool = False
+) -> Array:
+    """Noncompiled version of :meth:`jac` for debugging."""
+    return jax.jacrev(f)(p) if rev else jax.jacfwd(f)(p)
+
+
+class ToG(G, ABC):
+    r"""
+    Adapts a pure function
+
+    .. math::
+        f: \mathbb{R}^{k} \to \mathbb{R}^{m}, \quad
+        p \mapsto f(p)
+
+    where :math:`k, m` are shapes (natural numbers or tuples
+    of natural numbers) to the function interface ``G``.
+    """
+
+    def __init__(
+        self,
+        f: Callable[[Array], Array],
+        rev: bool = True,
+        jit: bool = True,
+    ):
+        """
+        Creates a new instance of this class.
+
+        :param f: The function :math:`f`.
+        :param rev: Use reverse mode for the Jacobian.
+        :param jit: Switches JIT compilation on and off (for debugging).
+        """
+        self._f = jax.jit(f) if jit else f
+        self._jit = jit
+        self._rev = rev
+
+    def eval(self, p: np.ndarray) -> np.ndarray:
+        p_ = jnp.asarray(p)
+        y_ = self._f(p_)
+        return np.asarray(y_)
+
+    def jac(self, p: np.ndarray) -> np.ndarray:
+        p_ = jnp.asarray(p)
+        g_ = (
+            jac(self._f, p_, self._rev)
+            if self._jit
+            else jac_no_jit(self._f, p_, self._rev)
+        )
+        return np.asarray(g_)
+
+    def lpu(
+        self, p: np.ndarray, u: np.ndarray, diag: bool = False
+    ) -> np.ndarray:
+        p_ = jnp.asarray(p)
+        u_ = jnp.asarray(u)
+        v_ = lpu(p_.ndim, jac(self._f, p_, self._rev), u_, diag)
+        return np.asarray(v_)
+
+    @property
+    def f(self) -> Callable[[Array, Array], Array]:
+        return self._f