mirror of
https://github.com/Richard-Sti/csiborgtools.git
synced 2024-12-22 12:28:03 +00:00
Fixing tidal field calculation (#68)
* Add import * Fix tidal calculation * Add env to paths * Improve plot routines * Add env classification
This commit is contained in:
Richard Stiskalek
GitHub
parent
ccbbbd24b4
commit
cbfd1cbc99
5 changed files with 233 additions and 34 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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.")
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
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)
|
||||
|
|
|
|||
Loading…
Reference in a new issue