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

Fixing tidal field calculation (#68)

Browse source 
* Add import

* Fix tidal calculation

* Add env to paths

* Improve plot routines

* Add env classification
This commit is contained in:
Richard Stiskalek 2023-06-16 18:31:43 +01:00 committed by GitHub
parent ccbbbd24b4
commit cbfd1cbc99
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 233 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
csiborgtools/field/__init__.py
View file

@ -17,7 +17,7 @@ from warnings import warn
try:
import MAS_library as MASL # noqa
from .density import DensityField, PotentialField, VelocityField # noqa
from .density import DensityField, PotentialField, VelocityField, TidalTensorField # noqa
from .interp import (evaluate_cartesian, evaluate_sky, field2rsp, # noqa
make_sky, fill_outside) # noqa
from .utils import nside2radec, smoothen_field # noqa

126
csiborgtools/field/density.py
View file

@ -391,17 +391,43 @@ class TidalTensorField(BaseField):
-------
eigvals : 3-dimensional array of shape `(grid, grid, grid)`
"""
# TODO needs to be checked further
grid = tidal_tensor.T00.shape[0]
eigvals = numpy.full((grid, grid, grid, 3), numpy.nan,
dtype=numpy.float32)
dummy = numpy.full((3, 3), numpy.nan, dtype=numpy.float32)
# FILL IN THESER ARGUMENTS
tidal_tensor_to_eigenvalues(eigvals, dummy, ...)
dummy_vector = numpy.full(3, numpy.nan, dtype=numpy.float32)
dummy_tensor = numpy.full((3, 3), numpy.nan, dtype=numpy.float32)
tidal_tensor_to_eigenvalues(eigvals, dummy_vector, dummy_tensor,
tidal_tensor.T00, tidal_tensor.T01,
tidal_tensor.T02, tidal_tensor.T11,
tidal_tensor.T12, tidal_tensor.T22)
return eigvals
@staticmethod
def eigvals_to_environment(eigvals, threshold=0.0):
"""
Calculate the environment of each grid cell based on the eigenvalues
of the tidal tensor field.
Parameters
----------
eigvals : 4-dimensional array of shape `(grid, grid, grid, 3)`
The eigenvalues of the tidal tensor field.
Returns
-------
environment : 3-dimensional array of shape `(grid, grid, grid)`
The environment of each grid cell. Possible values are:
- 0: void
- 1: sheet
- 2: filament
- 3: knot
"""
environment = numpy.full(eigvals.shape[:-1], numpy.nan,
dtype=numpy.float32)
eigenvalues_to_environment(environment, eigvals, threshold)
return environment
def __call__(self, overdensity_field):
"""
Calculate the tidal tensor field.
@ -422,20 +448,90 @@ class TidalTensorField(BaseField):
@jit(nopython=True)
def tidal_tensor_to_eigenvalues(eigvals, dummy, T00, T01, T02, T11, T12, T22):
def tidal_tensor_to_eigenvalues(eigvals, dummy_vector, dummy_tensor,
T00, T01, T02, T11, T12, T22):
"""
TODO: needs to be checked further.
Calculate eigenvalues of the tidal tensor field, sorted in decreasing
absolute value order. JIT implementation to speed up the work.
Parameters
----------
eigvals : 3-dimensional array of shape `(grid, grid, grid)`
Array to store the eigenvalues.
dummy_vector : 1-dimensional array of shape `(3,)`
Dummy vector to store the eigenvalues.
dummy_tensor : 2-dimensional array of shape `(3, 3)`
Dummy tensor to store the tidal tensor.
T00 : 3-dimensional array of shape `(grid, grid, grid)`
Tidal tensor component :math:`T_{00}`.
T01 : 3-dimensional array of shape `(grid, grid, grid)`
Tidal tensor component :math:`T_{01}`.
T02 : 3-dimensional array of shape `(grid, grid, grid)`
Tidal tensor component :math:`T_{02}`.
T11 : 3-dimensional array of shape `(grid, grid, grid)`
Tidal tensor component :math:`T_{11}`.
T12 : 3-dimensional array of shape `(grid, grid, grid)`
Tidal tensor component :math:`T_{12}`.
T22 : 3-dimensional array of shape `(grid, grid, grid)`
Tidal tensor component :math:`T_{22}`.
Returns
-------
eigvals : 3-dimensional array of shape `(grid, grid, grid)`
"""
grid = T00.shape[0]
for i in range(grid):
for j in range(grid):
for k in range(grid):
dummy[0, 0] = T00[i, j, k]
dummy[0, 1] = T01[i, j, k]
dummy[0, 2] = T02[i, j, k]
dummy[1, 1] = T11[i, j, k]
dummy[1, 2] = T12[i, j, k]
dummy[2, 2] = T22[i, j, k]
eigvals[i, j, k, :] = numpy.linalg.eigvalsh(dummy, 'U')
eigvals[i, j, k, :] = numpy.sort(eigvals[i, j, k, :])
dummy_tensor[0, 0] = T00[i, j, k]
dummy_tensor[1, 1] = T11[i, j, k]
dummy_tensor[2, 2] = T22[i, j, k]
dummy_tensor[0, 1] = T01[i, j, k]
dummy_tensor[1, 0] = T01[i, j, k]
dummy_tensor[0, 2] = T02[i, j, k]
dummy_tensor[2, 0] = T02[i, j, k]
dummy_tensor[1, 2] = T12[i, j, k]
dummy_tensor[2, 1] = T12[i, j, k]
dummy_vector[:] = numpy.linalg.eigvalsh(dummy_tensor)
eigvals[i, j, k, :] = dummy_vector[
numpy.argsort(numpy.abs(dummy_vector))][::-1]
return eigvals
@jit(nopython=True)
def eigenvalues_to_environment(environment, eigvals, th):
"""
Classify the environment of each grid cell based on the eigenvalues of the
tidal tensor field.
Parameters
----------
environment : 3-dimensional array of shape `(grid, grid, grid)`
Array to store the environment.
eigvals : 4-dimensional array of shape `(grid, grid, grid, 3)`
The eigenvalues of the tidal tensor field.
th : float
Threshold value to classify the environment.
Returns
-------
environment : 3-dimensional array of shape `(grid, grid, grid)`
"""
grid = eigvals.shape[0]
for i in range(grid):
for j in range(grid):
for k in range(grid):
lmbda1, lmbda2, lmbda3 = eigvals[i, j, k, :]
if lmbda1 < th and lmbda2 < th and lmbda3 < th:
environment[i, j, k] = 0
elif lmbda1 < th and lmbda2 < th:
environment[i, j, k] = 1
elif lmbda1 < th:
environment[i, j, k] = 2
else:
environment[i, j, k] = 3
return environment

5
csiborgtools/read/paths.py
View file

@ -364,7 +364,7 @@ class Paths:
----------
kind : str
Field type. Must be one of: `density`, `velocity`, `potential`,
`radvel`.
`radvel`, `environment`.
MAS : str
Mass-assignment scheme.
grid : int
@ -379,7 +379,8 @@ class Paths:
path : str
"""
fdir = join(self.postdir, "environment")
assert kind in ["density", "velocity", "potential", "radvel"]
assert kind in ["density", "velocity", "potential", "radvel",
"environment"]
if not isdir(fdir):
makedirs(fdir)
warn(f"Created directory `{fdir}`.", UserWarning, stacklevel=1)

51
scripts/field_prop.py
View file

@ -19,6 +19,7 @@ simulations' final snapshot.
from argparse import ArgumentParser
from datetime import datetime
from distutils.util import strtobool
from gc import collect
import numpy
from mpi4py import MPI
@ -130,6 +131,51 @@ def radvel_field(nsim, parser_args):
numpy.save(fout, field)
###############################################################################
# Environment classification #
###############################################################################
def environment_field(nsim, parser_args):
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)
density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
gen = csiborgtools.field.TidalTensorField(box, parser_args.MAS)
# Load the real space overdensity field
if parser_args.verbose:
print(f"{datetime.now()}: loading density field.")
rho = numpy.load(paths.field("density", parser_args.MAS, parser_args.grid,
nsim, in_rsp=False))
rho = density_gen.overdensity_field(rho)
# Calculate the real space tidal tensor field, delete overdensity.
if parser_args.verbose:
print(f"{datetime.now()}: calculating tidal tensor field.")
tensor_field = gen(rho)
del rho
collect()
# Calculate the eigenvalues of the tidal tensor field, delete tensor field.
if parser_args.verbose:
print(f"{datetime.now()}: calculating eigenvalues.")
eigvals = gen.tensor_field_eigvals(tensor_field)
del tensor_field
collect()
# Classify the environment based on the eigenvalues.
if parser_args.verbose:
print(f"{datetime.now()}: classifying environment.")
env = gen.eigvals_to_environment(eigvals)
del eigvals
collect()
fout = paths.field("environment", parser_args.MAS, parser_args.grid,
nsim, parser_args.in_rsp)
print(f"{datetime.now()}: saving output to `{fout}`.")
numpy.save(fout, env)
###############################################################################
# Command line interface #
###############################################################################
@ -140,7 +186,8 @@ if __name__ == "__main__":
parser.add_argument("--nsims", type=int, nargs="+", default=None,
help="IC realisations. `-1` for all simulations.")
parser.add_argument("--kind", type=str,
choices=["density", "velocity", "radvel", "potential"],
choices=["density", "velocity", "radvel", "potential",
"environment"],
help="What derived field to calculate?")
parser.add_argument("--MAS", type=str,
choices=["NGP", "CIC", "TSC", "PCS"])
@ -163,6 +210,8 @@ if __name__ == "__main__":
radvel_field(nsim, parser_args)
elif parser_args.kind == "potential":
potential_field(nsim, parser_args)
elif parser_args.kind == "environment":
environment_field(nsim, parser_args)
else:
raise RuntimeError(f"Field {parser_args.kind} is not implemented.")

83
scripts_plots/plot_data.py
View file

@ -240,9 +240,9 @@ def load_field(kind, nsim, grid, MAS, in_rsp=False):
def plot_projected_field(kind, nsim, grid, in_rsp, MAS="PCS",
highres_only=True, pdf=False):
highres_only=True, slice_find=None, pdf=False):
"""
Plot the mean projected field.
Plot the mean projected field, however can also plot a single slice.
Parameters
----------
@ -258,6 +258,8 @@ def plot_projected_field(kind, nsim, grid, in_rsp, MAS="PCS",
Mass assignment scheme.
highres_only : bool, optional
Whether to only plot the high-resolution region.
slice_find : float, optional
Which slice to plot in fractional units (i.e. 1. is the last slice)
pdf : bool, optional
Whether to save the figure as a PDF.
@ -284,23 +286,69 @@ def plot_projected_field(kind, nsim, grid, in_rsp, MAS="PCS",
end = round(field.shape[0] * 0.73)
field = field[start:end, start:end, start:end]
if kind != "environment":
cmap = "viridis"
else:
cmap = "brg"
labels = [r"$y-z$", r"$x-z$", r"$x-y$"]
with plt.style.context(plt_utils.mplstyle):
fig, ax = plt.subplots(figsize=(3.5 * 2, 2.625), ncols=3, sharey=True,
sharex=True)
sharex="col")
fig.subplots_adjust(hspace=0, wspace=0)
for i in range(3):
img = numpy.nanmean(field, axis=i)
if slice_find is None:
img = numpy.nanmean(field, axis=i)
else:
ii = int(field.shape[i] * slice_find)
img = numpy.take(field, ii, axis=i)
if i == 0:
vmin, vmax = numpy.nanpercentile(img, [1, 99])
im = ax[i].imshow(numpy.nanmean(field, axis=i), vmin=vmin,
vmax=vmax)
im = ax[i].imshow(img, vmin=vmin, vmax=vmax, cmap=cmap)
else:
ax[i].imshow(numpy.nanmean(field, axis=i), vmin=vmin,
vmax=vmax)
ax[i].imshow(img, vmin=vmin, vmax=vmax, cmap=cmap)
if not highres_only:
theta = numpy.linspace(0, 2 * numpy.pi, 100)
rad = 155.5 / 677.7 * grid
ax[i].plot(rad * numpy.cos(theta) + grid // 2,
rad * numpy.sin(theta) + grid // 2,
lw=plt.rcParams["lines.linewidth"], zorder=1,
c="red", ls="--")
ax[i].set_title(labels[i])
if highres_only:
ncells = end - start
size = ncells / grid * 677.7
else:
ncells = grid
size = 677.7
# Get beautiful ticks
yticks = numpy.linspace(0, ncells, 6).astype(int)
yticks = numpy.append(yticks, ncells // 2)
ax[0].set_yticks(yticks)
ax[0].set_yticklabels((yticks * size / ncells - size / 2).astype(int))
ax[0].set_ylabel(r"$x_i ~ [\mathrm{Mpc} / h]$")
for i in range(3):
xticks = numpy.linspace(0, ncells, 6).astype(int)
xticks = numpy.append(xticks, ncells // 2)
xticks = numpy.sort(xticks)
if i < 2:
xticks = xticks[:-1]
ax[i].set_xticks(xticks)
ax[i].set_xticklabels(
(xticks * size / ncells - size / 2).astype(int))
ax[i].set_xlabel(r"$x_j ~ [\mathrm{Mpc} / h]$")
cbar_ax = fig.add_axes([1.0, 0.1, 0.025, 0.8])
fig.colorbar(im, cax=cbar_ax, label="Mean projected field")
if slice_find is None:
clabel = "Mean projected field"
else:
clabel = "Sliced field"
fig.colorbar(im, cax=cbar_ax, label=clabel)
fig.tight_layout(h_pad=0, w_pad=0)
for ext in ["png"] if pdf is False else ["png", "pdf"]:
@ -310,6 +358,7 @@ def plot_projected_field(kind, nsim, grid, in_rsp, MAS="PCS",
fig.savefig(fout, dpi=plt_utils.dpi, bbox_inches="tight")
plt.close()
###############################################################################
# Sky distribution #
###############################################################################
@ -461,12 +510,16 @@ if __name__ == "__main__":
plot_hmf(pdf=False)
if False:
plot_sky_distribution("radvel", 7444, 256, nside=64,
plot_groups=False, dmin=50, dmax=100,
kind = "environment"
grid = 256
plot_sky_distribution(kind, 7444, grid, nside=64,
plot_groups=False, dmin=0, dmax=25,
plot_halos=5e13, volume_weight=False)
if True:
plot_projected_field("overdensity", 7444, 1024, in_rsp=True,
highres_only=False)
plot_projected_field("overdensity", 7444, 1024, in_rsp=False,
highres_only=False)
kind = "environment"
grid = 256
# plot_projected_field("overdensity", 7444, grid, in_rsp=True,
# highres_only=False)
plot_projected_field(kind, 7444, grid, in_rsp=False,
slice_find=0.5, highres_only=False)