Merge pull request #5 from entity-toolkit/0.6.0rc

Support for plotting with ghost cells and multidomain

Author: haykh

README.md

@@ -27,6 +27,8 @@ data.particles  # < dict[int : xr.Dataset]
 data.spectra    # < xr.Dataset
 ```
 
+> If using Jupyter notebook, you can quickly preview the loaded metadata by simply running a cell with just `data` in it (or in regular python, by doing `print(data)`).
+
 #### Examples
 
 Plot a field (in cartesian space) at a specific time (or output step):
@@ -96,7 +98,42 @@ def plot(t, data):
 data.makeMovie(plot)
 ```
 
-> If using Jupyter notebook, you can quickly preview the loaded metadata by simply running a cell with just `data` in it (or in regular python, by doing `print(data)`).
+You may also access the movie-making functionality directly in case you want to use it for other things:
+
+```python
+import nt2.export as nt2e
+
+def plot(t):
+  ...
+
+#             this will be the array of `t`-s passed to `plot`
+#                           |
+#                           V
+nt2e.makeFrames(plot, np.arange(100), "myAnim")
+nt2e.makeMovie(
+    input="myAnim/", output="myAnim.mp4", number=5, overwrite=True
+)
+
+# or combined together
+nt2e.makeFramesAndMovie(
+    name="myAnim", plot=plot, times=np.arange(100)
+)
+```
+
+#### Plots for debugging
+
+If the simulation also outputs the ghost cells, `nt2py` will interpret the fields differently, and instead of reading the physical coordinates, will build the coordinates based on the number of cells (including ghost cells). In particular, instead of, e.g., `data.fields.x` it will contain `data.fields.i1`. The data will also contain information about the meshblock decomposition. For instance, if you have `Nx` meshblocks in the `x` direction, each having `nx` cells, the coordinates `data.fields.i1` will go from `0` to `nx * NX + 2 * NGHOSTS * Nx`.
+
+You can overplot both the coordinate grid as well as the active zones of the meshblocks using the following:
+
+```python
+ax = plt.gca()
+data.fields.Ex.isel(t=ti).plot(ax=ax)
+data.plotGrid(ax=ax)
+data.plotDomains(ax=ax)
+```
+
+> Keep in mind, that by default `Entity` converts all quantities to tetrad basis (or contravariant in GR) and interpolates to cell centers before outputting (except for the ghost cells). So when doing plots for debugging, make sure to also set `as_is = true` (together with `ghosts = true`) in the `[output.debug]` section of the `toml` input file. This will ensure the fields are being output as is, with no conversion or interpolation. This does not apply to particle moments, which are never stored in the code and are computed only during the output.
 
 ### Dashboard
 
@@ -119,4 +156,7 @@ This will output the port where the dashboard server is running, e.g., `Dashboar
 
 - [ ] Unit tests
 - [ ] Plugins for other simulation data formats
+- [ ] Support for multiple runs
+- [ ] Interactive regime (`hvplot`, `bokeh`, `panel`)
+- [x] Ghost cells support
 - [x] Usage examples

dist/nt2py-0.6.0-py3-none-any.whl

@@ -1,4 +1,4 @@
-__version__ = "0.5.3"
+__version__ = "0.6.0"
 
 from nt2.data import Data as Data
 from nt2.dashboard import Dashboard as Dashboard

dist/nt2py-0.6.0.tar.gz

@@ -28,9 +28,6 @@ class Container:
 
     Kwargs
     ------
-    single_file : bool, optional
-        Whether the data is stored in a single file. Default is False.
-
     pickle : bool, optional
         Whether to use pickle for reading the data. Default is True.
 
@@ -67,17 +64,15 @@ class Container:
 
     """
 
-    def __init__(
-        self, path, single_file=False, pickle=True, greek=False, dask_props={}
-    ):
+    def __init__(self, path, pickle=True, greek=False, dask_props={}):
         super(Container, self).__init__()
 
+        self.path = path
         self.configs: dict[str, Any] = {
-            "single_file": single_file,
+            "single_file": self.path.endswith(".h5"),
             "use_pickle": pickle,
             "use_greek": greek,
         }
-        self.path = path
         self.metadata = {}
         self.mesh = None
         if self.configs["single_file"]:
@@ -95,6 +90,8 @@ def __init__(
 
         self.attrs = _read_attribs_SingleFile(self.master_file)
 
+        assert self.master_file is not None, "Master file not found"
+
         self.configs["ngh"] = int(self.master_file.attrs.get("NGhosts", 0))
         self.configs["layout"] = (
             "right" if self.master_file.attrs.get("LayoutRight", 1) == 1 else "left"
@@ -106,16 +103,31 @@ def __init__(
         if self.configs["coordinates"] == "qsph":
             self.configs["coordinates"] = "sph"
 
+        if self.isDebug():
+            self.configs["coordinates"] = "cart"
+
         if not self.configs["single_file"]:
             self.master_file.close()
             self.master_file = None
 
+    def isDebug(self):
+        return self.configs["ngh"] > 0
+
     def __del__(self):
         if self.master_file is not None:
             self.master_file.close()
 
     def plotGrid(self, ax, **kwargs):
-        from matplotlib import patches
+        try:
+            assert self.mesh is not None, "Mesh not found"
+        except AttributeError:
+            raise AttributeError("Mesh not found")
+        except AssertionError:
+            raise AssertionError("Mesh not found")
+
+        assert len(self.mesh["xc"]) == 2, "Data must be 2D for plotGrid to work"
+
+        from matplotlib import patches as mpatches
 
         xlim, ylim = ax.get_xlim(), ax.get_ylim()
         options = {
@@ -125,15 +137,23 @@ def plotGrid(self, ax, **kwargs):
         }
         options.update(kwargs)
 
+        x1_emin, x2_emin = list(self.mesh["xe_min"].keys())
+        x1_emax, x2_emax = list(self.mesh["xe_max"].keys())
+        x1_e = list(self.mesh["xe_min"][x1_emin][1]) + [
+            self.mesh["xe_max"][x1_emax][1][-1]
+        ]
+        x2_e = list(self.mesh["xe_min"][x2_emin][1]) + [
+            self.mesh["xe_max"][x2_emax][1][-1]
+        ]
         if self.configs["coordinates"] == "cart":
-            for x in self.attrs["X1"]:
-                ax.plot([x, x], [self.attrs["X2Min"], self.attrs["X2Max"]], **options)
-            for y in self.attrs["X2"]:
-                ax.plot([self.attrs["X1Min"], self.attrs["X1Max"]], [y, y], **options)
+            for x1 in x1_e:
+                ax.plot([x1, x1], [x2_e[0], x2_e[-1]], **options)
+            for x2 in x2_e:
+                ax.plot([x1_e[0], x1_e[-1]], [x2, x2], **options)
         else:
-            for r in self.attrs["X1"]:
+            for r in x1_e:
                 ax.add_patch(
-                    patches.Arc(
+                    mpatches.Arc(
                         (0, 0),
                         2 * r,
                         2 * r,
@@ -143,15 +163,15 @@ def plotGrid(self, ax, **kwargs):
                         **options,
                     )
                 )
-            for th in self.attrs["X2"]:
+            for th in x2_e:
                 ax.plot(
                     [
-                        self.attrs["X1Min"] * np.sin(th),
-                        self.attrs["X1Max"] * np.sin(th),
+                        x1_e[0] * np.sin(th),
+                        x1_e[-1] * np.sin(th),
                     ],
                     [
-                        self.attrs["X1Min"] * np.cos(th),
-                        self.attrs["X1Max"] * np.cos(th),
+                        x1_e[0] * np.cos(th),
+                        x1_e[-1] * np.cos(th),
                     ],
                     **options,
                 )

nt2/__init__.py

@@ -6,7 +6,9 @@
 from nt2.containers.utils import (
     _read_category_metadata,
     _read_coordinates,
+    _preload_domain_shapes,
     _preload_field,
+    _preload_field_with_ghosts,
 )
 
 from nt2.plotters.polar import (
@@ -103,9 +105,12 @@ def __init__(self, **kwargs):
                     False, "f", self.fields_files
                 )
 
-        coords = list(CoordinateDict[self.configs["coordinates"]].values())[::-1][
-            -self.configs["dimension"] :
-        ]
+        if not self.isDebug():
+            coords = list(CoordinateDict[self.configs["coordinates"]].values())[::-1][
+                -self.configs["dimension"] :
+            ]
+        else:
+            coords = ["i3", "i2", "i1"][-self.configs["dimension"] :]
 
         if self.configs["single_file"]:
             assert self.master_file is not None, "Master file not found"
@@ -116,30 +121,74 @@ def __init__(self, **kwargs):
         self._fields = xr.Dataset()
 
         if "fields" in self.metadata and len(self.metadata["fields"]["outsteps"]) > 0:
-            for k in self.metadata["fields"]["quantities"]:
-                name, dset = _preload_field(
-                    single_file=self.configs["single_file"],
-                    k=k,
-                    dim=self.configs["dimension"],
-                    ngh=self.configs["ngh"],
-                    outsteps=self.metadata["fields"]["outsteps"],
-                    times=self.metadata["fields"]["times"],
-                    steps=self.metadata["fields"]["steps"],
-                    coords=coords,
-                    xc_coords=self.mesh["x_c"],
-                    xe_min_coords=self.mesh["x_emin"],
-                    xe_max_coords=self.mesh["x_emax"],
-                    coord_replacements=list(
-                        CoordinateDict[self.configs["coordinates"]].items()
-                    ),
-                    field_replacements=list(QuantityDict.items()),
-                    layout=self.configs["layout"],
-                    file=(
-                        self.master_file
-                        if self.configs["single_file"] and self.master_file is not None
-                        else self.fields_files
-                    ),
+            self.domains = xr.Dataset()
+            for i in range(self.configs["dimension"]):
+                self.domains[f"x{i+1}"], self.domains[f"sx{i+1}"] = (
+                    _preload_domain_shapes(
+                        single_file=self.configs["single_file"],
+                        k=f"N{i+1}l",
+                        outsteps=self.metadata["fields"]["outsteps"],
+                        times=self.metadata["fields"]["times"],
+                        steps=self.metadata["fields"]["steps"],
+                        file=(
+                            self.master_file
+                            if self.configs["single_file"]
+                            and self.master_file is not None
+                            else self.fields_files
+                        ),
+                    )
                 )
+
+            for k in self.metadata["fields"]["quantities"]:
+                if not self.isDebug():
+                    name, dset = _preload_field(
+                        single_file=self.configs["single_file"],
+                        k=k,
+                        outsteps=self.metadata["fields"]["outsteps"],
+                        times=self.metadata["fields"]["times"],
+                        steps=self.metadata["fields"]["steps"],
+                        coords=coords,
+                        xc_coords=self.mesh["xc"],
+                        xe_min_coords=self.mesh["xe_min"],
+                        xe_max_coords=self.mesh["xe_max"],
+                        coord_replacements=list(
+                            CoordinateDict[self.configs["coordinates"]].items()
+                        ),
+                        field_replacements=list(QuantityDict.items()),
+                        layout=self.configs["layout"],
+                        file=(
+                            self.master_file
+                            if self.configs["single_file"]
+                            and self.master_file is not None
+                            else self.fields_files
+                        ),
+                    )
+                else:
+                    (
+                        name,
+                        dset,
+                        self.mesh["xc"],
+                        self.mesh["xe_min"],
+                        self.mesh["xe_max"],
+                    ) = _preload_field_with_ghosts(
+                        single_file=self.configs["single_file"],
+                        k=k,
+                        outsteps=self.metadata["fields"]["outsteps"],
+                        times=self.metadata["fields"]["times"],
+                        steps=self.metadata["fields"]["steps"],
+                        coords=coords,
+                        coord_replacements=list(
+                            CoordinateDict[self.configs["coordinates"]].items()
+                        ),
+                        field_replacements=list(QuantityDict.items()),
+                        layout=self.configs["layout"],
+                        file=(
+                            self.master_file
+                            if self.configs["single_file"]
+                            and self.master_file is not None
+                            else self.fields_files
+                        ),
+                    )
                 self.fields[name] = dset
 
     @property
@@ -190,3 +239,55 @@ def compactify(lst):
             string += "Fields: empty\n"
 
         return string
+
+    def plotDomains(self, ax, ti=None, t=None, **kwargs):
+        if self.domains is None:
+            raise AttributeError("Domains not found")
+
+        assert len(self.domains.data_vars) == 4, "Data must be 2D for plotGrid to work"
+
+        import matplotlib.patches as mpatches
+
+        ngh = self.configs["ngh"]
+
+        xlim, ylim = ax.get_xlim(), ax.get_ylim()
+        options = {
+            "lw": 2,
+            "color": "r",
+            "ls": "-",
+        }
+        options.update(kwargs)
+
+        for dom in self.domains.dom:
+            selection = self.domains.sel(dom=dom)
+            if ti is not None:
+                selection = selection.sel(t=ti)
+            elif t is not None:
+                selection = selection.sel(t=t, method="nearest")
+            else:
+                selection = selection.isel(t=0)
+
+            x1c, sx1 = selection.x1.values[()], selection.sx1.values[()]
+            x2c, sx2 = selection.x2.values[()], selection.sx2.values[()]
+
+            # add rectangle
+            ax.add_patch(
+                mpatches.Rectangle(
+                    (x1c + ngh, x2c + ngh),
+                    sx1 - 2 * ngh,
+                    sx2 - 2 * ngh,
+                    fill=None,
+                    **options,
+                )
+            )
+
+            # ax.plot(
+            #     self.domains[x1][j],
+            #     self.domains[x2][j],
+            #     **options,
+            # )
+            # ax.plot(
+            #     self.domains[x1_e][j],
+            #     self.domains[x2_e][j],
+            #     **options,
+            # )