Aquila-Consortium/csiborgtools
0
0
Fork 0
mirror of https://github.com/Richard-Sti/csiborgtools.git synced 2024-10-18 16:25:06 +02:00
Code Issues Actions Packages Projects Releases Wiki Activity

Moving environment to RS (#70)

Browse Source 
* 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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

51
csiborgtools/field/interp.py
View File

@ -168,17 +168,17 @@ def divide_nonzero(field0, field1):
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):
"""
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
positions reconstructs back the field on a grid by NGP interpolation.
Parameters
----------
field : 3-dimensional array of shape `(grid, grid, grid)`
Real space field to be evolved to redshift space.
fields : (list of) 3-dimensional array of shape `(grid, grid, grid)`
Real space fields to be evolved to redshift space.
parts : 2-dimensional array of shape `(n_parts, 6)`
Particle positions and velocities in real space. Must be organised as
`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_field = numpy.full(field.shape, init_value, dtype=numpy.float32)
cellcounts = numpy.zeros(rsp_field.shape, dtype=numpy.float32)
# We iterate over the fields and in the inner loop over the particles. This
# is slower than iterating over the particles and in the inner loop over
# the fields, but it is more memory efficient. Typically we will only have
# one field.
nfields = len(fields)
# Check that all fields have the same shape.
if nfields > 1:
assert all(fields[0].shape == fields[i].shape
for i in range(1, nfields))
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]
batch_size = nparts // nbatch
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)
pos = parts[start:end]
pos, vel = pos[:, :3], pos[:, 3:6]
if flip_partsxz:
pos[:, [0, 2]] = pos[:, [2, 0]]
vel[:, [0, 2]] = vel[:, [2, 0]]
# Evaluate the field at the particle positions in real space.
values = evaluate_cartesian(field, pos=pos)
# Move particles to redshift space.
# Then move the particles to redshift space.
rsp_pos = real2redshift(pos, vel, [0.5, 0.5, 0.5], box,
in_box_units=True, periodic_wrap=True,
make_copy=True)
# Assign particles' values to the grid.
MASL.MA(rsp_pos, rsp_field, 1., "NGP", W=values)
# Count the number of particles in each grid cell.
# ... and count the number of particles in each grid cell.
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:
break
start = end
# Finally divide by the number of particles in each cell and smooth.
divide_nonzero(rsp_field, cellcounts)
return rsp_field
# We divide by the number of particles in each cell.
for i in range(len(fields)):
divide_nonzero(rsp_fields[i], cellcounts)
if len(fields) == 1:
return rsp_fields[0]
return rsp_fields
@jit(nopython=True)

21
scripts/field_prop.py
View File

@ -193,7 +193,7 @@ def potential_field(nsim, parser_args, to_save=True):
if parser_args.in_rsp:
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)
if to_save:
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
"""
if parser_args.in_rsp:
raise NotImplementedError("Env. field in RSP not implemented.")
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim))
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
@ -292,7 +290,22 @@ def environment_field(nsim, parser_args, to_save=True):
del rho
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.
if parser_args.verbose:

59
scripts/match_all.py
View File

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

17
scripts_plots/plot_data.py
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)
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)
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
R = 155.5 / 677.7 * grid
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,
rad * numpy.sin(theta) + grid // 2,
lw=0.75 * plt.rcParams["lines.linewidth"], zorder=1,
c="red", ls="--")
ax[i].set_title(labels[i])
if highres_only:
@ -551,7 +560,7 @@ if __name__ == "__main__":
plot_halos=5e13, volume_weight=False)
if True:
kind = "environment"
kind = "density"
grid = 256
smooth_scale = 0
# plot_projected_field("overdensity", 7444, grid, in_rsp=True,

4
scripts_plots/plot_knn.py
View File

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