Fixing tidal field calculation (#68)

* Add import

* Fix tidal calculation

* Add env to paths

* Improve plot routines

* Add env classification

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

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)
+        dummy_vector = numpy.full(3, numpy.nan, dtype=numpy.float32)
-
+        dummy_tensor = numpy.full((3, 3), numpy.nan, dtype=numpy.float32)
-        # FILL IN THESER ARGUMENTS
-        tidal_tensor_to_eigenvalues(eigvals, dummy, ...)
 
+        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_tensor[0, 0] = T00[i, j, k]
-                dummy[0, 1] = T01[i, j, k]
+                dummy_tensor[1, 1] = T11[i, j, k]
-                dummy[0, 2] = T02[i, j, k]
+                dummy_tensor[2, 2] = T22[i, j, k]
-                dummy[1, 1] = T11[i, j, k]
+
-                dummy[1, 2] = T12[i, j, k]
+                dummy_tensor[0, 1] = T01[i, j, k]
-                dummy[2, 2] = T22[i, j, k]
+                dummy_tensor[1, 0] = T01[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, 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

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)

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.")
 

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,
+                im = ax[i].imshow(img, vmin=vmin, vmax=vmax, cmap=cmap)
-                                  vmax=vmax)
             else:
-                ax[i].imshow(numpy.nanmean(field, axis=i), vmin=vmin,
+                ax[i].imshow(img, vmin=vmin, vmax=vmax, cmap=cmap)
-                             vmax=vmax)
+
+            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,
+        kind = "environment"
-                              plot_groups=False, dmin=50, dmax=100,
+        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,
+        kind = "environment"
-                             highres_only=False)
+        grid = 256
-        plot_projected_field("overdensity", 7444, 1024, in_rsp=False,
+        # plot_projected_field("overdensity", 7444, grid, in_rsp=True,
-                             highres_only=False)
+        #                      highres_only=False)
+        plot_projected_field(kind, 7444, grid, in_rsp=False,
+                             slice_find=0.5, highres_only=False)