From 209eb45d64a191bdd80b5b8b1bce10b551af3758 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 2 Oct 2023 19:43:28 -0400 Subject: [PATCH 1/2] [pre-commit.ci] pre-commit autoupdate (#114) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit updates: - [github.com/astral-sh/ruff-pre-commit: v0.0.287 → v0.0.292](https://github.com/astral-sh/ruff-pre-commit/compare/v0.0.287...v0.0.292) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ab6b7e97..a4133ee7 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -13,7 +13,7 @@ exclude: | repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: "v0.0.287" + rev: "v0.0.292" hooks: - id: ruff # TODO: turning off line length check From a31ffd1e340c2ecd39e1561c510a1421e9112387 Mon Sep 17 00:00:00 2001 From: Lehman Garrison Date: Mon, 9 Oct 2023 12:20:14 -0400 Subject: [PATCH 2/2] power_spectrum: flesh out dtype and paste arg handling --- abacusnbody/analysis/power_spectrum.py | 40 +++++++++++++++----------- 1 file changed, 24 insertions(+), 16 deletions(-) diff --git a/abacusnbody/analysis/power_spectrum.py b/abacusnbody/analysis/power_spectrum.py index c44b73d1..fdd7c317 100644 --- a/abacusnbody/analysis/power_spectrum.py +++ b/abacusnbody/analysis/power_spectrum.py @@ -589,7 +589,7 @@ def calc_pk_from_deltak(field_fft, Lbox, k_bin_edges, mu_bin_edges, field2_fft=N return binned_pk, Nmode, binned_poles, N_mode_poles -def get_field(pos, Lbox, nmesh, paste, w=None, d=0., nthread=MAX_THREADS): +def get_field(pos, Lbox, nmesh, paste, w=None, d=0., nthread=MAX_THREADS, dtype=np.float32): r""" Construct real-space 3D field given particle positions. @@ -609,6 +609,8 @@ def get_field(pos, Lbox, nmesh, paste, w=None, d=0., nthread=MAX_THREADS): uniform shift to particle positions. nthread : int, optional Number of numba threads to use + dtype : np.dtype, optional + Data type of the field Returns ------- @@ -618,21 +620,19 @@ def get_field(pos, Lbox, nmesh, paste, w=None, d=0., nthread=MAX_THREADS): # check if weights are requested if w is not None: assert pos.shape[0] == len(w) - pos = pos.astype(np.float32, copy=False) + field = np.zeros((nmesh, nmesh, nmesh), dtype=dtype) paste = paste.upper() if paste == 'TSC': - if d != 0.: - field = tsc_parallel(pos, nmesh, Lbox, weights=w, nthread=nthread, offset=d) - else: - field = tsc_parallel(pos, nmesh, Lbox, weights=w, nthread=nthread) + tsc_parallel(pos, field, Lbox, weights=w, nthread=nthread, offset=d) elif paste == 'CIC': - field = np.zeros((nmesh, nmesh, nmesh), dtype=np.float32) warnings.warn("Note that currently CIC pasting, unlike TSC, supports only a non-parallel implementation.") if d != 0.: cic_serial(pos + d, field, Lbox, weights=w) else: cic_serial(pos, field, Lbox, weights=w) + else: + raise ValueError(f"Unknown pasting method: {paste}") if w is None: # in the zcv code the weights are already normalized, so don't normalize here # TODO assuming normalized weights is fragile # same as passing "Value" to nbodykit (1+delta)(x) V(x) @@ -777,7 +777,7 @@ def get_interlaced_field_fft(pos, Lbox, nmesh, paste, w, nthread=MAX_THREADS, ve return field_fft -def get_field_fft(pos, Lbox, nmesh, paste, w, W, compensated, interlaced, nthread=MAX_THREADS, verbose=False): +def get_field_fft(pos, Lbox, nmesh, paste, w, W, compensated, interlaced, nthread=MAX_THREADS, verbose=False, dtype=np.float32): r""" Calculate field from particle positions and return 3D Fourier field. @@ -799,6 +799,10 @@ def get_field_fft(pos, Lbox, nmesh, paste, w, W, compensated, interlaced, nthrea want to apply interlacing? nthread : int, optional Number of numba threads to use + verbose : bool, optional + Print out debugging info + dtype : np.dtype, optional + Data type of the field Returns ------- @@ -806,19 +810,17 @@ def get_field_fft(pos, Lbox, nmesh, paste, w, W, compensated, interlaced, nthrea interlaced 3D Fourier field. """ - # get field in real space - field = get_field(pos, Lbox, nmesh, paste, w, nthread=nthread) - if verbose: - print("field, pos", field.dtype, pos.dtype) if interlaced: # get interlaced field field_fft = get_interlaced_field_fft(pos, Lbox, nmesh, paste, w, nthread=nthread) else: # get field in real space - field = get_field(pos, Lbox, nmesh, paste, w, nthread=nthread) + field = get_field(pos, Lbox, nmesh, paste, w, nthread=nthread, dtype=dtype) + if verbose: + print("field, pos", field.dtype, pos.dtype) # get Fourier modes from skewers grid - inv_size = np.float32(1 / field.size) + inv_size = dtype(1 / field.size) field_fft = rfftn(field, overwrite_x=True, workers=nthread) _normalize(field_fft, inv_size, nthread=nthread) @@ -868,11 +870,14 @@ def get_W_compensated(Lbox, nmesh, paste, interlaced): k = (fftfreq(nmesh, d=d) * 2. * np.pi).astype(np.float32) # h/Mpc # apply deconvolution + paste = paste.upper() if interlaced: if paste == 'TSC': p = 3. elif paste == 'CIC': p = 2. + else: + raise ValueError(f"Unknown pasting method {paste}") W = np.sinc(0.5*k/kN)**p # sinc def else: # first order correction of interlacing (aka aliasing) s = np.sin(0.5 * np.pi * k/kN)**2 @@ -899,6 +904,7 @@ def calc_power(pos, w2 = None, poles = None, nthread = MAX_THREADS, + dtype = np.float32, ): r""" Compute the 3D power spectrum given particle positions by first painting them on a @@ -945,6 +951,8 @@ def calc_power(pos, Default of None gives the monopole. nthread : int, optional Number of numba threads to use + dtype : np.dtype, optional + Data type of the field Returns ------- @@ -987,12 +995,12 @@ def calc_power(pos, W = None # convert to fourier space - field_fft = get_field_fft(pos, Lbox, nmesh, paste, w, W, compensated, interlaced, nthread=nthread) + field_fft = get_field_fft(pos, Lbox, nmesh, paste, w, W, compensated, interlaced, nthread=nthread, dtype=dtype) # if second field provided if pos2 is not None: # convert to fourier space - field2_fft = get_field_fft(pos2, Lbox, nmesh, paste, w2, W, compensated, interlaced, nthread=nthread) + field2_fft = get_field_fft(pos2, Lbox, nmesh, paste, w2, W, compensated, interlaced, nthread=nthread, dtype=dtype) else: field2_fft = None