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 2994e9bd..5614abe3 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 @@ -51,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 @@ -64,23 +68,34 @@ 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.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] # pad image, psf and weight map, get FFTs, apply cen_phases @@ -90,18 +105,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( - 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 # @@ -128,17 +142,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( @@ -151,8 +165,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: @@ -225,19 +242,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 @@ -278,7 +303,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): @@ -382,22 +407,92 @@ def _zero_pad_and_compute_fft(im, cen_row, cen_col, target_dim, ap_rad): return kpim, pad_cen_row, pad_cen_col -def _deconvolve_im_psf_inplace(kim, kpsf_im, max_amp, min_psf_frac=1e-5): +# 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(tuple(ii) for ii in im), + float(cen_row), float(cen_col), int(target_dim), float(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 _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_col = 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) def _ksigma_kernels( dim, kernel_size, @@ -505,6 +600,7 @@ def _ksigma_kernels( ) +@functools.lru_cache(maxsize=128) def _gauss_kernels( dim, kernel_size, @@ -596,7 +692,16 @@ def _gauss_kernels( ) -def _check_obs_and_get_psf_obs(obs, no_psf): +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, +): if not isinstance(obs, Observation): raise ValueError("input obs must be an Observation") @@ -608,14 +713,19 @@ def _check_obs_and_get_psf_obs(obs, no_psf): raise RuntimeError("The PSF must be set to measure a pre-PSF moment!") if not no_psf: - psf_obs = obs.get_psf() - - if psf_obs.jacobian.get_galsim_wcs() != obs.jacobian.get_galsim_wcs(): - raise RuntimeError( - "The PSF and observation must have the same WCS " - "Jacobian for measuring pre-PSF moments." - ) + conv_psfs = extra_conv_psfs or [] + deconv_psfs = [obs.get_psf()] + (extra_deconv_psfs or []) else: - psf_obs = None + 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 psf_obs + return conv_psfs, deconv_psfs diff --git a/ngmix/tests/test_prepsfmom.py b/ngmix/tests/test_prepsfmom.py index d24a1ff6..0e2d96f8 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 @@ -13,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( @@ -30,6 +65,188 @@ 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] + ) + 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 + + 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( + 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 + + 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, + 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] + + no_psf = sdata["psf_im"] is None + if sdata["extra_psf_im"] 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_im"] is None: + obs = Observation( + image=_im, + weight=sdata["wgt"], + jacobian=sdata["jac"], + psf=( + Observation(image=sdata["psf_im"], jacobian=sdata["psf_jac"]) + if not no_psf + else None + ), + ) + + _res = fitter.go(obs=obs, no_psf=no_psf) + 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' @@ -57,6 +274,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) @@ -94,36 +328,127 @@ 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" +@pytest.mark.parametrize("cls", [KSigmaMom, PGaussMom]) +def test_prepsfmom_raises_badjacob_extra_conv(cls): + fitter = cls(1.2) - if len(v.shape) == 1: - return (dstr, v.shape[0]) - else: - return (dstr, v.shape) - else: - raise RuntimeError("cannot interpret dtype of '%s'" % v) + 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) - 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 + psf_jac = Jacobian( + y=0, x=0, + dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, + dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy*2) - return d + 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(): + mom_fwhm = 2 + + rng = np.random.RandomState(seed=100) + + 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, + ) + + # 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( + fwhm=mom_fwhm + 1e-3, + ) + + im = sdata["im"] + rng.normal(size=sdata["im"].shape, scale=sdata["noise"]) + obs = Observation( + image=im, + weight=sdata["wgt"], + jacobian=sdata["jac"], + psf=Observation(image=sdata["psf_im"] + 1e-8, jacobian=sdata["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 + fitter = PGaussMom( + fwhm=mom_fwhm, + ) + + obs = Observation( + image=im, + weight=sdata["wgt"], + jacobian=sdata["jac"], + psf=Observation(image=sdata["psf_im"], jacobian=sdata["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 + + # now we test with full caching + 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]) @@ -134,99 +459,36 @@ 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) - 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 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) - obs = Observation( - image=_im, - weight=wgt, - jacobian=jac, - psf=Observation(image=psf_im, jacobian=psf_jac), - ) - - _res = fitter.go(obs=obs) - 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") @@ -263,99 +525,38 @@ 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. """ 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] + 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, ) - 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=53, - ny=53, - wcs=gs_wcs - ).array 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) - obs = Observation( - image=_im, - weight=wgt, - jacobian=jac, - psf=Observation(image=psf_im, jacobian=psf_jac), - ) - - _res = fitter.go(obs=obs) - 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) @@ -417,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) @@ -581,69 +740,33 @@ 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)]) @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( - pad_factor, psf_image_size, image_size, fwhm, psf_fwhm, pixel_scale, cls +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 - 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 fitter = cls( fwhm=mom_fwhm, @@ -651,28 +774,69 @@ def test_prepsfmom_gauss_true_flux( ) # 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"] + T_true = res["T"] + g1_true = res["e1"] + g2_true = res["e2"] 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=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"] + 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( + 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"] + # 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( + 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=sdata["psf_im"], jacobian=sdata["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) + 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])