Skip to content

Commit

Permalink
[pre-commit.ci] auto fixes from pre-commit.com hooks
Browse files Browse the repository at this point in the history 
for more information, see https://pre-commit.ci
  • Loading branch information
@pre-commit-ci @lgarrison
pre-commit-ci[bot] authored and lgarrison committed Feb 3, 2025
1 parent 2fc1c47 commit 6947516
Show file tree
Hide file tree
Showing 10 changed files with 153 additions and 157 deletions.
2 changes: 1 addition & 1 deletion abacusnbody/analysis/tsc.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ def tsc_parallel(
if npartition > n1d // 3 and npartition != n1d // 2 and nthread > 1:
raise ValueError(
f'npartition {npartition} must be less than'
f' ngrid//3 = {n1d//3} or equal to ngrid//2 = {n1d//2}'
f' ngrid//3 = {n1d // 3} or equal to ngrid//2 = {n1d // 2}'
)
if npartition > 1 and npartition % 2 != 0 and nthread > 1:
raise ValueError(f'npartition {npartition} not divisible by 2')
Expand Down
4 changes: 2 additions & 2 deletions abacusnbody/data/pipe_asdf.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -155,11 +155,11 @@ def unpack_to_pipe(
tot_time = timer() - start_time
if verbose:
print(
f'[pipe_asdf.py] Read + decompressed {nbytes_tot/1e6:.3g} MB in {read_time:.3g} s at {nbytes_tot/1e6/read_time:.3g} MB/s',
f'[pipe_asdf.py] Read + decompressed {nbytes_tot / 1e6:.3g} MB in {read_time:.3g} s at {nbytes_tot / 1e6 / read_time:.3g} MB/s',
file=sys.stderr,
)
print(
f'[pipe_asdf.py] Processed {nbytes_tot/1e6:.3g} MB in {tot_time:.3g} s at {nbytes_tot/1e6/tot_time:.3g} MB/s',
f'[pipe_asdf.py] Processed {nbytes_tot / 1e6:.3g} MB in {tot_time:.3g} s at {nbytes_tot / 1e6 / tot_time:.3g} MB/s',
file=sys.stderr,
)

Expand Down
2 changes: 1 addition & 1 deletion abacusnbody/data/read_abacus.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,7 @@ def read_asdf(fn, load=None, colname=None, dtype=np.float32, verbose=True, **kwa
header['SubsampleFraction'] = SubsampleFraction
if verbose:
print(
f'Loading "{basename(fn)}", which contains the A and B subsamples ({int(SubsampleFraction*100):d}% total)'
f'Loading "{basename(fn)}", which contains the A and B subsamples ({int(SubsampleFraction * 100):d}% total)'
)

table = Table(meta=header)
Expand Down
2 changes: 1 addition & 1 deletion abacusnbody/hod/GRAND_HOD.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1713,7 +1713,7 @@ def gen_gal_cat(
else:
ascii.write(
table,
outdir / (f"{tracer}s_chunk{params['chunk']:d}.dat"),
outdir / (f'{tracer}s_chunk{params["chunk"]:d}.dat'),
overwrite=True,
format='ecsv',
)
Expand Down
164 changes: 80 additions & 84 deletions abacusnbody/hod/abacus_hod.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -185,9 +185,9 @@ def __init__(
# setting up chunking
self.chunk = chunk
self.n_chunks = n_chunks
assert (
self.chunk < self.n_chunks
), 'Total number of chunks needs to be larger than current chunk index'
assert self.chunk < self.n_chunks, (
'Total number of chunks needs to be larger than current chunk index'
)

if not skip_staging:
# load the subsample particles
Expand Down Expand Up @@ -1400,28 +1400,26 @@ def apply_zcv(self, mock_dict, config, load_presaved=False):

# compute real space and redshift space
# assert config['HOD_params']['want_rsd'], "Currently want_rsd=False not implemented"
assert (
len(mock_dict.keys()) == 1
), (
assert len(mock_dict.keys()) == 1, (
'Currently implemented only a single tracer'
) # should make a dict of dicts, but need cross
assert (
len(config['power_params']['poles']) <= 3
), 'Currently implemented only multipoles 0, 2, 4; need to change ZeNBu'
assert (
config['power_params']['nbins_mu'] == 1
), 'Currently wedges are not implemented; need to change ZeNBu'
assert len(config['power_params']['poles']) <= 3, (
'Currently implemented only multipoles 0, 2, 4; need to change ZeNBu'
)
assert config['power_params']['nbins_mu'] == 1, (
'Currently wedges are not implemented; need to change ZeNBu'
)
if 'nmesh' not in config['power_params'].keys():
config['power_params']['nmesh'] = config['zcv_params']['nmesh']
assert (
config['zcv_params']['nmesh'] == config['power_params']['nmesh']
), '`nmesh` in `power_params` and `zcv_params` should match.'
assert config['zcv_params']['nmesh'] == config['power_params']['nmesh'], (
'`nmesh` in `power_params` and `zcv_params` should match.'
)

# create save directory
save_dir = (
Path(config['zcv_params']['zcv_dir']) / config['sim_params']['sim_name']
)
save_z_dir = save_dir / f"z{config['sim_params']['z_mock']:.3f}"
save_z_dir = save_dir / f'z{config["sim_params"]["z_mock"]:.3f}'
rsd_str = '_rsd' if config['HOD_params']['want_rsd'] else ''

# define bins
Expand All @@ -1444,34 +1442,34 @@ def apply_zcv(self, mock_dict, config, load_presaved=False):
if config['power_params']['nbins_k'] == config['zcv_params']['nmesh'] // 2:
power_rsd_tr_fn = (
save_z_dir
/ f"power{rsd_str}_tr_nmesh{config['zcv_params']['nmesh']}.asdf"
/ f'power{rsd_str}_tr_nmesh{config["zcv_params"]["nmesh"]}.asdf'
)
power_rsd_ij_fn = (
save_z_dir
/ f"power{rsd_str}_ij_nmesh{config['zcv_params']['nmesh']}.asdf"
/ f'power{rsd_str}_ij_nmesh{config["zcv_params"]["nmesh"]}.asdf'
)
power_tr_fn = (
save_z_dir / f"power_tr_nmesh{config['zcv_params']['nmesh']}.asdf"
save_z_dir / f'power_tr_nmesh{config["zcv_params"]["nmesh"]}.asdf'
)
power_ij_fn = (
save_z_dir / f"power_ij_nmesh{config['zcv_params']['nmesh']}.asdf"
save_z_dir / f'power_ij_nmesh{config["zcv_params"]["nmesh"]}.asdf'
)
else:
power_rsd_tr_fn = (
save_z_dir
/ f"power{rsd_str}_tr_nmesh{config['zcv_params']['nmesh']}_dk{dk:.3f}.asdf"
/ f'power{rsd_str}_tr_nmesh{config["zcv_params"]["nmesh"]}_dk{dk:.3f}.asdf'
)
power_rsd_ij_fn = (
save_z_dir
/ f"power{rsd_str}_ij_nmesh{config['zcv_params']['nmesh']}_dk{dk:.3f}.asdf"
/ f'power{rsd_str}_ij_nmesh{config["zcv_params"]["nmesh"]}_dk{dk:.3f}.asdf'
)
power_tr_fn = (
save_z_dir
/ f"power_tr_nmesh{config['zcv_params']['nmesh']}_dk{dk:.3f}.asdf"
/ f'power_tr_nmesh{config["zcv_params"]["nmesh"]}_dk{dk:.3f}.asdf'
)
power_ij_fn = (
save_z_dir
/ f"power_ij_nmesh{config['zcv_params']['nmesh']}_dk{dk:.3f}.asdf"
/ f'power_ij_nmesh{config["zcv_params"]["nmesh"]}_dk{dk:.3f}.asdf'
)
pk_fns = [power_rsd_tr_fn, power_rsd_ij_fn, power_tr_fn, power_ij_fn]
for fn in pk_fns:
Expand All @@ -1485,33 +1483,33 @@ def apply_zcv(self, mock_dict, config, load_presaved=False):
if load_presaved:
pk_rsd_tr_dict = asdf.open(power_rsd_tr_fn)['data']
pk_rsd_ij_dict = asdf.open(power_rsd_ij_fn)['data']
assert np.allclose(
k_binc, pk_rsd_tr_dict['k_binc']
), f'Mismatching file: {str(power_rsd_tr_fn)}'
assert np.allclose(
k_binc, pk_rsd_ij_dict['k_binc']
), f'Mismatching file: {str(power_rsd_ij_fn)}'
assert np.allclose(
mu_binc, pk_rsd_tr_dict['mu_binc']
), f'Mismatching file: {str(power_rsd_tr_fn)}'
assert np.allclose(
mu_binc, pk_rsd_ij_dict['mu_binc']
), f'Mismatching file: {str(power_rsd_ij_fn)}'
assert np.allclose(k_binc, pk_rsd_tr_dict['k_binc']), (
f'Mismatching file: {str(power_rsd_tr_fn)}'
)
assert np.allclose(k_binc, pk_rsd_ij_dict['k_binc']), (
f'Mismatching file: {str(power_rsd_ij_fn)}'
)
assert np.allclose(mu_binc, pk_rsd_tr_dict['mu_binc']), (
f'Mismatching file: {str(power_rsd_tr_fn)}'
)
assert np.allclose(mu_binc, pk_rsd_ij_dict['mu_binc']), (
f'Mismatching file: {str(power_rsd_ij_fn)}'
)
if config['HOD_params']['want_rsd']:
pk_tr_dict = asdf.open(power_tr_fn)['data']
pk_ij_dict = asdf.open(power_ij_fn)['data']
assert np.allclose(
k_binc, pk_tr_dict['k_binc']
), f'Mismatching file: {str(power_tr_fn)}'
assert np.allclose(
k_binc, pk_ij_dict['k_binc']
), f'Mismatching file: {str(power_ij_fn)}'
assert np.allclose(
mu_binc, pk_tr_dict['mu_binc']
), f'Mismatching file: {str(power_tr_fn)}'
assert np.allclose(
mu_binc, pk_ij_dict['mu_binc']
), f'Mismatching file: {str(power_ij_fn)}'
assert np.allclose(k_binc, pk_tr_dict['k_binc']), (
f'Mismatching file: {str(power_tr_fn)}'
)
assert np.allclose(k_binc, pk_ij_dict['k_binc']), (
f'Mismatching file: {str(power_ij_fn)}'
)
assert np.allclose(mu_binc, pk_tr_dict['mu_binc']), (
f'Mismatching file: {str(power_tr_fn)}'
)
assert np.allclose(mu_binc, pk_ij_dict['mu_binc']), (
f'Mismatching file: {str(power_ij_fn)}'
)
else:
pk_tr_dict, pk_ij_dict = None, None

Expand All @@ -1532,12 +1530,12 @@ def apply_zcv(self, mock_dict, config, load_presaved=False):
tracer_pos, config['HOD_params']['want_rsd'], config
)
pk_rsd_ij_dict = asdf.open(power_rsd_ij_fn)['data']
assert np.allclose(
k_binc, pk_rsd_ij_dict['k_binc']
), f'Mismatching file: {str(power_rsd_ij_fn)}'
assert np.allclose(
mu_binc, pk_rsd_ij_dict['mu_binc']
), f'Mismatching file: {str(power_rsd_ij_fn)}'
assert np.allclose(k_binc, pk_rsd_ij_dict['k_binc']), (
f'Mismatching file: {str(power_rsd_ij_fn)}'
)
assert np.allclose(mu_binc, pk_rsd_ij_dict['mu_binc']), (
f'Mismatching file: {str(power_rsd_ij_fn)}'
)
# run version without rsd if rsd was requested
if config['HOD_params']['want_rsd']:
mock_dict = self.run_hod(
Expand All @@ -1564,12 +1562,12 @@ def apply_zcv(self, mock_dict, config, load_presaved=False):
tracer_pos, want_rsd=False, config=config
)
pk_ij_dict = asdf.open(power_ij_fn)['data']
assert np.allclose(
k_binc, pk_ij_dict['k_binc']
), f'Mismatching file: {str(power_ij_fn)}'
assert np.allclose(
mu_binc, pk_ij_dict['mu_binc']
), f'Mismatching file: {str(power_ij_fn)}'
assert np.allclose(k_binc, pk_ij_dict['k_binc']), (
f'Mismatching file: {str(power_ij_fn)}'
)
assert np.allclose(mu_binc, pk_ij_dict['mu_binc']), (
f'Mismatching file: {str(power_ij_fn)}'
)
else:
pk_tr_dict, pk_ij_dict = None, None

Expand All @@ -1590,31 +1588,29 @@ def apply_zcv_xi(self, mock_dict, config, load_presaved=False):
from ..analysis.power_spectrum import pk_to_xi

# compute real space and redshift space
assert config['HOD_params'][
'want_rsd'
], 'Currently want_rsd=False not implemented'
assert (
len(mock_dict.keys()) == 1
), (
assert config['HOD_params']['want_rsd'], (
'Currently want_rsd=False not implemented'
)
assert len(mock_dict.keys()) == 1, (
'Currently implemented only a single tracer'
) # should make a dict of dicts, but need cross
assert (
len(config['power_params']['poles']) <= 3
), 'Currently implemented only multipoles 0, 2, 4; need to change ZeNBu'
assert (
config['power_params']['nbins_mu'] == 1
), 'Currently wedges are not implemented; need to change ZeNBu'
assert len(config['power_params']['poles']) <= 3, (
'Currently implemented only multipoles 0, 2, 4; need to change ZeNBu'
)
assert config['power_params']['nbins_mu'] == 1, (
'Currently wedges are not implemented; need to change ZeNBu'
)
if 'nmesh' not in config['power_params'].keys():
config['power_params']['nmesh'] = config['zcv_params']['nmesh']
assert (
config['zcv_params']['nmesh'] == config['power_params']['nmesh']
), '`nmesh` in `power_params` and `zcv_params` should match.'
assert config['zcv_params']['nmesh'] == config['power_params']['nmesh'], (
'`nmesh` in `power_params` and `zcv_params` should match.'
)

# create save directory
save_dir = (
Path(config['zcv_params']['zcv_dir']) / config['sim_params']['sim_name']
)
save_z_dir = save_dir / f"z{config['sim_params']['z_mock']:.3f}"
save_z_dir = save_dir / f'z{config["sim_params"]["z_mock"]:.3f}'
rsd_str = '_rsd' if config['HOD_params']['want_rsd'] else ''

# construct names of files based on fields
Expand All @@ -1627,30 +1623,30 @@ def apply_zcv_xi(self, mock_dict, config, load_presaved=False):
pk_ij_fns = []
pk_rsd_tr_fns.append(
save_z_dir
/ f"power{rsd_str}_tr_tr_nmesh{config['zcv_params']['nmesh']:d}.asdf"
/ f'power{rsd_str}_tr_tr_nmesh{config["zcv_params"]["nmesh"]:d}.asdf'
)
pk_tr_fns.append(
save_z_dir / f"power_tr_tr_nmesh{config['zcv_params']['nmesh']:d}.asdf"
save_z_dir / f'power_tr_tr_nmesh{config["zcv_params"]["nmesh"]:d}.asdf'
)
for i in range(len(keynames)):
pk_rsd_tr_fns.append(
save_z_dir
/ f"power{rsd_str}_{keynames[i]}_tr_nmesh{config['zcv_params']['nmesh']:d}.asdf"
/ f'power{rsd_str}_{keynames[i]}_tr_nmesh{config["zcv_params"]["nmesh"]:d}.asdf'
)
pk_tr_fns.append(
save_z_dir
/ f"power_{keynames[i]}_tr_nmesh{config['zcv_params']['nmesh']:d}.asdf"
/ f'power_{keynames[i]}_tr_nmesh{config["zcv_params"]["nmesh"]:d}.asdf'
)
for j in range(len(keynames)):
if i < j:
continue
pk_rsd_ij_fns.append(
save_z_dir
/ f"power{rsd_str}_{keynames[i]}_{keynames[j]}_nmesh{config['zcv_params']['nmesh']:d}.asdf"
/ f'power{rsd_str}_{keynames[i]}_{keynames[j]}_nmesh{config["zcv_params"]["nmesh"]:d}.asdf'
)
pk_ij_fns.append(
save_z_dir
/ f"power_{keynames[i]}_{keynames[j]}_nmesh{config['zcv_params']['nmesh']:d}.asdf"
/ f'power_{keynames[i]}_{keynames[j]}_nmesh{config["zcv_params"]["nmesh"]:d}.asdf'
)

if not load_presaved:
Expand Down Expand Up @@ -1719,11 +1715,11 @@ def apply_zcv_xi(self, mock_dict, config, load_presaved=False):
r_bins = np.linspace(0.0, 200.0, 201)
pk_rsd_tr_fns = [
save_z_dir
/ f"power{rsd_str}_tr_tr_nmesh{config['zcv_params']['nmesh']:d}.asdf"
/ f'power{rsd_str}_tr_tr_nmesh{config["zcv_params"]["nmesh"]:d}.asdf'
] # TODO: same as other (could check that we have this if presaved)
power_cv_tr_fn = (
save_z_dir
/ f"power{rsd_str}_ZCV_tr_nmesh{config['zcv_params']['nmesh']:d}.asdf"
/ f'power{rsd_str}_ZCV_tr_nmesh{config["zcv_params"]["nmesh"]:d}.asdf'
) # TODO: should be an output (could check that we have this if presaved; run_zcv too)
r_binc, binned_poles_zcv, Npoles = pk_to_xi(
asdf.open(power_cv_tr_fn)['data']['P_k3D_tr_tr_zcv'],
Expand Down
54 changes: 27 additions & 27 deletions abacusnbody/hod/zcv/advect_fields.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -183,27 +183,27 @@ def main(
f = asdf.open(fields_fft_fn[i])
fields_fft.append(f['data'])
header = f['header']
assert (
header['sim_name'] == sim_name
), f'Mismatch in the files: {str(fields_fft_fn[i])}'
assert np.isclose(
header['Lbox'], Lbox
), f'Mismatch in the file: {str(fields_fft_fn[i])}'
assert (
header['nmesh'] == nmesh
), f'Mismatch in the file: {str(fields_fft_fn[i])}'
assert np.isclose(
header['kcut'], kcut
), f'Mismatch in the file: {str(fields_fft_fn[i])}'
assert (
header['compensated'] == compensated
), f'Mismatch in the file: {str(fields_fft_fn[i])}'
assert (
header['interlaced'] == interlaced
), f'Mismatch in the file: {str(fields_fft_fn[i])}'
assert (
header['paste'] == paste
), f'Mismatch in the file: {str(fields_fft_fn[i])}'
assert header['sim_name'] == sim_name, (
f'Mismatch in the files: {str(fields_fft_fn[i])}'
)
assert np.isclose(header['Lbox'], Lbox), (
f'Mismatch in the file: {str(fields_fft_fn[i])}'
)
assert header['nmesh'] == nmesh, (
f'Mismatch in the file: {str(fields_fft_fn[i])}'
)
assert np.isclose(header['kcut'], kcut), (
f'Mismatch in the file: {str(fields_fft_fn[i])}'
)
assert header['compensated'] == compensated, (
f'Mismatch in the file: {str(fields_fft_fn[i])}'
)
assert header['interlaced'] == interlaced, (
f'Mismatch in the file: {str(fields_fft_fn[i])}'
)
assert header['paste'] == paste, (
f'Mismatch in the file: {str(fields_fft_fn[i])}'
)
else:
# load density field and displacements
f = asdf.open(ic_fn)
Expand Down Expand Up @@ -249,12 +249,12 @@ def main(
else:
f = asdf.open(fields_fn)
header = f['header']
assert (
header['nmesh'] == nmesh
), 'Mismatch in the file: {str(fields_fn)}'
assert np.isclose(
header['kcut'], kcut
), 'Mismatch in the file: {str(fields_fn)}'
assert header['nmesh'] == nmesh, (
'Mismatch in the file: {str(fields_fn)}'
)
assert np.isclose(header['kcut'], kcut), (
'Mismatch in the file: {str(fields_fn)}'
)
w = f['data'][keynames[i]][:, :, :].flatten()
f.close()
field_fft = get_field_fft(
Expand Down
Loading

0 comments on commit 6947516

Please sign in to comment.