py-pluto 1.1.4__py3-none-any.whl

pyPLUTO/loadfuncs/offsetfluid.py

@@ -0,0 +1,408 @@
+"""Docstring for pyPLUTO.loadfuncs.offsetfluid module."""
+
+import mmap
+import struct
+import warnings
+from pathlib import Path
+
+import h5py
+import numpy as np
+
+from pyPLUTO.loadfuncs.readgridalone import GridManager
+from pyPLUTO.loadmixin import LoadMixin
+from pyPLUTO.loadstate import LoadState
+
+
+class OffsetFluid(LoadMixin):
+    """Class that computes the fluid offsets in single_file format."""
+
+    from pyPLUTO.amr import _DataScanHDF5, _inspect_hdf5, _read_gridfile
+    from pyPLUTO.loadfuncs.readgridout import _split_gridfile
+    from pyPLUTO.toolfuncs.transform import _congrid
+
+    def __init__(self, state: LoadState) -> None:
+        self.state = state
+        self.varoffset, self.varshape = ({}, {})
+        self.GridAloneManager = GridManager(state)
+
+    def offset_bin(
+        self, _i: int, var: str | None, exout: int, _mm: mmap.mmap
+    ) -> None:
+        """Compute the offset and shape of the variables to be loaded.
+
+        The routine, knowing the grid shape, computes the offset and stores the
+        shape dependng on wether the variable is staggered or not.
+
+        Returns
+        -------
+        - None
+
+        Parameters
+        ----------
+        - i (not optional): int
+            The index of the file to be loaded.
+        - var (not optional): str
+            The variable to be loaded.
+
+        ----
+
+        Examples
+        --------
+        - Example #1: Load all the variables
+
+            >>> _offset_bin(0, True)
+
+        """
+        off_start = 0
+
+        grid_sizes: dict[str, tuple[int | tuple[int, ...] | None, int]] = {
+            "Bx1s": (self.nshp_st1, self.gridsize_st1),
+            "Ex1s": (self.nshp_st1, self.gridsize_st1),
+            "Bx2s": (self.nshp_st2, self.gridsize_st2),
+            "Ex2s": (self.nshp_st2, self.gridsize_st2),
+            "Bx3s": (self.nshp_st3, self.gridsize_st3),
+            "Ex3s": (self.nshp_st3, self.gridsize_st3),
+        }
+
+        # Loop over the variables to be loaded (None for single files)
+        varloop = self.d_info["varslist"][exout] if var is None else [var]
+
+        for eachvar in varloop:
+            # Get the grid shape and size (centered or staggered)
+            grid_size = grid_sizes.get(eachvar, [self.nshp, self.gridsize])
+            self.varshape[eachvar] = grid_size[0]
+            # Assign the offset
+            self.varoffset[eachvar] = off_start
+            # Move to next variable
+            if isinstance(grid_size[1], int):
+                off_start += grid_size[1] * self.charsize
+            else:
+                raise ValueError("Grid size must be an integer.")
+
+        # End of function
+
+    def offset_h5(
+        self, i: int, _var: str | None, exout: int, _mm: mmap.mmap
+    ) -> None:
+        """Compute the offset and shape of the variable in hdf5 format.
+
+        The routine, knowing the grid shape, computes the offset and stores the
+        shape dependng on wether the variable is staggered or not.
+
+        Returns
+        -------
+        - None
+
+        Parameters
+        ----------
+        - i (not optional): int
+            The index of the file to be loaded.
+        - var (not optional): str
+            The variable to be loaded.
+
+        ----
+
+        Examples
+        --------
+        - Example #1: Load all the variables
+
+            >>> _offset_h5(0, True)
+
+        """
+        # Open the file with the h5py library
+        h5file = h5py.File(str(self.filepath), "r")
+
+        # Selects the binformat
+        self.d_info["binformat"][exout] = (
+            "d" if self.format == "dbl.h5" else "f"
+        )
+
+        # Safely access the timestep group and its sub-items to avoid
+        # treating an h5py.Datatype as a subscriptable object.
+        timestep_key = f"Timestep_{exout}"
+        timestep = h5file.get(timestep_key, None)
+
+        if isinstance(timestep, h5py.Group):
+            timeattr = timestep.attrs["Time"]
+            idx = np.searchsorted(self.outlist, exout)
+            self.timelist[idx] = float(timeattr)
+            idx = np.searchsorted(self.noutlist, exout)
+            self.ntimelist[idx] = float(timeattr)
+            cellvs = timestep.get("vars", {})
+            stagvs = timestep.get("stag_vars", {})
+            cellvs = (
+                {}
+                if cellvs is None or isinstance(cellvs, h5py.Datatype)
+                else cellvs
+            )
+
+            stagvs = (
+                {}
+                if stagvs is None or isinstance(stagvs, h5py.Datatype)
+                else stagvs
+            )
+        else:
+            # Timestep group not present or not a group; treat as empty
+            cellvs = {}
+            stagvs = {}
+
+        # If standalone file, finds the variables to be loaded, else
+        # remove variables in the .out file that are not present in the actual
+        # file
+        if (
+            self.alone is True
+            and isinstance(cellvs, h5py.Group | dict)
+            and isinstance(stagvs, h5py.Group | dict)
+        ):
+            self.d_info["varslist"][exout] = set(cellvs.keys()) | set(
+                stagvs.keys()
+            )
+        elif not isinstance(cellvs, h5py.Group | dict) or not isinstance(
+            stagvs, h5py.Group | dict
+        ):
+            raise ValueError(
+                "Error: Variables group not found in the HDF5 file."
+            )
+
+        # Loop over the variables and store the offset and shape
+        for j in self.d_info["varslist"][exout]:
+            if j in cellvs:
+                obj = cellvs[j]
+            elif j in stagvs:
+                obj = stagvs[j]
+            else:
+                obj = None
+                warnings.warn(
+                    f"Warning: Variable {j} not found in the HDF5 file.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+            # Only Dataset objects provide shape and a file offset
+            if isinstance(obj, h5py.Dataset):
+                self.varoffset[j] = obj.id.get_offset()
+                self.varshape[j] = obj.shape
+            elif obj is not None:
+                raise ValueError(
+                    f"Error: Variable {j} in the HDF5 file is not a dataset."
+                )
+
+        if self.alone is True and self.infogrid is True:
+            self.x1 = h5file["cell_coords"]["X"][:].T
+            self.x2 = h5file["cell_coords"]["Y"][:].T
+            self.x3 = h5file["cell_coords"]["Z"][:].T
+            self.x1r = h5file["node_coords"]["X"][:].T
+            self.x2r = h5file["node_coords"]["Y"][:].T
+            self.x3r = h5file["node_coords"]["Z"][:].T
+            self.GridAloneManager.readgridh5()
+        self.infogrid = False
+
+        # Close the file
+        h5file.close()
+
+    def offset_vtk(
+        self, i: int, var: str | None, exout: int, mm: mmap.mmap
+    ) -> None:
+        """Compute the offset and shape of the variables to be loaded.
+
+        The routine, knowing the grid shape, computes the offset and stores the
+        shape dependng on wether the variable is staggered or not.
+
+        Returns
+        -------
+        - None
+
+        Parameters
+        ----------
+        - i (not optional): int
+            The index of the file to be loaded.
+        - var (not optional): str
+            The variable to be loaded.
+
+        ----
+
+        Examples
+        --------
+        - Example #1: Load all the variables
+
+            >>> _offset_vtk(0, True)
+
+        """
+        dir_map: dict[str, str] = {}
+        gridvars: list[str] = []
+
+        self.d_info["endianess"][exout] = (
+            ">" if self.endian is None else self.d_info["endianess"][exout]
+        )
+        if self.d_info["endianess"][exout] is None:
+            raise ValueError("Error: Wrong endianess in vtk file.")
+
+        if self.alone is True:
+            self.d_info["binformat"][exout] = (
+                f"{self.d_info['endianess'][exout]}f{self.charsize}"
+            )
+        search_pos = 0
+        while True:
+            line_end = mm.find(b"\n", search_pos)
+            if line_end == -1:
+                break  # No more occurrences found
+            line = mm[search_pos:line_end]
+
+            if line.startswith(b"SCALARS"):
+                parts = line.split()
+                namevar = parts[1].decode()
+
+                # Move to the start of the scalar data
+                lookup_table_pos = mm.find(b"LOOKUP_TABLE default", line_end)
+                offset = mm.find(b"\n", lookup_table_pos) + 1
+
+                self.varoffset[namevar] = offset
+                self.varshape[namevar] = self.nshp
+                if var is not None:
+                    break
+
+                search_pos = offset + self.gridsize * self.charsize
+
+            elif line.startswith(b"VECTORS"):
+                # Handle VECTORS data
+                warnings.warn(
+                    "Warning: VECTORS data is not supported.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+                search_pos = line_end + 1
+            elif line.startswith(b"DIMENSIONS") and self.infogrid is True:
+                self.nx1, self.nx2, self.nx3 = [
+                    max(int(i) - 1, 1) for i in line.split()[1:4]
+                ]
+                if self.nx3 == 1 and self.nx2 == 1:
+                    self.dim = 1
+                    self.nshp = self.nx1
+                    self.gridsize = self.nx1
+                    nshp_grid = self.nx1 + 1
+                    gridvars = ["x1r", "x2", "x3"]
+                elif self.nx3 == 1:
+                    self.dim = 2
+                    self.nshp = (self.nx2, self.nx1)
+                    self.gridsize = self.nx1 * self.nx2
+                    nshp_grid = (self.nx2 + 1, self.nx1 + 1)
+                    gridvars = ["x1r", "x2r", "x3"]
+                else:
+                    self.dim = 3
+                    self.nshp = (self.nx3, self.nx2, self.nx1)
+                    self.gridsize = self.nx1 * self.nx2 * self.nx3
+                    nshp_grid = (self.nx3 + 1, self.nx2 + 1, self.nx1 + 1)
+                    gridvars = ["x1r", "x2r", "x3r"]
+                dir_map = {"X": gridvars[0], "Y": gridvars[1], "Z": gridvars[2]}
+                search_pos = line_end + 1
+            elif line.startswith(b"TIME") and self.alone is True:
+                try:
+                    parts = line.split()
+                    binf = 8 if parts[3].decode() == "double" else 4
+                    raw = mm[line_end + 1 : line_end + 1 + binf]
+                    scrh = struct.unpack(
+                        self.d_info["endianess"][exout] + "d", raw
+                    )[0]
+                    self.timelist[exout] = scrh
+                    idx = np.searchsorted(self.noutlist, exout)
+                    self.ntimelist[idx] = scrh
+                except Exception:
+                    binf = 0
+                search_pos = line_end + 1 + binf
+
+            elif line[1:].startswith(b"_COORDINATES") and self.infogrid is True:
+                self.geom = "CARTESIAN"
+                linesplit = line.split()
+                var_sel = linesplit[0].decode()[0]
+                binf = (
+                    self.d_info["endianess"][exout] + "d"
+                    if linesplit[2].decode() == "double"
+                    else self.d_info["endianess"][exout] + "f"
+                )
+                scrh = np.ndarray(
+                    shape=int(linesplit[1]),
+                    dtype=binf,
+                    buffer=mm,
+                    offset=line_end + 1,
+                    order="C",
+                ).T
+                setattr(self, dir_map[var_sel], scrh)
+                search_pos = line_end + 1
+            else:
+                search_pos = line_end + 1
+
+        if (
+            self.d_info["typefile"][exout] == "single_file"
+            and self.alone is True
+        ):
+            self.d_info["varslist"][exout] = list(self.varoffset.keys())
+        if self.infogrid is True:
+            self.GridAloneManager.readgridvtk(gridvars)
+            self.infogrid = False
+
+    def offset_hdf5(
+        self, i: int, _var: str | None, exout: int, _mm: mmap.mmap
+    ) -> None:
+        """Load AMR data from PLUTO/Chombo HDF5 files.
+
+        This method is intentionally isolated from other formats.
+        """
+        # Bridge AMR helpers (legacy naming) to the Newload state fields.
+        self._filepath = self.filepath
+        self._pathgrid = self.pathdir / Path("grid.out")
+        if not hasattr(self, "level"):
+            self.level = 0
+
+        # Keep the original varslist structure used by Newload internals.
+        varslist_table = self.d_info.get("varslist", None)
+        self._inspect_hdf5(i, exout)
+
+        # AMR reader stores a flat variable list; normalize as plain strings.
+        loaded_vars = []
+        for v in self.d_info.get("varslist", []):
+            loaded_vars.append(v.decode() if isinstance(v, bytes) else str(v))
+
+        if (
+            isinstance(varslist_table, list)
+            and len(varslist_table) > exout
+            and isinstance(varslist_table[exout], list | set | np.ndarray)
+        ):
+            self.d_info["varslist"] = varslist_table
+            self.d_info["varslist"][exout] = loaded_vars
+        else:
+            self.d_info["varslist"] = loaded_vars
+
+        self.load_vars = loaded_vars
+
+        # Persist AMR metadata to state so Newload exposes it.
+        for amr_key in ("AMRLevel", "AMRBoxes", "Dt", "n1", "n2", "n3"):
+            if hasattr(self, amr_key):
+                setattr(self.state, amr_key, getattr(self, amr_key))
+
+        # Keep times aligned with nout/out tables when available.
+        try:
+            idx = np.searchsorted(self.outlist, exout)
+            self.timelist[idx] = float(self.ntime)
+        except Exception:
+            pass
+        try:
+            idx = np.searchsorted(self.noutlist, exout)
+            self.ntimelist[idx] = float(self.ntime)
+        except Exception:
+            pass
+
+        # Populate d_vars so Newload can expose loaded AMR variables normally.
+        for varname in loaded_vars:
+            if not hasattr(self, varname):
+                continue
+            data = getattr(self, varname)
+            self.varshape[varname] = np.shape(data)
+            self.varoffset[varname] = 0
+
+            if self.lennout != 1:
+                if varname not in self.d_vars:
+                    self.d_vars[varname] = {}
+                self.d_vars[varname][exout] = data
+            else:
+                self.d_vars[varname] = data
+
+        self.infogrid = False

pyPLUTO/loadfuncs/read_files.py

@@ -0,0 +1,213 @@
+import warnings
+from typing import Any
+
+import h5py
+import numpy as np
+
+
+def read_vtk(self) -> None:
+    """Read the data from a VTK file.
+
+    Returns
+    -------
+    - None
+
+    Parameters
+    ----------
+    - None
+
+    ----
+
+    Examples
+    --------
+    - Example #1: Read the data from a VTK file
+
+        >>> read_vtk()
+
+    """
+    raise NotImplementedError("read_vtk() is not yet implemented.")
+
+
+def _read_h5(self, filename: str, **kwargs: Any) -> None:
+    """Read the data from a HDF5 file.
+
+    Returns
+    -------
+    - the data in a dictionary
+
+    Parameters
+    ----------
+    - filename: str
+        The name of the file to be read.
+    - kwargs: Any
+        Any additional arguments.
+
+    ----
+
+    Examples
+    --------
+    - Example #1: Read the data from a HDF5 file
+
+        >>> read_h5("filename.h5")
+
+    """
+    # Create the path to the HDF5 file
+    try:
+        self._pathh5 = self.pathdir / filename
+    except FileNotFoundError:
+        self._pathh5 = filename
+
+    # Open the HDF5 file
+    data_dict = {}
+    with h5py.File(self._pathh5, "r") as f:
+        for key in f.keys():
+            data_dict[key] = f[key][()]  # Store data in the dictionary
+
+    return data_dict  # Return the dictionary
+
+
+def read_tab(self) -> None:
+    """Read the data from a tab file.
+
+    Returns
+    -------
+    - None
+
+    Parameters
+    ----------
+    - None
+
+    ----
+
+    Examples
+    --------
+    - Example #1: Read the data from a tab file
+
+        >>> read_tab()
+
+    """
+    raise NotImplementedError("read_tab() is not yet implemented.")
+
+
+def read_bin(self) -> None:
+    """Read the data from a binary file.
+
+    Returns
+    -------
+    - None
+
+    Parameters
+    ----------
+    - None
+
+    ----
+
+    Examples
+    --------
+    - Example #1: Read the data from a binary file
+
+        >>> read_bin()
+
+    """
+    raise NotImplementedError("read_bin() is not yet implemented.")
+
+
+def _read_dat(self, filename: str, **kwargs: Any) -> None:
+    """Read the data from a dat file.
+
+    Returns
+    -------
+    - The file columns as a dicionary
+
+    Parameters
+    ----------
+    - filename (not optional): str
+        The name of the file to be read.
+    - kwargs: Any
+        Any additional arguments.
+    - names: bool, default True
+        If True, checks for column names in the file.
+    - skip: int, default 0
+        The number of rows to skip.
+
+    ----
+
+    Examples
+    --------
+    - Example #1: Read the data from a dat file
+
+        >>> read_dat("filename.dat")
+
+    """
+    # Create the path to the HDF5 file
+    try:
+        self._pathh5 = self.pathdir / filename
+    except FileNotFoundError:
+        self._pathh5 = filename
+
+    names = kwargs.get("names", True)
+    skip = kwargs.get("skip", 0)
+
+    # Open the dat file with np.genfromtext
+    data = np.genfromtxt(self._pathh5, names=names, skip_header=skip)
+    if names:
+        loop = data.dtype.names
+        analysis = {name: data[name] for name in loop}
+    else:
+        loop = range(data.shape[1])
+        analysis = {f"col_{i}": data[:, i] for i in loop}
+
+    return analysis
+
+
+def read_file(
+    self, filename: str, datatype: str | None = None, **kwargs: Any
+) -> Any:
+    """Read the data from the output files.
+
+    Returns
+    -------
+    - the data as a dicionary
+
+    Parameters
+    ----------
+    - filename (not optional): str
+        The name of the file to be read.
+    - datatype: str
+        The type of the file.
+    - **kwargs: Any
+        Any additional arguments.
+
+    ----
+
+    Examples
+    --------
+    - Example #1: Read the data from the output files
+
+        >>> read_files()
+
+    """
+    # Check the datatype of the input data
+    if datatype is None:
+        datatype = filename.rsplit(".", maxsplit=1)[-1]
+    poss_types = {"dbl", "flt", "vtk", "h5", "tab", "dat"}
+    if datatype not in poss_types:
+        warn = f"Invalid datatype: {datatype}. Resetting to 'h5'"
+        warnings.warn(warn)
+        datatype = "h5"
+
+    # Check the format of the output files
+    if datatype == "h5":
+        res = _read_h5(self, filename, **kwargs)
+    elif datatype == "dat":
+        res = _read_dat(self, filename, **kwargs)
+    else:
+        warn = (
+            f"Invalid datatype: {datatype}, not implemented yet! "
+            "Resetting to 'h5'"
+        )
+        warnings.warn(warn)
+        pass
+
+    # End of the function
+    return res