From 08204d300a504c3d3a9f0fa6153ddeb757b6216b Mon Sep 17 00:00:00 2001 From: beckermr Date: Wed, 24 Aug 2022 13:53:20 -0500 Subject: [PATCH 1/8] PERF add O(N) phase shift computation --- ngmix/prepsfmom.py | 46 +++++++++++++++++++++++++++++++++++++++------- 1 file changed, 39 insertions(+), 7 deletions(-) diff --git a/ngmix/prepsfmom.py b/ngmix/prepsfmom.py index 4604e1f7..fc5cdb6d 100644 --- a/ngmix/prepsfmom.py +++ b/ngmix/prepsfmom.py @@ -276,6 +276,23 @@ def _measure_moments_fft( cen_phase = _compute_cen_phase_shift(drow, dcol, dim, msk=msk) kim *= cen_phase + # we only sum where the kernel is nonzero + fkf = kernels["fkf"] + fkr = kernels["fkr"] + fkp = kernels["fkp"] + fkc = kernels["fkc"] + + mom_norm = kernels["fk00"] + + return _measure_moments_fft_numba( + kim, kpsf_im, dim, eff_pad_factor, fkf, fkr, fkp, fkc, mom_norm, tot_var, + ) + + +@njit +def _measure_moments_fft_numba( + kim, kpsf_im, dim, eff_pad_factor, fkf, fkr, fkp, fkc, mom_norm, tot_var, +): # build the flux, radial, plus and cross kernels / moments # the inverse FFT in our convention has a factor of 1/n per dimension # the sums below are inverse FFTs but only computing the values at the @@ -285,13 +302,6 @@ def _measure_moments_fft( df2 = df * df df4 = df2 * df2 - # we only sum where the kernel is nonzero - fkf = kernels["fkf"] - fkr = kernels["fkr"] - fkp = kernels["fkp"] - fkc = kernels["fkc"] - - mom_norm = kernels["fk00"] mf = np.sum((kim * fkf).real) * df2 mr = np.sum((kim * fkr).real) * df2 mp = np.sum((kim * fkp).real) * df2 @@ -389,6 +399,28 @@ def _zero_pad_image(im, target_dim): def _compute_cen_phase_shift(cen_row, cen_col, dim, msk=None): """computes exp(i*2*pi*k*cen) for shifting the phases of FFTS. + If you feed the centroid of a profile, then this factor times the raw FFT + of that profile will result in an FFT centered at the profile. + """ + f = fft.fftfreq(dim) * (2.0 * np.pi) + # this reshaping makes sure the arrays broadcast nicely into a grid + fx = f.reshape(1, -1) + fy = f.reshape(-1, 1) + kcen_x = fx*cen_col + kcen_y = fy*cen_row + px = np.cos(kcen_x) + 1j*np.sin(kcen_x) + py = np.cos(kcen_y) + 1j*np.sin(kcen_y) + + pxy = px * py + if msk is not None: + pxy = pxy[msk] + + return pxy + + +def _compute_cen_phase_shift_orig(cen_row, cen_col, dim, msk=None): + """computes exp(i*2*pi*k*cen) for shifting the phases of FFTS. + If you feed the centroid of a profile, then this factor times the raw FFT of that profile will result in an FFT centered at the profile. """ From ea1b51a38a4796553d625365ec7979f4aba13731 Mon Sep 17 00:00:00 2001 From: beckermr Date: Wed, 24 Aug 2022 14:51:31 -0500 Subject: [PATCH 2/8] PERF more numba --- ngmix/prepsfmom.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/ngmix/prepsfmom.py b/ngmix/prepsfmom.py index fc5cdb6d..107087f6 100644 --- a/ngmix/prepsfmom.py +++ b/ngmix/prepsfmom.py @@ -403,6 +403,16 @@ def _compute_cen_phase_shift(cen_row, cen_col, dim, msk=None): of that profile will result in an FFT centered at the profile. """ f = fft.fftfreq(dim) * (2.0 * np.pi) + pxy = _compute_cen_phase_shift_numba(f, cen_row, cen_col) + + if msk is not None: + pxy = pxy[msk] + + return pxy + + +@njit +def _compute_cen_phase_shift_numba(f, cen_row, cen_col): # this reshaping makes sure the arrays broadcast nicely into a grid fx = f.reshape(1, -1) fy = f.reshape(-1, 1) @@ -410,11 +420,7 @@ def _compute_cen_phase_shift(cen_row, cen_col, dim, msk=None): kcen_y = fy*cen_row px = np.cos(kcen_x) + 1j*np.sin(kcen_x) py = np.cos(kcen_y) + 1j*np.sin(kcen_y) - pxy = px * py - if msk is not None: - pxy = pxy[msk] - return pxy From f7812028425f55c226d3166c0c0c7b98f308c05b Mon Sep 17 00:00:00 2001 From: beckermr Date: Wed, 24 Aug 2022 15:03:13 -0500 Subject: [PATCH 3/8] TST run more trials to get a failure --- ngmix/tests/test_leastsqbound.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ngmix/tests/test_leastsqbound.py b/ngmix/tests/test_leastsqbound.py index f4a5f53e..7ddbe79a 100644 --- a/ngmix/tests/test_leastsqbound.py +++ b/ngmix/tests/test_leastsqbound.py @@ -115,7 +115,7 @@ def test_leastsqbound_smoke(use_prior): def test_leastsqbound_bounds(fracdev_bounds): rng = np.random.RandomState(2830) - ntrial = 10 + ntrial = 100 fit_model = 'bd' scale = 0.263 From 89aec4310b0cbbb7f37992899d50878cfdd77d2a Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Wed, 24 Aug 2022 15:08:29 -0500 Subject: [PATCH 4/8] Update ngmix/prepsfmom.py --- ngmix/prepsfmom.py | 1 - 1 file changed, 1 deletion(-) diff --git a/ngmix/prepsfmom.py b/ngmix/prepsfmom.py index 107087f6..664b0a8c 100644 --- a/ngmix/prepsfmom.py +++ b/ngmix/prepsfmom.py @@ -276,7 +276,6 @@ def _measure_moments_fft( cen_phase = _compute_cen_phase_shift(drow, dcol, dim, msk=msk) kim *= cen_phase - # we only sum where the kernel is nonzero fkf = kernels["fkf"] fkr = kernels["fkr"] fkp = kernels["fkp"] From 88b695c7928f4c674d2cf91a331965594c98ad53 Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Wed, 24 Aug 2022 15:10:23 -0500 Subject: [PATCH 5/8] Update ngmix/prepsfmom.py --- ngmix/prepsfmom.py | 1 + 1 file changed, 1 insertion(+) diff --git a/ngmix/prepsfmom.py b/ngmix/prepsfmom.py index 664b0a8c..788659e2 100644 --- a/ngmix/prepsfmom.py +++ b/ngmix/prepsfmom.py @@ -301,6 +301,7 @@ def _measure_moments_fft_numba( df2 = df * df df4 = df2 * df2 + # we only sum where the kernel is nonzero mf = np.sum((kim * fkf).real) * df2 mr = np.sum((kim * fkr).real) * df2 mp = np.sum((kim * fkp).real) * df2 From 1c7063e2d25c4ad421d7b9c7efcfe6a7b15a826a Mon Sep 17 00:00:00 2001 From: beckermr Date: Wed, 24 Aug 2022 15:23:16 -0500 Subject: [PATCH 6/8] PROD prep for version release --- CHANGES.md | 3 ++- ngmix/_version.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 971cbf30..29456b82 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -1,8 +1,9 @@ -## unreleased +## v2.2.0 ### new features - Added function to regularize moments results. + - Added numba to pre-PSF moments to reduce runtime. ## v2.1.0 diff --git a/ngmix/_version.py b/ngmix/_version.py index d1ba3084..8acd0c49 100644 --- a/ngmix/_version.py +++ b/ngmix/_version.py @@ -1 +1 @@ -__version__ = '2.1.0' # noqa +__version__ = '2.2.0' # noqa From 8a6fe225aaa4f5f46ad0a3e5aa49e4db4c8f9610 Mon Sep 17 00:00:00 2001 From: beckermr Date: Wed, 24 Aug 2022 15:40:10 -0500 Subject: [PATCH 7/8] TST added test for optimized function --- ngmix/tests/test_prepsfmom.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/ngmix/tests/test_prepsfmom.py b/ngmix/tests/test_prepsfmom.py index 7b4cdd77..49a9ccae 100644 --- a/ngmix/tests/test_prepsfmom.py +++ b/ngmix/tests/test_prepsfmom.py @@ -11,6 +11,8 @@ PrePSFMom, _gauss_kernels, _zero_pad_and_compute_fft_cached_impl, + _compute_cen_phase_shift, + _compute_cen_phase_shift_orig, ) from ngmix import Jacobian from ngmix import Observation @@ -18,6 +20,25 @@ import ngmix.flags +@pytest.mark.parametrize("row", [-0.4, 0, 1.2, 4.5]) +@pytest.mark.parametrize("col", [-0.32434, 0, 1.43232, 4.56775]) +@pytest.mark.parametrize("dim,msk", [ + (100, None), + (453, None), + (3, np.array([[True, False, True], [True, True, False], [True, True, True]])), + (4, np.array([ + [False, True, False, True], + [False, True, True, False], + [True, True, True, True], + [False, False, True, False]])), +]) +def test_cen_phase_shift(row, col, msk, dim): + np.testing.assert_allclose( + _compute_cen_phase_shift(row, col, dim, msk=msk), + _compute_cen_phase_shift_orig(row, col, dim, msk=msk) + ) + + def _report_info(s, arr, mn, err): if mn is not None and err is not None: print( From 8de09b70d9c4f5fa9d9149f3572f556b31cc3e9f Mon Sep 17 00:00:00 2001 From: beckermr Date: Wed, 24 Aug 2022 15:42:11 -0500 Subject: [PATCH 8/8] REF make this clearer --- ngmix/prepsfmom.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ngmix/prepsfmom.py b/ngmix/prepsfmom.py index 788659e2..7ff229cf 100644 --- a/ngmix/prepsfmom.py +++ b/ngmix/prepsfmom.py @@ -402,7 +402,7 @@ def _compute_cen_phase_shift(cen_row, cen_col, dim, msk=None): If you feed the centroid of a profile, then this factor times the raw FFT of that profile will result in an FFT centered at the profile. """ - f = fft.fftfreq(dim) * (2.0 * np.pi) + f = fft.fftfreq(dim) pxy = _compute_cen_phase_shift_numba(f, cen_row, cen_col) if msk is not None: @@ -416,8 +416,8 @@ def _compute_cen_phase_shift_numba(f, cen_row, cen_col): # this reshaping makes sure the arrays broadcast nicely into a grid fx = f.reshape(1, -1) fy = f.reshape(-1, 1) - kcen_x = fx*cen_col - kcen_y = fy*cen_row + kcen_x = fx * (2.0 * np.pi * cen_col) + kcen_y = fy * (2.0 * np.pi * cen_row) px = np.cos(kcen_x) + 1j*np.sin(kcen_x) py = np.cos(kcen_y) + 1j*np.sin(kcen_y) pxy = px * py