From f67f10bc6ac919e67361af3db23c8cd9e2165b45 Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Fri, 18 Mar 2022 15:14:54 -0500 Subject: [PATCH 01/19] ENH add caching to FFTs --- ngmix/prepsfmom.py | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/ngmix/prepsfmom.py b/ngmix/prepsfmom.py index 2994e9bd..1819cdd7 100644 --- a/ngmix/prepsfmom.py +++ b/ngmix/prepsfmom.py @@ -1,4 +1,5 @@ import logging +import functools import numpy as np import scipy.fft as fft @@ -91,7 +92,7 @@ def _meas(self, obs, psf_obs, return_kernels): fft_dim = kim.shape[0] if psf_obs is not None: - kpsf_im, psf_im_row, psf_im_col = _zero_pad_and_compute_fft( + kpsf_im, psf_im_row, psf_im_col = _zero_pad_and_compute_fft_cached( psf_obs.image, psf_obs.jacobian.row0, psf_obs.jacobian.col0, target_dim, @@ -382,6 +383,29 @@ def _zero_pad_and_compute_fft(im, cen_row, cen_col, target_dim, ap_rad): return kpim, pad_cen_row, pad_cen_col +# see https://stackoverflow.com/a/52332109 for how this works +@functools.lru_cache(maxsize=128) +def _zero_pad_and_compute_fft_cached_impl( + im_tuple, cen_row, cen_col, target_dim, ap_rad +): + return _zero_pad_and_compute_fft( + np.array(im_tuple), cen_row, cen_col, target_dim, ap_rad + ) + + +@functools.wraps(_zero_pad_and_compute_fft) +def _zero_pad_and_compute_fft_cached(im, cen_row, cen_col, target_dim, ap_rad): + return _zero_pad_and_compute_fft_cached_impl( + tuple(im), cen_row, cen_col, target_dim, ap_rad + ) + + +_zero_pad_and_compute_fft_cached.cache_info \ + = _zero_pad_and_compute_fft_cached_impl.cache_info +_zero_pad_and_compute_fft_cached.cache_clear \ + = _zero_pad_and_compute_fft_cached_impl.cache_clear + + def _deconvolve_im_psf_inplace(kim, kpsf_im, max_amp, min_psf_frac=1e-5): """deconvolve the PSF from an image in place. @@ -398,6 +422,7 @@ def _deconvolve_im_psf_inplace(kim, kpsf_im, max_amp, min_psf_frac=1e-5): return kim, kpsf_im, msk +@functools.lru_cache(maxsize=128) def _ksigma_kernels( dim, kernel_size, @@ -505,6 +530,7 @@ def _ksigma_kernels( ) +@functools.lru_cache(maxsize=128) def _gauss_kernels( dim, kernel_size, From 48ee1d7e144d76740b6d529af782b80b43036ceb Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Fri, 18 Mar 2022 15:50:40 -0500 Subject: [PATCH 02/19] ENH add support for fourier images --- ngmix/prepsfmom.py | 169 +++++++++++++++++++++++++++++++-------------- 1 file changed, 118 insertions(+), 51 deletions(-) diff --git a/ngmix/prepsfmom.py b/ngmix/prepsfmom.py index 1819cdd7..93fde642 100644 --- a/ngmix/prepsfmom.py +++ b/ngmix/prepsfmom.py @@ -52,7 +52,10 @@ def __init__(self, fwhm, kernel, pad_factor=4, ap_rad=1.5): "The kernel '%s' for PrePSFMom is not recognized!" % self.kernel ) - def go(self, obs, return_kernels=False, no_psf=False): + def go( + self, obs, return_kernels=False, no_psf=False, + extra_conv_psfs=None, extra_deconv_psfs=None, + ): """Measure the pre-PSF ksigma moments. Parameters @@ -65,23 +68,32 @@ def go(self, obs, return_kernels=False, no_psf=False): no_psf : bool, optional If True, allow inputs without a PSF observation. Defaults to False so that any input observation without a PSF will raise an error. + extra_conv_psfs : list of ngmix obs, optional + If specified, these PSFs will be convolved into the image before the moments + are measured. + extra_deconv_psfs : list of ngmix obs, optional + If specified, these PSFs will be deconvolved from the image before the + moments are measured. Returns ------- result dictionary """ - psf_obs = _check_obs_and_get_psf_obs(obs, no_psf) - return self._meas(obs, psf_obs, return_kernels) - - def _meas(self, obs, psf_obs, return_kernels): - # pick the larger size - if psf_obs is not None: - if obs.image.shape[0] > psf_obs.image.shape[0]: - target_dim = int(obs.image.shape[0] * self.pad_factor) - else: - target_dim = int(psf_obs.image.shape[0] * self.pad_factor) - else: - target_dim = int(obs.image.shape[0] * self.pad_factor) + conv_psf_obs_list, deconv_psf_obs_list = _check_obs_and_get_psf_obs( + obs, no_psf, + extra_conv_psfs=extra_conv_psfs, + extra_deconv_psfs=extra_deconv_psfs, + ) + return self._meas(obs, conv_psf_obs_list, deconv_psf_obs_list, return_kernels) + + def _meas(self, obs, conv_psf_obs_list, deconv_psf_obs_list, return_kernels): + # pick the largest size + target_dim = max( + max([o.shape[0] for o in conv_psf_obs_list]), + max([o.shape[0] for o in deconv_psf_obs_list]), + obs.shape[0], + ) + target_dim = int(target_dim * self.pad_factor) eff_pad_factor = target_dim / obs.image.shape[0] # pad image, psf and weight map, get FFTs, apply cen_phases @@ -91,18 +103,17 @@ def _meas(self, obs, psf_obs, return_kernels): ) fft_dim = kim.shape[0] - if psf_obs is not None: - kpsf_im, psf_im_row, psf_im_col = _zero_pad_and_compute_fft_cached( - psf_obs.image, - psf_obs.jacobian.row0, psf_obs.jacobian.col0, - target_dim, - 0, # we do not apodize PSF stamps since it should not be needed - ) - else: - # delta function in k-space - kpsf_im = np.ones_like(kim, dtype=np.complex128) - psf_im_row = 0.0 - psf_im_col = 0.0 + ( + deconv_kpsf_im, + deconv_psf_im_row, + deconv_psf_im_col, + ) = _zero_pad_and_compute_fft_cached_list(deconv_psf_obs_list, kim, target_dim) + + ( + conv_kpsf_im, + conv_psf_im_row, + conv_psf_im_col, + ) = _zero_pad_and_compute_fft_cached_list(conv_psf_obs_list, kim, target_dim) # the final, deconvolved image we want is # @@ -129,17 +140,17 @@ def _meas(self, obs, psf_obs, return_kernels): # now build the kernels if self.kernel == "ksigma": kernels = _ksigma_kernels( - target_dim, - self.fwhm, - obs.jacobian.dvdrow, obs.jacobian.dvdcol, - obs.jacobian.dudrow, obs.jacobian.dudcol, + int(target_dim), + float(self.fwhm), + float(obs.jacobian.dvdrow), float(obs.jacobian.dvdcol), + float(obs.jacobian.dudrow), float(obs.jacobian.dudcol), ) elif self.kernel in ["gauss", "pgauss"]: kernels = _gauss_kernels( - target_dim, - self.fwhm, - obs.jacobian.dvdrow, obs.jacobian.dvdcol, - obs.jacobian.dudrow, obs.jacobian.dudcol, + int(target_dim), + float(self.fwhm), + float(obs.jacobian.dvdrow), float(obs.jacobian.dvdcol), + float(obs.jacobian.dudrow), float(obs.jacobian.dudcol), ) else: raise ValueError( @@ -152,8 +163,11 @@ def _meas(self, obs, psf_obs, return_kernels): # run the actual measurements and return mom, mom_cov = _measure_moments_fft( - kim, kpsf_im, tot_var, eff_pad_factor, kernels, - im_row - psf_im_row, im_col - psf_im_col, + kim, + deconv_kpsf_im, conv_kpsf_im, + tot_var, eff_pad_factor, kernels, + im_row - deconv_psf_im_row + conv_psf_im_row, + im_col - deconv_psf_im_col + conv_psf_im_col, ) res = make_mom_result(mom, mom_cov) if res['flags'] != 0: @@ -226,19 +240,27 @@ def __init__(self, fwhm, pad_factor=4, ap_rad=1.5): PrePSFGaussMom = PGaussMom -def _measure_moments_fft(kim, kpsf_im, tot_var, eff_pad_factor, kernels, drow, dcol): +def _measure_moments_fft( + kim, deconv_kpsf_im, conv_kpsf_im, tot_var, eff_pad_factor, kernels, + drow, dcol +): # we only need to do things where the kernel is non-zero # this saves a bunch of CPU cycles msk = kernels["msk"] dim = kim.shape[0] # deconvolve PSF - kim, kpsf_im, _ = _deconvolve_im_psf_inplace( - kim[msk], - kpsf_im[msk], - # max amplitude is flux which is 0,0 in the standard FFT convention - np.abs(kpsf_im[0, 0]), - ) + if deconv_kpsf_im is not None or conv_kpsf_im is not None: + kim, inv_kpsf_im, _ = _deconvolve_im_psf_inplace( + kim[msk], + deconv_kpsf_im[msk] if deconv_kpsf_im is not None else None, + # max amplitude is flux which is 0,0 in the standard FFT convention + np.abs(deconv_kpsf_im[0, 0]) if deconv_kpsf_im is not None else None, + conv_kpsf_im[msk] if conv_kpsf_im is not None else None, + ) + else: + kim = kim[msk] + inv_kpsf_im = np.ones_like(kim, dtype=np.complex128) # put in phase shift as described above # the sin and cos are expensive so we only compute them where we will @@ -279,7 +301,7 @@ def _measure_moments_fft(kim, kpsf_im, tot_var, eff_pad_factor, kernels, drow, d m_cov[1, 1] = 1 tot_var *= eff_pad_factor**2 tot_var_df4 = tot_var * df4 - kerns = [fkp / kpsf_im, fkc / kpsf_im, fkr / kpsf_im, fkf / kpsf_im] + kerns = [fkp * inv_kpsf_im, fkc * inv_kpsf_im, fkr * inv_kpsf_im, fkf * inv_kpsf_im] conj_kerns = [np.conj(k) for k in kerns] for i in range(2, 6): for j in range(i, 6): @@ -396,7 +418,7 @@ def _zero_pad_and_compute_fft_cached_impl( @functools.wraps(_zero_pad_and_compute_fft) def _zero_pad_and_compute_fft_cached(im, cen_row, cen_col, target_dim, ap_rad): return _zero_pad_and_compute_fft_cached_impl( - tuple(im), cen_row, cen_col, target_dim, ap_rad + tuple(im), float(cen_row), float(cen_col), int(target_dim), float(ap_rad) ) @@ -406,20 +428,65 @@ def _zero_pad_and_compute_fft_cached(im, cen_row, cen_col, target_dim, ap_rad): = _zero_pad_and_compute_fft_cached_impl.cache_clear -def _deconvolve_im_psf_inplace(kim, kpsf_im, max_amp, min_psf_frac=1e-5): +def _zero_pad_and_compute_fft_cached_list(psf_obs_list, kim, target_dim): + if len(psf_obs_list) == 0: + # delta function in k-space + kpsf_im = None + psf_im_row = 0.0 + psf_im_col = 0.0 + elif len(psf_obs_list) == 1: + return _zero_pad_and_compute_fft_cached( + psf_obs_list[0].image, + psf_obs_list[0].jacobian.row0, psf_obs_list[0].jacobian.col0, + target_dim, + 0, # we do not apodize PSF stamps since it should not be needed + ) + else: + fft_res = [ + _zero_pad_and_compute_fft_cached( + psf_obs.image, + psf_obs.jacobian.row0, psf_obs.jacobian.col0, + target_dim, + 0, # we do not apodize PSF stamps since it should not be needed + ) + for psf_obs in psf_obs_list + ] + kpsf_im = np.prod([f[0] for f in fft_res], axis=0) + psf_im_row = sum([f[1] for f in fft_res]) + psf_im_row = sum([f[2] for f in fft_res]) + + return kpsf_im, psf_im_row, psf_im_col + + +def _deconvolve_im_psf_inplace( + kim, deconv_kpsf_im, max_amp, conv_kpsf_im, min_psf_frac=1e-5 +): """deconvolve the PSF from an image in place. Returns the deconvolved image, the kpsf_im used, and a bool mask marking PSF modes that were truncated """ min_amp = min_psf_frac * max_amp - abs_kpsf_im = np.abs(kpsf_im) - msk = abs_kpsf_im <= min_amp - if np.any(msk): - kpsf_im[msk] = kpsf_im[msk] / abs_kpsf_im[msk] * min_amp + if deconv_kpsf_im is not None: + abs_kpsf_im = np.abs(deconv_kpsf_im) + msk = abs_kpsf_im <= min_amp + if np.any(msk): + deconv_kpsf_im[msk] = deconv_kpsf_im[msk] / abs_kpsf_im[msk] * min_amp + + if conv_kpsf_im is not None: + inv_kpsf_im = conv_kpsf_im / deconv_kpsf_im + else: + inv_kpsf_im = 1.0 / deconv_kpsf_im + else: + msk = np.ones_like(kim, dtype=bool) + + if conv_kpsf_im is not None: + inv_kpsf_im = conv_kpsf_im + else: + inv_kpsf_im = np.ones_like(kim, dtype=np.complex128) - kim /= kpsf_im - return kim, kpsf_im, msk + kim *= inv_kpsf_im + return kim, inv_kpsf_im, msk @functools.lru_cache(maxsize=128) From c0f8ba77a4eda5a7a1094bacd3132052ec667f00 Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Fri, 18 Mar 2022 16:09:47 -0500 Subject: [PATCH 03/19] ENH try it --- ngmix/prepsfmom.py | 34 +++++++++++++++++++++++++--------- 1 file changed, 25 insertions(+), 9 deletions(-) diff --git a/ngmix/prepsfmom.py b/ngmix/prepsfmom.py index 93fde642..2e820236 100644 --- a/ngmix/prepsfmom.py +++ b/ngmix/prepsfmom.py @@ -89,9 +89,11 @@ def go( def _meas(self, obs, conv_psf_obs_list, deconv_psf_obs_list, return_kernels): # pick the largest size target_dim = max( - max([o.shape[0] for o in conv_psf_obs_list]), - max([o.shape[0] for o in deconv_psf_obs_list]), - obs.shape[0], + max([o.image.shape[0] for o in conv_psf_obs_list]) + if conv_psf_obs_list else -1, + max([o.image.shape[0] for o in deconv_psf_obs_list]) + if deconv_psf_obs_list else -1, + obs.image.shape[0], ) target_dim = int(target_dim * self.pad_factor) eff_pad_factor = target_dim / obs.image.shape[0] @@ -418,7 +420,8 @@ def _zero_pad_and_compute_fft_cached_impl( @functools.wraps(_zero_pad_and_compute_fft) def _zero_pad_and_compute_fft_cached(im, cen_row, cen_col, target_dim, ap_rad): return _zero_pad_and_compute_fft_cached_impl( - tuple(im), float(cen_row), float(cen_col), int(target_dim), float(ap_rad) + tuple(tuple(ii) for ii in im), + float(cen_row), float(cen_col), int(target_dim), float(ap_rad) ) @@ -689,7 +692,9 @@ def _gauss_kernels( ) -def _check_obs_and_get_psf_obs(obs, no_psf): +def _check_obs_and_get_psf_obs( + obs, no_psf, extra_conv_psfs=None, extra_deconv_psfs=None, +): if not isinstance(obs, Observation): raise ValueError("input obs must be an Observation") @@ -700,15 +705,26 @@ def _check_obs_and_get_psf_obs(obs, no_psf): if not obs.has_psf() and not no_psf: raise RuntimeError("The PSF must be set to measure a pre-PSF moment!") + if no_psf and (extra_conv_psfs or extra_deconv_psfs): + raise RuntimeError( + "You can only use extra conv. or deconv. PSFs with observations " + "that have a PSF!" + ) + if not no_psf: - psf_obs = obs.get_psf() + conv_psfs = extra_conv_psfs or [] + deconv_psfs = [obs.get_psf()] + (extra_deconv_psfs or []) - if psf_obs.jacobian.get_galsim_wcs() != obs.jacobian.get_galsim_wcs(): + if any( + psf_obs.jacobian.get_galsim_wcs() != obs.jacobian.get_galsim_wcs() + for psf_obs in conv_psfs + deconv_psfs + ): raise RuntimeError( "The PSF and observation must have the same WCS " "Jacobian for measuring pre-PSF moments." ) else: - psf_obs = None + conv_psfs = [] + deconv_psfs = [] - return psf_obs + return conv_psfs, deconv_psfs From 8578f5cfdded9cf28f8bb1055a1167fa44328441 Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Fri, 18 Mar 2022 16:30:57 -0500 Subject: [PATCH 04/19] TST added test for speed and caching --- ngmix/tests/test_prepsfmom.py | 131 ++++++++++++++++++++++++++++++++++ 1 file changed, 131 insertions(+) diff --git a/ngmix/tests/test_prepsfmom.py b/ngmix/tests/test_prepsfmom.py index d24a1ff6..34a9040a 100644 --- a/ngmix/tests/test_prepsfmom.py +++ b/ngmix/tests/test_prepsfmom.py @@ -1,11 +1,14 @@ import galsim import numpy as np import pytest +import time from ngmix.prepsfmom import ( KSigmaMom, PGaussMom, _build_square_apodization_mask, PrePSFMom, + _gauss_kernels, + _zero_pad_and_compute_fft_cached_impl, ) from ngmix import Jacobian from ngmix import Observation @@ -126,6 +129,134 @@ def _get_dtype(v): return d +def test_prepsfmom_speed_and_cache(): + image_size = 48 + psf_image_size = 53 + pixel_scale = 0.263 + fwhm = 0.9 + psf_fwhm = 0.9 + snr = 20 + mom_fwhm = 2 + + rng = np.random.RandomState(seed=100) + + cen = (image_size - 1)/2 + psf_cen = (psf_image_size - 1)/2 + gs_wcs = galsim.ShearWCS( + pixel_scale, galsim.Shear(g1=-0.1, g2=0.06)).jacobian() + scale = np.sqrt(gs_wcs.pixelArea()) + shift = rng.uniform(low=-scale/2, high=scale/2, size=2) + psf_shift = rng.uniform(low=-scale/2, high=scale/2, size=2) + xy = gs_wcs.toImage(galsim.PositionD(shift)) + psf_xy = gs_wcs.toImage(galsim.PositionD(psf_shift)) + + jac = Jacobian( + y=cen + xy.y, x=cen + xy.x, + dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, + dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy) + + psf_jac = Jacobian( + y=psf_cen + psf_xy.y, x=psf_cen + psf_xy.x, + dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, + dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy) + + gal = galsim.Gaussian( + fwhm=fwhm + ).shear( + g1=-0.1, g2=0.2 + ).withFlux( + 400 + ).shift( + dx=shift[0], dy=shift[1] + ) + psf = galsim.Gaussian( + fwhm=psf_fwhm + ).shear( + g1=0.3, g2=-0.15 + ) + im = galsim.Convolve([gal, psf]).drawImage( + nx=image_size, + ny=image_size, + wcs=gs_wcs + ).array + noise = np.sqrt(np.sum(im**2)) / snr + wgt = np.ones_like(im) / noise**2 + + psf_im = psf.shift( + dx=psf_shift[0], dy=psf_shift[1] + ).drawImage( + nx=psf_image_size, + ny=psf_image_size, + wcs=gs_wcs + ).array + + # now we test the speed + caching + + # start by clearing the caches + _gauss_kernels.cache_clear() + _zero_pad_and_compute_fft_cached_impl.cache_clear() + + # the first fit will do numba stuff, so we exclude it + # we also perturb the various inputs to fool our caches + fitter = PGaussMom( + fwhm=mom_fwhm + 1e-3, + ) + + im += rng.normal(size=im.shape, scale=noise) + obs = Observation( + image=im, + weight=wgt, + jacobian=jac, + psf=Observation(image=psf_im + 1e-8, jacobian=psf_jac), + ) + + dt = time.time() + fitter.go(obs=obs) + dt = time.time() - dt + print("\n%0.4f ms for first fit" % (dt*1000)) + + # we miss once here + assert _gauss_kernels.cache_info().misses == 1 + assert _zero_pad_and_compute_fft_cached_impl.cache_info().misses == 1 + + # the second fit will have numba cached, but not the other kernel and FFT caches + # we also perturb the various inputs to fool our caches + fitter = PGaussMom( + fwhm=mom_fwhm, + ) + + im += rng.normal(size=im.shape, scale=noise) + obs = Observation( + image=im, + weight=wgt, + jacobian=jac, + psf=Observation(image=psf_im, jacobian=psf_jac), + ) + + dt = time.time() + fitter.go(obs=obs) + dt = time.time() - dt + print("%0.4f ms for second fit" % (dt*1000)) + + # we miss twice since we changed the moments width and psf slightly + assert _gauss_kernels.cache_info().misses == 2 + assert _zero_pad_and_compute_fft_cached_impl.cache_info().misses == 2 + + # finally, we test with full caching + # we also perturb the various inputs to fool our caches + nfit = 1000 + dt = time.time() + for _ in range(nfit): + fitter.go(obs=obs) + dt = time.time() - dt + + print("%0.4f ms per fit" % (dt/nfit*1000)) + + # we should never miss again for the calls above + assert _gauss_kernels.cache_info().misses == 2 + assert _zero_pad_and_compute_fft_cached_impl.cache_info().misses == 2 + + @pytest.mark.parametrize("cls", [KSigmaMom, PGaussMom]) @pytest.mark.parametrize('snr', [1e1, 1e3]) @pytest.mark.parametrize('pixel_scale', [0.125, 0.25]) From e053658fdc822e2a1b6537191a9416ca015d9545 Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Fri, 18 Mar 2022 16:40:41 -0500 Subject: [PATCH 05/19] TST added test for mixing no_psf with extra conv or deconv PSFs --- ngmix/tests/test_prepsfmom.py | 123 ++++++++++++++-------------------- 1 file changed, 50 insertions(+), 73 deletions(-) diff --git a/ngmix/tests/test_prepsfmom.py b/ngmix/tests/test_prepsfmom.py index 34a9040a..f3f9aacc 100644 --- a/ngmix/tests/test_prepsfmom.py +++ b/ngmix/tests/test_prepsfmom.py @@ -7,8 +7,6 @@ KSigmaMom, PGaussMom, _build_square_apodization_mask, PrePSFMom, - _gauss_kernels, - _zero_pad_and_compute_fft_cached_impl, ) from ngmix import Jacobian from ngmix import Observation @@ -60,6 +58,23 @@ def test_prepsfmom_raises_nopsf(cls): fitter.go(obs, no_psf=True) +@pytest.mark.parametrize("cls", [KSigmaMom, PGaussMom]) +def test_prepsfmom_raises_nopsf_with_extra(cls): + fitter = cls(20) + obs = Observation(image=np.zeros((1000, 1000))) + with pytest.raises(RuntimeError) as e: + fitter.go(obs, extra_deconv_psfs=[10], no_psf=True) + assert "You can only use extra conv." in str(e.value) + + with pytest.raises(RuntimeError) as e: + fitter.go(obs, extra_conv_psfs=[10], no_psf=True) + assert "You can only use extra conv." in str(e.value) + + with pytest.raises(RuntimeError) as e: + fitter.go(obs, extra_conv_psfs=[10], extra_deconv_psfs=[11], no_psf=True) + assert "You can only use extra conv." in str(e.value) + + @pytest.mark.parametrize("cls", [KSigmaMom, PGaussMom]) def test_prepsfmom_raises_nonsquare(cls): fitter = cls(20) @@ -97,38 +112,6 @@ def test_prepsfmom_raises_badjacob(cls): assert "same WCS Jacobia" in str(e.value) -def _stack_list_of_dicts(res): - def _get_dtype(v): - if isinstance(v, float): - return ('f8',) - elif isinstance(v, int): - return ('i4',) - elif isinstance(v, str): - return ('U256',) - elif hasattr(v, "dtype") and hasattr(v, "shape"): - if "float" in str(v.dtype): - dstr = "f8" - else: - dstr = "i8" - - if len(v.shape) == 1: - return (dstr, v.shape[0]) - else: - return (dstr, v.shape) - else: - raise RuntimeError("cannot interpret dtype of '%s'" % v) - - dtype = [] - for k, v in res[0].items(): - dtype.append((k,) + _get_dtype(v)) - d = np.zeros(len(res), dtype=dtype) - for i in range(len(res)): - for k, v in res[i].items(): - d[k][i] = v - - return d - - def test_prepsfmom_speed_and_cache(): image_size = 48 psf_image_size = 53 @@ -191,36 +174,7 @@ def test_prepsfmom_speed_and_cache(): ).array # now we test the speed + caching - - # start by clearing the caches - _gauss_kernels.cache_clear() - _zero_pad_and_compute_fft_cached_impl.cache_clear() - # the first fit will do numba stuff, so we exclude it - # we also perturb the various inputs to fool our caches - fitter = PGaussMom( - fwhm=mom_fwhm + 1e-3, - ) - - im += rng.normal(size=im.shape, scale=noise) - obs = Observation( - image=im, - weight=wgt, - jacobian=jac, - psf=Observation(image=psf_im + 1e-8, jacobian=psf_jac), - ) - - dt = time.time() - fitter.go(obs=obs) - dt = time.time() - dt - print("\n%0.4f ms for first fit" % (dt*1000)) - - # we miss once here - assert _gauss_kernels.cache_info().misses == 1 - assert _zero_pad_and_compute_fft_cached_impl.cache_info().misses == 1 - - # the second fit will have numba cached, but not the other kernel and FFT caches - # we also perturb the various inputs to fool our caches fitter = PGaussMom( fwhm=mom_fwhm, ) @@ -236,14 +190,9 @@ def test_prepsfmom_speed_and_cache(): dt = time.time() fitter.go(obs=obs) dt = time.time() - dt - print("%0.4f ms for second fit" % (dt*1000)) - - # we miss twice since we changed the moments width and psf slightly - assert _gauss_kernels.cache_info().misses == 2 - assert _zero_pad_and_compute_fft_cached_impl.cache_info().misses == 2 + print("\n%0.4f ms for first fit" % (dt*1000)) - # finally, we test with full caching - # we also perturb the various inputs to fool our caches + # we test with full caching nfit = 1000 dt = time.time() for _ in range(nfit): @@ -252,9 +201,37 @@ def test_prepsfmom_speed_and_cache(): print("%0.4f ms per fit" % (dt/nfit*1000)) - # we should never miss again for the calls above - assert _gauss_kernels.cache_info().misses == 2 - assert _zero_pad_and_compute_fft_cached_impl.cache_info().misses == 2 + +def _stack_list_of_dicts(res): + def _get_dtype(v): + if isinstance(v, float): + return ('f8',) + elif isinstance(v, int): + return ('i4',) + elif isinstance(v, str): + return ('U256',) + elif hasattr(v, "dtype") and hasattr(v, "shape"): + if "float" in str(v.dtype): + dstr = "f8" + else: + dstr = "i8" + + if len(v.shape) == 1: + return (dstr, v.shape[0]) + else: + return (dstr, v.shape) + else: + raise RuntimeError("cannot interpret dtype of '%s'" % v) + + dtype = [] + for k, v in res[0].items(): + dtype.append((k,) + _get_dtype(v)) + d = np.zeros(len(res), dtype=dtype) + for i in range(len(res)): + for k, v in res[i].items(): + d[k][i] = v + + return d @pytest.mark.parametrize("cls", [KSigmaMom, PGaussMom]) From 2a1fe9ec0ec616c790e1f7d957275202fd0f5000 Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Fri, 18 Mar 2022 16:44:39 -0500 Subject: [PATCH 06/19] BUG add bag odd thing --- ngmix/tests/test_prepsfmom.py | 36 ++++++++++++++++++++++++++++++++++- 1 file changed, 35 insertions(+), 1 deletion(-) diff --git a/ngmix/tests/test_prepsfmom.py b/ngmix/tests/test_prepsfmom.py index f3f9aacc..bb13c5d5 100644 --- a/ngmix/tests/test_prepsfmom.py +++ b/ngmix/tests/test_prepsfmom.py @@ -7,6 +7,8 @@ KSigmaMom, PGaussMom, _build_square_apodization_mask, PrePSFMom, + _gauss_kernels, + _zero_pad_and_compute_fft_cached_impl, ) from ngmix import Jacobian from ngmix import Observation @@ -175,6 +177,30 @@ def test_prepsfmom_speed_and_cache(): # now we test the speed + caching # the first fit will do numba stuff, so we exclude it + # we also perturb the various inputs to fool our caches + fitter = PGaussMom( + fwhm=mom_fwhm + 1e-3, + ) + + im += rng.normal(size=im.shape, scale=noise) + obs = Observation( + image=im, + weight=wgt, + jacobian=jac, + psf=Observation(image=psf_im + 1e-8, jacobian=psf_jac), + ) + + dt = time.time() + fitter.go(obs=obs) + dt = time.time() - dt + print("\n%0.4f ms for first fit" % (dt*1000)) + + # we miss once here + assert _gauss_kernels.cache_info().misses == 1 + assert _zero_pad_and_compute_fft_cached_impl.cache_info().misses == 1 + + # the second fit will have numba cached, but not the other kernel and FFT caches + # we also perturb the various inputs to fool our caches fitter = PGaussMom( fwhm=mom_fwhm, ) @@ -190,7 +216,11 @@ def test_prepsfmom_speed_and_cache(): dt = time.time() fitter.go(obs=obs) dt = time.time() - dt - print("\n%0.4f ms for first fit" % (dt*1000)) + print("%0.4f ms for second fit" % (dt*1000)) + + # we miss twice since we changed the moments width and psf slightly + assert _gauss_kernels.cache_info().misses == 2 + assert _zero_pad_and_compute_fft_cached_impl.cache_info().misses == 2 # we test with full caching nfit = 1000 @@ -201,6 +231,10 @@ def test_prepsfmom_speed_and_cache(): print("%0.4f ms per fit" % (dt/nfit*1000)) + # we should never miss again for the calls above + assert _gauss_kernels.cache_info().misses == 2 + assert _zero_pad_and_compute_fft_cached_impl.cache_info().misses == 2 + def _stack_list_of_dicts(res): def _get_dtype(v): From 80c70b59d8e9f561f83f7cbb4ad0b885af5e6d35 Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Fri, 18 Mar 2022 16:45:46 -0500 Subject: [PATCH 07/19] DOC update comments --- ngmix/tests/test_prepsfmom.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ngmix/tests/test_prepsfmom.py b/ngmix/tests/test_prepsfmom.py index bb13c5d5..2493b9f7 100644 --- a/ngmix/tests/test_prepsfmom.py +++ b/ngmix/tests/test_prepsfmom.py @@ -200,7 +200,6 @@ def test_prepsfmom_speed_and_cache(): assert _zero_pad_and_compute_fft_cached_impl.cache_info().misses == 1 # the second fit will have numba cached, but not the other kernel and FFT caches - # we also perturb the various inputs to fool our caches fitter = PGaussMom( fwhm=mom_fwhm, ) @@ -222,7 +221,7 @@ def test_prepsfmom_speed_and_cache(): assert _gauss_kernels.cache_info().misses == 2 assert _zero_pad_and_compute_fft_cached_impl.cache_info().misses == 2 - # we test with full caching + # now we test with full caching nfit = 1000 dt = time.time() for _ in range(nfit): From 5ed7293915100eeb6e5cff0d917b05ae54497c2c Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Fri, 18 Mar 2022 16:49:28 -0500 Subject: [PATCH 08/19] TST added test for all psf jacobians --- ngmix/tests/test_prepsfmom.py | 45 +++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/ngmix/tests/test_prepsfmom.py b/ngmix/tests/test_prepsfmom.py index 2493b9f7..4acda891 100644 --- a/ngmix/tests/test_prepsfmom.py +++ b/ngmix/tests/test_prepsfmom.py @@ -114,6 +114,51 @@ def test_prepsfmom_raises_badjacob(cls): assert "same WCS Jacobia" in str(e.value) +@pytest.mark.parametrize("cls", [KSigmaMom, PGaussMom]) +def test_prepsfmom_raises_badjacob_extra_conv(cls): + fitter = cls(1.2) + + gs_wcs = galsim.ShearWCS( + 0.2, galsim.Shear(g1=-0.1, g2=0.06)).jacobian() + jac = Jacobian( + y=0, x=0, + dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, + dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy) + + psf_jac = Jacobian( + y=0, x=0, + dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, + dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy*2) + + obs = Observation( + image=np.zeros((10, 10)), + jacobian=jac, + psf=Observation(image=np.zeros((10, 10)), jacobian=jac), + ) + + with pytest.raises(RuntimeError) as e: + fitter.go( + obs, + extra_conv_psfs=[Observation(image=np.zeros((10, 10)), jacobian=psf_jac)], + ) + assert "same WCS Jacobia" in str(e.value) + + with pytest.raises(RuntimeError) as e: + fitter.go( + obs, + extra_deconv_psfs=[Observation(image=np.zeros((10, 10)), jacobian=psf_jac)], + ) + assert "same WCS Jacobia" in str(e.value) + + with pytest.raises(RuntimeError) as e: + fitter.go( + obs, + extra_conv_psfs=[Observation(image=np.zeros((10, 10)), jacobian=psf_jac)], + extra_deconv_psfs=[Observation(image=np.zeros((10, 10)), jacobian=psf_jac)], + ) + assert "same WCS Jacobia" in str(e.value) + + def test_prepsfmom_speed_and_cache(): image_size = 48 psf_image_size = 53 From 75772fa8eb397a21c2011917e7ab237258f41512 Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Fri, 18 Mar 2022 16:53:44 -0500 Subject: [PATCH 09/19] REF make sure close to machine precision --- ngmix/prepsfmom.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/ngmix/prepsfmom.py b/ngmix/prepsfmom.py index 2e820236..9987bee4 100644 --- a/ngmix/prepsfmom.py +++ b/ngmix/prepsfmom.py @@ -692,6 +692,13 @@ def _gauss_kernels( ) +def _jacobian_close(jac1, jac2): + return np.allclose( + [jac1.dudcol, jac1.dudrow, jac1.dvdcol, jac1.dvdrow], + [jac2.dudcol, jac2.dudrow, jac2.dvdcol, jac2.dvdrow] + ) + + def _check_obs_and_get_psf_obs( obs, no_psf, extra_conv_psfs=None, extra_deconv_psfs=None, ): @@ -716,7 +723,7 @@ def _check_obs_and_get_psf_obs( deconv_psfs = [obs.get_psf()] + (extra_deconv_psfs or []) if any( - psf_obs.jacobian.get_galsim_wcs() != obs.jacobian.get_galsim_wcs() + not _jacobian_close(psf_obs.jacobian, obs.jacobian) for psf_obs in conv_psfs + deconv_psfs ): raise RuntimeError( From 81097ca7dd51da19ec0328b008de4416ba9578d9 Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Fri, 18 Mar 2022 17:05:34 -0500 Subject: [PATCH 10/19] BUG clear the caches in the test first --- ngmix/tests/test_prepsfmom.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ngmix/tests/test_prepsfmom.py b/ngmix/tests/test_prepsfmom.py index 4acda891..0a05c08f 100644 --- a/ngmix/tests/test_prepsfmom.py +++ b/ngmix/tests/test_prepsfmom.py @@ -221,6 +221,9 @@ def test_prepsfmom_speed_and_cache(): ).array # now we test the speed + caching + _gauss_kernels.cache_clear() + _zero_pad_and_compute_fft_cached_impl.cache_clear() + # the first fit will do numba stuff, so we exclude it # we also perturb the various inputs to fool our caches fitter = PGaussMom( From 6858c51759988ce5015418c499a98ed799d57910 Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Fri, 18 Mar 2022 18:15:07 -0500 Subject: [PATCH 11/19] TST added accuracy test of deconv and conv handling --- ngmix/prepsfmom.py | 2 +- ngmix/tests/test_prepsfmom.py | 74 ++++++++++++++++++++++++++++++----- 2 files changed, 65 insertions(+), 11 deletions(-) diff --git a/ngmix/prepsfmom.py b/ngmix/prepsfmom.py index 9987bee4..cce624d8 100644 --- a/ngmix/prepsfmom.py +++ b/ngmix/prepsfmom.py @@ -456,7 +456,7 @@ def _zero_pad_and_compute_fft_cached_list(psf_obs_list, kim, target_dim): ] kpsf_im = np.prod([f[0] for f in fft_res], axis=0) psf_im_row = sum([f[1] for f in fft_res]) - psf_im_row = sum([f[2] for f in fft_res]) + psf_im_col = sum([f[2] for f in fft_res]) return kpsf_im, psf_im_row, psf_im_col diff --git a/ngmix/tests/test_prepsfmom.py b/ngmix/tests/test_prepsfmom.py index 0a05c08f..542e920c 100644 --- a/ngmix/tests/test_prepsfmom.py +++ b/ngmix/tests/test_prepsfmom.py @@ -770,6 +770,7 @@ def test_moments_make_mom_result_flags(): assert res["T_flagstr"] == "" +@pytest.mark.parametrize("extra_psf_fwhm", [None, 0.8, 1.1]) @pytest.mark.parametrize("cls", [PGaussMom, KSigmaMom]) @pytest.mark.parametrize('pixel_scale', [0.125, 0.25]) @pytest.mark.parametrize('fwhm,psf_fwhm', [(0.6, 0.9)]) @@ -777,7 +778,8 @@ def test_moments_make_mom_result_flags(): @pytest.mark.parametrize('psf_image_size', [33, 34]) @pytest.mark.parametrize('pad_factor', [4, 3.5]) def test_prepsfmom_gauss_true_flux( - pad_factor, psf_image_size, image_size, fwhm, psf_fwhm, pixel_scale, cls + pad_factor, psf_image_size, image_size, fwhm, psf_fwhm, pixel_scale, + cls, extra_psf_fwhm ): rng = np.random.RandomState(seed=100) @@ -834,6 +836,27 @@ def test_prepsfmom_gauss_true_flux( wcs=gs_wcs ).array + if extra_psf_fwhm is not None: + extra_psf = galsim.Gaussian( + fwhm=extra_psf_fwhm + ).shear( + g1=-0.2, g2=0.3 + ) + extra_psf_shift = rng.uniform(low=-scale/2, high=scale/2, size=2) + extra_psf_xy = gs_wcs.toImage(galsim.PositionD(extra_psf_shift)) + extra_psf_im = extra_psf.shift( + dx=extra_psf_shift[0], dy=extra_psf_shift[1] + ).drawImage( + nx=psf_image_size, + ny=psf_image_size, + wcs=gs_wcs + ).array + extra_psf_jac = Jacobian( + y=psf_cen + extra_psf_xy.y, x=psf_cen + extra_psf_xy.x, + dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, + dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy, + ) + fitter = cls( fwhm=mom_fwhm, pad_factor=pad_factor, @@ -851,17 +874,48 @@ def test_prepsfmom_gauss_true_flux( ) res = fitter.go(obs=obs, no_psf=True) flux_true = res["flux"] + T_true = res["T"] assert np.allclose(flux_true, 400, atol=0, rtol=5e-3) - obs = Observation( - image=im, - weight=wgt, - jacobian=jac, - psf=Observation(image=psf_im, jacobian=psf_jac), - ) - res = fitter.go(obs=obs) - flux_true = res["flux"] - assert np.allclose(flux_true, 400, atol=0, rtol=5e-3) + if extra_psf_fwhm is None: + obs = Observation( + image=im, + weight=wgt, + jacobian=jac, + psf=Observation(image=psf_im, jacobian=psf_jac), + ) + res = fitter.go(obs=obs) + flux_true = res["flux"] + assert np.allclose(flux_true, 400, atol=0, rtol=5e-3) + assert np.allclose(res["T"], T_true, atol=0, rtol=5e-4) + else: + # this should fail since it is the wrong PSF + obs = Observation( + image=im, + weight=wgt, + jacobian=jac, + psf=Observation(image=extra_psf_im, jacobian=extra_psf_jac), + ) + res = fitter.go(obs=obs) + flux_true = res["flux"] + # we use a big relative difference since the T should be way off + assert not np.allclose(res["T"], T_true, atol=0, rtol=1e-1) + + # this should work since we have the PSF correction in there + obs = Observation( + image=im, + weight=wgt, + jacobian=jac, + psf=Observation(image=extra_psf_im, jacobian=extra_psf_jac), + ) + res = fitter.go( + obs=obs, + extra_deconv_psfs=[Observation(image=psf_im, jacobian=psf_jac)], + extra_conv_psfs=[obs.psf], + ) + flux_true = res["flux"] + assert np.allclose(flux_true, 400, atol=0, rtol=5e-3) + assert np.allclose(res["T"], T_true, atol=0, rtol=5e-4) @pytest.mark.parametrize('pixel_scale', [0.25, 0.125]) From 8bee261cfb25c5f578c0e28dbbcf2e21c1c23740 Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Fri, 18 Mar 2022 18:20:04 -0500 Subject: [PATCH 12/19] TST rename test since it works on T too --- ngmix/tests/test_prepsfmom.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ngmix/tests/test_prepsfmom.py b/ngmix/tests/test_prepsfmom.py index 542e920c..048f42e5 100644 --- a/ngmix/tests/test_prepsfmom.py +++ b/ngmix/tests/test_prepsfmom.py @@ -777,7 +777,7 @@ def test_moments_make_mom_result_flags(): @pytest.mark.parametrize('image_size', [250]) @pytest.mark.parametrize('psf_image_size', [33, 34]) @pytest.mark.parametrize('pad_factor', [4, 3.5]) -def test_prepsfmom_gauss_true_flux( +def test_prepsfmom_gauss_true_flux_T( pad_factor, psf_image_size, image_size, fwhm, psf_fwhm, pixel_scale, cls, extra_psf_fwhm ): From 3270a22a16b58f09f755c20565a5425a1ec52d0d Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Fri, 18 Mar 2022 20:18:59 -0500 Subject: [PATCH 13/19] TST added tests of extra PSF stuff for errors --- ngmix/tests/test_prepsfmom.py | 105 +++++++++++++++++++++++++++++----- 1 file changed, 90 insertions(+), 15 deletions(-) diff --git a/ngmix/tests/test_prepsfmom.py b/ngmix/tests/test_prepsfmom.py index 048f42e5..f1922947 100644 --- a/ngmix/tests/test_prepsfmom.py +++ b/ngmix/tests/test_prepsfmom.py @@ -323,9 +323,10 @@ def _get_dtype(v): @pytest.mark.parametrize('image_size', [57, 58]) @pytest.mark.parametrize('psf_image_size', [33, 34]) @pytest.mark.parametrize('pad_factor', [3.5, 2]) +@pytest.mark.parametrize("extra_psf_fwhm", [None, 0.8, 1.2]) def test_prepsfmom_gauss( pad_factor, image_size, psf_image_size, fwhm, psf_fwhm, pixel_scale, snr, mom_fwhm, - cls, + cls, extra_psf_fwhm, ): """fast test at a range of parameters to check that things come out ok""" rng = np.random.RandomState(seed=100) @@ -380,6 +381,27 @@ def test_prepsfmom_gauss( wcs=gs_wcs ).array + if extra_psf_fwhm is not None: + extra_psf = galsim.Gaussian( + fwhm=extra_psf_fwhm + ).shear( + g1=-0.2, g2=0.3 + ) + extra_psf_shift = rng.uniform(low=-scale/2, high=scale/2, size=2) + extra_psf_xy = gs_wcs.toImage(galsim.PositionD(extra_psf_shift)) + extra_psf_im = extra_psf.shift( + dx=extra_psf_shift[0], dy=extra_psf_shift[1] + ).drawImage( + nx=53, + ny=53, + wcs=gs_wcs + ).array + extra_psf_jac = Jacobian( + y=26 + extra_psf_xy.y, x=26 + extra_psf_xy.x, + dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, + dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy, + ) + fitter = cls( fwhm=mom_fwhm, pad_factor=pad_factor, @@ -404,14 +426,29 @@ def test_prepsfmom_gauss( res = [] for _ in range(100): _im = im + rng.normal(size=im.shape, scale=noise) - obs = Observation( - image=_im, - weight=wgt, - jacobian=jac, - psf=Observation(image=psf_im, jacobian=psf_jac), - ) + if extra_psf_fwhm is None: + obs = Observation( + image=_im, + weight=wgt, + jacobian=jac, + psf=Observation(image=psf_im, jacobian=psf_jac), + ) + + _res = fitter.go(obs=obs) + else: + obs = Observation( + image=_im, + weight=wgt, + jacobian=jac, + psf=Observation(image=extra_psf_im, jacobian=extra_psf_jac), + ) + + _res = fitter.go( + obs=obs, + extra_deconv_psfs=[Observation(image=psf_im, jacobian=psf_jac)], + extra_conv_psfs=[obs.psf], + ) - _res = fitter.go(obs=obs) if _res['flags'] == 0: res.append(_res) @@ -452,8 +489,10 @@ def test_prepsfmom_gauss( @pytest.mark.parametrize('pad_factor', [ 1.5, ]) +@pytest.mark.parametrize("extra_psf_fwhm", [None, 0.8, 1.2]) def test_prepsfmom_mn_cov( pad_factor, image_size, fwhm, psf_fwhm, pixel_scale, snr, mom_fwhm, cls, + extra_psf_fwhm, ): """Slower test to make sure means and errors are right w/ tons of monte carlo samples. @@ -509,6 +548,27 @@ def test_prepsfmom_mn_cov( wcs=gs_wcs ).array + if extra_psf_fwhm is not None: + extra_psf = galsim.Gaussian( + fwhm=extra_psf_fwhm + ).shear( + g1=-0.2, g2=0.3 + ) + extra_psf_shift = rng.uniform(low=-scale/2, high=scale/2, size=2) + extra_psf_xy = gs_wcs.toImage(galsim.PositionD(extra_psf_shift)) + extra_psf_im = extra_psf.shift( + dx=extra_psf_shift[0], dy=extra_psf_shift[1] + ).drawImage( + nx=53, + ny=53, + wcs=gs_wcs + ).array + extra_psf_jac = Jacobian( + y=26 + extra_psf_xy.y, x=26 + extra_psf_xy.x, + dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, + dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy, + ) + fitter = cls( fwhm=mom_fwhm, pad_factor=pad_factor, @@ -533,14 +593,29 @@ def test_prepsfmom_mn_cov( res = [] for _ in range(10_000): _im = im + rng.normal(size=im.shape, scale=noise) - obs = Observation( - image=_im, - weight=wgt, - jacobian=jac, - psf=Observation(image=psf_im, jacobian=psf_jac), - ) + if extra_psf_fwhm is None: + obs = Observation( + image=_im, + weight=wgt, + jacobian=jac, + psf=Observation(image=psf_im, jacobian=psf_jac), + ) + + _res = fitter.go(obs=obs) + else: + obs = Observation( + image=_im, + weight=wgt, + jacobian=jac, + psf=Observation(image=extra_psf_im, jacobian=extra_psf_jac), + ) + + _res = fitter.go( + obs=obs, + extra_deconv_psfs=[Observation(image=psf_im, jacobian=psf_jac)], + extra_conv_psfs=[obs.psf], + ) - _res = fitter.go(obs=obs) if _res['flags'] == 0: res.append(_res) From 808708a1638b92c555ce6b2144c896d51933fa52 Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Sat, 19 Mar 2022 08:00:01 -0500 Subject: [PATCH 14/19] REF combine a bunch of repeated code --- ngmix/tests/test_prepsfmom.py | 556 ++++++++++++++-------------------- 1 file changed, 233 insertions(+), 323 deletions(-) diff --git a/ngmix/tests/test_prepsfmom.py b/ngmix/tests/test_prepsfmom.py index f1922947..0ca84f51 100644 --- a/ngmix/tests/test_prepsfmom.py +++ b/ngmix/tests/test_prepsfmom.py @@ -16,6 +16,38 @@ import ngmix.flags +def _stack_list_of_dicts(res): + def _get_dtype(v): + if isinstance(v, float): + return ('f8',) + elif isinstance(v, int): + return ('i4',) + elif isinstance(v, str): + return ('U256',) + elif hasattr(v, "dtype") and hasattr(v, "shape"): + if "float" in str(v.dtype): + dstr = "f8" + else: + dstr = "i8" + + if len(v.shape) == 1: + return (dstr, v.shape[0]) + else: + return (dstr, v.shape) + else: + raise RuntimeError("cannot interpret dtype of '%s'" % v) + + dtype = [] + for k, v in res[0].items(): + dtype.append((k,) + _get_dtype(v)) + d = np.zeros(len(res), dtype=dtype) + for i in range(len(res)): + for k, v in res[i].items(): + d[k][i] = v + + return d + + def _report_info(s, arr, mn, err): if mn is not None and err is not None: print( @@ -33,6 +65,161 @@ def _report_info(s, arr, mn, err): ) +def _make_prepsfmom_sim( + *, image_size, psf_image_size, pixel_scale, rng, fwhm, psf_fwhm, + snr, extra_psf_fwhm +): + cen = (image_size - 1)/2 + psf_cen = (psf_image_size - 1)/2 + gs_wcs = galsim.ShearWCS( + pixel_scale, galsim.Shear(g1=-0.1, g2=0.06)).jacobian() + scale = np.sqrt(gs_wcs.pixelArea()) + shift = rng.uniform(low=-scale/2, high=scale/2, size=2) + psf_shift = rng.uniform(low=-scale/2, high=scale/2, size=2) + xy = gs_wcs.toImage(galsim.PositionD(shift)) + psf_xy = gs_wcs.toImage(galsim.PositionD(psf_shift)) + + jac = Jacobian( + y=cen + xy.y, x=cen + xy.x, + dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, + dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy) + + psf_jac = Jacobian( + y=psf_cen + psf_xy.y, x=psf_cen + psf_xy.x, + dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, + dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy) + + gal = galsim.Gaussian( + fwhm=fwhm + ).shear( + g1=-0.1, g2=0.2 + ).withFlux( + 400 + ).shift( + dx=shift[0], dy=shift[1] + ) + psf = galsim.Gaussian( + fwhm=psf_fwhm + ).shear( + g1=0.3, g2=-0.15 + ) + im = galsim.Convolve([gal, psf]).drawImage( + nx=image_size, + ny=image_size, + wcs=gs_wcs + ).array + noise = np.sqrt(np.sum(im**2)) / snr + wgt = np.ones_like(im) / noise**2 + + psf_im = psf.shift( + dx=psf_shift[0], dy=psf_shift[1] + ).drawImage( + nx=psf_image_size, + ny=psf_image_size, + wcs=gs_wcs + ).array + + if extra_psf_fwhm is not None: + extra_psf = galsim.Gaussian( + fwhm=extra_psf_fwhm + ).shear( + g1=-0.2, g2=0.3 + ) + extra_psf_shift = rng.uniform(low=-scale/2, high=scale/2, size=2) + extra_psf_xy = gs_wcs.toImage(galsim.PositionD(extra_psf_shift)) + extra_psf_im = extra_psf.shift( + dx=extra_psf_shift[0], dy=extra_psf_shift[1] + ).drawImage( + nx=image_size, + ny=image_size, + wcs=gs_wcs + ).array + extra_psf_jac = Jacobian( + y=cen + extra_psf_xy.y, x=cen + extra_psf_xy.x, + dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, + dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy, + ) + else: + extra_psf_im = None + extra_psf_jac = None + + im_true = gal.drawImage( + nx=image_size, + ny=image_size, + wcs=gs_wcs, + method='no_pixel').array + + return dict( + im=im, + im_true=im_true, + jac=jac, + psf_jac=psf_jac, + psf_im=psf_im, + noise=noise, + wgt=wgt, + extra_psf_im=extra_psf_im, + extra_psf_jac=extra_psf_jac, + ) + + +def _run_prepsfmom_sims(sdata, fitter, rng, nitr): + # get true flux + obs = Observation( + image=sdata["im_true"], + jacobian=sdata["jac"], + ) + res = fitter.go(obs=obs, no_psf=True) + flux_true = res["flux"] + T_true = res["T"] + g1_true = res["e"][0] + g2_true = res["e"][1] + + res = [] + for _ in range(nitr): + _im = sdata["im"] + rng.normal(size=sdata["im"].shape, scale=sdata["noise"]) + if sdata["extra_psf_fwhm"] is None: + obs = Observation( + image=_im, + weight=sdata["wgt"], + jacobian=sdata["jac"], + psf=Observation(image=sdata["psf_im"], jacobian=sdata["psf_jac"]), + ) + + _res = fitter.go(obs=obs) + else: + obs = Observation( + image=_im, + weight=sdata["wgt"], + jacobian=sdata["jac"], + psf=Observation( + image=sdata["extra_psf_im"], + jacobian=sdata["extra_psf_jac"], + ), + ) + + _res = fitter.go( + obs=obs, + extra_deconv_psfs=[Observation( + image=sdata["psf_im"], + jacobian=sdata["psf_jac"] + )], + extra_conv_psfs=[obs.psf], + ) + + if _res['flags'] == 0: + res.append(_res) + + res = _stack_list_of_dicts(res) + + return dict( + flux_true=flux_true, + T_true=T_true, + g1_true=g1_true, + g2_true=g2_true, + res=res, + ) + + def test_prepsfmom_kind(): fitter = PrePSFMom(2.0, 'gauss') assert fitter.kind == 'pgauss' @@ -160,65 +347,20 @@ def test_prepsfmom_raises_badjacob_extra_conv(cls): def test_prepsfmom_speed_and_cache(): - image_size = 48 - psf_image_size = 53 - pixel_scale = 0.263 - fwhm = 0.9 - psf_fwhm = 0.9 - snr = 20 mom_fwhm = 2 rng = np.random.RandomState(seed=100) - cen = (image_size - 1)/2 - psf_cen = (psf_image_size - 1)/2 - gs_wcs = galsim.ShearWCS( - pixel_scale, galsim.Shear(g1=-0.1, g2=0.06)).jacobian() - scale = np.sqrt(gs_wcs.pixelArea()) - shift = rng.uniform(low=-scale/2, high=scale/2, size=2) - psf_shift = rng.uniform(low=-scale/2, high=scale/2, size=2) - xy = gs_wcs.toImage(galsim.PositionD(shift)) - psf_xy = gs_wcs.toImage(galsim.PositionD(psf_shift)) - - jac = Jacobian( - y=cen + xy.y, x=cen + xy.x, - dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, - dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy) - - psf_jac = Jacobian( - y=psf_cen + psf_xy.y, x=psf_cen + psf_xy.x, - dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, - dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy) - - gal = galsim.Gaussian( - fwhm=fwhm - ).shear( - g1=-0.1, g2=0.2 - ).withFlux( - 400 - ).shift( - dx=shift[0], dy=shift[1] + sdata = _make_prepsfmom_sim( + image_size=48, + psf_image_size=53, + pixel_scale=0.263, + fwhm=0.9, + psf_fwhm=0.9, + snr=20, + rng=rng, + extra_psf_fwhm=None, ) - psf = galsim.Gaussian( - fwhm=psf_fwhm - ).shear( - g1=0.3, g2=-0.15 - ) - im = galsim.Convolve([gal, psf]).drawImage( - nx=image_size, - ny=image_size, - wcs=gs_wcs - ).array - noise = np.sqrt(np.sum(im**2)) / snr - wgt = np.ones_like(im) / noise**2 - - psf_im = psf.shift( - dx=psf_shift[0], dy=psf_shift[1] - ).drawImage( - nx=psf_image_size, - ny=psf_image_size, - wcs=gs_wcs - ).array # now we test the speed + caching _gauss_kernels.cache_clear() @@ -230,12 +372,12 @@ def test_prepsfmom_speed_and_cache(): fwhm=mom_fwhm + 1e-3, ) - im += rng.normal(size=im.shape, scale=noise) + im = sdata["im"] + rng.normal(size=sdata["im"].shape, scale=sdata["noise"]) obs = Observation( image=im, - weight=wgt, - jacobian=jac, - psf=Observation(image=psf_im + 1e-8, jacobian=psf_jac), + weight=sdata["wgt"], + jacobian=sdata["jac"], + psf=Observation(image=sdata["psf_im"] + 1e-8, jacobian=sdata["psf_jac"]), ) dt = time.time() @@ -252,12 +394,11 @@ def test_prepsfmom_speed_and_cache(): fwhm=mom_fwhm, ) - im += rng.normal(size=im.shape, scale=noise) obs = Observation( image=im, - weight=wgt, - jacobian=jac, - psf=Observation(image=psf_im, jacobian=psf_jac), + weight=sdata["wgt"], + jacobian=sdata["jac"], + psf=Observation(image=sdata["psf_im"], jacobian=sdata["psf_jac"]), ) dt = time.time() @@ -283,38 +424,6 @@ def test_prepsfmom_speed_and_cache(): assert _zero_pad_and_compute_fft_cached_impl.cache_info().misses == 2 -def _stack_list_of_dicts(res): - def _get_dtype(v): - if isinstance(v, float): - return ('f8',) - elif isinstance(v, int): - return ('i4',) - elif isinstance(v, str): - return ('U256',) - elif hasattr(v, "dtype") and hasattr(v, "shape"): - if "float" in str(v.dtype): - dstr = "f8" - else: - dstr = "i8" - - if len(v.shape) == 1: - return (dstr, v.shape[0]) - else: - return (dstr, v.shape) - else: - raise RuntimeError("cannot interpret dtype of '%s'" % v) - - dtype = [] - for k, v in res[0].items(): - dtype.append((k,) + _get_dtype(v)) - d = np.zeros(len(res), dtype=dtype) - for i in range(len(res)): - for k, v in res[i].items(): - d[k][i] = v - - return d - - @pytest.mark.parametrize("cls", [KSigmaMom, PGaussMom]) @pytest.mark.parametrize('snr', [1e1, 1e3]) @pytest.mark.parametrize('pixel_scale', [0.125, 0.25]) @@ -331,128 +440,28 @@ def test_prepsfmom_gauss( """fast test at a range of parameters to check that things come out ok""" rng = np.random.RandomState(seed=100) - cen = (image_size - 1)/2 - psf_cen = (psf_image_size - 1)/2 - gs_wcs = galsim.ShearWCS( - pixel_scale, galsim.Shear(g1=-0.1, g2=0.06)).jacobian() - scale = np.sqrt(gs_wcs.pixelArea()) - shift = rng.uniform(low=-scale/2, high=scale/2, size=2) - psf_shift = rng.uniform(low=-scale/2, high=scale/2, size=2) - xy = gs_wcs.toImage(galsim.PositionD(shift)) - psf_xy = gs_wcs.toImage(galsim.PositionD(psf_shift)) - - jac = Jacobian( - y=cen + xy.y, x=cen + xy.x, - dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, - dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy) - - psf_jac = Jacobian( - y=psf_cen + psf_xy.y, x=psf_cen + psf_xy.x, - dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, - dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy) - - gal = galsim.Gaussian( - fwhm=fwhm - ).shear( - g1=-0.1, g2=0.2 - ).withFlux( - 400 - ).shift( - dx=shift[0], dy=shift[1] + sdata = _make_prepsfmom_sim( + image_size=image_size, + psf_image_size=psf_image_size, + pixel_scale=pixel_scale, + fwhm=fwhm, + psf_fwhm=psf_fwhm, + snr=snr, + rng=rng, + extra_psf_fwhm=extra_psf_fwhm, ) - psf = galsim.Gaussian( - fwhm=psf_fwhm - ).shear( - g1=0.3, g2=-0.15 - ) - im = galsim.Convolve([gal, psf]).drawImage( - nx=image_size, - ny=image_size, - wcs=gs_wcs - ).array - noise = np.sqrt(np.sum(im**2)) / snr - wgt = np.ones_like(im) / noise**2 - - psf_im = psf.shift( - dx=psf_shift[0], dy=psf_shift[1] - ).drawImage( - nx=psf_image_size, - ny=psf_image_size, - wcs=gs_wcs - ).array - - if extra_psf_fwhm is not None: - extra_psf = galsim.Gaussian( - fwhm=extra_psf_fwhm - ).shear( - g1=-0.2, g2=0.3 - ) - extra_psf_shift = rng.uniform(low=-scale/2, high=scale/2, size=2) - extra_psf_xy = gs_wcs.toImage(galsim.PositionD(extra_psf_shift)) - extra_psf_im = extra_psf.shift( - dx=extra_psf_shift[0], dy=extra_psf_shift[1] - ).drawImage( - nx=53, - ny=53, - wcs=gs_wcs - ).array - extra_psf_jac = Jacobian( - y=26 + extra_psf_xy.y, x=26 + extra_psf_xy.x, - dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, - dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy, - ) fitter = cls( fwhm=mom_fwhm, pad_factor=pad_factor, ) - # get true flux - im_true = gal.drawImage( - nx=image_size, - ny=image_size, - wcs=gs_wcs, - method='no_pixel').array - obs = Observation( - image=im_true, - jacobian=jac, - ) - res = cls(fwhm=mom_fwhm, pad_factor=pad_factor).go(obs=obs, no_psf=True) - flux_true = res["flux"] - T_true = res["T"] - g1_true = res["e"][0] - g2_true = res["e"][1] - - res = [] - for _ in range(100): - _im = im + rng.normal(size=im.shape, scale=noise) - if extra_psf_fwhm is None: - obs = Observation( - image=_im, - weight=wgt, - jacobian=jac, - psf=Observation(image=psf_im, jacobian=psf_jac), - ) - - _res = fitter.go(obs=obs) - else: - obs = Observation( - image=_im, - weight=wgt, - jacobian=jac, - psf=Observation(image=extra_psf_im, jacobian=extra_psf_jac), - ) - - _res = fitter.go( - obs=obs, - extra_deconv_psfs=[Observation(image=psf_im, jacobian=psf_jac)], - extra_conv_psfs=[obs.psf], - ) - - if _res['flags'] == 0: - res.append(_res) - - res = _stack_list_of_dicts(res) + sres = _run_prepsfmom_sims(sdata, fitter, rng, 100) + flux_true = sres["flux_true"] + T_true = sres["T_true"] + g1_true = sres["g1_true"] + g2_true = sres["g2_true"] + res = sres["res"] if np.mean(res["flux"])/np.mean(res["flux_err"]) > 7: print("\n") @@ -499,127 +508,28 @@ def test_prepsfmom_mn_cov( """ rng = np.random.RandomState(seed=100) - cen = (image_size - 1)/2 - gs_wcs = galsim.ShearWCS( - pixel_scale, galsim.Shear(g1=-0.1, g2=0.06)).jacobian() - scale = np.sqrt(gs_wcs.pixelArea()) - shift = rng.uniform(low=-scale/2, high=scale/2, size=2) - psf_shift = rng.uniform(low=-scale/2, high=scale/2, size=2) - xy = gs_wcs.toImage(galsim.PositionD(shift)) - psf_xy = gs_wcs.toImage(galsim.PositionD(psf_shift)) - - jac = Jacobian( - y=cen + xy.y, x=cen + xy.x, - dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, - dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy) - - psf_jac = Jacobian( - y=26 + psf_xy.y, x=26 + psf_xy.x, - dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, - dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy) - - gal = galsim.Gaussian( - fwhm=fwhm - ).shear( - g1=-0.1, g2=0.2 - ).withFlux( - 400 - ).shift( - dx=shift[0], dy=shift[1] - ) - psf = galsim.Gaussian( - fwhm=psf_fwhm - ).shear( - g1=0.3, g2=-0.15 + sdata = _make_prepsfmom_sim( + image_size=image_size, + psf_image_size=53, + pixel_scale=pixel_scale, + fwhm=fwhm, + psf_fwhm=psf_fwhm, + snr=snr, + rng=rng, + extra_psf_fwhm=extra_psf_fwhm, ) - im = galsim.Convolve([gal, psf]).drawImage( - nx=image_size, - ny=image_size, - wcs=gs_wcs - ).array - noise = np.sqrt(np.sum(im**2)) / snr - wgt = np.ones_like(im) / noise**2 - - psf_im = psf.shift( - dx=psf_shift[0], dy=psf_shift[1] - ).drawImage( - nx=53, - ny=53, - wcs=gs_wcs - ).array - - if extra_psf_fwhm is not None: - extra_psf = galsim.Gaussian( - fwhm=extra_psf_fwhm - ).shear( - g1=-0.2, g2=0.3 - ) - extra_psf_shift = rng.uniform(low=-scale/2, high=scale/2, size=2) - extra_psf_xy = gs_wcs.toImage(galsim.PositionD(extra_psf_shift)) - extra_psf_im = extra_psf.shift( - dx=extra_psf_shift[0], dy=extra_psf_shift[1] - ).drawImage( - nx=53, - ny=53, - wcs=gs_wcs - ).array - extra_psf_jac = Jacobian( - y=26 + extra_psf_xy.y, x=26 + extra_psf_xy.x, - dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, - dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy, - ) fitter = cls( fwhm=mom_fwhm, pad_factor=pad_factor, ) - # get true flux - im_true = gal.drawImage( - nx=image_size, - ny=image_size, - wcs=gs_wcs, - method='no_pixel').array - obs = Observation( - image=im_true, - jacobian=jac, - ) - res = cls(fwhm=mom_fwhm, pad_factor=pad_factor).go(obs=obs, no_psf=True) - flux_true = res["flux"] - T_true = res["T"] - g1_true = res["e"][0] - g2_true = res["e"][1] - - res = [] - for _ in range(10_000): - _im = im + rng.normal(size=im.shape, scale=noise) - if extra_psf_fwhm is None: - obs = Observation( - image=_im, - weight=wgt, - jacobian=jac, - psf=Observation(image=psf_im, jacobian=psf_jac), - ) - - _res = fitter.go(obs=obs) - else: - obs = Observation( - image=_im, - weight=wgt, - jacobian=jac, - psf=Observation(image=extra_psf_im, jacobian=extra_psf_jac), - ) - - _res = fitter.go( - obs=obs, - extra_deconv_psfs=[Observation(image=psf_im, jacobian=psf_jac)], - extra_conv_psfs=[obs.psf], - ) - - if _res['flags'] == 0: - res.append(_res) - - res = _stack_list_of_dicts(res) + sres = _run_prepsfmom_sims(sdata, fitter, rng, 10_000) + flux_true = sres["flux_true"] + T_true = sres["T_true"] + g1_true = sres["g1_true"] + g2_true = sres["g2_true"] + res = sres["res"] print("\n") _report_info("snr", np.mean(res["flux"])/np.mean(res["flux_err"]), None, None) From e8987366536590cf09c2871355056efd62af2552 Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Sat, 19 Mar 2022 08:10:10 -0500 Subject: [PATCH 15/19] REF combine a bit more --- ngmix/tests/test_prepsfmom.py | 149 +++++++++++++++------------------- 1 file changed, 67 insertions(+), 82 deletions(-) diff --git a/ngmix/tests/test_prepsfmom.py b/ngmix/tests/test_prepsfmom.py index 0ca84f51..10666e9f 100644 --- a/ngmix/tests/test_prepsfmom.py +++ b/ngmix/tests/test_prepsfmom.py @@ -98,26 +98,37 @@ def _make_prepsfmom_sim( ).shift( dx=shift[0], dy=shift[1] ) - psf = galsim.Gaussian( - fwhm=psf_fwhm - ).shear( - g1=0.3, g2=-0.15 - ) - im = galsim.Convolve([gal, psf]).drawImage( - nx=image_size, - ny=image_size, - wcs=gs_wcs - ).array + if psf_fwhm is not None: + psf = galsim.Gaussian( + fwhm=psf_fwhm + ).shear( + g1=0.3, g2=-0.15 + ) + im = galsim.Convolve([gal, psf]).drawImage( + nx=image_size, + ny=image_size, + wcs=gs_wcs + ).array + else: + im = gal.drawImage( + nx=image_size, + ny=image_size, + wcs=gs_wcs + ).array + noise = np.sqrt(np.sum(im**2)) / snr wgt = np.ones_like(im) / noise**2 - psf_im = psf.shift( - dx=psf_shift[0], dy=psf_shift[1] - ).drawImage( - nx=psf_image_size, - ny=psf_image_size, - wcs=gs_wcs - ).array + if psf_fwhm is not None: + psf_im = psf.shift( + dx=psf_shift[0], dy=psf_shift[1] + ).drawImage( + nx=psf_image_size, + ny=psf_image_size, + wcs=gs_wcs + ).array + else: + psf_im = None if extra_psf_fwhm is not None: extra_psf = galsim.Gaussian( @@ -143,11 +154,19 @@ def _make_prepsfmom_sim( extra_psf_im = None extra_psf_jac = None - im_true = gal.drawImage( - nx=image_size, - ny=image_size, - wcs=gs_wcs, - method='no_pixel').array + if psf_fwhm is not None: + # true image has pixel removed if we deconvolve the PSF + im_true = gal.drawImage( + nx=image_size, + ny=image_size, + wcs=gs_wcs, + method='no_pixel').array + else: + im_true = gal.drawImage( + nx=image_size, + ny=image_size, + wcs=gs_wcs + ).array return dict( im=im, @@ -174,18 +193,26 @@ def _run_prepsfmom_sims(sdata, fitter, rng, nitr): g1_true = res["e"][0] g2_true = res["e"][1] + no_psf = sdata["psf_im"] is None + if sdata["extra_psf_fwhm"] is not None: + assert not no_psf + res = [] for _ in range(nitr): _im = sdata["im"] + rng.normal(size=sdata["im"].shape, scale=sdata["noise"]) - if sdata["extra_psf_fwhm"] is None: + if sdata["extra_psf_im"] is None: obs = Observation( image=_im, weight=sdata["wgt"], jacobian=sdata["jac"], - psf=Observation(image=sdata["psf_im"], jacobian=sdata["psf_jac"]), + psf=( + Observation(image=sdata["psf_im"], jacobian=sdata["psf_jac"]) + if not no_psf + else None + ), ) - _res = fitter.go(obs=obs) + _res = fitter.go(obs=obs, no_psf=no_psf) else: obs = Observation( image=_im, @@ -591,70 +618,28 @@ def test_prepsfmom_mn_cov_nopsf( """ rng = np.random.RandomState(seed=100) - cen = (image_size - 1)/2 - gs_wcs = galsim.ShearWCS( - pixel_scale, galsim.Shear(g1=-0.1, g2=0.06)).jacobian() - scale = np.sqrt(gs_wcs.pixelArea()) - shift = rng.uniform(low=-scale/2, high=scale/2, size=2) - xy = gs_wcs.toImage(galsim.PositionD(shift)) - - jac = Jacobian( - y=cen + xy.y, x=cen + xy.x, - dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, - dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy) - - gal = galsim.Gaussian( - fwhm=fwhm - ).shear( - g1=-0.1, g2=0.2 - ).withFlux( - 400 - ).shift( - dx=shift[0], dy=shift[1] + sdata = _make_prepsfmom_sim( + image_size=image_size, + psf_image_size=53, + pixel_scale=pixel_scale, + fwhm=fwhm, + psf_fwhm=None, + snr=snr, + rng=rng, + extra_psf_fwhm=None, ) - im = gal.drawImage( - nx=image_size, - ny=image_size, - wcs=gs_wcs - ).array - noise = np.sqrt(np.sum(im**2)) / snr - wgt = np.ones_like(im) / noise**2 fitter = cls( fwhm=mom_fwhm, pad_factor=pad_factor, ) - # get true flux - im_true = gal.drawImage( - nx=image_size, - ny=image_size, - wcs=gs_wcs, - ).array - obs = Observation( - image=im_true, - jacobian=jac, - ) - res = cls(fwhm=mom_fwhm, pad_factor=pad_factor).go(obs=obs, no_psf=True) - flux_true = res["flux"] - T_true = res["T"] - g1_true = res["e"][0] - g2_true = res["e"][1] - - res = [] - for _ in range(10_000): - _im = im + rng.normal(size=im.shape, scale=noise) - obs = Observation( - image=_im, - weight=wgt, - jacobian=jac, - ) - - _res = fitter.go(obs=obs, no_psf=True) - if _res['flags'] == 0: - res.append(_res) - - res = _stack_list_of_dicts(res) + sres = _run_prepsfmom_sims(sdata, fitter, rng, 10_000) + flux_true = sres["flux_true"] + T_true = sres["T_true"] + g1_true = sres["g1_true"] + g2_true = sres["g2_true"] + res = sres["res"] print("\n") _report_info("snr", np.mean(res["flux"])/np.mean(res["flux_err"]), None, None) From 5db830d35e628c288297d3d5b5328766e3da3881 Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Sat, 19 Mar 2022 08:17:45 -0500 Subject: [PATCH 16/19] REF combine a biut more --- ngmix/tests/test_prepsfmom.py | 121 +++++++++------------------------- 1 file changed, 31 insertions(+), 90 deletions(-) diff --git a/ngmix/tests/test_prepsfmom.py b/ngmix/tests/test_prepsfmom.py index 10666e9f..e6b9df68 100644 --- a/ngmix/tests/test_prepsfmom.py +++ b/ngmix/tests/test_prepsfmom.py @@ -752,80 +752,19 @@ def test_prepsfmom_gauss_true_flux_T( cls, extra_psf_fwhm ): rng = np.random.RandomState(seed=100) - snr = 1e8 mom_fwhm = 15.0 - cen = (image_size - 1)/2 - psf_cen = (psf_image_size - 1)/2 - gs_wcs = galsim.ShearWCS( - pixel_scale, galsim.Shear(g1=-0.1, g2=0.06)).jacobian() - scale = np.sqrt(gs_wcs.pixelArea()) - shift = rng.uniform(low=-scale/2, high=scale/2, size=2) - psf_shift = rng.uniform(low=-scale/2, high=scale/2, size=2) - xy = gs_wcs.toImage(galsim.PositionD(shift)) - psf_xy = gs_wcs.toImage(galsim.PositionD(psf_shift)) - - jac = Jacobian( - y=cen + xy.y, x=cen + xy.x, - dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, - dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy) - - psf_jac = Jacobian( - y=psf_cen + psf_xy.y, x=psf_cen + psf_xy.x, - dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, - dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy) - - gal = galsim.Gaussian( - fwhm=fwhm - ).shear( - g1=-0.1, g2=0.2 - ).withFlux( - 400 - ).shift( - dx=shift[0], dy=shift[1] - ) - psf = galsim.Gaussian( - fwhm=psf_fwhm - ).shear( - g1=0.3, g2=-0.15 + sdata = _make_prepsfmom_sim( + image_size=image_size, + psf_image_size=psf_image_size, + pixel_scale=pixel_scale, + fwhm=fwhm, + psf_fwhm=psf_fwhm, + snr=snr, + rng=rng, + extra_psf_fwhm=extra_psf_fwhm, ) - im = galsim.Convolve([gal, psf]).drawImage( - nx=image_size, - ny=image_size, - wcs=gs_wcs - ).array - noise = np.sqrt(np.sum(im**2)) / snr - wgt = np.ones_like(im) / noise**2 - - psf_im = psf.shift( - dx=psf_shift[0], dy=psf_shift[1] - ).drawImage( - nx=psf_image_size, - ny=psf_image_size, - wcs=gs_wcs - ).array - - if extra_psf_fwhm is not None: - extra_psf = galsim.Gaussian( - fwhm=extra_psf_fwhm - ).shear( - g1=-0.2, g2=0.3 - ) - extra_psf_shift = rng.uniform(low=-scale/2, high=scale/2, size=2) - extra_psf_xy = gs_wcs.toImage(galsim.PositionD(extra_psf_shift)) - extra_psf_im = extra_psf.shift( - dx=extra_psf_shift[0], dy=extra_psf_shift[1] - ).drawImage( - nx=psf_image_size, - ny=psf_image_size, - wcs=gs_wcs - ).array - extra_psf_jac = Jacobian( - y=psf_cen + extra_psf_xy.y, x=psf_cen + extra_psf_xy.x, - dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, - dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy, - ) fitter = cls( fwhm=mom_fwhm, @@ -833,14 +772,9 @@ def test_prepsfmom_gauss_true_flux_T( ) # get true flux - im_true = gal.drawImage( - nx=image_size, - ny=image_size, - wcs=gs_wcs, - method='no_pixel').array obs = Observation( - image=im_true, - jacobian=jac, + image=sdata["im_true"], + jacobian=sdata["jac"], ) res = fitter.go(obs=obs, no_psf=True) flux_true = res["flux"] @@ -849,10 +783,10 @@ def test_prepsfmom_gauss_true_flux_T( if extra_psf_fwhm is None: obs = Observation( - image=im, - weight=wgt, - jacobian=jac, - psf=Observation(image=psf_im, jacobian=psf_jac), + image=sdata["im"], + weight=sdata["wgt"], + jacobian=sdata["jac"], + psf=Observation(image=sdata["psf_im"], jacobian=sdata["psf_jac"]), ) res = fitter.go(obs=obs) flux_true = res["flux"] @@ -861,10 +795,12 @@ def test_prepsfmom_gauss_true_flux_T( else: # this should fail since it is the wrong PSF obs = Observation( - image=im, - weight=wgt, - jacobian=jac, - psf=Observation(image=extra_psf_im, jacobian=extra_psf_jac), + image=sdata["im"], + weight=sdata["wgt"], + jacobian=sdata["jac"], + psf=Observation( + image=sdata["extra_psf_im"], jacobian=sdata["extra_psf_jac"] + ), ) res = fitter.go(obs=obs) flux_true = res["flux"] @@ -873,14 +809,19 @@ def test_prepsfmom_gauss_true_flux_T( # this should work since we have the PSF correction in there obs = Observation( - image=im, - weight=wgt, - jacobian=jac, - psf=Observation(image=extra_psf_im, jacobian=extra_psf_jac), + image=sdata["im"], + weight=sdata["wgt"], + jacobian=sdata["jac"], + psf=Observation( + image=sdata["extra_psf_im"], + jacobian=sdata["extra_psf_jac"] + ), ) res = fitter.go( obs=obs, - extra_deconv_psfs=[Observation(image=psf_im, jacobian=psf_jac)], + extra_deconv_psfs=[Observation( + image=sdata["psf_im"], jacobian=sdata["psf_jac"], + )], extra_conv_psfs=[obs.psf], ) flux_true = res["flux"] From afa3bd8421b833624a6b1acbb32384219e9f25b4 Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Sat, 19 Mar 2022 08:27:40 -0500 Subject: [PATCH 17/19] ENH test the shape too --- ngmix/tests/test_prepsfmom.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/ngmix/tests/test_prepsfmom.py b/ngmix/tests/test_prepsfmom.py index e6b9df68..3053fa89 100644 --- a/ngmix/tests/test_prepsfmom.py +++ b/ngmix/tests/test_prepsfmom.py @@ -751,6 +751,8 @@ def test_prepsfmom_gauss_true_flux_T( pad_factor, psf_image_size, image_size, fwhm, psf_fwhm, pixel_scale, cls, extra_psf_fwhm ): + """This test ensures the kernels are normalize properly so + we have the correct total flux, etc.""" rng = np.random.RandomState(seed=100) snr = 1e8 mom_fwhm = 15.0 @@ -779,6 +781,8 @@ def test_prepsfmom_gauss_true_flux_T( res = fitter.go(obs=obs, no_psf=True) flux_true = res["flux"] T_true = res["T"] + g1_true = res["e1"] + g2_true = res["e2"] assert np.allclose(flux_true, 400, atol=0, rtol=5e-3) if extra_psf_fwhm is None: @@ -792,6 +796,8 @@ def test_prepsfmom_gauss_true_flux_T( flux_true = res["flux"] assert np.allclose(flux_true, 400, atol=0, rtol=5e-3) assert np.allclose(res["T"], T_true, atol=0, rtol=5e-4) + assert np.allclose(res["e1"], g1_true, atol=5e-3, rtol=5e-3) + assert np.allclose(res["e2"], g2_true, atol=5e-3, rtol=5e-3) else: # this should fail since it is the wrong PSF obs = Observation( @@ -806,6 +812,8 @@ def test_prepsfmom_gauss_true_flux_T( flux_true = res["flux"] # we use a big relative difference since the T should be way off assert not np.allclose(res["T"], T_true, atol=0, rtol=1e-1) + assert not np.allclose(res["e1"], g1_true, atol=0, rtol=1e-1) + assert not np.allclose(res["e2"], g2_true, atol=0, rtol=1e-1) # this should work since we have the PSF correction in there obs = Observation( @@ -827,6 +835,8 @@ def test_prepsfmom_gauss_true_flux_T( flux_true = res["flux"] assert np.allclose(flux_true, 400, atol=0, rtol=5e-3) assert np.allclose(res["T"], T_true, atol=0, rtol=5e-4) + assert np.allclose(res["e1"], g1_true, atol=5e-3, rtol=5e-3) + assert np.allclose(res["e2"], g2_true, atol=5e-3, rtol=5e-3) @pytest.mark.parametrize('pixel_scale', [0.25, 0.125]) From e16eb081d1715ff9133f0f396f339990748a16a5 Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Sat, 19 Mar 2022 08:28:51 -0500 Subject: [PATCH 18/19] BUG wrong key --- ngmix/tests/test_prepsfmom.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ngmix/tests/test_prepsfmom.py b/ngmix/tests/test_prepsfmom.py index 3053fa89..0e2d96f8 100644 --- a/ngmix/tests/test_prepsfmom.py +++ b/ngmix/tests/test_prepsfmom.py @@ -194,7 +194,7 @@ def _run_prepsfmom_sims(sdata, fitter, rng, nitr): g2_true = res["e"][1] no_psf = sdata["psf_im"] is None - if sdata["extra_psf_fwhm"] is not None: + if sdata["extra_psf_im"] is not None: assert not no_psf res = [] From 08064133d63a2ad282eccbe31636baf69ac8797f Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Sat, 9 Apr 2022 07:16:15 -0500 Subject: [PATCH 19/19] cache metacal; allow psf corr for PSF only --- ngmix/metacal/metacal.py | 34 +++++++++++++++++++++++----------- ngmix/prepsfmom.py | 28 +++++++++++----------------- 2 files changed, 34 insertions(+), 28 deletions(-) diff --git a/ngmix/metacal/metacal.py b/ngmix/metacal/metacal.py index a79f3d74..9838fb9f 100644 --- a/ngmix/metacal/metacal.py +++ b/ngmix/metacal/metacal.py @@ -6,6 +6,8 @@ """ import copy import logging +from functools import lru_cache + import numpy as np from ..gexceptions import GMixRangeError, BootPSFFailure from ..shape import Shape @@ -20,6 +22,14 @@ logger = logging.getLogger(__name__) +@lru_cache(maxsize=128) +def _cached_galsim_stuff(img, wcs_repr, xinterp): + import galsim + image = galsim.Image(np.array(img), wcs=eval(wcs_repr)) + image_int = galsim.InterpolatedImage(image, x_interpolant=xinterp) + return image, image_int + + class MetacalDilatePSF(object): """ Create manipulated images for use in metacalibration @@ -355,22 +365,24 @@ def _set_data(self): # these would share data with the original numpy arrays, make copies # to be sure they don't get modified # - self.image = galsim.Image(obs.image.copy(), - wcs=self.get_wcs()) - - self.psf_image = galsim.Image(obs.psf.image.copy(), - wcs=self.get_psf_wcs()) + image, image_int = _cached_galsim_stuff( + tuple(tuple(ii) for ii in obs.image.copy()), + repr(self.get_wcs()), + self.interp, + ) + self.image = image + self.image_int = image_int - # interpolated psf image - psf_int = galsim.InterpolatedImage(self.psf_image, - x_interpolant=self.interp) + psf_image, psf_int = _cached_galsim_stuff( + tuple(tuple(ii) for ii in obs.psf.image.copy()), + repr(self.get_psf_wcs()), + self.interp, + ) + self.psf_image = psf_image # this can be used to deconvolve the psf from the galaxy image psf_int_inv = galsim.Deconvolve(psf_int) - self.image_int = galsim.InterpolatedImage(self.image, - x_interpolant=self.interp) - # deconvolved galaxy image, psf+pixel removed self.image_int_nopsf = galsim.Convolve(self.image_int, psf_int_inv) diff --git a/ngmix/prepsfmom.py b/ngmix/prepsfmom.py index cce624d8..5614abe3 100644 --- a/ngmix/prepsfmom.py +++ b/ngmix/prepsfmom.py @@ -712,26 +712,20 @@ def _check_obs_and_get_psf_obs( if not obs.has_psf() and not no_psf: raise RuntimeError("The PSF must be set to measure a pre-PSF moment!") - if no_psf and (extra_conv_psfs or extra_deconv_psfs): - raise RuntimeError( - "You can only use extra conv. or deconv. PSFs with observations " - "that have a PSF!" - ) - if not no_psf: conv_psfs = extra_conv_psfs or [] deconv_psfs = [obs.get_psf()] + (extra_deconv_psfs or []) - - if any( - not _jacobian_close(psf_obs.jacobian, obs.jacobian) - for psf_obs in conv_psfs + deconv_psfs - ): - raise RuntimeError( - "The PSF and observation must have the same WCS " - "Jacobian for measuring pre-PSF moments." - ) else: - conv_psfs = [] - deconv_psfs = [] + conv_psfs = extra_conv_psfs or [] + deconv_psfs = extra_deconv_psfs or [] + + if any( + not _jacobian_close(psf_obs.jacobian, obs.jacobian) + for psf_obs in conv_psfs + deconv_psfs + ): + raise RuntimeError( + "The PSF and observation must have the same WCS " + "Jacobian for measuring pre-PSF moments." + ) return conv_psfs, deconv_psfs