Add user-supplied domain interior function and fix float roundoff bug in gp-1d demo notebook

UncertainSCI/gp/__init__.py

@@ -2,6 +2,8 @@
 Build Gaussian processes.
 """
 
+from collections.abc import Callable
+
 import equinox as eqx
 import jax
 import jax.numpy as jnp
@@ -392,6 +394,7 @@ def get_sample_point(
         n: int = 1000,
         optim = DEFAULT_OPTIM,
         optim_kwargs: dict = DEFAULT_OPTIM_KWARGS,
+        inside_domain: Callable[[jax.Array], bool | jax.Array] | None = None,
     ) -> jax.Array:
         """
         Optimize a coordinate ``x`` to maximize posterior variance.
@@ -400,19 +403,25 @@ def get_sample_point(
             x (jax.Array or numpy.ndarray):
                 Candidate coordinate from which to start optimization.
             hull (scipy.spatial.ConvexHull, optional):
-                Convex hull of coordinate domain,
-                must be supplied if ``ranges`` is not.
+                Convex hull of coordinate domain. Mutually exclusive with
+                ``ranges`` and ``inside_domain``.
             ranges (jax.Array or numpy.ndarray, optional):
-                Array of shape (2, dim) of corners of prism of coordinate domain.
-                Must be supplied if ``hull`` is not.
+                Array of shape (2, dim) containing the lower and upper corners
+                of the coordinate domain. Mutually exclusive with ``hull`` and
+                ``inside_domain``.
             tol (float):
-                Tolerance of hull inclusion criterion.
+                Tolerance of the inclusion criterion constructed from ``hull``
+                or ``ranges``.
             n (int):
                 Number of steps to take in optimization.
             optim (callable):
                 Optax or Optax-style optimizer.
             optim_kwargs (dict):
                 kwargs for ``optim``.
+            inside_domain (callable, optional):
+                Predicate that accepts a candidate array of shape ``(1, dim)``
+                and returns whether it is in the coordinate domain. Mutually
+                exclusive with ``hull`` and ``ranges``.
         """
         if not hasattr(self, 'train_x') or not hasattr(self, 'train_cov_factor'):
             raise ValueError('Gaussian process must be conditioned before sampling.')
@@ -426,50 +435,41 @@ def get_sample_point(
             raise ValueError
         x = x.reshape((1, self.dim))
 
-        if not ((hull is not None) ^ (ranges is not None)):
-            raise ValueError
-
-        if self.dim > 1:
-            if hull is None:
-                ranges = jnp.asarray(ranges)
-                if ranges.ndim != 2:
-                    raise ValueError
-                if ranges.shape != (2, self.dim):
-                    raise ValueError
-
-                hull = spatial.ConvexHull(
-                    jnp.stack(
-                        jnp.meshgrid(
-                            *ranges.T,
-                            indexing='ij'
-                        ),
-                        axis=-1
-                    ).reshape((-1, self.dim))
-                )
+        domain_specs = (hull, ranges, inside_domain)
+        if sum(spec is not None for spec in domain_specs) != 1:
+            raise ValueError(
+                'Supply exactly one of hull, ranges, or inside_domain.'
+            )
 
-            A = hull.equations[:, :-1]
-            b = hull.equations[:, -1]
-            def inside_domain(x_cand: jax.Array):
-                return jnp.all(jnp.dot(x_cand, A.T) + b <= tol)
+        if inside_domain is not None:
+            domain_contains = inside_domain
+        elif hull is not None:
+            equations = jnp.asarray(hull.equations)
 
+            def domain_contains(x_cand: jax.Array) -> jax.Array:
+                return jnp.all(
+                    x_cand @ equations[:, :-1].T + equations[:, -1] <= tol
+                )
         else:
-            if ranges is None:
-                raise ValueError
-
             ranges = jnp.asarray(ranges)
-            def inside_domain(x_cand: jax.Array):
+            if ranges.shape != (2, self.dim):
+                raise ValueError(
+                    f'ranges must have shape (2, {self.dim}).'
+                )
+
+            def domain_contains(x_cand: jax.Array) -> jax.Array:
                 return jnp.all(
-                    (x_cand >= ranges[0, 0] - tol) &
-                    (x_cand <= ranges[1, 0] + tol)
+                    (x_cand >= ranges[0] - tol) &
+                    (x_cand <= ranges[1] + tol)
                 )
 
-        if not inside_domain(x):
-            raise ValueError
+        if not domain_contains(x):
+            raise ValueError('Initial coordinate is outside the domain.')
 
         optim = optim(**optim_kwargs)
         optim_state = optim.init(x)
 
-        for i in range(n):
+        for _ in range(n):
             optim_state, x_cand, _variance = GaussianProcess._step_sample_point(
                 optim,
                 optim_state,
@@ -479,12 +479,11 @@ def inside_domain(x_cand: jax.Array):
                 x
             )
 
-            if inside_domain(x_cand):
-                x = x_cand
-            else:
-                return jnp.asarray(x)
+            if not domain_contains(x_cand):
+                break
+            x = x_cand
 
-        return jnp.asarray(x)
+        return x
 
     @staticmethod
     @jax.jit(static_argnames=('optim',))

demos/gp-1d.ipynb

@@ -190,8 +190,8 @@
    "source": [
     "N_INIT = 10\n",
     "\n",
-    "train_x = jnp.linspace(0, 2 * jnp.pi, N_INIT).reshape(-1, DIM)\n",
-    "train_y, train_s = f_noisy(train_x)    \n",
+    "train_x = jnp.linspace(*DOMAIN, N_INIT).reshape(-1, DIM)\n",
+    "train_y, train_s = f_noisy(train_x)\n",
     "\n",
     "g.condition(train_x, train_y, train_s)\n"
    ]
@@ -294,9 +294,10 @@
     "N_TRAIN = 10\n",
     "N_SAMPLE = 10\n",
     "\n",
-    "x_sample = jnp.linspace(0, 2 * jnp.pi, N_SAMPLE).reshape(-1, DIM)\n",
     "\n",
     "for i in range(N_TRAIN):\n",
+    "    x_sample = jnp.array(rng.uniform(*DOMAIN, N_SAMPLE)).reshape(-1, DIM)\n",
+    "    \n",
     "    # Greedily sample posterior variance:\n",
     "    _v = jnp.diag(g.posterior_covariance(x_sample, x_sample))\n",
     "    _x = x_sample[jnp.argmax(_v)].reshape(-1, DIM)\n",

tests/test_gp.py

@@ -53,6 +53,72 @@ def test_get_sample_point_optimizes_posterior_variance(self):
         self.assertLessEqual(float(x_sample[0, 0]), 2.0 + 1e-6)
         self.assertGreater(float(sample_variance), float(start_variance))
 
+    def test_get_sample_point_accepts_domain_predicate(self):
+        """Test posterior-variance sampling with a custom domain predicate."""
+        mu = gp.mean.Affine(
+            dim=1,
+            cdim=1,
+            a=jnp.zeros((1, 1)),
+            b=0.0,
+            a_is_static=True,
+            b_is_static=True
+        )
+        k = gp.kernel.Gaussian(
+            dim=1,
+            cdim=1,
+            D=jnp.eye(1),
+            D_is_static=True
+        )
+        g = gp.GaussianProcess(dim=1, cdim=1, mu=mu, k=k, seed=0)
+        g.condition(jnp.array([[0.0]]), jnp.array([[0.0]]), 1e-8)
+
+        shapes = []
+
+        def inside_domain(x):
+            shapes.append(x.shape)
+            return jnp.all((x >= 0.0) & (x <= 0.75))
+
+        x_sample = g.get_sample_point(
+            jnp.array([0.25]),
+            inside_domain=inside_domain,
+            n=100,
+            optim_kwargs={'learning_rate': 0.05}
+        )
+
+        self.assertTrue(shapes)
+        self.assertTrue(all(shape == (1, 1) for shape in shapes))
+        self.assertGreater(float(x_sample[0, 0]), 0.25)
+        self.assertLessEqual(float(x_sample[0, 0]), 0.75)
+
+    def test_get_sample_point_requires_one_domain_specification(self):
+        """Test that domain specifications are mutually exclusive."""
+        mu = gp.mean.Affine(
+            dim=1,
+            cdim=1,
+            a=jnp.zeros((1, 1)),
+            b=0.0,
+            a_is_static=True,
+            b_is_static=True
+        )
+        k = gp.kernel.Gaussian(
+            dim=1,
+            cdim=1,
+            D=jnp.eye(1),
+            D_is_static=True
+        )
+        g = gp.GaussianProcess(dim=1, cdim=1, mu=mu, k=k, seed=0)
+        g.condition(jnp.array([[0.0]]), jnp.array([[0.0]]), 1e-8)
+
+        with self.assertRaisesRegex(ValueError, 'exactly one'):
+            g.get_sample_point(jnp.array([0.25]), n=1)
+        with self.assertRaisesRegex(ValueError, 'exactly one'):
+            g.get_sample_point(
+                jnp.array([0.25]),
+                ranges=jnp.array([[0.0], [1.0]]),
+                inside_domain=lambda x: True,
+                n=1
+            )
+
     def test_kronecker_scalar_noise_factor_matches_dense_solve(self):
         """Test Kronecker structured solves with scalar observation noise."""
         x = jnp.array([[0.0], [0.5], [1.0]])