diff --git a/tests/test_empirical_2pcf.py b/tests/test_empirical_2pcf.py index db456e8..e812fe2 100644 --- a/tests/test_empirical_2pcf.py +++ b/tests/test_empirical_2pcf.py @@ -1,6 +1,7 @@ import numpy as np import treegp import copy +import warnings from treegp_test_helper import timer from treegp_test_helper import get_correlation_length_matrix @@ -14,6 +15,7 @@ _apod, _adaptive_moments, ) +from treegp.grid_gp import _upsample_map, _symmetrize_kernel_image def make_elliptical_gaussian(npix, size, g1, g2): @@ -28,7 +30,7 @@ def make_elliptical_gaussian(npix, size, g1, g2): return np.exp(-0.5 * arg) -def make_gp(npoints=2000, noise=0.3, white_noise=0.0, seed=42): +def make_gp(npoints=2000, noise=0.3, white_noise=0.0, seed=42, **kwargs): """Generate a 2d GRF with a known anisotropic kernel and return an initialized and solved GPInterpolation using empirical-2pcf.""" L = get_correlation_length_matrix(2.0, 0.2, 0.2) @@ -41,6 +43,7 @@ def make_gp(npoints=2000, noise=0.3, white_noise=0.0, seed=42): white_noise=white_noise, max_sep=6.0, pixel_size=0.5, + **kwargs, ) gp.initialize(X, y, y_err=y_err) gp.solve() @@ -78,14 +81,14 @@ def test_empirical_2pcf_gp(): @timer -def test_empirical_2pcf_eigenvalue_clipping(): +def test_empirical_2pcf_psd_repair(): # The tabulated kernel is not guaranteed to be positive - # semi-definite between arbitrary points. By default the negative - # eigenvalues of the covariance matrix are clipped to zero - # (equivalent to the singular value clipping of Gomes et al. 2025); - # without the clipping, the Cholesky decomposition fails on this - # data set. - gp, X, y, y_err, noise = make_gp() + # semi-definite between arbitrary points. Calling the kernel with + # clip_eigenvalues=True (the default) clips the negative eigenvalues + # of the covariance matrix to zero (equivalent to the singular value + # clipping of Gomes et al. 2025), and the raw covariance of this + # data set is indeed indefinite. + gp, X, y, y_err, noise = make_gp(solve_method="direct") K = gp.kernel(X) eigenvalues = np.linalg.eigvalsh(K) assert np.all(eigenvalues > -1e-10) @@ -93,8 +96,29 @@ def test_empirical_2pcf_eigenvalue_clipping(): gp.kernel.clip_eigenvalues = False K_raw = gp.kernel(X) assert np.min(np.linalg.eigvalsh(K_raw)) < 0.0 + + # The direct solve path never pays for the eigenvalue clipping: it + # factors the raw covariance and repairs it with an escalating + # diagonal jitter, with a warning, and the prediction quality is + # preserved. gp._alpha = None - np.testing.assert_raises(np.linalg.LinAlgError, gp.predict, X) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + y_predict = gp.predict(X) + assert any("jitter" in str(wi.message) for wi in w) + assert np.var(y - y_predict) < 0.5 * np.var(y) + + # The spectral solve path is positive semi-definite by construction + # (the negative Fourier modes of the padded kernel image are clipped + # to zero once), so it needs no repair at all: no jitter warning. + gp_s, X_s, y_s, _, _ = make_gp(solve_method="spectral") + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + y_predict_s = gp_s.predict(X_s) + assert not any("jitter" in str(wi.message) for wi in w) + assert gp_s._engine.spectrum_min < 0.0 + assert np.all(gp_s._engine._spectrum >= 0.0) + assert np.var(y_s - y_predict_s) < 0.5 * np.var(y_s) @timer @@ -458,9 +482,154 @@ def test_empirical_2pcf_validation(): assert np.var(y - y_predict) < 0.5 * np.var(y) +@timer +def test_spectral_vs_direct(): + # The spectral and direct solve methods differ only by the tent + # smoothing of the kernel (one fine-grid pixel wide) and by the PSD + # repair (Fourier-space clipping vs diagonal jitter): predictions + # agree at a small fraction of the signal, and both catch the field. + gp_s, X, y, y_err, noise = make_gp(solve_method="spectral") + gp_d, _, _, _, _ = make_gp(solve_method="direct") + pred_s = gp_s.predict(X) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + pred_d = gp_d.predict(X) + assert np.std(pred_s - pred_d) < 0.15 * np.std(pred_d) + assert np.var(y - pred_s) < 0.5 * np.var(y) + assert np.var(y - pred_d) < 0.5 * np.var(y) + # The conjugate gradient converged and its diagnostics are filled. + assert gp_s._engine.n_iterations > 0 + assert gp_s._engine.xi0_eff > 0.0 + + # Solving is lazy (first predict) and survives a re-initialize with + # new values: the engine solution is reset and recomputed. + gp_s.initialize(X, y + 1.0, y_err=y_err) + assert gp_s._engine._alpha is None + pred_s2 = gp_s.predict(X) + # The agreement is limited by the conjugate gradient tolerance. + np.testing.assert_allclose(pred_s2, pred_s + 1.0, atol=1e-5) + + +@timer +def test_spectral_per_ccd(): + # Predicting per detector (many small calls) samples the same + # precomputed mean field as one big call: results are identical + # and the per-call cost does not depend on the training set size. + gp, X, y, y_err, noise = make_gp(solve_method="spectral") + rng = np.random.default_rng(5) + X_test = rng.uniform(np.min(X), np.max(X), (300, 2)) + full = gp.predict(X_test) + chunks = [gp.predict(X_test[i : i + 37]) for i in range(0, len(X_test), 37)] + np.testing.assert_allclose(np.concatenate(chunks), full, atol=1e-12) + + +@timer +def test_spectral_orientation(): + # The anisotropy of the tabulated kernel survives the spectral + # representation: the pairwise covariance implied by the engine + # (W G W^T) matches the tabulated kernel, including its + # orientation, up to the tent smoothing. + npix = 24 + pixel_size = 0.5 + lag = (np.arange(npix) - npix // 2) * pixel_size + dx, dy = np.meshgrid(lag, lag) + xi_grid = 4.0 * np.exp(-0.5 * (dx**2 / 4.0 + dy**2 / 1.0)) + kernel = treegp.EmpiricalCorrelationKernel(lag, lag, xi_grid) + + X = np.array([[1.5, 0.0], [0.0, 0.0], [0.0, 1.5]]) + engine = treegp.GridConvolutionGP(xi_grid, pixel_size, upsample=4) + engine._setup_geometry(X) + w = engine._spread_matrix(X) + K_eng = np.empty((3, 3)) + for j in range(3): + field = (w.T @ np.eye(3)[j]).reshape(engine._ny, engine._nx) + K_eng[:, j] = w @ engine._convolve(field).ravel() + K_tab = kernel(X) + np.testing.assert_allclose(K_eng, K_tab, atol=0.15) + # Long correlation length along x: the x-separated pair is more + # correlated than the y-separated one. + assert K_eng[0, 1] > 2.0 * K_eng[2, 1] + + +@timer +def test_grid_gp_helpers(): + # _upsample_map is exact on the input nodes and doubles the + # sampling of a band-limited map. + npix = 16 + lag = np.arange(npix) - npix // 2 + dx, dy = np.meshgrid(lag, lag) + xi = np.exp(-0.5 * (dx**2 / 9.0 + dy**2 / 4.0)) + for upsample in [1, 2, 3]: + up = _upsample_map(xi, upsample) + assert up.shape == (npix * upsample, npix * upsample) + np.testing.assert_allclose(up[::upsample, ::upsample], xi, atol=1e-12) + + # _symmetrize_kernel_image returns an exactly point-symmetric map + # with the unpaired edge lags halved (the grid-space equivalent of + # the 0.5 * (K + K.T) symmetrization of the dense path). + sym = _symmetrize_kernel_image(xi) + assert sym.shape == (npix + 1, npix + 1) + np.testing.assert_allclose(sym, sym[::-1, ::-1], atol=1e-14) + np.testing.assert_allclose(sym[npix, 1:npix], 0.5 * xi[0, :0:-1], atol=1e-14) + + # The engine solves (K + diag(y_err^2)) alpha = y: check the + # residual of the linear system through the engine's own matvec. + rng = np.random.default_rng(2) + X = rng.uniform(-6.0, 6.0, (400, 2)) + y = rng.normal(size=400) + y_err = np.full(400, 0.5) + engine = treegp.GridConvolutionGP(xi, 1.0, upsample=2, cg_rtol=1e-10) + engine.solve(X, y, y_err) + field = (engine._w_train.T @ engine._alpha).reshape(engine._ny, engine._nx) + resid = ( + engine._w_train @ engine._convolve(field).ravel() + y_err**2 * engine._alpha - y + ) + assert np.linalg.norm(resid) < 1e-8 * np.linalg.norm(y) + + # Prediction at the training points equals K alpha, and far from + # the data (beyond the kernel support) the mean field is zero. + pred = engine.predict(X) + np.testing.assert_allclose( + pred, engine._w_train @ engine._convolve(field).ravel(), atol=1e-12 + ) + np.testing.assert_allclose( + engine.predict(np.array([[100.0, 100.0], [-50.0, 3.0]])), 0.0, atol=1e-14 + ) + + # predict before solve raises, and invalid inputs are rejected. + engine.reset() + np.testing.assert_raises(RuntimeError, engine.predict, X) + np.testing.assert_raises(ValueError, treegp.GridConvolutionGP, xi[:, :-1], 1.0) + np.testing.assert_raises(ValueError, treegp.GridConvolutionGP, xi[:-1, :-1], 1.0) + np.testing.assert_raises(ValueError, treegp.GridConvolutionGP, xi, 1.0, upsample=0) + np.testing.assert_raises( + ValueError, treegp.GridConvolutionGP, xi, 1.0, upsample=1.5 + ) + + # Zero measurement errors trigger the conditioning warning. + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + engine.solve(X, y, np.zeros(400)) + assert any("zero error" in str(wi.message) for wi in w) + + # solve_method validation in GPInterpolation. + np.testing.assert_raises( + ValueError, + treegp.GPInterpolation, + optimizer="empirical-2pcf", + solve_method="woodbury", + ) + np.testing.assert_raises( + ValueError, + treegp.GPInterpolation, + optimizer="anisotropic", + solve_method="spectral", + ) + + if __name__ == "__main__": test_empirical_2pcf_gp() - test_empirical_2pcf_eigenvalue_clipping() + test_empirical_2pcf_psd_repair() test_empirical_2pcf_extrapolation() test_empirical_2pcf_introspection() test_empirical_kernel_orientation() @@ -468,3 +637,7 @@ def test_empirical_2pcf_validation(): test_adaptive_moments() test_empirical_2pcf_anisotropic_apod() test_empirical_2pcf_validation() + test_spectral_vs_direct() + test_spectral_per_ccd() + test_spectral_orientation() + test_grid_gp_helpers() diff --git a/treegp/__init__.py b/treegp/__init__.py index 44a1678..e94b5b8 100644 --- a/treegp/__init__.py +++ b/treegp/__init__.py @@ -11,6 +11,7 @@ from .two_pcf import two_pcf from .log_likelihood import log_likelihood from .empirical_2pcf import empirical_2pcf +from .grid_gp import GridConvolutionGP from .kernels import AnisotropicRBF from .kernels import VonKarman @@ -30,6 +31,7 @@ "two_pcf", "log_likelihood", "empirical_2pcf", + "GridConvolutionGP", "AnisotropicRBF", "VonKarman", "AnisotropicVonKarman", diff --git a/treegp/gp_interp.py b/treegp/gp_interp.py index 61ccfc3..1e5b8f1 100644 --- a/treegp/gp_interp.py +++ b/treegp/gp_interp.py @@ -5,9 +5,11 @@ import treegp import numpy as np import copy +import warnings from .kernels import eval_kernel from .kernels import EmpiricalCorrelationKernel +from .grid_gp import GridConvolutionGP from sklearn.neighbors import KNeighborsRegressor from scipy.linalg import cholesky, cho_solve @@ -53,10 +55,15 @@ class GPInterpolation(object): hyperparameters, so it is rejected by the fitting optimizers and can only be used with "empirical-2pcf" or "none". As the tabulated kernel is not guaranteed to be - positive semi-definite between arbitrary points, the negative - eigenvalues of the covariance matrix are clipped to zero - (equivalent to the singular value clipping of Gomes et al. 2025). - If the Cholesky decomposition still fails with a LinAlgError, + positive semi-definite between arbitrary points, the + covariance is repaired: the spectral solve method is + positive semi-definite by construction (the negative + Fourier modes of the padded kernel image are clipped + to zero), and the direct solve method adds an + escalating diagonal jitter until the Cholesky + factorization succeeds (both equivalent in spirit to + the singular value clipping of Gomes et al. 2025). + If the factorization still fails with a LinAlgError, increase white_noise. :param normalize: Whether to normalize the interpolation parameters to have a mean of 0. Normally, the parameters being interpolated are not mean 0, so you would @@ -116,6 +123,27 @@ class GPInterpolation(object): elliptical window when apod_anisotropy="auto", to soften (< 1) or exaggerate (> 1) the anisotropy of the taper. Used only by the "empirical-2pcf" optimizer. [default: 1.] + :param solve_method: How to solve for the gaussian process weights and evaluate + predictions with the "empirical-2pcf" optimizer (other + optimizers only accept "direct"). "spectral" exploits the grid + structure of the empirical kernel (see GridConvolutionGP): the + training covariance is never built, the solve is a + preconditioned conjugate gradient where each product costs one + FFT, and predictions sample a precomputed mean field, so time + and memory stay far below the O(N^3) / O(N^2) of the dense + path (minutes and several GB at N ~ 10^4). "direct" is the + dense path: covariance matrix plus Cholesky factorization. + Predictions with return_cov=True always use the dense path. + If None, defaults to "spectral" for the "empirical-2pcf" + optimizer and "direct" otherwise. [default: None] + :param spread_upsample: Integer upsampling factor of the spreading grid used by + solve_method="spectral", with respect to the correlation + function grid pixels. Controls the (second order) smoothing + of the kernel by the bilinear spreading. [default: 2] + :param cg_rtol: Relative tolerance of the conjugate gradient solve of + solve_method="spectral". [default: 1e-7] + :param cg_maxiter: Maximum number of conjugate gradient iterations of + solve_method="spectral". [default: 500] """ def __init__( @@ -138,6 +166,10 @@ def __init__( apod_radius=None, apod_anisotropy=None, apod_g_scale=1.0, + solve_method=None, + spread_upsample=2, + cg_rtol=1e-7, + cg_maxiter=500, ): self.normalize = normalize self.optimizer = optimizer @@ -153,6 +185,23 @@ def __init__( self.apod_radius = apod_radius self.apod_anisotropy = apod_anisotropy self.apod_g_scale = apod_g_scale + if solve_method is None: + solve_method = "spectral" if optimizer == "empirical-2pcf" else "direct" + if solve_method not in ["spectral", "direct"]: + raise ValueError( + "Only spectral and direct are supported for solve_method. " + "Current value: %s" % (solve_method) + ) + if solve_method == "spectral" and optimizer != "empirical-2pcf": + raise ValueError( + "solve_method='spectral' is only available for the " + "empirical-2pcf optimizer. Current optimizer: %s" % (optimizer) + ) + self.solve_method = solve_method + self.spread_upsample = spread_upsample + self.cg_rtol = cg_rtol + self.cg_maxiter = cg_maxiter + self._engine = None if self.optimizer == "anisotropic": self.robust_fit = True @@ -223,6 +272,7 @@ def _fit(self, kernel, X, y, y_err): """ self._alpha = None self._fact = None + self._engine = None if self.optimizer != "none": # Hyperparameters estimation using 2-point correlation # function information. @@ -258,6 +308,14 @@ def _fit(self, kernel, X, y, y_err): apod_g_scale=self.apod_g_scale, ) kernel = self._optimizer.optimizer(kernel) + if self.solve_method == "spectral": + self._engine = GridConvolutionGP( + self._optimizer._xi_clean, + self._optimizer.pixel_size, + upsample=self.spread_upsample, + cg_rtol=self.cg_rtol, + cg_maxiter=self.cg_maxiter, + ) # Hyperparameters estimation using maximum likelihood fit. if self.optimizer == "log-likelihood": self._optimizer = treegp.log_likelihood(X, y, y_err) @@ -273,6 +331,23 @@ def predict(self, X, return_cov=False): y_init = copy.deepcopy(self._y) y_err = copy.deepcopy(self._y_err) + if self._engine is not None and not return_cov: + # Spectral solve and predict (see GridConvolutionGP): the + # weights are solved once by conjugate gradient, the mean + # field is computed once as a grid convolution, and each + # predict call only samples it, so its cost is independent + # of the number of training points. return_cov=True falls + # through to the dense path below. + if self._engine._alpha is None: + self._engine.solve( + self._X, + y_init - self._mean - self._spatial_average, + y_err, + ) + y_interp = self._engine.predict(X) + y_interp += self._mean + self._build_average_meanify(X) + return y_interp + y_interp, y_cov = self.return_gp_predict( y_init - self._mean - self._spatial_average, self._X, @@ -300,15 +375,59 @@ def return_gp_predict(self, y, X1, X2, kernel, y_err, return_cov=False): """ HT = kernel.__call__(X2, Y=X1) if self._alpha is None: - K = kernel.__call__(X1) + np.eye(len(y)) * y_err**2 + if isinstance(kernel, EmpiricalCorrelationKernel): + # The PSD repair of the tabulated kernel is done at the + # factorization level below (escalating diagonal + # jitter), which is much cheaper than clipping the + # eigenvalues of the covariance matrix, so bypass the + # clipping when building the training covariance. + clip = kernel.clip_eigenvalues + kernel.clip_eigenvalues = False + try: + K = kernel.__call__(X1) + finally: + kernel.clip_eigenvalues = clip + else: + K = kernel.__call__(X1) + K[np.diag_indices_from(K)] += y_err**2 try: self._fact = cholesky(K, lower=True) except np.linalg.LinAlgError as e: - raise np.linalg.LinAlgError( - "Cholesky decomposition of the covariance matrix failed " - "(%s). The kernel might not be positive definite; " - "increasing white_noise can help." % (str(e)) + if not _kernel_contains(kernel, EmpiricalCorrelationKernel): + raise np.linalg.LinAlgError( + "Cholesky decomposition of the covariance matrix failed " + "(%s). The kernel might not be positive definite; " + "increasing white_noise can help." % (str(e)) + ) + # A tabulated correlation function is not guaranteed to + # be positive semi-definite between arbitrary points: + # add an escalating diagonal jitter until the + # factorization succeeds. + scale = np.mean(np.diagonal(K)) + jitter = 1e-6 * scale + added = 0.0 + fact = None + for _ in range(8): + K[np.diag_indices_from(K)] += jitter - added + added = jitter + try: + fact = cholesky(K, lower=True) + break + except np.linalg.LinAlgError: + jitter *= 10.0 + if fact is None: + raise np.linalg.LinAlgError( + "Cholesky decomposition of the covariance matrix " + "failed even with a diagonal jitter of %.3e (%s). " + "Increasing white_noise can help." % (added, str(e)) + ) + warnings.warn( + "The covariance matrix of the tabulated kernel is not " + "positive semi-definite: a diagonal jitter of %.3e " + "(%.2f%% of its mean diagonal) was added to make the " + "Cholesky factorization succeed." % (added, 100.0 * added / scale) ) + self._fact = fact self._alpha = cho_solve((self._fact, True), y) y_predict = np.dot(HT, self._alpha.reshape((len(self._alpha), 1))).T[0] if return_cov: @@ -351,6 +470,8 @@ def initialize(self, X, y, y_err=None): # input data. self._alpha = None self._fact = None + if self._engine is not None: + self._engine.reset() def _build_average_meanify(self, X): """Compute spatial average from meanify output for a given coordinate using KN interpolation. diff --git a/treegp/grid_gp.py b/treegp/grid_gp.py new file mode 100644 index 0000000..fb34a39 --- /dev/null +++ b/treegp/grid_gp.py @@ -0,0 +1,321 @@ +""" +.. module:: grid_gp +""" + +import warnings +import numpy as np + +from scipy import fft +from scipy import sparse +from scipy.sparse.linalg import LinearOperator, cg + + +def _upsample_map(xi, upsample): + """Upsample a 2d map by an integer factor using exact trigonometric + (Fourier zero-padding) interpolation. The input map is assumed to be + square with even side length and centered on pixel N//2 (treecorr + TwoD layout after _shift_and_bin); the output is centered on pixel + (N * upsample) // 2 and matches the input exactly on the input + nodes. Exact for band-limited maps, which the cleaned correlation + function is by construction (its Fourier power spectrum was + thresholded on the same grid). + + :param xi: 2d map, zero lag at pixel N//2. (N, N) ndarray + :param upsample: Integer upsampling factor. [required] + """ + if upsample == 1: + return xi.copy() + n = xi.shape[0] + m = n * upsample + # Spectrum with the frequency origin at the center, frequencies + # running from -n/2 to n/2 - 1 in each axis. + f_shift = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(xi))) + # Split the (unpaired) Nyquist row and column between -n/2 and + # +n/2 so that the interpolant of a real map is real. + f_ext = np.zeros((n + 1, n + 1), dtype=complex) + f_ext[:n, :n] = f_shift + f_ext[n, :n] = f_shift[0, :] + f_ext[:n, n] = f_shift[:, 0] + f_ext[n, n] = f_shift[0, 0] + f_ext[0, :] *= 0.5 + f_ext[n, :] *= 0.5 + f_ext[:, 0] *= 0.5 + f_ext[:, n] *= 0.5 + # Embed into the fine-grid spectrum, frequencies -m/2 to m/2 - 1. + pad = m // 2 - n // 2 + f_big = np.zeros((m, m), dtype=complex) + f_big[pad : pad + n + 1, pad : pad + n + 1] = f_ext + out = np.fft.fftshift(np.fft.ifft2(np.fft.ifftshift(f_big))) + return out.real * upsample**2 + + +def _symmetrize_kernel_image(xi): + """Return a point-symmetric copy of a kernel image, extended by one + row and column. The input is a square even-sized map with zero lag + at pixel N//2, whose first row and column (the most negative lag) + have no positive counterpart on the grid; the output has odd side + length N + 1, zero lag exactly at its center, and satisfies + out[c + i, c + j] == out[c - i, c - j]. This is the grid-space + equivalent of the 0.5 * (K + K.T) symmetrization used for the dense + covariance matrices (the unpaired edge lags end up halved, exactly + as in the dense case where their mirror image falls outside the + grid and evaluates to zero), and makes the Fourier transform of the + image exactly real. + + :param xi: 2d kernel image, zero lag at pixel N//2. (N, N) ndarray + """ + n = xi.shape[0] + ext = np.zeros((n + 1, n + 1)) + ext[:n, :n] = xi + return 0.5 * (ext + ext[::-1, ::-1]) + + +class GridConvolutionGP(object): + """Fast, low-memory solve and predict engine for the empirical-2pcf + (Gomes et al. 2025) kernel, exploiting its grid structure instead of + dense linear algebra. + + The training covariance is represented as K = W G W^T, where W is + the sparse bilinear (tent) spreading matrix of the points onto a + uniform fine grid and G is the circular convolution by the kernel + image on a zero-padded grid, diagonalized by FFT. The negative + Fourier modes of the padded kernel image are clipped to zero once, + in Fourier space, which makes the effective kernel + (tent * xi * tent, with a non-negative spectrum) continuously + positive semi-definite: the covariance of ANY point set is PSD by + construction, replacing the O(N^3) per-point-set eigenvalue clipping + of the dense path. The padded grid is large enough that the circular + wrap-around never reaches a real pair separation. + + (K + diag(y_err^2)) alpha = y is solved by Jacobi-preconditioned + conjugate gradient, where each matvec costs one FFT pair on the + padded grid: O(n_grid log n_grid + n_points) time and memory instead + of O(N^3) time and O(N^2) memory. The posterior mean is then a + single field, ŷ(x) = sum_j alpha_j k(x - x_j), computed once as a + convolution on the grid; each predict call only bilinearly samples + that precomputed field, so its cost is independent of the number of + training points (useful when predicting per CCD on a full visit). + + Compared to the dense path, the kernel is smoothed by the tent + spreading (one fine-grid pixel wide, i.e. pixel_size / upsample) and + the PSD projection is done on the padded grid instead of on the + point-set covariance; both effects are second order and shrink with + upsample. + + :param xi_grid: Cleaned 2d correlation function in treecorr TwoD + layout ([iy, ix]), zero lag at pixel npix//2, as + produced by empirical_2pcf. (npix, npix) ndarray + :param pixel_size: Pixel size of xi_grid, in the same units as the + coordinates of the field. + :param upsample: Integer upsampling factor of the spreading grid + with respect to the xi_grid pixels. The kernel + image is upsampled exactly (Fourier zero-padding, + xi_grid is band-limited by construction), so + upsample only controls the tent smoothing scale + and the grid memory. [default: 2] + :param cg_rtol: Relative tolerance of the conjugate gradient + solve. [default: 1e-7] + :param cg_maxiter: Maximum number of conjugate gradient iterations. + [default: 500] + """ + + def __init__(self, xi_grid, pixel_size, upsample=2, cg_rtol=1e-7, cg_maxiter=500): + xi_grid = np.asarray(xi_grid, dtype=float) + if xi_grid.ndim != 2 or xi_grid.shape[0] != xi_grid.shape[1]: + raise ValueError( + "xi_grid must be square. Current shape: %s" % (str(xi_grid.shape)) + ) + if xi_grid.shape[0] % 2 != 0: + raise ValueError( + "xi_grid side length must be even. Current shape: %s" + % (str(xi_grid.shape)) + ) + if not isinstance(upsample, (int, np.integer)) or upsample < 1: + raise ValueError( + "upsample must be a positive integer. Current value: %s" % (upsample) + ) + self.pixel_size = float(pixel_size) + self.upsample = int(upsample) + self.cg_rtol = float(cg_rtol) + self.cg_maxiter = int(cg_maxiter) + # Half width of the kernel support, from the zero lag to the + # grid edge. + self.max_sep = (xi_grid.shape[0] // 2) * self.pixel_size + # Fine grid spacing and exactly-upsampled, symmetrized kernel + # image (odd side length, zero lag at its center). + self.h = self.pixel_size / self.upsample + self._kernel_image = _symmetrize_kernel_image( + _upsample_map(xi_grid, self.upsample) + ) + self.reset() + + def reset(self): + """Forget the current solution (and the training-set geometry), + so that the next predict triggers a new solve.""" + self._alpha = None + self._mean_field = None + self._nx = None + self._ny = None + self.n_iterations = None + self.spectrum_min = None + self.xi0_eff = None + + def _setup_geometry(self, X): + """Build the fine spreading grid covering the training points + plus the kernel support, the padded FFT grid, and the clipped + non-negative kernel spectrum on it.""" + ks = self._kernel_image.shape[0] + margin = self.max_sep + self.h + self._x0 = np.min(X[:, 0]) - margin + self._y0 = np.min(X[:, 1]) - margin + # Number of grid nodes covering the domain (bilinear spreading + # needs one node beyond the last point in each axis). + self._nx = int(np.ceil((np.max(X[:, 0]) - self._x0) / self.h)) + 2 + self._ny = int(np.ceil((np.max(X[:, 1]) - self._y0) / self.h)) + 2 + # Padded FFT grid: large enough for a linear (wrap-free) + # convolution of the field with the kernel image. + self._lx = fft.next_fast_len(self._nx + ks - 1, real=True) + self._ly = fft.next_fast_len(self._ny + ks - 1, real=True) + # Kernel image with its zero lag rolled to pixel (0, 0) of the + # padded grid, and its spectrum, exactly real thanks to the + # point symmetry of the image. Clipping the negative modes to + # zero is the one-time PSD projection. + buf = np.zeros((self._ly, self._lx)) + buf[:ks, :ks] = self._kernel_image + buf = np.roll(buf, (-(ks // 2), -(ks // 2)), axis=(0, 1)) + spectrum = fft.rfft2(buf).real + self.spectrum_min = float(np.min(spectrum)) + self._spectrum = np.clip(spectrum, 0.0, None) + # Effective zero-lag variance of the clipped kernel. + image_clip = fft.irfft2(self._spectrum, s=(self._ly, self._lx)) + self.xi0_eff = float(image_clip[0, 0]) + + def _spread_matrix(self, X, name="X"): + """Sparse bilinear (tent) spreading matrix of the given points + onto the fine grid: W[i, iy * nx + ix] holds the weight of point + i on grid node (iy, ix). Points must be inside the grid.""" + gx = (X[:, 0] - self._x0) / self.h + gy = (X[:, 1] - self._y0) / self.h + if np.any(gx < 0) or np.any(gx > self._nx - 2): + raise ValueError("Some %s coordinates fall outside the grid." % (name)) + if np.any(gy < 0) or np.any(gy > self._ny - 2): + raise ValueError("Some %s coordinates fall outside the grid." % (name)) + ix = np.floor(gx).astype(np.intp) + iy = np.floor(gy).astype(np.intp) + tx = gx - ix + ty = gy - iy + n = len(X) + rows = np.repeat(np.arange(n), 4) + cols = np.empty((n, 4), dtype=np.intp) + cols[:, 0] = iy * self._nx + ix + cols[:, 1] = iy * self._nx + ix + 1 + cols[:, 2] = (iy + 1) * self._nx + ix + cols[:, 3] = (iy + 1) * self._nx + ix + 1 + vals = np.empty((n, 4)) + vals[:, 0] = (1.0 - tx) * (1.0 - ty) + vals[:, 1] = tx * (1.0 - ty) + vals[:, 2] = (1.0 - tx) * ty + vals[:, 3] = tx * ty + return sparse.csr_matrix( + (vals.ravel(), (rows, cols.ravel())), + shape=(n, self._nx * self._ny), + ) + + def _convolve(self, field): + """Convolve a field given on the (ny, nx) grid nodes with the + clipped kernel image, through the zero-padded FFT grid.""" + buf = np.zeros((self._ly, self._lx)) + buf[: self._ny, : self._nx] = field + out = fft.irfft2(fft.rfft2(buf) * self._spectrum, s=(self._ly, self._lx)) + return out[: self._ny, : self._nx] + + def solve(self, X, y, y_err): + """Solve (K + diag(y_err^2)) alpha = y by preconditioned + conjugate gradient, where K is the PSD-by-construction grid + kernel evaluated between the training points. + + :param X: Coordinates of the field. (n_samples, 2) + :param y: Values of the field, already centered (mean and + spatial average subtracted). (n_samples) + :param y_err: Error of y. (n_samples) + """ + X = np.asarray(X, dtype=float) + y = np.asarray(y, dtype=float) + y_err = np.asarray(y_err, dtype=float) + self.reset() + self._setup_geometry(X) + self._w_train = self._spread_matrix(X, name="training") + noise = y_err**2 + if np.min(noise) <= 0.0: + warnings.warn( + "Some points have zero error: K + diag(y_err^2) is only " + "positive SEMI-definite and the conjugate gradient solve " + "may struggle to converge; setting white_noise > 0 can help." + ) + + w = self._w_train + + def matvec(v): + field = (w.T @ v).reshape(self._ny, self._nx) + return w @ self._convolve(field).ravel() + noise * v + + n = len(y) + operator = LinearOperator((n, n), matvec=matvec, dtype=float) + precond_diag = 1.0 / (self.xi0_eff + noise) + precond = LinearOperator((n, n), matvec=lambda v: precond_diag * v, dtype=float) + iterations = [0] + + def callback(xk): + iterations[0] += 1 + + alpha, info = cg( + operator, + y, + rtol=self.cg_rtol, + atol=0.0, + maxiter=self.cg_maxiter, + M=precond, + callback=callback, + ) + if info > 0: + resid = np.linalg.norm(matvec(alpha) - y) / np.linalg.norm(y) + warnings.warn( + "Conjugate gradient did not converge to rtol=%.1e in %i " + "iterations (relative residual: %.1e). The solution is " + "used anyway; increasing white_noise or cg_maxiter can " + "help." % (self.cg_rtol, self.cg_maxiter, resid) + ) + self.n_iterations = iterations[0] + self._alpha = alpha + + def predict(self, X): + """Evaluate the posterior mean ŷ(x) = sum_j alpha_j k(x - x_j) + at the given coordinates by bilinear sampling of the mean field, + which is computed once (a single convolution of the spread alpha + with the kernel image) and cached until the next solve. Points + beyond the kernel support of every training point return 0. + + :param X: The coordinates at which to interpolate. (n_samples, 2) + """ + if self._alpha is None: + raise RuntimeError("solve() must be called before predict().") + if self._mean_field is None: + field = (self._w_train.T @ self._alpha).reshape(self._ny, self._nx) + self._mean_field = self._convolve(field) + X = np.asarray(X, dtype=float) + gx = (X[:, 0] - self._x0) / self.h + gy = (X[:, 1] - self._y0) / self.h + inside = (gx >= 0.0) & (gx <= self._nx - 1) & (gy >= 0.0) & (gy <= self._ny - 1) + ix = np.clip(np.floor(gx).astype(np.intp), 0, self._nx - 2) + iy = np.clip(np.floor(gy).astype(np.intp), 0, self._ny - 2) + tx = gx - ix + ty = gy - iy + f = self._mean_field + out = ( + (1.0 - tx) * (1.0 - ty) * f[iy, ix] + + tx * (1.0 - ty) * f[iy, ix + 1] + + (1.0 - tx) * ty * f[iy + 1, ix] + + tx * ty * f[iy + 1, ix + 1] + ) + out[~inside] = 0.0 + return out diff --git a/treegp/kernels.py b/treegp/kernels.py index 26f01ff..bc76749 100644 --- a/treegp/kernels.py +++ b/treegp/kernels.py @@ -451,6 +451,11 @@ class EmpiricalCorrelationKernel(StationaryKernelMixin, Kernel): [default: True] """ + # Maximum number of pair lags interpolated at once by __call__, to + # bound the size of the temporary arrays (a few 100 MB) instead of + # scaling with n_samples^2. + _chunk_size = 4_000_000 + def __init__(self, x_grid, y_grid, xi_grid, clip_eigenvalues=True): self.x_grid = np.asarray(x_grid) self.y_grid = np.asarray(y_grid) @@ -461,21 +466,82 @@ def __init__(self, x_grid, y_grid, xi_grid, clip_eigenvalues=True): "xi_grid shape %s does not match (len(y_grid), len(x_grid)) = %s" % (str(self.xi_grid.shape), str((len(self.y_grid), len(self.x_grid)))) ) - # xi_grid is indexed [iy, ix], while the interpolator axes are + # xi_grid is indexed [iy, ix], while the interpolation axes are # (x lag, y lag), hence the transpose. - self._interp = RegularGridInterpolator( - (self.x_grid, self.y_grid), - self.xi_grid.T, - method="linear", - bounds_error=False, - fill_value=0.0, + self._z = np.ascontiguousarray(self.xi_grid.T) + # Uniform grids (the empirical_2pcf solver always produces them) + # take a direct vectorized bilinear interpolation, much faster + # than the generic RegularGridInterpolator, which is kept as a + # fallback for non-uniform grids. + steps_x = np.diff(self.x_grid) + steps_y = np.diff(self.y_grid) + uniform = ( + len(steps_x) > 0 + and len(steps_y) > 0 + and np.all(steps_x > 0) + and np.all(steps_y > 0) + and np.allclose(steps_x, steps_x[0]) + and np.allclose(steps_y, steps_y[0]) + ) + if uniform: + self._interp = None + self._hx = steps_x[0] + self._hy = steps_y[0] + else: + self._interp = RegularGridInterpolator( + (self.x_grid, self.y_grid), + self._z, + method="linear", + bounds_error=False, + fill_value=0.0, + ) + + def _eval_lags(self, dx, dy): + """Bilinear interpolation of the tabulated correlation function + at the given (dx, dy) lags, zero outside the grid. + + :param dx, dy: Lags where to evaluate the correlation function. + ndarrays of matching shape. + """ + if self._interp is not None: + return self._interp(np.stack([dx, dy], axis=-1)) + nx = len(self.x_grid) + ny = len(self.y_grid) + gx = (dx - self.x_grid[0]) / self._hx + gy = (dy - self.y_grid[0]) / self._hy + inside = (gx >= 0.0) & (gx <= nx - 1) & (gy >= 0.0) & (gy <= ny - 1) + ix = np.clip(np.floor(gx).astype(np.intp), 0, nx - 2) + iy = np.clip(np.floor(gy).astype(np.intp), 0, ny - 2) + tx = gx - ix + ty = gy - iy + z = self._z + out = ( + (1.0 - tx) * (1.0 - ty) * z[ix, iy] + + tx * (1.0 - ty) * z[ix + 1, iy] + + (1.0 - tx) * ty * z[ix, iy + 1] + + tx * ty * z[ix + 1, iy + 1] ) + return np.where(inside, out, 0.0) + + def _pair_covariance(self, X, Y): + """Evaluate K[i, j] = xi(X_i - Y_j) in chunks of rows, so that + the temporary arrays never scale with len(X) * len(Y).""" + n1 = len(X) + n2 = len(Y) + K = np.empty((n1, n2)) + chunk = max(1, self._chunk_size // max(n2, 1)) + for start in range(0, n1, chunk): + end = min(start + chunk, n1) + dx = X[start:end, 0][:, np.newaxis] - Y[np.newaxis, :, 0] + dy = X[start:end, 1][:, np.newaxis] - Y[np.newaxis, :, 1] + K[start:end] = self._eval_lags(dx, dy) + return K @property def xi0(self): """Zero-lag value of the correlation function, i.e. the variance of the field.""" - return float(self._interp(np.zeros((1, 2)))[0]) + return float(self._eval_lags(np.zeros(1), np.zeros(1))[0]) def __call__(self, X, Y=None, eval_gradient=False): if eval_gradient: @@ -487,8 +553,7 @@ def __call__(self, X, Y=None, eval_gradient=False): "Current ndim: %i" % (np.shape(X)[1]) ) if Y is None: - d = X[:, np.newaxis, :] - X[np.newaxis, :, :] - K = self._interp(d) + K = self._pair_covariance(X, X) # The tabulated correlation function is point-symmetric except # for its first row/column (the most negative lag has no # positive counterpart on the grid), so symmetrize to get an @@ -503,8 +568,7 @@ def __call__(self, X, Y=None, eval_gradient=False): K = 0.5 * (K + K.T) else: Y = np.atleast_2d(Y) - d = X[:, np.newaxis, :] - Y[np.newaxis, :, :] - K = self._interp(d) + K = self._pair_covariance(X, Y) return K def diag(self, X):