From cbfd1cbc99ff8fe6bdac74ff81db09e0bc2a4d3e Mon Sep 17 00:00:00 2001 From: Richard Stiskalek Date: Fri, 16 Jun 2023 18:31:43 +0100 Subject: [PATCH] Fixing tidal field calculation (#68) * Add import * Fix tidal calculation * Add env to paths * Improve plot routines * Add env classification --- csiborgtools/field/__init__.py | 2 +- csiborgtools/field/density.py | 126 +++++++++++++++++++++++++++++---- csiborgtools/read/paths.py | 5 +- scripts/field_prop.py | 51 ++++++++++++- scripts_plots/plot_data.py | 83 ++++++++++++++++++---- 5 files changed, 233 insertions(+), 34 deletions(-) diff --git a/csiborgtools/field/__init__.py b/csiborgtools/field/__init__.py index 7957ffc..dcd3773 100644 --- a/csiborgtools/field/__init__.py +++ b/csiborgtools/field/__init__.py @@ -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 diff --git a/csiborgtools/field/density.py b/csiborgtools/field/density.py index 9bcb893..78efc1c 100644 --- a/csiborgtools/field/density.py +++ b/csiborgtools/field/density.py @@ -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 diff --git a/csiborgtools/read/paths.py b/csiborgtools/read/paths.py index a64d319..f032c5d 100644 --- a/csiborgtools/read/paths.py +++ b/csiborgtools/read/paths.py @@ -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) diff --git a/scripts/field_prop.py b/scripts/field_prop.py index fe2f6db..7cf73df 100644 --- a/scripts/field_prop.py +++ b/scripts/field_prop.py @@ -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.") diff --git a/scripts_plots/plot_data.py b/scripts_plots/plot_data.py index 6bfd731..3d8c9b7 100644 --- a/scripts_plots/plot_data.py +++ b/scripts_plots/plot_data.py @@ -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)