From 8e175a14a77e057bdc692ef3433f0500fc65e5fe Mon Sep 17 00:00:00 2001 From: PFLeget Date: Thu, 20 Aug 2026 11:05:26 -0400 Subject: [PATCH 1/4] Implement Gomes 2025 witht the help of Claude/Fable5. --- README.rst | 8 +- docs/treegp_gp_interp.rst | 31 ++++ tests/test_empirical_2pcf.py | 303 +++++++++++++++++++++++++++++++++++ treegp/__init__.py | 4 + treegp/empirical_2pcf.py | 281 ++++++++++++++++++++++++++++++++ treegp/gp_interp.py | 130 ++++++++++++++- treegp/kernels.py | 96 +++++++++++ 7 files changed, 845 insertions(+), 8 deletions(-) create mode 100644 tests/test_empirical_2pcf.py create mode 100644 treegp/empirical_2pcf.py diff --git a/README.rst b/README.rst index 70f1662..c48636b 100644 --- a/README.rst +++ b/README.rst @@ -20,12 +20,14 @@ Overview ``treegp`` has some special features compared to other available Gaussian Processes codes: * Hyperparameters estimation will scale in O(N log(N)) with the the 2-points correlation function estimation compared to O(N^3) with the classical maximum likelihood. - + * Gaussian process interpolation can be performed around a mean function - + * A tool is provided to compute the mean function (``meanify``) -``treegp`` was originally developed for Point Spread Function interpolation within `Piff `_. There is a specific article that describes the math used in ``treegp`` in the context of modelling astrometric shifts of the Subaru Telescope due to atmospheric turbulences. This article can be found +* The measured anisotropic 2D 2-points correlation function can be used directly as the kernel, without fitting any hyperparameters (``optimizer="empirical-2pcf"``), implementing `Gomes et al. (2025) `_. The other optimizers (``"two-pcf"``, ``"anisotropic"``) implement `Léget et al. (2021) `_. + +``treegp`` was originally developed for Point Spread Function interpolation within `Piff `_. There is a specific article that describes the math used in ``treegp`` in the context of modelling astrometric shifts of the Subaru Telescope due to atmospheric turbulences. This article can be found `here `_. diff --git a/docs/treegp_gp_interp.rst b/docs/treegp_gp_interp.rst index c8642ec..f4a7ac3 100644 --- a/docs/treegp_gp_interp.rst +++ b/docs/treegp_gp_interp.rst @@ -118,6 +118,37 @@ To do the gaussian process interpolation with ``treegp`` it follow this API: .. image:: fitted.png +Using the measured 2-points correlation function as the kernel +============================================================== + +Instead of fitting the hyperparameters of a parametric kernel on the +measured 2-points correlation function (`Léget et al 2021 +`_), the measured anisotropic 2D +2-points correlation function can be used directly as the kernel, following +`Gomes et al 2025 `_. The measured +correlation function is cleaned by apodization and thresholding of its +Fourier power spectrum, tabulated on a grid, and interpolated at the pair +separations. No hyperparameters are fitted, so the ``kernel`` argument is +ignored; the grid is controlled by ``max_sep`` (half width) and +``pixel_size``: + +.. code:: ipython3 + + gp = treegp.GPInterpolation( + optimizer="empirical-2pcf", + normalize=True, + max_sep=6.0, + pixel_size=0.5, + ) + gp.initialize(x, y, y_err=y_err) + # measure, clean, and tabulate the 2-points correlation function + gp.solve() + y_test, y_test_cov = gp.predict(x_test, return_cov=True) + + # measured and cleaned correlation functions can be inspected with: + xi, xi_clean, distance, pixel_size = gp.return_empirical_2pcf() + + diff --git a/tests/test_empirical_2pcf.py b/tests/test_empirical_2pcf.py new file mode 100644 index 0000000..a228826 --- /dev/null +++ b/tests/test_empirical_2pcf.py @@ -0,0 +1,303 @@ +import numpy as np +import treegp +import copy + +from treegp_test_helper import timer +from treegp_test_helper import get_correlation_length_matrix +from treegp_test_helper import make_2d_grf + +from treegp.empirical_2pcf import ( + _shift_and_bin, + _threshold, + _corr2power, + _power2corr, + _apod, +) + + +def make_gp(npoints=2000, noise=0.3, white_noise=0.0, seed=42): + """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) + invLam = np.linalg.inv(L) + kernel = 2.0**2 * treegp.AnisotropicRBF(invLam=invLam) + X, y, y_err = make_2d_grf(kernel, noise=noise, seed=seed, npoints=npoints) + gp = treegp.GPInterpolation( + optimizer="empirical-2pcf", + normalize=True, + white_noise=white_noise, + max_sep=6.0, + pixel_size=0.5, + ) + gp.initialize(X, y, y_err=y_err) + gp.solve() + return gp, X, y, y_err, noise + + +@timer +def test_empirical_2pcf_gp(): + gp, X, y, y_err, noise = make_gp() + + # No hyperparameters are fitted: the kernel is a tabulated + # correlation function with an empty theta. + assert isinstance(gp.kernel, treegp.EmpiricalCorrelationKernel) + assert len(gp.kernel.theta) == 0 + + # The zero lag of the cleaned correlation function is the + # variance of the field (up to noise in the measured 2-pcf). + np.testing.assert_allclose(gp.kernel.xi0, np.var(y) - noise**2, rtol=3e-1) + + y_predict, y_cov = gp.predict(X, return_cov=True) + y_std = np.sqrt(np.diag(y_cov)) + + # Check that the GP interpolation catches a good fraction of + # the field variance. + residuals = y - y_predict + assert np.var(residuals) < 0.5 * np.var(y) + + # Pull distribution should have a mean of 0 and a std < 1 + # (as the interpolation is better than the noise). + pull = residuals / np.sqrt(gp._y_err**2 + y_std**2) + mean_pull = np.mean(pull) + std_pull = np.std(pull) + assert np.abs(mean_pull) < 3.0 * std_pull / np.sqrt(len(y)) + assert std_pull < 1.0 + + +@timer +def test_empirical_2pcf_eigenvalue_clipping(): + # 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() + K = gp.kernel(X) + eigenvalues = np.linalg.eigvalsh(K) + assert np.all(eigenvalues > -1e-10) + + gp.kernel.clip_eigenvalues = False + K_raw = gp.kernel(X) + assert np.min(np.linalg.eigvalsh(K_raw)) < 0.0 + gp._alpha = None + np.testing.assert_raises(np.linalg.LinAlgError, gp.predict, X) + + +@timer +def test_empirical_2pcf_extrapolation(): + gp, X, y, y_err, noise = make_gp() + + # Far from the data (beyond max_sep), the kernel is zero, so the GP + # returns the mean of the field with a variance equal to the zero + # lag of the correlation function. + np.random.seed(42) + X_far = np.random.uniform( + np.max(X) + 6.0 * gp._optimizer.max_sep, + np.max(X) + 12.0 * gp._optimizer.max_sep, + size=20, + ).reshape((10, 2)) + y_far, y_cov_far = gp.predict(X_far, return_cov=True) + np.testing.assert_allclose(y_far, np.mean(y), atol=1e-10) + np.testing.assert_allclose( + np.sqrt(np.diag(y_cov_far)), np.sqrt(gp.kernel.xi0), atol=1e-10 + ) + + +@timer +def test_empirical_2pcf_introspection(): + gp, X, y, y_err, noise = make_gp() + + npix = gp._optimizer.npix + xi, xi_clean, distance, pixel_size = gp.return_empirical_2pcf() + assert xi.shape == (npix, npix) + assert xi_clean.shape == (npix, npix) + assert distance.shape == (npix * npix, 2) + assert pixel_size == 0.5 + np.testing.assert_allclose(xi.flatten(), gp._optimizer._2pcf, atol=1e-10) + np.testing.assert_allclose(xi_clean.flatten(), gp._optimizer._2pcf_fit, atol=1e-10) + + # plot_fitted_kernel uses _2pcf, _2pcf_fit and _2pcf_dist, so it + # works for the empirical-2pcf optimizer. + import matplotlib + + matplotlib.use("Agg") + gp.plot_fitted_kernel() + + # return_2pcf is only meaningful for the two-pcf and anisotropic + # optimizers. + np.testing.assert_raises(NotImplementedError, gp.return_2pcf) + + # return_empirical_2pcf is only available for the empirical-2pcf + # optimizer. + gp_iso = treegp.GPInterpolation(optimizer="two-pcf") + np.testing.assert_raises(NotImplementedError, gp_iso.return_empirical_2pcf) + + +@timer +def test_empirical_kernel_orientation(): + # Build an analytic anisotropic correlation function, with a + # correlation length larger along x than along y, on a grid in + # treecorr TwoD layout, i.e. indexed [iy, ix]. + npix = 20 + pixel_size = 1.0 + lag = (np.arange(npix) - npix // 2) * pixel_size + dx, dy = np.meshgrid(lag, lag) + xi_grid = np.exp(-0.5 * (dx**2 / 16.0 + dy**2 / 1.0)) + + kernel = treegp.EmpiricalCorrelationKernel(lag, lag, xi_grid) + + # A separation along x should be evaluated with the long + # correlation length, a separation along y with the short one. + X = np.array([[3.0, 0.0], [0.0, 0.0], [0.0, 3.0]]) + K = kernel(X) + np.testing.assert_allclose(K[0, 1], np.exp(-0.5 * 9.0 / 16.0), atol=1e-10) + np.testing.assert_allclose(K[2, 1], np.exp(-0.5 * 9.0 / 1.0), atol=1e-10) + np.testing.assert_allclose(K, K.T, atol=1e-10) + np.testing.assert_allclose(np.diag(K), kernel.diag(X), atol=1e-10) + np.testing.assert_allclose(kernel.xi0, 1.0, atol=1e-10) + + # Cross covariance between two sets of points. + X2 = np.array([[1.0, 0.0], [0.0, 1.0], [2.0, 2.0], [10.0, -3.0]]) + HT = kernel(X2, Y=X) + assert HT.shape == (len(X2), len(X)) + np.testing.assert_allclose(HT[0, 1], np.exp(-0.5 * 1.0 / 16.0), atol=1e-10) + np.testing.assert_allclose(HT[1, 1], np.exp(-0.5 * 1.0 / 1.0), atol=1e-10) + + # Beyond the grid, the kernel is zero. + X_far = np.array([[100.0, 100.0]]) + np.testing.assert_allclose(kernel(X_far, Y=X), 0.0, atol=1e-10) + + # sklearn kernel API with an empty theta. + assert len(kernel.theta) == 0 + kernel_clone = kernel.clone_with_theta(kernel.theta) + np.testing.assert_allclose(kernel_clone(X), K, atol=1e-10) + kernel_copy = copy.deepcopy(kernel) + np.testing.assert_allclose(kernel_copy(X), K, atol=1e-10) + assert kernel.is_stationary() + + # Only 2d coordinates are supported. + np.testing.assert_raises(ValueError, kernel, np.array([[1.0], [2.0]])) + + +@timer +def test_empirical_2pcf_helpers(): + # _shift_and_bin conserves the sum (up to the 1/4 normalization) + # and moves the zero lag from the intersection of the 4 central + # pixels to the center of pixel N//2. + np.random.seed(42) + raw = np.random.uniform(size=(40, 40)) + binned = _shift_and_bin(raw) + assert binned.shape == (20, 20) + np.testing.assert_allclose(np.sum(binned), np.sum(raw) / 4.0, atol=1e-10) + raw_peak = np.zeros((40, 40)) + raw_peak[19:21, 19:21] = 1.0 + binned_peak = _shift_and_bin(raw_peak) + assert binned_peak[10, 10] == 1.0 + np.testing.assert_allclose(np.sum(binned_peak), 1.0, atol=1e-10) + np.testing.assert_raises(ValueError, _shift_and_bin, np.zeros((39, 39))) + np.testing.assert_raises(ValueError, _shift_and_bin, np.zeros((40, 20))) + + # _power2corr is the inverse of _corr2power for a symmetric map. + npix = 20 + lag = np.arange(npix) - npix // 2 + dx, dy = np.meshgrid(lag, lag) + xi = np.exp(-0.5 * (dx**2 + dy**2) / 9.0) + np.testing.assert_allclose(_power2corr(_corr2power(xi)), xi, atol=1e-10) + + # _threshold zeroes the elements below the threshold and keeps + # the ones above. + p = np.ones((20, 20)) + p[10, 10] = 1e4 + p_thresh = _threshold(p, n_sigma=3.0) + assert p_thresh[10, 10] == 1e4 + assert np.sum(p_thresh != 0.0) == 1 + # An all-zero power spectrum is returned unchanged. + np.testing.assert_allclose( + _threshold(np.zeros((20, 20))), np.zeros((20, 20)), atol=1e-10 + ) + + # _apod is 1 at zero lag and goes to zero at the edge of the grid + # (the 4-term Blackman-Harris window is 6e-5 at its edge). + window = _apod(xi) + np.testing.assert_allclose(window[npix // 2, npix // 2], 1.0, atol=1e-10) + np.testing.assert_allclose(window[npix // 2, 0], 0.0, atol=1e-4) + + +@timer +def test_empirical_2pcf_validation(): + # Only 2d fields are supported. + X = np.random.uniform(-10, 10, 100).reshape((100, 1)) + y = np.random.normal(size=100) + np.testing.assert_raises(ValueError, treegp.empirical_2pcf, X, y, np.zeros_like(y)) + + # max_sep must span at least 2 pixels. + X = np.random.uniform(-10, 10, 200).reshape((100, 2)) + np.testing.assert_raises( + ValueError, + treegp.empirical_2pcf, + X, + y, + np.zeros_like(y), + 1.0, + 2.0, + ) + + # Unknown optimizer is rejected. + np.testing.assert_raises(ValueError, treegp.GPInterpolation, optimizer="gomes25") + + # The empirical-2pcf optimizer builds its own kernel: passing one + # is rejected. + np.testing.assert_raises( + ValueError, + treegp.GPInterpolation, + kernel="2.0**2 * AnisotropicVonKarman(scale_length=[1.0, 1.0])", + optimizer="empirical-2pcf", + ) + + # An EmpiricalCorrelationKernel has no hyperparameters: it is + # rejected by the fitting optimizers (also inside a composite + # kernel), but allowed with optimizer="none". + kernel_str = ( + "EmpiricalCorrelationKernel(array([-1., 0., 1.]), array([-1., 0., 1.]), " + "array([[0., 0., 0.], [0., 1., 0.], [0., 0., 0.]]))" + ) + for opt in ["two-pcf", "anisotropic", "log-likelihood"]: + np.testing.assert_raises( + ValueError, treegp.GPInterpolation, kernel=kernel_str, optimizer=opt + ) + np.testing.assert_raises( + ValueError, + treegp.GPInterpolation, + kernel="2.0**2 * " + kernel_str, + optimizer=opt, + ) + gp_none = treegp.GPInterpolation(kernel=kernel_str, optimizer="none") + assert isinstance(gp_none.kernel_template, treegp.EmpiricalCorrelationKernel) + + # y_err = 0 runs (unweighted 2-point correlation function). + L = get_correlation_length_matrix(2.0, 0.2, 0.2) + invLam = np.linalg.inv(L) + kernel = 2.0**2 * treegp.AnisotropicRBF(invLam=invLam) + X, y, _ = make_2d_grf(kernel, noise=None, seed=42, npoints=1000) + gp = treegp.GPInterpolation( + optimizer="empirical-2pcf", + normalize=True, + white_noise=0.7, + max_sep=6.0, + pixel_size=0.5, + ) + gp.initialize(X, y) + gp.solve() + y_predict = gp.predict(X) + assert np.var(y - y_predict) < 0.5 * np.var(y) + + +if __name__ == "__main__": + test_empirical_2pcf_gp() + test_empirical_2pcf_eigenvalue_clipping() + test_empirical_2pcf_extrapolation() + test_empirical_2pcf_introspection() + test_empirical_kernel_orientation() + test_empirical_2pcf_helpers() + test_empirical_2pcf_validation() diff --git a/treegp/__init__.py b/treegp/__init__.py index 14f93e7..44a1678 100644 --- a/treegp/__init__.py +++ b/treegp/__init__.py @@ -10,10 +10,12 @@ from .two_pcf import two_pcf from .log_likelihood import log_likelihood +from .empirical_2pcf import empirical_2pcf from .kernels import AnisotropicRBF from .kernels import VonKarman from .kernels import AnisotropicVonKarman +from .kernels import EmpiricalCorrelationKernel from .kernels import eval_kernel from .meanify import meanify, MeanifyStream @@ -27,9 +29,11 @@ "GPInterpolation", "two_pcf", "log_likelihood", + "empirical_2pcf", "AnisotropicRBF", "VonKarman", "AnisotropicVonKarman", + "EmpiricalCorrelationKernel", "eval_kernel", "meanify", "MeanifyStream", diff --git a/treegp/empirical_2pcf.py b/treegp/empirical_2pcf.py new file mode 100644 index 0000000..324f789 --- /dev/null +++ b/treegp/empirical_2pcf.py @@ -0,0 +1,281 @@ +""" +.. module:: empirical_2pcf +""" + +import copy +import numpy as np +import treecorr + +from scipy import fft + +from .kernels import EmpiricalCorrelationKernel + + +def _blackman_harris(r, r_max=1.0): + """Return Blackman-Harris function of variable r with + peak at 1.0 and going to zero at r_max. + + :param r: Radii where to evaluate the window. (ndarray) + :param r_max: Radius where the window reaches zero. [default: 1.] + """ + upi = np.pi * r / r_max + out = ( + 0.35875 + + 0.48829 * np.cos(upi) + + 0.14128 * np.cos(2 * upi) + + 0.01168 * np.cos(3 * upi) + ) + return np.where(np.abs(upi) <= np.pi, out, 0.0) + + +def _apod(corr): + """Return a Blackman-Harris apodization window with the same shape + as the given 2d correlation function, equal to 1 at zero lag + (pixel N//2) and going to zero at N//2 pixels from it. + + :param corr: 2d correlation function, zero lag at pixel N//2. (N, N) ndarray + """ + s = corr.shape[0] // 2 + yx = np.indices(corr.shape) + ctr = np.array(corr.shape) // 2 + yx -= ctr[:, np.newaxis, np.newaxis] + rad = np.hypot(yx[0], yx[1]) + return _blackman_harris(rad, s) + + +def _shift_and_bin(corr_func): + """Bin a 2d correlation function 2x2 in such a way that the zero + lag, located at the intersection of the 4 central pixels of the + input (treecorr TwoD convention), ends up at the center of pixel + N//2 of the output, where N is the output side length. + + :param corr_func: 2d correlation function from treecorr TwoD binning, + with even side length. (2N, 2N) ndarray + """ + if corr_func.shape[0] != corr_func.shape[1]: + raise ValueError( + "corr_func must be square. Current shape: %s" % (str(corr_func.shape)) + ) + if corr_func.shape[0] % 2 != 0: + raise ValueError( + "corr_func side length must be even. Current shape: %s" + % (str(corr_func.shape)) + ) + # Roll points astride zero lag into origin corner + s = corr_func.shape[0] // 2 + corr = np.roll( + corr_func, -(s - 1), axis=0 + ) # This places the 2 mirror-image low-f pixels at 0,1 + corr = np.sum(corr.reshape(s, 2, -1), axis=1) # Bin by 2 in y + corr = np.roll(corr, -(s - 1), axis=1) + corr = np.sum(corr.reshape(-1, s, 2), axis=2) # Bin by 2 in x + # Roll DC back to center from its current position at (0,0) + s = corr.shape[0] // 2 + corr = np.roll(corr, s, axis=0) + corr = np.roll(corr, s, axis=1) + return corr / 4.0 + + +def _threshold(p, n_sigma=3.0): + """Keep only the elements of the power spectrum p that are above + n_sigma times the noise, where the noise is estimated from the + 16-50-84 percentiles of the non-zero elements of p. Elements below + the threshold are set to zero. + + :param p: 2d power spectrum. (ndarray) + :param n_sigma: Multiple of the noise below which elements + are zeroed out. [default: 3.] + """ + nonzero = p[p != 0.0] + if len(nonzero) == 0: + return p + hml = np.percentile(nonzero, (16.0, 50.0, 84.0)) + thresh = hml[1] + n_sigma * 0.5 * (hml[2] - hml[0]) + return np.where(p > thresh, p, 0.0) + + +def _corr2power(corr): + """Get power spectrum from a (square, even side length) 2d correlation + function. Both have their origin at pixel N//2 along the first axis. + + :param corr: 2d correlation function, zero lag at pixel N//2. (N, N) ndarray + """ + s = corr.shape[0] // 2 + tmp = np.roll(corr, -s, axis=0) + tmp = np.roll(tmp, -s, axis=1) + pk = fft.rfft2(tmp) + return np.roll(pk.real, s, axis=0) + + +def _power2corr(pk): + """Get 2d correlation function from a power spectrum, origin at + pixel N//2 in the power spectrum first axis. Inverse of _corr2power. + + :param pk: 2d power spectrum from _corr2power. (N, N//2+1) ndarray + """ + s = pk.shape[0] // 2 + corr = fft.irfft2(np.roll(pk, -s, axis=0)) + corr = np.roll(corr, s, axis=0) + corr = np.roll(corr, s, axis=1) + return corr + + +class empirical_2pcf(object): + """ + Build a gaussian process kernel directly from the measured anisotropic + 2d 2-point correlation function, following Gomes et al. (2025) + (AJ 170:361, doi:10.3847/1538-3881/ae1a7b). No hyperparameters are + fitted: the measured correlation function, cleaned by apodization and + thresholding of its Fourier power spectrum, is tabulated and used + as the kernel. + + :param X: Coordinates of the field. (n_samples, 2) + :param y: Values of the field. (n_samples) + :param y_err: Error of y. (n_samples) + :param max_sep: Maximum separation in each coordinate of the + 2d correlation function grid (half width of the + grid), in the same units as X. Rounded up to an + integer number of pixels. Computed automatically + (half of the field diagonal) if not given. + [default: None] + :param pixel_size: Pixel size of the 2d correlation function grid, + in the same units as X. Computed automatically + (twice the mean separation between points) if + not given. [default: None] + :param power_threshold: Signal-to-noise threshold below which Fourier + modes of the measured correlation function are + set to zero. [default: 2.5] + :param apodize: Whether to apodize the measured correlation + function with a Blackman-Harris window before + taking its Fourier transform. [default: True] + """ + + def __init__( + self, + X, + y, + y_err, + max_sep=None, + pixel_size=None, + power_threshold=2.5, + apodize=True, + ): + self.ndim = np.shape(X)[1] + if self.ndim != 2: + raise ValueError( + "empirical-2pcf supports only 2d modeling. Current ndim: %i" + % (self.ndim) + ) + self.X = X + self.y = y + self.y_err = y_err + self.power_threshold = power_threshold + self.apodize = apodize + + size_x = np.max(X[:, 0]) - np.min(X[:, 0]) + size_y = np.max(X[:, 1]) - np.min(X[:, 1]) + rho = float(len(X[:, 0])) / (size_x * size_y) + # if max_sep is None, set max_sep to half of the size of the + # given field. + if max_sep is None: + max_sep = np.sqrt(size_x**2 + size_y**2) / 2.0 + # if pixel_size is None, set it to twice the mean separation + # between points, so treecorr bins (half a pixel) match the + # mean point separation. + if pixel_size is None: + pixel_size = 2.0 * np.sqrt(1.0 / rho) + + half_npix = int(np.ceil(max_sep / pixel_size)) + if half_npix < 2: + raise ValueError( + "max_sep must span at least 2 pixels. " + "Current max_sep: %f, pixel_size: %f" % (max_sep, pixel_size) + ) + # Final grid is (npix, npix) with zero lag at the center of + # pixel npix//2 in each axis, covering [-max_sep, max_sep]. + self.npix = 2 * half_npix + self.pixel_size = pixel_size + self.max_sep = half_npix * pixel_size + + def comp_2pcf(self, X, y, y_err): + """ + Estimate the anisotropic 2d 2-point correlation function + using TreeCorr. + + TreeCorr TwoD binning puts zero lag at the intersection of the 4 + central pixels, so the correlation function is measured at half + the requested pixel size and then binned 2x2 so that zero lag + ends up at the center of pixel npix//2. + + :param X: Coordinates of the field. (n_samples, 2) + :param y: Values of the field. (n_samples) + :param y_err: Error of y. (n_samples) + """ + if np.sum(y_err) == 0: + w = None + else: + w = 1.0 / y_err**2 + + cat = treecorr.Catalog(x=X[:, 0], y=X[:, 1], k=(y - np.mean(y)), w=w) + kk = treecorr.KKCorrelation( + max_sep=self.max_sep, + nbins=2 * self.npix, + bin_type="TwoD", + bin_slop=0, + ) + kk.process(cat) + return _shift_and_bin(kk.xi) + + def clean(self, xi): + """ + Clean the measured 2d 2-point correlation function by apodizing + it (if requested) and keeping only the Fourier modes above + power_threshold times the noise. The surviving power is + positive, so the cleaned correlation function is positive + semi-definite on its grid. + + :param xi: Measured 2d correlation function, zero lag at + pixel npix//2. (npix, npix) ndarray + """ + if self.apodize: + pk = _corr2power(xi * _apod(xi)) + else: + pk = _corr2power(xi) + pk = _threshold(pk, n_sigma=self.power_threshold) + if np.all(pk == 0.0): + raise RuntimeError( + "All Fourier modes of the measured 2-point correlation " + "function are below the power threshold. The field might " + "not have significant correlations; try lowering " + "power_threshold (current value: %f)." % (self.power_threshold) + ) + return _power2corr(pk) + + def optimizer(self, kernel): + """ + Build the gaussian process kernel from the measured 2d 2-point + correlation function. Contrary to the other optimizers, no + hyperparameters are fitted and the given kernel is ignored: the + cleaned measured correlation function is returned as a tabulated + EmpiricalCorrelationKernel. + + :param kernel: sklearn.gaussian_process kernel. (ignored) + """ + xi = self.comp_2pcf(self.X, self.y, self.y_err) + xi_clean = self.clean(xi) + + lag = (np.arange(self.npix) - self.npix // 2) * self.pixel_size + kernel_out = EmpiricalCorrelationKernel(lag, lag, xi_clean) + + # Keep diagnostics around. dx varies along the columns of the + # treecorr TwoD maps, so the flattened distances match the + # flattened correlation functions. + dx, dy = np.meshgrid(lag, lag) + self._xi = xi + self._xi_clean = xi_clean + self._2pcf = xi.flatten() + self._2pcf_fit = xi_clean.flatten() + self._2pcf_dist = np.array([dx.flatten(), dy.flatten()]).T + self._2pcf_mask = np.ones(self.npix**2, dtype=bool) + self._kernel = copy.deepcopy(kernel_out) + return kernel_out diff --git a/treegp/gp_interp.py b/treegp/gp_interp.py index c136906..5917300 100644 --- a/treegp/gp_interp.py +++ b/treegp/gp_interp.py @@ -7,11 +7,23 @@ import copy from .kernels import eval_kernel +from .kernels import EmpiricalCorrelationKernel from sklearn.neighbors import KNeighborsRegressor from scipy.linalg import cholesky, cho_solve +def _kernel_contains(kernel, kernel_class): + """Return whether the given (possibly composite) sklearn kernel is, + or contains, an instance of kernel_class.""" + if isinstance(kernel, kernel_class): + return True + return any( + isinstance(param, kernel_class) + for param in kernel.get_params(deep=True).values() + ) + + class GPInterpolation(object): """ An interpolator that uses 2-point correlation function informations @@ -22,12 +34,30 @@ class GPInterpolation(object): sklearn.gaussian_process.kernels.Kernel object. The reprs of sklearn.gaussian_process.kernels will work, as well as the repr of a custom treegp VonKarman object. [default: 'RBF(1)'] - :param optimizer: Indicates which techniques to use for optimizing the kernel. Three options + :param optimizer: Indicates which techniques to use for optimizing the kernel. Five options are available. "none" does not optimize hyperparameters and used the one given in the kernel. "two-pcf" optimize the kernel on the 1d 2-point correlation function estimate by treecorr. "anisotropic" optimize the kernel on the 2d 2-point correlation function estimate by treecorr. + "two-pcf" and "anisotropic" implement Leget et al. 2021 + (A&A 650, A81, arXiv:2103.09881). "log-likelihood" used the classical maximum likelihood method. + "empirical-2pcf" implements Gomes et al. 2025 (AJ 170:361, + doi:10.3847/1538-3881/ae1a7b): the measured 2d 2-point correlation + function, cleaned by apodization and thresholding of its Fourier + power spectrum, is used directly as a tabulated kernel; no + hyperparameters are fitted: the kernel argument must be left to + its default (an error is raised otherwise), and min_sep and + nbins are ignored (the grid is controlled by max_sep and + pixel_size). Conversely, an EmpiricalCorrelationKernel has no + 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, + 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 want this to be True, but if your parameters have an a priori mean of 0, @@ -49,6 +79,16 @@ class GPInterpolation(object): :param average_fits: A fits file that have the spatial average functions of the interpolated parameter build in it. Build using meanify output across different exposures. See meanify documentation. [default: None] + :param pixel_size: Pixel size of the 2d correlation function grid used by the + "empirical-2pcf" optimizer, in the same units as the coordinates + of the field. Computed automatically (twice the mean separation + between points) if it is not given. [default: None] + :param power_threshold: Signal-to-noise threshold below which Fourier modes of the + measured 2-point correlation function are set to zero. Used only + by the "empirical-2pcf" optimizer. [default: 2.5] + :param apodize: Whether to apodize the measured 2-point correlation function with + a Blackman-Harris window before taking its Fourier transform. + Used only by the "empirical-2pcf" optimizer. [default: True] """ def __init__( @@ -64,6 +104,9 @@ def __init__( nbins=20, min_sep=None, max_sep=None, + pixel_size=None, + power_threshold=2.5, + apodize=True, ): self.normalize = normalize self.optimizer = optimizer @@ -72,6 +115,9 @@ def __init__( self.nbins = nbins self.min_sep = min_sep self.max_sep = max_sep + self.pixel_size = pixel_size + self.power_threshold = power_threshold + self.apodize = apodize if self.optimizer == "anisotropic": self.robust_fit = True @@ -88,10 +134,33 @@ def __init__( "kernel should be a string a list or a numpy.ndarray of string" ) - if self.optimizer not in ["anisotropic", "two-pcf", "log-likelihood", "none"]: + if self.optimizer not in [ + "anisotropic", + "two-pcf", + "empirical-2pcf", + "log-likelihood", + "none", + ]: raise ValueError( - "Only anisotropic, two-pcf, log-likelihood and none are supported for optimizer. Current value: %s" - % (self.optimizer) + "Only anisotropic, two-pcf, empirical-2pcf, log-likelihood and none " + "are supported for optimizer. Current value: %s" % (self.optimizer) + ) + + if self.optimizer == "empirical-2pcf" and kernel != "RBF(1)": + raise ValueError( + "The empirical-2pcf optimizer builds its own " + "EmpiricalCorrelationKernel from the measured 2-point " + "correlation function, so the kernel argument is ignored and " + "should be left to its default value. Current value: %s" % (kernel) + ) + if self.optimizer != "none" and _kernel_contains( + self.kernel_template, EmpiricalCorrelationKernel + ): + raise ValueError( + "EmpiricalCorrelationKernel has no hyperparameters to fit, so " + "it can only be used with optimizer='empirical-2pcf' (where it " + "is built automatically) or optimizer='none'. " + "Current optimizer: %s" % (self.optimizer) ) if average_fits is not None: @@ -136,6 +205,20 @@ def _fit(self, kernel, X, y, y_err): p0=self.p0_robust_fit, ) kernel = self._optimizer.optimizer(kernel) + # Kernel built directly from the measured 2d 2-point correlation + # function (Gomes et al. 2025). No hyperparameters are fitted and + # the given kernel is ignored. + if self.optimizer == "empirical-2pcf": + self._optimizer = treegp.empirical_2pcf( + X, + y, + y_err, + max_sep=self.max_sep, + pixel_size=self.pixel_size, + power_threshold=self.power_threshold, + apodize=self.apodize, + ) + kernel = self._optimizer.optimizer(kernel) # Hyperparameters estimation using maximum likelihood fit. if self.optimizer == "log-likelihood": self._optimizer = treegp.log_likelihood(X, y, y_err) @@ -179,7 +262,14 @@ 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 - self._fact = cholesky(K, lower=True) + 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)) + ) 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: @@ -258,6 +348,11 @@ def return_2pcf(self): """ Return 2-point correlation function and its variance using Bootstrap. """ + if self.optimizer == "empirical-2pcf": + raise NotImplementedError( + "return_2pcf is not available for the empirical-2pcf optimizer. " + "Use return_empirical_2pcf instead." + ) anisotropic = self.optimizer == "anisotropic" pcf = treegp.two_pcf( self._X, @@ -271,6 +366,31 @@ def return_2pcf(self): xi, xi_weight, distance, coord, mask = pcf.return_2pcf() return xi, xi_weight, distance, coord, mask + def return_empirical_2pcf(self): + """ + Return the measured and cleaned 2d 2-point correlation functions + used as kernel by the empirical-2pcf optimizer, the lag coordinates + of the grid, and the pixel size of the grid. + + Returns xi, xi_clean, distance, pixel_size where xi and xi_clean + are (npix, npix) arrays with zero lag at pixel npix//2, and + distance is a (npix*npix, 2) array of the (dx, dy) lags matching + the flattened correlation functions. + """ + if self.optimizer != "empirical-2pcf": + raise NotImplementedError( + "return_empirical_2pcf is only available for the " + "empirical-2pcf optimizer. Current optimizer: %s" % (self.optimizer) + ) + if not hasattr(self, "_optimizer"): + raise RuntimeError("solve() must be called before return_empirical_2pcf.") + return ( + self._optimizer._xi, + self._optimizer._xi_clean, + self._optimizer._2pcf_dist, + self._optimizer.pixel_size, + ) + def return_log_likelihood(self, theta=None): """ Return of log likehood of gaussian process diff --git a/treegp/kernels.py b/treegp/kernels.py index 888b5c1..26f01ff 100644 --- a/treegp/kernels.py +++ b/treegp/kernels.py @@ -5,6 +5,7 @@ import numpy as np from scipy.spatial.distance import pdist, cdist, squareform from scipy import special +from scipy.interpolate import RegularGridInterpolator from sklearn.gaussian_process.kernels import ( StationaryKernelMixin, NormalizedKernelMixin, @@ -418,3 +419,98 @@ def __repr__(self): @property def bounds(self): return self._bounds + + +class EmpiricalCorrelationKernel(StationaryKernelMixin, Kernel): + """A tabulated stationary anisotropic kernel built from a measured + 2d 2-point correlation function, following Gomes et al. (2025) + (AJ 170:361, doi:10.3847/1538-3881/ae1a7b). + + The kernel has no free hyperparameters (theta is empty): it evaluates + K(X, Y)[i, j] = xi(X_i - Y_j) by bilinear interpolation of the given + correlation function grid, and is zero beyond the grid (compact + support). + + A tabulated correlation function is not guaranteed to be positive + semi-definite between arbitrary points, so by default the covariance + matrices built when Y is None are projected onto the closest positive + semi-definite matrix by clipping their negative eigenvalues to zero. + This is equivalent to the singular value clipping used by + Gomes et al. (2025). + + Input is expected to be 2-dimensional, i.e. X.shape = (n_samples, 2). + + :param x_grid: Lag coordinates of the grid columns, zero lag + at index len(x_grid)//2. (nx,) ndarray + :param y_grid: Lag coordinates of the grid rows, zero lag + at index len(y_grid)//2. (ny,) ndarray + :param xi_grid: 2d correlation function in treecorr TwoD layout, + i.e. indexed [iy, ix]. (ny, nx) ndarray + :param clip_eigenvalues: Whether to clip the negative eigenvalues of + the covariance matrices built when Y is None. + [default: True] + """ + + 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) + self.xi_grid = np.asarray(xi_grid) + self.clip_eigenvalues = clip_eigenvalues + if self.xi_grid.shape != (len(self.y_grid), len(self.x_grid)): + raise ValueError( + "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 + # (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, + ) + + @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]) + + def __call__(self, X, Y=None, eval_gradient=False): + if eval_gradient: + raise ValueError("Gradient can not be evaluated.") + X = np.atleast_2d(X) + if np.shape(X)[1] != 2: + raise ValueError( + "EmpiricalCorrelationKernel supports only 2d coordinates. " + "Current ndim: %i" % (np.shape(X)[1]) + ) + if Y is None: + d = X[:, np.newaxis, :] - X[np.newaxis, :, :] + K = self._interp(d) + # 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 + # exactly symmetric covariance matrix. + K = 0.5 * (K + K.T) + if self.clip_eigenvalues: + eigenvalues, eigenvectors = np.linalg.eigh(K) + if np.any(eigenvalues < 0.0): + K = (eigenvectors * np.clip(eigenvalues, 0.0, None)).dot( + eigenvectors.T + ) + K = 0.5 * (K + K.T) + else: + Y = np.atleast_2d(Y) + d = X[:, np.newaxis, :] - Y[np.newaxis, :, :] + K = self._interp(d) + return K + + def diag(self, X): + return np.full(len(X), self.xi0) + + def __repr__(self): + return "{0}(npix={1!r}, xi0={2:.4g})".format( + self.__class__.__name__, self.xi_grid.shape, self.xi0 + ) From 297ad0c72ab472e355ba9665e52ee6abca6760d4 Mon Sep 17 00:00:00 2001 From: PFLeget Date: Fri, 21 Aug 2026 16:29:09 -0400 Subject: [PATCH 2/4] expose radius of blackman-harris filter and add Hans filtering --- tests/test_empirical_2pcf.py | 61 ++++++++++++++++++++++++++++ treegp/empirical_2pcf.py | 77 +++++++++++++++++++++++++++++++----- treegp/gp_interp.py | 22 +++++++++-- 3 files changed, 147 insertions(+), 13 deletions(-) diff --git a/tests/test_empirical_2pcf.py b/tests/test_empirical_2pcf.py index a228826..900adfe 100644 --- a/tests/test_empirical_2pcf.py +++ b/tests/test_empirical_2pcf.py @@ -223,6 +223,29 @@ def test_empirical_2pcf_helpers(): np.testing.assert_allclose(window[npix // 2, npix // 2], 1.0, atol=1e-10) np.testing.assert_allclose(window[npix // 2, 0], 0.0, atol=1e-4) + # The Hann window is exactly zero at its edge, 1/2 at half radius, + # and gentler (larger) than Blackman-Harris in the interior. (Right + # at the edge Blackman-Harris is larger: it floors at 6e-5 while + # Hann goes to zero quadratically.) + window_hann = _apod(xi, window="hann") + np.testing.assert_allclose(window_hann[npix // 2, npix // 2], 1.0, atol=1e-10) + np.testing.assert_allclose(window_hann[npix // 2, 0], 0.0, atol=1e-10) + np.testing.assert_allclose( + window_hann[npix // 2, npix // 2 + npix // 4], 0.5, atol=1e-10 + ) + interior = _apod(xi, r_max=0.9 * (npix // 2)) > 0.0 + assert np.all(window_hann[interior] >= window[interior]) + + # A smaller apodization radius reaches zero earlier. + window_small = _apod(xi, r_max=npix // 4) + np.testing.assert_allclose(window_small[npix // 2, npix // 2], 1.0, atol=1e-10) + np.testing.assert_allclose( + window_small[npix // 2, npix // 2 + npix // 4 + 1 :], 0.0, atol=1e-10 + ) + # A larger radius leaves the window non-zero at the grid edge. + window_large = _apod(xi, r_max=npix, window="hann") + assert window_large[npix // 2, 0] > 0.1 + @timer def test_empirical_2pcf_validation(): @@ -246,6 +269,26 @@ def test_empirical_2pcf_validation(): # Unknown optimizer is rejected. np.testing.assert_raises(ValueError, treegp.GPInterpolation, optimizer="gomes25") + # Unknown apodization window and non-positive apodization radius + # are rejected. + X2d = np.random.uniform(-10, 10, 200).reshape((100, 2)) + np.testing.assert_raises( + ValueError, + treegp.empirical_2pcf, + X2d, + y, + np.zeros_like(y), + apod_window="tukey", + ) + np.testing.assert_raises( + ValueError, + treegp.empirical_2pcf, + X2d, + y, + np.zeros_like(y), + apod_radius=-1.0, + ) + # The empirical-2pcf optimizer builds its own kernel: passing one # is rejected. np.testing.assert_raises( @@ -292,6 +335,24 @@ def test_empirical_2pcf_validation(): y_predict = gp.predict(X) assert np.var(y - y_predict) < 0.5 * np.var(y) + # The apodization window and radius are threaded through + # GPInterpolation, and the interpolation still works with them. + gp = treegp.GPInterpolation( + optimizer="empirical-2pcf", + normalize=True, + white_noise=0.7, + max_sep=6.0, + pixel_size=0.5, + apod_window="hann", + apod_radius=4.0, + ) + gp.initialize(X, y) + gp.solve() + assert gp._optimizer.apod_window == "hann" + assert gp._optimizer.apod_radius == 4.0 + y_predict = gp.predict(X) + assert np.var(y - y_predict) < 0.5 * np.var(y) + if __name__ == "__main__": test_empirical_2pcf_gp() diff --git a/treegp/empirical_2pcf.py b/treegp/empirical_2pcf.py index 324f789..48aa41c 100644 --- a/treegp/empirical_2pcf.py +++ b/treegp/empirical_2pcf.py @@ -28,19 +28,47 @@ def _blackman_harris(r, r_max=1.0): return np.where(np.abs(upi) <= np.pi, out, 0.0) -def _apod(corr): - """Return a Blackman-Harris apodization window with the same shape - as the given 2d correlation function, equal to 1 at zero lag - (pixel N//2) and going to zero at N//2 pixels from it. +def _hann(r, r_max=1.0): + """Return Hann function of variable r with peak at 1.0 and going + to zero at r_max. Gentler taper (less bias on the correlation + function) than Blackman-Harris, at the price of more spectral + leakage. - :param corr: 2d correlation function, zero lag at pixel N//2. (N, N) ndarray + :param r: Radii where to evaluate the window. (ndarray) + :param r_max: Radius where the window reaches zero. [default: 1.] """ - s = corr.shape[0] // 2 + out = 0.5 * (1.0 + np.cos(np.pi * r / r_max)) + return np.where(np.abs(r) <= r_max, out, 0.0) + + +APOD_WINDOWS = { + "blackman-harris": _blackman_harris, + "hann": _hann, +} + + +def _apod(corr, r_max=None, window="blackman-harris"): + """Return an apodization window with the same shape as the given + 2d correlation function, equal to 1 at zero lag (pixel N//2) and + going to zero at r_max pixels from it. + + :param corr: 2d correlation function, zero lag at pixel N//2. + (N, N) ndarray + :param r_max: Radius (in pixels) where the window reaches zero. + N//2 (the grid edge) if not given. A radius beyond + the grid edge gives a gentler taper but leaves the + window non-zero at the edge, reintroducing some + spectral leakage. [default: None] + :param window: Name of the window function, one of "blackman-harris" + or "hann". [default: "blackman-harris"] + """ + if r_max is None: + r_max = corr.shape[0] // 2 yx = np.indices(corr.shape) ctr = np.array(corr.shape) // 2 yx -= ctr[:, np.newaxis, np.newaxis] rad = np.hypot(yx[0], yx[1]) - return _blackman_harris(rad, s) + return APOD_WINDOWS[window](rad, r_max) def _shift_and_bin(corr_func): @@ -146,8 +174,20 @@ class empirical_2pcf(object): modes of the measured correlation function are set to zero. [default: 2.5] :param apodize: Whether to apodize the measured correlation - function with a Blackman-Harris window before - taking its Fourier transform. [default: True] + function before taking its Fourier transform. + [default: True] + :param apod_window: Name of the apodization window, one of + "blackman-harris" or "hann". Hann is a gentler + taper (less bias on the correlation function) + at the price of more spectral leakage. + [default: "blackman-harris"] + :param apod_radius: Radius where the apodization window reaches + zero, in the same units as X. max_sep (the + grid edge) if not given. A radius beyond + max_sep gives a gentler taper but leaves the + window non-zero at the grid edge, + reintroducing some spectral leakage. + [default: None] """ def __init__( @@ -159,6 +199,8 @@ def __init__( pixel_size=None, power_threshold=2.5, apodize=True, + apod_window="blackman-harris", + apod_radius=None, ): self.ndim = np.shape(X)[1] if self.ndim != 2: @@ -166,11 +208,22 @@ def __init__( "empirical-2pcf supports only 2d modeling. Current ndim: %i" % (self.ndim) ) + if apod_window not in APOD_WINDOWS: + raise ValueError( + "Only %s are supported for apod_window. Current value: %s" + % (sorted(APOD_WINDOWS), apod_window) + ) + if apod_radius is not None and apod_radius <= 0: + raise ValueError( + "apod_radius must be positive. Current value: %s" % (apod_radius) + ) self.X = X self.y = y self.y_err = y_err self.power_threshold = power_threshold self.apodize = apodize + self.apod_window = apod_window + self.apod_radius = apod_radius size_x = np.max(X[:, 0]) - np.min(X[:, 0]) size_y = np.max(X[:, 1]) - np.min(X[:, 1]) @@ -238,7 +291,11 @@ def clean(self, xi): pixel npix//2. (npix, npix) ndarray """ if self.apodize: - pk = _corr2power(xi * _apod(xi)) + if self.apod_radius is None: + r_max = None + else: + r_max = self.apod_radius / self.pixel_size + pk = _corr2power(xi * _apod(xi, r_max=r_max, window=self.apod_window)) else: pk = _corr2power(xi) pk = _threshold(pk, n_sigma=self.power_threshold) diff --git a/treegp/gp_interp.py b/treegp/gp_interp.py index 5917300..ed1d72d 100644 --- a/treegp/gp_interp.py +++ b/treegp/gp_interp.py @@ -86,9 +86,19 @@ class GPInterpolation(object): :param power_threshold: Signal-to-noise threshold below which Fourier modes of the measured 2-point correlation function are set to zero. Used only by the "empirical-2pcf" optimizer. [default: 2.5] - :param apodize: Whether to apodize the measured 2-point correlation function with - a Blackman-Harris window before taking its Fourier transform. - Used only by the "empirical-2pcf" optimizer. [default: True] + :param apodize: Whether to apodize the measured 2-point correlation function + before taking its Fourier transform. Used only by the + "empirical-2pcf" optimizer. [default: True] + :param apod_window: Name of the apodization window, one of "blackman-harris" or + "hann". Hann is a gentler taper (less bias on the correlation + function) at the price of more spectral leakage. Used only by + the "empirical-2pcf" optimizer. [default: "blackman-harris"] + :param apod_radius: Radius where the apodization window reaches zero, in the same + units as the coordinates of the field. max_sep (the grid edge) + if it is not given. A radius beyond max_sep gives a gentler + taper but leaves the window non-zero at the grid edge, + reintroducing some spectral leakage. Used only by the + "empirical-2pcf" optimizer. [default: None] """ def __init__( @@ -107,6 +117,8 @@ def __init__( pixel_size=None, power_threshold=2.5, apodize=True, + apod_window="blackman-harris", + apod_radius=None, ): self.normalize = normalize self.optimizer = optimizer @@ -118,6 +130,8 @@ def __init__( self.pixel_size = pixel_size self.power_threshold = power_threshold self.apodize = apodize + self.apod_window = apod_window + self.apod_radius = apod_radius if self.optimizer == "anisotropic": self.robust_fit = True @@ -217,6 +231,8 @@ def _fit(self, kernel, X, y, y_err): pixel_size=self.pixel_size, power_threshold=self.power_threshold, apodize=self.apodize, + apod_window=self.apod_window, + apod_radius=self.apod_radius, ) kernel = self._optimizer.optimizer(kernel) # Hyperparameters estimation using maximum likelihood fit. From 5f3383fc235d19cc72892bf5bbedb195ab764a63 Mon Sep 17 00:00:00 2001 From: PFLeget Date: Fri, 21 Aug 2026 17:34:36 -0400 Subject: [PATCH 3/4] add adaptative anisotropic window filtering. --- tests/test_empirical_2pcf.py | 106 ++++++++++++++++++ treegp/empirical_2pcf.py | 206 ++++++++++++++++++++++++++++++++--- treegp/gp_interp.py | 23 ++++ 3 files changed, 322 insertions(+), 13 deletions(-) diff --git a/tests/test_empirical_2pcf.py b/tests/test_empirical_2pcf.py index 900adfe..db456e8 100644 --- a/tests/test_empirical_2pcf.py +++ b/tests/test_empirical_2pcf.py @@ -12,9 +12,22 @@ _corr2power, _power2corr, _apod, + _adaptive_moments, ) +def make_elliptical_gaussian(npix, size, g1, g2): + """Elliptical gaussian map in treecorr TwoD layout ([iy, ix]), + centered on pixel npix//2, using the Leget et al. 2021 shear + parametrization (size = major axis, in pixels).""" + L = get_correlation_length_matrix(size, g1, g2) + invL = np.linalg.inv(L) + lag = np.arange(npix) - npix // 2 + dx, dy = np.meshgrid(lag, lag) + arg = invL[0, 0] * dx**2 + 2.0 * invL[0, 1] * dx * dy + invL[1, 1] * dy**2 + return np.exp(-0.5 * arg) + + def make_gp(npoints=2000, noise=0.3, white_noise=0.0, seed=42): """Generate a 2d GRF with a known anisotropic kernel and return an initialized and solved GPInterpolation using empirical-2pcf.""" @@ -236,6 +249,14 @@ def test_empirical_2pcf_helpers(): interior = _apod(xi, r_max=0.9 * (npix // 2)) > 0.0 assert np.all(window_hann[interior] >= window[interior]) + # An elliptical window with g1 > 0 (major axis along x) is wider + # along x than along y, and reduces exactly to the isotropic + # window at g1 = g2 = 0. + window_ell = _apod(xi, window="hann", g1=0.4, g2=0.0) + d = npix // 4 + assert window_ell[npix // 2, npix // 2 + d] > window_ell[npix // 2 + d, npix // 2] + np.testing.assert_allclose(_apod(xi, g1=0.0, g2=0.0), _apod(xi), atol=1e-14) + # A smaller apodization radius reaches zero earlier. window_small = _apod(xi, r_max=npix // 4) np.testing.assert_allclose(window_small[npix // 2, npix // 2], 1.0, atol=1e-10) @@ -247,6 +268,70 @@ def test_empirical_2pcf_helpers(): assert window_large[npix // 2, 0] > 0.1 +@timer +def test_adaptive_moments(): + # Adaptive moments recover the shear of an analytic elliptical + # gaussian in the get_correlation_length_matrix convention. + for g1_true, g2_true in [(0.0, 0.0), (0.3, 0.0), (0.0, -0.2), (0.2, 0.2)]: + f = make_elliptical_gaussian(64, 6.0, g1_true, g2_true) + g1, g2 = _adaptive_moments(f) + np.testing.assert_allclose([g1, g2], [g1_true, g2_true], atol=1e-2) + + # An all-zero (or all-negative, clipped to zero) map returns no + # anisotropy. + assert _adaptive_moments(np.zeros((32, 32))) == (0.0, 0.0) + assert _adaptive_moments(-np.ones((32, 32))) == (0.0, 0.0) + + +@timer +def test_empirical_2pcf_anisotropic_apod(): + L = get_correlation_length_matrix(2.0, 0.2, 0.2) + invLam = np.linalg.inv(L) + kernel = 2.0**2 * treegp.AnisotropicRBF(invLam=invLam) + X, y, y_err = make_2d_grf(kernel, noise=0.3, seed=42, npoints=2000) + + def run(**kwargs): + gp = treegp.GPInterpolation( + optimizer="empirical-2pcf", + normalize=True, + max_sep=6.0, + pixel_size=0.5, + **kwargs, + ) + gp.initialize(X, y, y_err=y_err) + gp.solve() + return gp + + # Auto mode: the field was generated with positive (g1, g2), so the + # measured anisotropy of its 2-pcf must have positive components, + # and the applied shear is the measured one times apod_g_scale. + gp = run(apod_anisotropy="auto", apod_g_scale=0.5) + g1_m, g2_m = gp._optimizer._apod_g_measured + assert g1_m > 0.0 + assert g2_m > 0.0 + np.testing.assert_allclose( + gp._optimizer._apod_g_applied, [0.5 * g1_m, 0.5 * g2_m], atol=1e-12 + ) + assert gp._optimizer._xi_clean_pass1 is not None + y_predict = gp.predict(X) + assert np.var(y - y_predict) < 0.5 * np.var(y) + + # apod_g_scale = 0 in auto mode applies an isotropic window, so the + # final map is identical to the first pass. + gp = run(apod_anisotropy="auto", apod_g_scale=0.0) + np.testing.assert_allclose( + gp._optimizer._xi_clean, gp._optimizer._xi_clean_pass1, atol=1e-14 + ) + + # Manual mode: the given shear is applied directly, nothing is + # measured. + gp = run(apod_anisotropy=(0.3, 0.1)) + assert gp._optimizer._apod_g_measured is None + np.testing.assert_allclose(gp._optimizer._apod_g_applied, [0.3, 0.1], atol=1e-12) + y_predict = gp.predict(X) + assert np.var(y - y_predict) < 0.5 * np.var(y) + + @timer def test_empirical_2pcf_validation(): # Only 2d fields are supported. @@ -289,6 +374,25 @@ def test_empirical_2pcf_validation(): apod_radius=-1.0, ) + # Invalid apod_anisotropy and apod_g_scale values are rejected. + for bad in ["hsm", (0.1, 0.2, 0.3), (0.8, 0.7)]: + np.testing.assert_raises( + ValueError, + treegp.empirical_2pcf, + X2d, + y, + np.zeros_like(y), + apod_anisotropy=bad, + ) + np.testing.assert_raises( + ValueError, + treegp.empirical_2pcf, + X2d, + y, + np.zeros_like(y), + apod_g_scale=-0.5, + ) + # The empirical-2pcf optimizer builds its own kernel: passing one # is rejected. np.testing.assert_raises( @@ -361,4 +465,6 @@ def test_empirical_2pcf_validation(): test_empirical_2pcf_introspection() test_empirical_kernel_orientation() test_empirical_2pcf_helpers() + test_adaptive_moments() + test_empirical_2pcf_anisotropic_apod() test_empirical_2pcf_validation() diff --git a/treegp/empirical_2pcf.py b/treegp/empirical_2pcf.py index 48aa41c..84097bc 100644 --- a/treegp/empirical_2pcf.py +++ b/treegp/empirical_2pcf.py @@ -3,12 +3,14 @@ """ import copy +import warnings import numpy as np import treecorr from scipy import fft from .kernels import EmpiricalCorrelationKernel +from .two_pcf import get_correlation_length_matrix def _blackman_harris(r, r_max=1.0): @@ -47,11 +49,18 @@ def _hann(r, r_max=1.0): } -def _apod(corr, r_max=None, window="blackman-harris"): +def _apod(corr, r_max=None, window="blackman-harris", g1=0.0, g2=0.0): """Return an apodization window with the same shape as the given 2d correlation function, equal to 1 at zero lag (pixel N//2) and going to zero at r_max pixels from it. + The window can be made elliptical using the (g1, g2) shear + parametrization of Leget et al. (2021) (same convention as + get_correlation_length_matrix): the window reaches zero at r_max + along the major axis, whose direction is phi = 0.5 arctan2(g2, g1) + from the x (column) axis, and at r_max * q along the minor axis, + with q = (1 - g) / (1 + g). At g1 = g2 = 0 the window is isotropic. + :param corr: 2d correlation function, zero lag at pixel N//2. (N, N) ndarray :param r_max: Radius (in pixels) where the window reaches zero. @@ -61,14 +70,96 @@ def _apod(corr, r_max=None, window="blackman-harris"): spectral leakage. [default: None] :param window: Name of the window function, one of "blackman-harris" or "hann". [default: "blackman-harris"] + :param g1, g2: Shear applied to the isotropic window. + [default: 0., 0.] """ if r_max is None: r_max = corr.shape[0] // 2 yx = np.indices(corr.shape) ctr = np.array(corr.shape) // 2 yx -= ctr[:, np.newaxis, np.newaxis] - rad = np.hypot(yx[0], yx[1]) - return APOD_WINDOWS[window](rad, r_max) + dy = yx[0] + dx = yx[1] + # Dimensionless elliptical radius, equal to 1 on the ellipse with + # semi-major axis r_max. Reduces to rad / r_max at g1 = g2 = 0. + L = get_correlation_length_matrix(r_max, g1, g2) + invL = np.linalg.inv(L) + r_ell = np.sqrt( + invL[0, 0] * dx**2 + 2.0 * invL[0, 1] * dx * dy + invL[1, 1] * dy**2 + ) + return APOD_WINDOWS[window](r_ell, 1.0) + + +def _adaptive_moments(xi_map, n_iter=30, tol=1e-6): + """Measure the anisotropy of a 2d correlation function map using + HSM-like adaptive weighted second moments: the weight function is an + elliptical gaussian that is iterated until it matches the measured + moments (Hirata & Seljak 2003). Negative values of the map are + clipped to zero, and the map is assumed to be centered on pixel + N//2 (no centroid iteration). + + Returns the measured anisotropy as a (g1, g2) shear, in the same + convention as get_correlation_length_matrix. + + :param xi_map: 2d correlation function, zero lag at pixel N//2. + (N, N) ndarray + :param n_iter: Maximum number of iterations. [default: 30] + :param tol: Relative tolerance on the moment matrix for + convergence. [default: 1e-6] + """ + f = np.clip(xi_map, 0.0, None) + if np.all(f == 0.0): + return 0.0, 0.0 + + yx = np.indices(f.shape) + ctr = np.array(f.shape) // 2 + yx = yx - ctr[:, np.newaxis, np.newaxis] + dy = yx[0].astype(float) + dx = yx[1].astype(float) + + # Start from an isotropic gaussian weight. + sigma = f.shape[0] / 8.0 + M = np.array([[sigma**2, 0.0], [0.0, sigma**2]]) + converged = False + for _ in range(n_iter): + invM = np.linalg.inv(M) + arg = invM[0, 0] * dx**2 + 2.0 * invM[0, 1] * dx * dy + invM[1, 1] * dy**2 + w = np.exp(-0.5 * arg) + norm = np.sum(w * f) + if norm <= 0.0: + return 0.0, 0.0 + mxx = np.sum(w * f * dx * dx) / norm + myy = np.sum(w * f * dy * dy) / norm + mxy = np.sum(w * f * dx * dy) / norm + # The factor 2 makes the iteration converge to the true + # covariance for a gaussian map: the weighted moments measure + # (C^-1 + M^-1)^-1, whose fixed point after doubling is M = C. + M_new = 2.0 * np.array([[mxx, mxy], [mxy, myy]]) + if np.max(np.abs(M_new - M)) < tol * np.trace(M_new): + M = M_new + converged = True + break + M = M_new + if not converged: + warnings.warn( + "Adaptive moments did not converge after %i iterations; " + "using the last iterate." % (n_iter) + ) + + # Moments give the distortion chi; convert it to the shear g used + # by get_correlation_length_matrix. + trace = M[0, 0] + M[1, 1] + chi1 = (M[0, 0] - M[1, 1]) / trace + chi2 = 2.0 * M[0, 1] / trace + chi = np.hypot(chi1, chi2) + if chi == 0.0: + return 0.0, 0.0 + if chi >= 1.0: + # Degenerate (essentially 1d) map; cap just below 1. + chi = 1.0 - 1e-12 + q = np.sqrt((1.0 - chi) / (1.0 + chi)) + g = (1.0 - q) / (1.0 + q) + return g * chi1 / chi, g * chi2 / chi def _shift_and_bin(corr_func): @@ -188,6 +279,28 @@ class empirical_2pcf(object): window non-zero at the grid edge, reintroducing some spectral leakage. [default: None] + :param apod_anisotropy: Anisotropy of the apodization window, using + the (g1, g2) shear parametrization of Leget + et al. (2021). None gives an isotropic + window. A (g1, g2) tuple applies the given + shear: the window reaches zero at + apod_radius along the major axis (direction + 0.5 arctan2(g2, g1) from the x axis) and at + apod_radius * q along the minor axis, with + q = (1 - g) / (1 + g). "auto" measures + (g1, g2) on the correlation function itself: + a first cleaning pass is done with the + isotropic window, the anisotropy of its + output is measured with adaptive weighted + second moments, and the raw correlation + function is re-cleaned with the matched + elliptical window. Ignored if apodize is + False. [default: None] + :param apod_g_scale: Factor multiplying the measured (g1, g2) + before building the elliptical window in + "auto" mode, to soften (< 1) or exaggerate + (> 1) the anisotropy of the taper. + [default: 1.] """ def __init__( @@ -201,6 +314,8 @@ def __init__( apodize=True, apod_window="blackman-harris", apod_radius=None, + apod_anisotropy=None, + apod_g_scale=1.0, ): self.ndim = np.shape(X)[1] if self.ndim != 2: @@ -217,6 +332,25 @@ def __init__( raise ValueError( "apod_radius must be positive. Current value: %s" % (apod_radius) ) + if apod_anisotropy is not None and not ( + isinstance(apod_anisotropy, str) and apod_anisotropy == "auto" + ): + apod_anisotropy = np.asarray(apod_anisotropy, dtype=float) + if apod_anisotropy.shape != (2,): + raise ValueError( + "apod_anisotropy must be None, 'auto', or a (g1, g2) pair. " + "Current value: %s" % (apod_anisotropy) + ) + if np.hypot(apod_anisotropy[0], apod_anisotropy[1]) >= 1.0: + raise ValueError( + "The norm of the apod_anisotropy (g1, g2) shear must be " + "lower than one. Current value: %s" % (apod_anisotropy) + ) + if not np.isfinite(apod_g_scale) or apod_g_scale < 0: + raise ValueError( + "apod_g_scale must be finite and non-negative. " + "Current value: %s" % (apod_g_scale) + ) self.X = X self.y = y self.y_err = y_err @@ -224,6 +358,8 @@ def __init__( self.apodize = apodize self.apod_window = apod_window self.apod_radius = apod_radius + self.apod_anisotropy = apod_anisotropy + self.apod_g_scale = apod_g_scale size_x = np.max(X[:, 0]) - np.min(X[:, 0]) size_y = np.max(X[:, 1]) - np.min(X[:, 1]) @@ -279,23 +415,23 @@ def comp_2pcf(self, X, y, y_err): kk.process(cat) return _shift_and_bin(kk.xi) - def clean(self, xi): - """ - Clean the measured 2d 2-point correlation function by apodizing - it (if requested) and keeping only the Fourier modes above - power_threshold times the noise. The surviving power is - positive, so the cleaned correlation function is positive - semi-definite on its grid. + def _clean_pass(self, xi, g1=0.0, g2=0.0): + """Single cleaning pass: apodize (if requested) with the given + window shear, threshold the Fourier power spectrum, and + transform back. - :param xi: Measured 2d correlation function, zero lag at - pixel npix//2. (npix, npix) ndarray + :param xi: Measured 2d correlation function, zero lag at + pixel npix//2. (npix, npix) ndarray + :param g1, g2: Shear applied to the apodization window. + [default: 0., 0.] """ if self.apodize: if self.apod_radius is None: r_max = None else: r_max = self.apod_radius / self.pixel_size - pk = _corr2power(xi * _apod(xi, r_max=r_max, window=self.apod_window)) + window = _apod(xi, r_max=r_max, window=self.apod_window, g1=g1, g2=g2) + pk = _corr2power(xi * window) else: pk = _corr2power(xi) pk = _threshold(pk, n_sigma=self.power_threshold) @@ -308,6 +444,50 @@ def clean(self, xi): ) return _power2corr(pk) + def clean(self, xi): + """ + Clean the measured 2d 2-point correlation function by apodizing + it (if requested) and keeping only the Fourier modes above + power_threshold times the noise. The surviving power is + positive, so the cleaned correlation function is positive + semi-definite on its grid. + + With apod_anisotropy="auto", a first pass is cleaned with the + isotropic window, the anisotropy of its output is measured with + adaptive weighted second moments, and the raw correlation + function is re-cleaned with the matched elliptical window + (scaled by apod_g_scale). The measured and applied shears are + stored in _apod_g_measured and _apod_g_applied, and the first + pass in _xi_clean_pass1. + + :param xi: Measured 2d correlation function, zero lag at + pixel npix//2. (npix, npix) ndarray + """ + self._apod_g_measured = None + self._xi_clean_pass1 = None + g1, g2 = 0.0, 0.0 + if self.apodize and self.apod_anisotropy is not None: + if isinstance(self.apod_anisotropy, str): + # "auto": isotropic pass, measure, elliptical re-clean. + xi_pass1 = self._clean_pass(xi) + g1_m, g2_m = _adaptive_moments(xi_pass1) + self._xi_clean_pass1 = xi_pass1 + self._apod_g_measured = (g1_m, g2_m) + g1 = self.apod_g_scale * g1_m + g2 = self.apod_g_scale * g2_m + g_norm = np.hypot(g1, g2) + if g_norm >= 0.9: + warnings.warn( + "The scaled apodization shear norm (%f) was capped " + "at 0.9." % (g_norm) + ) + g1 *= 0.9 / g_norm + g2 *= 0.9 / g_norm + else: + g1, g2 = self.apod_anisotropy + self._apod_g_applied = (g1, g2) + return self._clean_pass(xi, g1=g1, g2=g2) + def optimizer(self, kernel): """ Build the gaussian process kernel from the measured 2d 2-point diff --git a/treegp/gp_interp.py b/treegp/gp_interp.py index ed1d72d..61ccfc3 100644 --- a/treegp/gp_interp.py +++ b/treegp/gp_interp.py @@ -99,6 +99,23 @@ class GPInterpolation(object): taper but leaves the window non-zero at the grid edge, reintroducing some spectral leakage. Used only by the "empirical-2pcf" optimizer. [default: None] + :param apod_anisotropy: Anisotropy of the apodization window, using the (g1, g2) + shear parametrization of Leget et al. 2021 (same convention as + get_correlation_length_matrix). None gives an isotropic window. + A (g1, g2) tuple applies the given shear: the window reaches + zero at apod_radius along the major axis (direction + 0.5 arctan2(g2, g1) from the x axis) and at apod_radius * q + along the minor axis, with q = (1 - g) / (1 + g). "auto" + measures (g1, g2) on the correlation function itself: a first + cleaning pass is done with the isotropic window, the anisotropy + of its output is measured with adaptive weighted second + moments, and the raw correlation function is re-cleaned with + the matched elliptical window. Ignored if apodize is False. + Used only by the "empirical-2pcf" optimizer. [default: None] + :param apod_g_scale: Factor multiplying the measured (g1, g2) before building the + 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.] """ def __init__( @@ -119,6 +136,8 @@ def __init__( apodize=True, apod_window="blackman-harris", apod_radius=None, + apod_anisotropy=None, + apod_g_scale=1.0, ): self.normalize = normalize self.optimizer = optimizer @@ -132,6 +151,8 @@ def __init__( self.apodize = apodize self.apod_window = apod_window self.apod_radius = apod_radius + self.apod_anisotropy = apod_anisotropy + self.apod_g_scale = apod_g_scale if self.optimizer == "anisotropic": self.robust_fit = True @@ -233,6 +254,8 @@ def _fit(self, kernel, X, y, y_err): apodize=self.apodize, apod_window=self.apod_window, apod_radius=self.apod_radius, + apod_anisotropy=self.apod_anisotropy, + apod_g_scale=self.apod_g_scale, ) kernel = self._optimizer.optimizer(kernel) # Hyperparameters estimation using maximum likelihood fit. From dae683e76f78b96f52c4450c49b83f8d5cbb2d5b Mon Sep 17 00:00:00 2001 From: PFLeget Date: Mon, 31 Aug 2026 15:15:10 -0400 Subject: [PATCH 4/4] oh my god this is fast (#38) --- tests/test_empirical_2pcf.py | 193 +++++++++++++++++++-- treegp/__init__.py | 2 + treegp/gp_interp.py | 139 ++++++++++++++- treegp/grid_gp.py | 321 +++++++++++++++++++++++++++++++++++ treegp/kernels.py | 88 ++++++++-- 5 files changed, 712 insertions(+), 31 deletions(-) create mode 100644 treegp/grid_gp.py 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):