Moving environment to RS (#70)

* Drag several fields at once

* add env in RSP

* Add docstrings

* Make __main__

* Fix bug

* Fix plotting little bug
This commit is contained in:
Richard Stiskalek 2023-06-18 11:42:21 +01:00 committed by GitHub
parent 35ccfb5c67
commit 27e1c181a2
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 96 additions and 56 deletions

View file

@ -168,17 +168,17 @@ def divide_nonzero(field0, field1):
field0[i, j, k] /= field1[i, j, k] field0[i, j, k] /= field1[i, j, k]
def field2rsp(field, parts, box, nbatch=30, flip_partsxz=True, init_value=0., def field2rsp(*fields, parts, box, nbatch=30, flip_partsxz=True, init_value=0.,
verbose=True): verbose=True):
""" """
Forward model real space scalar field to redshift space. Attaches field Forward model real space scalar fields to redshift space. Attaches field
values to particles, those are then moved to redshift space and from their values to particles, those are then moved to redshift space and from their
positions reconstructs back the field on a grid by NGP interpolation. positions reconstructs back the field on a grid by NGP interpolation.
Parameters Parameters
---------- ----------
field : 3-dimensional array of shape `(grid, grid, grid)` fields : (list of) 3-dimensional array of shape `(grid, grid, grid)`
Real space field to be evolved to redshift space. Real space fields to be evolved to redshift space.
parts : 2-dimensional array of shape `(n_parts, 6)` parts : 2-dimensional array of shape `(n_parts, 6)`
Particle positions and velocities in real space. Must be organised as Particle positions and velocities in real space. Must be organised as
`x, y, z, vx, vy, vz`. `x, y, z, vx, vy, vz`.
@ -199,39 +199,50 @@ def field2rsp(field, parts, box, nbatch=30, flip_partsxz=True, init_value=0.,
------- -------
rsp_fields : (list of) 3-dimensional array of shape `(grid, grid, grid)` rsp_fields : (list of) 3-dimensional array of shape `(grid, grid, grid)`
""" """
rsp_field = numpy.full(field.shape, init_value, dtype=numpy.float32) nfields = len(fields)
cellcounts = numpy.zeros(rsp_field.shape, dtype=numpy.float32) # Check that all fields have the same shape.
# We iterate over the fields and in the inner loop over the particles. This if nfields > 1:
# is slower than iterating over the particles and in the inner loop over assert all(fields[0].shape == fields[i].shape
# the fields, but it is more memory efficient. Typically we will only have for i in range(1, nfields))
# one field.
rsp_fields = [numpy.full(field.shape, init_value, dtype=numpy.float32)
for field in fields]
cellcounts = numpy.zeros(rsp_fields[0].shape, dtype=numpy.float32)
nparts = parts.shape[0] nparts = parts.shape[0]
batch_size = nparts // nbatch batch_size = nparts // nbatch
start = 0 start = 0
for k in trange(nbatch + 1) if verbose else range(nbatch + 1): for __ in trange(nbatch + 1) if verbose else range(nbatch + 1):
# We first load the batch of particles into memory and flip x and z.
end = min(start + batch_size, nparts) end = min(start + batch_size, nparts)
pos = parts[start:end] pos = parts[start:end]
pos, vel = pos[:, :3], pos[:, 3:6] pos, vel = pos[:, :3], pos[:, 3:6]
if flip_partsxz: if flip_partsxz:
pos[:, [0, 2]] = pos[:, [2, 0]] pos[:, [0, 2]] = pos[:, [2, 0]]
vel[:, [0, 2]] = vel[:, [2, 0]] vel[:, [0, 2]] = vel[:, [2, 0]]
# Evaluate the field at the particle positions in real space. # Then move the particles to redshift space.
values = evaluate_cartesian(field, pos=pos)
# Move particles to redshift space.
rsp_pos = real2redshift(pos, vel, [0.5, 0.5, 0.5], box, rsp_pos = real2redshift(pos, vel, [0.5, 0.5, 0.5], box,
in_box_units=True, periodic_wrap=True, in_box_units=True, periodic_wrap=True,
make_copy=True) make_copy=True)
# Assign particles' values to the grid. # ... and count the number of particles in each grid cell.
MASL.MA(rsp_pos, rsp_field, 1., "NGP", W=values)
# Count the number of particles in each grid cell.
MASL.MA(rsp_pos, cellcounts, 1., "NGP") MASL.MA(rsp_pos, cellcounts, 1., "NGP")
# Now finally we evaluate the field at the particle positions in real
# space and then assign the values to the grid in redshift space.
for i in range(nfields):
values = evaluate_cartesian(fields[i], pos=pos)
MASL.MA(rsp_pos, rsp_fields[i], 1., "NGP", W=values)
if end == nparts: if end == nparts:
break break
start = end start = end
# Finally divide by the number of particles in each cell and smooth. # We divide by the number of particles in each cell.
divide_nonzero(rsp_field, cellcounts) for i in range(len(fields)):
return rsp_field divide_nonzero(rsp_fields[i], cellcounts)
if len(fields) == 1:
return rsp_fields[0]
return rsp_fields
@jit(nopython=True) @jit(nopython=True)

View file

@ -193,7 +193,7 @@ def potential_field(nsim, parser_args, to_save=True):
if parser_args.in_rsp: if parser_args.in_rsp:
parts = csiborgtools.read.read_h5(paths.particles(nsim))["particles"] parts = csiborgtools.read.read_h5(paths.particles(nsim))["particles"]
field = csiborgtools.field.field2rsp(field, parts=parts, box=box, field = csiborgtools.field.field2rsp(*field, parts=parts, box=box,
verbose=parser_args.verbose) verbose=parser_args.verbose)
if to_save: if to_save:
fout = paths.field(parser_args.kind, parser_args.MAS, parser_args.grid, fout = paths.field(parser_args.kind, parser_args.MAS, parser_args.grid,
@ -268,8 +268,6 @@ def environment_field(nsim, parser_args, to_save=True):
------- -------
env : 3-dimensional array env : 3-dimensional array
""" """
if parser_args.in_rsp:
raise NotImplementedError("Env. field in RSP not implemented.")
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim)) nsnap = max(paths.get_snapshots(nsim))
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths) box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
@ -292,7 +290,22 @@ def environment_field(nsim, parser_args, to_save=True):
del rho del rho
collect() collect()
# TODO: Optionally drag the field to RSP. # Optionally drag the field to RSP.
if parser_args.in_rsp:
parts = csiborgtools.read.read_h5(paths.particles(nsim))["particles"]
fields = (tensor_field.T00, tensor_field.T11, tensor_field.T22,
tensor_field.T01, tensor_field.T02, tensor_field.T12)
T00, T11, T22, T01, T02, T12 = csiborgtools.field.field2rsp(
*fields, parts=parts, box=box, verbose=parser_args.verbose)
tensor_field.T00[...] = T00
tensor_field.T11[...] = T11
tensor_field.T22[...] = T22
tensor_field.T01[...] = T01
tensor_field.T02[...] = T02
tensor_field.T12[...] = T12
del T00, T11, T22, T01, T02, T12
collect()
# Calculate the eigenvalues of the tidal tensor field, delete tensor field. # Calculate the eigenvalues of the tidal tensor field, delete tensor field.
if parser_args.verbose: if parser_args.verbose:

View file

@ -16,12 +16,14 @@ Script to match all pairs of CSiBORG simulations. Mathches main haloes whose
mass is above 1e12 solar masses. mass is above 1e12 solar masses.
""" """
from argparse import ArgumentParser from argparse import ArgumentParser
from datetime import datetime
from distutils.util import strtobool from distutils.util import strtobool
from itertools import combinations from itertools import combinations
from random import Random from random import Random
from mpi4py import MPI from mpi4py import MPI
from taskmaster import work_delegation
from match_singlematch import pair_match
try: try:
import csiborgtools import csiborgtools
@ -31,28 +33,16 @@ except ModuleNotFoundError:
sys.path.append("../") sys.path.append("../")
import csiborgtools import csiborgtools
from taskmaster import master_process, worker_process
from match_singlematch import pair_match
# Argument parser
parser = ArgumentParser()
parser.add_argument("--sigma", type=float, default=None)
parser.add_argument("--smoothen", type=lambda x: bool(strtobool(x)),
default=None)
parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
default=False)
args = parser.parse_args()
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
nproc = comm.Get_size()
def get_combs(): def get_combs():
""" """
Get the list of all pairs of simulations, then permute them with a known Get the list of all pairs of simulations, then permute them with a known
seed to minimise loading the same files simultaneously. seed to minimise loading the same files simultaneously.
Returns
-------
combs : list
List of pairs of simulations.
""" """
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
ics = paths.get_ics("csiborg") ics = paths.get_ics("csiborg")
@ -62,18 +52,31 @@ def get_combs():
def do_work(comb): def do_work(comb):
"""
Match a pair of simulations.
Parameters
----------
comb : tuple
Pair of simulations.
Returns
-------
None
"""
nsim0, nsimx = comb nsim0, nsimx = comb
pair_match(nsim0, nsimx, args.sigma, args.smoothen, args.verbose) pair_match(nsim0, nsimx, args.sigma, args.smoothen, args.verbose)
if nproc > 1: if __name__ == "__main__":
if rank == 0: parser = ArgumentParser()
parser.add_argument("--sigma", type=float, default=None)
parser.add_argument("--smoothen", type=lambda x: bool(strtobool(x)),
default=None)
parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
default=False)
args = parser.parse_args()
comm = MPI.COMM_WORLD
combs = get_combs() combs = get_combs()
master_process(combs, comm, verbose=True) work_delegation(do_work, combs, comm, master_verbose=True)
else:
worker_process(do_work, comm, verbose=False)
else:
combs = get_combs()
for comb in combs:
print(f"{datetime.now()}: completing task `{comb}`.", flush=True)
do_work(comb)

View file

@ -322,15 +322,24 @@ def plot_projected_field(kind, nsim, grid, in_rsp, smooth_scale, MAS="PCS",
ax[i].imshow(img, vmin=vmin, vmax=vmax, cmap=cmap) ax[i].imshow(img, vmin=vmin, vmax=vmax, cmap=cmap)
frad = 155.5 / 677.7 frad = 155.5 / 677.7
if not highres_only and 0.5 - frad < slice_find < 0.5 + frad:
theta = numpy.linspace(0, 2 * numpy.pi, 100)
z = (slice_find - 0.5) * grid
R = 155.5 / 677.7 * grid R = 155.5 / 677.7 * grid
if slice_find is None:
rad = R
plot_circle = True
elif (not highres_only and 0.5 - frad < slice_find < 0.5 + frad):
z = (slice_find - 0.5) * grid
rad = R * numpy.sqrt(1 - z**2 / R**2) rad = R * numpy.sqrt(1 - z**2 / R**2)
plot_circle = True
else:
plot_circle = False
if not highres_only and plot_circle:
theta = numpy.linspace(0, 2 * numpy.pi, 100)
ax[i].plot(rad * numpy.cos(theta) + grid // 2, ax[i].plot(rad * numpy.cos(theta) + grid // 2,
rad * numpy.sin(theta) + grid // 2, rad * numpy.sin(theta) + grid // 2,
lw=0.75 * plt.rcParams["lines.linewidth"], zorder=1, lw=0.75 * plt.rcParams["lines.linewidth"], zorder=1,
c="red", ls="--") c="red", ls="--")
ax[i].set_title(labels[i]) ax[i].set_title(labels[i])
if highres_only: if highres_only:
@ -551,7 +560,7 @@ if __name__ == "__main__":
plot_halos=5e13, volume_weight=False) plot_halos=5e13, volume_weight=False)
if True: if True:
kind = "environment" kind = "density"
grid = 256 grid = 256
smooth_scale = 0 smooth_scale = 0
# plot_projected_field("overdensity", 7444, grid, in_rsp=True, # plot_projected_field("overdensity", 7444, grid, in_rsp=True,

View file

@ -41,6 +41,10 @@ def plot_knn(runname):
Parameters Parameters
---------- ----------
runname : str runname : str
Returns
-------
None
""" """
print(f"Plotting kNN CDF for {runname}.") print(f"Plotting kNN CDF for {runname}.")
cols = plt.rcParams["axes.prop_cycle"].by_key()["color"] cols = plt.rcParams["axes.prop_cycle"].by_key()["color"]