diff --git a/UncertainSCI/gp/__init__.py b/UncertainSCI/gp/__init__.py index d9478e1..ecf5307 100644 --- a/UncertainSCI/gp/__init__.py +++ b/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',)) diff --git a/demos/gp-1d.ipynb b/demos/gp-1d.ipynb index a31bc31..c53a1b2 100644 --- a/demos/gp-1d.ipynb +++ b/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", diff --git a/tests/test_gp.py b/tests/test_gp.py index 92deda0..6c0ff64 100644 --- a/tests/test_gp.py +++ b/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]])