new data reading layout

.gitignore

@@ -151,11 +151,6 @@ dmypy.json
 # Cython debug symbols
 cython_debug/
 
-# PyCharm
-#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
-#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
-#  and can be added to the global gitignore or merged into this file.  For a more nuclear
-#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
 test/
-temp/
+temp/
+*.bak

nt2/__init__.py

@@ -1 +1,4 @@
-__version__ = "0.4.1"
+__version__ = "0.5.0"
+
+from nt2.read import Data as Data
+from nt2.plotters import polarplot as polarplot

nt2/containers/container.py

@@ -0,0 +1,148 @@
+import h5py
+import numpy as np
+from typing import Any
+from dask.distributed import Client
+
+
+def _read_attribs_SingleFile(file: h5py.File):
+    attribs = {}
+    for k in file.attrs.keys():
+        attr = file.attrs[k]
+        if type(attr) is bytes or type(attr) is np.bytes_:
+            attribs[k] = attr.decode("UTF-8")
+        else:
+            attribs[k] = attr
+    return attribs
+
+
+class Container:
+    def __init__(
+        self, path, single_file=False, pickle=True, greek=False, dask_props={}
+    ):
+        super(Container, self).__init__()
+
+        self.client = Client(**dask_props)
+        if self.client.status == "running":
+            print("Dask client launched:")
+            print(self.client)
+
+        self.configs: dict[str, Any] = {
+            "single_file": single_file,
+            "use_pickle": pickle,
+            "use_greek": greek,
+        }
+        self.path = path
+        self.metadata = {}
+        self.mesh = None
+        if self.configs["single_file"]:
+            try:
+                self.master_file: h5py.File | None = h5py.File(self.path, "r")
+            except OSError:
+                raise OSError(f"Could not open file {self.path}")
+        else:
+            self.master_file: h5py.File | None = None
+            raise NotImplementedError("Multiple files not yet supported")
+
+        self.attrs = _read_attribs_SingleFile(self.master_file)
+
+        if self.configs["single_file"]:
+            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"
+            )
+            self.configs["dimension"] = int(self.master_file.attrs.get("Dimension", 1))
+            self.configs["coordinates"] = self.master_file.attrs.get(
+                "Coordinates", b"cart"
+            ).decode("UTF-8")
+            if self.configs["coordinates"] == "qsph":
+                self.configs["coordinates"] = "sph"
+            # if coordinates == "sph":
+            #     self.metric = SphericalMetric()
+            # else:
+            #     self.metric = MinkowskiMetric()
+
+    def plotGrid(self, ax, **kwargs):
+        from matplotlib import patches
+
+        xlim, ylim = ax.get_xlim(), ax.get_ylim()
+        options = {
+            "lw": 1,
+            "color": "k",
+            "ls": "-",
+        }
+        options.update(kwargs)
+
+        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)
+        else:
+            for r in self.attrs["X1"]:
+                ax.add_patch(
+                    patches.Arc(
+                        (0, 0),
+                        2 * r,
+                        2 * r,
+                        theta1=-90,
+                        theta2=90,
+                        fill=False,
+                        **options,
+                    )
+                )
+            for th in self.attrs["X2"]:
+                ax.plot(
+                    [
+                        self.attrs["X1Min"] * np.sin(th),
+                        self.attrs["X1Max"] * np.sin(th),
+                    ],
+                    [
+                        self.attrs["X1Min"] * np.cos(th),
+                        self.attrs["X1Max"] * np.cos(th),
+                    ],
+                    **options,
+                )
+        ax.set(xlim=xlim, ylim=ylim)
+
+    def print_container(self) -> str:
+        return f"Client {self.client}\n"
+
+    #
+    # def makeMovie(self, plot, makeframes=True, **kwargs):
+    #     """
+    #     Makes a movie from a plot function
+    #
+    #     Parameters
+    #     ----------
+    #     plot : function
+    #         The plot function to use; accepts output timestep and dataset as arguments.
+    #     makeframes : bool, optional
+    #         Whether to make the frames, or just proceed to making the movie. Default is True.
+    #     num_cpus : int, optional
+    #         The number of CPUs to use for making the frames. Default is None.
+    #     **kwargs :
+    #         Additional keyword arguments passed to `ffmpeg`.
+    #     """
+    #     import numpy as np
+    #
+    #     if makeframes:
+    #         makemovie = all(
+    #             exp.makeFrames(
+    #                 plot,
+    #                 np.arange(len(self.t)),
+    #                 f"{self.attrs['simulation.name']}/frames",
+    #                 data=self,
+    #                 num_cpus=kwargs.pop("num_cpus", None),
+    #             )
+    #         )
+    #     else:
+    #         makemovie = True
+    #     if makemovie:
+    #         exp.makeMovie(
+    #             input=f"{self.attrs['simulation.name']}/frames/",
+    #             overwrite=True,
+    #             output=f"{self.attrs['simulation.name']}.mp4",
+    #             number=5,
+    #             **kwargs,
+    #         )
+    #         return True

nt2/containers/fields.py

@@ -0,0 +1,229 @@
+import h5py
+import xarray as xr
+import numpy as np
+from dask.array.core import from_array
+from dask.array.core import stack
+
+from nt2.containers.container import Container
+from nt2.containers.utils import _read_category_metadata_SingleFile
+
+
+def _read_coordinates_SingleFile(coords: list[str], file: h5py.File):
+    for st in file:
+        group = file[st]
+        if isinstance(group, h5py.Group):
+            if any([k.startswith("X") for k in group if k is not None]):
+                # cell-centered coords
+                xc = {
+                    c: (
+                        np.asarray(xi[:])
+                        if isinstance(xi := group[f"X{i+1}"], h5py.Dataset) and xi
+                        else None
+                    )
+                    for i, c in enumerate(coords[::-1])
+                }
+                # cell edges
+                xe_min = {
+                    f"{c}_1": (
+                        c,
+                        (
+                            np.asarray(xi[:-1])
+                            if isinstance((xi := group[f"X{i+1}e"]), h5py.Dataset)
+                            else None
+                        ),
+                    )
+                    for i, c in enumerate(coords[::-1])
+                }
+                xe_max = {
+                    f"{c}_2": (
+                        c,
+                        (
+                            np.asarray(xi[1:])
+                            if isinstance((xi := group[f"X{i+1}e"]), h5py.Dataset)
+                            else None
+                        ),
+                    )
+                    for i, c in enumerate(coords[::-1])
+                }
+                return {"x_c": xc, "x_emin": xe_min, "x_emax": xe_max}
+        else:
+            raise ValueError(f"Unexpected type {type(file[st])}")
+    raise ValueError("Could not find coordinates in file")
+
+
+def _preload_field_SingleFile(
+    k: str,
+    dim: int,
+    ngh: int,
+    outsteps: list[int],
+    times: list[float],
+    steps: list[int],
+    coords: list[str],
+    xc_coords: dict[str, str],
+    xe_min_coords: dict[str, str],
+    xe_max_coords: dict[str, str],
+    coord_replacements: list[tuple[str, str]],
+    field_replacements: list[tuple[str, str]],
+    layout: str,
+    file: h5py.File,
+):
+    if dim == 1:
+        noghosts = slice(ngh, -ngh) if ngh > 0 else slice(None)
+    elif dim == 2:
+        noghosts = (slice(ngh, -ngh), slice(ngh, -ngh)) if ngh > 0 else slice(None)
+    elif dim == 3:
+        noghosts = (
+            (slice(ngh, -ngh), slice(ngh, -ngh), slice(ngh, -ngh))
+            if ngh > 0
+            else slice(None)
+        )
+    else:
+        raise ValueError("Invalid dimension")
+
+    dask_arrays = []
+    for s in outsteps:
+        dset = file[f"{s}/{k}"]
+        if isinstance(dset, h5py.Dataset):
+            array = from_array(np.transpose(dset) if layout == "right" else dset)
+            dask_arrays.append(array[noghosts])
+        else:
+            raise ValueError(f"Unexpected type {type(dset)}")
+
+    k_ = k[1:]
+    for c in coord_replacements:
+        if "_" not in k_:
+            k_ = k_.replace(c[0], c[1])
+        else:
+            k_ = "_".join([k_.split("_")[0].replace(c[0], c[1])] + k_.split("_")[1:])
+    for f in field_replacements:
+        k_ = k_.replace(*f)
+
+    return k_, xr.DataArray(
+        stack(dask_arrays, axis=0),
+        dims=["t", *coords],
+        name=k_,
+        coords={
+            "t": times,
+            "s": ("t", steps),
+            **xc_coords,
+            **xe_min_coords,
+            **xe_max_coords,
+        },
+    )
+
+
+class FieldsContainer(Container):
+    def __init__(self, **kwargs):
+        super(FieldsContainer, self).__init__(**kwargs)
+        QuantityDict = {
+            "Ttt": "E",
+            "Ttx": "Px",
+            "Tty": "Py",
+            "Ttz": "Pz",
+        }
+        CoordinateDict = {
+            "cart": {"x": "x", "y": "y", "z": "z", "1": "x", "2": "y", "3": "z"},
+            "sph": {
+                "r": "r",
+                "theta": "θ" if self.configs["use_greek"] else "th",
+                "phi": "φ" if self.configs["use_greek"] else "ph",
+                "1": "r",
+                "2": "θ" if self.configs["use_greek"] else "th",
+                "3": "φ" if self.configs["use_greek"] else "ph",
+            },
+        }
+        if self.configs["single_file"]:
+            assert self.master_file is not None, "Master file not found"
+            self.metadata["fields"] = _read_category_metadata_SingleFile(
+                "f", self.master_file
+            )
+        else:
+            try:
+                raise NotImplementedError("Multiple files not yet supported")
+            except OSError:
+                raise OSError(f"Could not open file {self.path}")
+
+        coords = list(CoordinateDict[self.configs["coordinates"]].values())[::-1][
+            -self.configs["dimension"] :
+        ]
+
+        if self.configs["single_file"]:
+            self.mesh = _read_coordinates_SingleFile(coords, self.master_file)
+        else:
+            raise NotImplementedError("Multiple files not yet supported")
+
+        self.fields = xr.Dataset()
+
+        if len(self.metadata["fields"]["outsteps"]) > 0:
+            if self.configs["single_file"]:
+                for k in self.metadata["fields"]["quantities"]:
+                    name, dset = _preload_field_SingleFile(
+                        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,
+                    )
+                    self.fields[name] = dset
+            else:
+                raise NotImplementedError("Multiple files not yet supported")
+
+    def __del__(self):
+        if self.configs["single_file"] and self.master_file is not None:
+            self.master_file.close()
+        else:
+            raise NotImplementedError("Multiple files not yet supported")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.configs["single_file"] and self.master_file is not None:
+            self.master_file.close()
+        else:
+            raise NotImplementedError("Multiple files not yet supported")
+
+    def print_fields(self) -> str:
+        def sizeof_fmt(num, suffix="B"):
+            for unit in ("", "K", "M", "G", "T", "P", "E", "Z"):
+                if abs(num) < 1e3:
+                    return f"{num:3.1f} {unit}{suffix}"
+                num /= 1e3
+            return f"{num:.1f} Y{suffix}"
+
+        def compactify(lst):
+            c = ""
+            cntr = 0
+            for l_ in lst:
+                if cntr > 5:
+                    c += "\n                "
+                    cntr = 0
+                c += l_ + ", "
+                cntr += 1
+            return c[:-2]
+
+        string = ""
+        field_keys = list(self.fields.data_vars.keys())
+
+        if len(field_keys) > 0:
+            string += "Fields:\n"
+            string += f"  - data axes: {compactify(self.fields.indexes.keys())}\n"
+            string += f"  - timesteps: {self.fields[field_keys[0]].shape[0]}\n"
+            string += f"  - shape: {self.fields[field_keys[0]].shape[1:]}\n"
+            string += f"  - quantities: {compactify(self.fields.data_vars.keys())}\n"
+            string += f"  - total size: {sizeof_fmt(self.fields.nbytes)}\n"
+        else:
+            string += "Fields: empty\n"
+
+        return string