py-pluto 1.1.4__py3-none-any.whl

pyPLUTO/loadfuncs/initload.py

@@ -0,0 +1,203 @@
+"""Module for initialization loading functions."""
+
+import warnings
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from pyPLUTO.baseloadmixin import BaseLoadMixin
+from pyPLUTO.baseloadstate import BaseLoadState
+from pyPLUTO.loadfuncs.codeselection import CodeManager
+from pyPLUTO.loadfuncs.descriptor import DescriptorManager
+from pyPLUTO.loadfuncs.findfiles import FindFilesManager
+from pyPLUTO.loadfuncs.findformat import FindFormat
+from pyPLUTO.loadfuncs.loadvars import LoadVariables
+from pyPLUTO.loadstate import LoadState
+from pyPLUTO.utils.inspector import track_kwargs
+
+
+@track_kwargs
+class InitLoadManager(BaseLoadMixin[BaseLoadState]):
+    """Class that handles the initialization loading process."""
+
+    def __init__(
+        self,
+        state: BaseLoadState,
+        nout: int | str | list[int | str] | None,
+        **kwargs: Any,  # Unpack[MyKwargs]
+    ) -> None:
+        """Initialize the InitLoadManager class."""
+        self.state = state
+
+        if nout is None:
+            warnings.warn("No output is loaded!", UserWarning, stacklevel=2)
+            return
+
+        # Check the endianess
+        self.check_endian(kwargs.get("endian", self.endian))
+
+        # Check the input multiple
+        self.multiple = kwargs.get("multiple", self.multiple)
+        if not isinstance(self.multiple, bool):
+            raise TypeError("Invalid data type. 'multiple' must be a boolean.")
+
+        # Check the path
+        self.check_path(kwargs.get("path", self.pathdir))
+
+        self.code = kwargs.get("code", self.code)
+        if self.code.lower() not in {"pluto", "gpluto"}:
+            self.CodeManager = CodeManager(state, nout, **kwargs)
+            return
+
+        self.FindFormat = FindFormat(state, **kwargs)
+
+        if isinstance(state, LoadState) and not self.alone:
+            try:
+                self.Descriptor = DescriptorManager(state, nout, **kwargs)
+            except UserWarning:
+                warnings.warn(
+                    "Failed to initialize descriptor manager.", stacklevel=2
+                )
+                self.alone = True
+
+        if not isinstance(state, LoadState) or self.alone:
+            self.findfile = FindFilesManager(state, nout, **kwargs)
+
+        loadvars = True
+        if kwargs.get("vars") is not None:
+            warnings.warn(
+                "'vars' argument is deprecated. Use 'var' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            loadvars = kwargs.get("vars", loadvars)
+        loadvars = kwargs.get("var", loadvars)
+
+        for i, exout in enumerate(self.noutlist):
+            LoadVariables(state, loadvars, i, exout)
+
+        for key in self.d_vars:
+            setattr(self.state, str(key), self.d_vars[key])
+
+        if isinstance(self.ntimelist, np.ndarray) and len(self.ntimelist) == 1:
+            self.ntime = self.ntimelist[0]
+            self.nout = self.noutlist[0]
+        else:
+            self.ntime = self.ntimelist
+            self.nout = self.noutlist
+
+    def check_endian(self, endian: str | None) -> None:
+        """Check the endian format.
+
+        If the endian is given, check if it is valid (either 'little' or
+        'big'). If not given, set it to 'little' by default.
+
+        Parameters
+        ----------
+        endian: str | None
+            The endian format. If None, it is set to 'little' by default.
+
+        Returns
+        -------
+        None
+
+        Examples
+        --------
+        - Example #1: Set endian to 'big'
+
+            >>> check_endian("big")
+
+        - Example #2: Set endian to 'little'
+
+            >>> check_endian("little")
+
+        - Example #3: Set endian to default ('little')
+
+            >>> check_endian(None)
+
+        - Example #4: Invalid endian value
+
+            >>> check_endian("invalid")
+            ValueError: Invalid endian value. Must be 'little' or 'big'.
+
+        """
+        # Check the input endianess
+        d_end = {
+            "big": ">",
+            "little": "<",
+            ">": ">",
+            "<": "<",
+            None: None,
+        }
+
+        self.endian = endian
+        if endian not in d_end:
+            error = f"Invalid endianess. Valid values are {d_end.keys()}"
+            raise ValueError(error)
+        self.endian = d_end[endian]
+
+    def check_path(self, path: str | Path) -> None:
+        """Check if the given path is consistent.
+
+        If the path is given through a non- empty string or set to the default
+        value, If the path is consistent, it is converted to a Path object.
+        Then, a check is performed to see if the path is a directory. The path
+        is stored in the class as a Path object self.pathdir.
+
+        Returns
+        -------
+        - None
+
+        Parameters
+        ----------
+        - path (not optional): str | Path
+            The path to the simulation directory.
+
+        ----
+
+        Examples
+        --------
+        - Example #1: path is a string
+
+            >>> _check_pathformat("path/to/simulation")
+
+        - Example #2: path is not a string
+
+            >>> _check_pathformat(1)
+            TypeError: Invalid data type. 'path' must be a non-empty string.
+
+        - Example #3: path is an empty string
+
+            >>> _check_pathformat("")
+            ValueError: 'path' cannot be an empty string.
+
+        - Example #4: path is not a directory
+
+            >>> _check_pathformat("path/to/simulation")
+            NotADirectoryError: Directory path/to/simulation not found.
+
+        """
+        # Check if the path is a non-empty string.
+        if not isinstance(path, str) and not isinstance(path, Path):
+            error = TypeError(
+                "Invalid data type. 'path' must be path or string"
+            )
+            raise TypeError(error)
+        # Check if the path is not empty
+        elif not isinstance(path, Path) and not path.strip():
+            raise ValueError("'path' cannot be an empty string.")
+        # Convert the path to a Path object and store it
+
+        self.pathdir = Path(path)
+
+        if not isinstance(self.pathdir, Path):
+            raise TypeError(
+                "Invalid data type. 'path' must be or converted to Path object."
+            )
+
+        # Check that the path is a directory
+        if not self.pathdir.is_dir():
+            raise NotADirectoryError(f"Directory {self.pathdir} not found.")
+
+        # End of the function

pyPLUTO/loadfuncs/loadvars.py

@@ -0,0 +1,227 @@
+"""Docstring for pyPLUTO.loadfuncs.loadvars module."""
+
+import mmap
+import warnings
+
+import numpy as np
+
+from pyPLUTO.baseloadmixin import BaseLoadMixin
+from pyPLUTO.baseloadstate import BaseLoadState
+from pyPLUTO.loadfuncs.offsetdata import OffsetData
+
+
+class LoadVariables(BaseLoadMixin[BaseLoadState]):
+    """Class that handles the loading of variables from PLUTO output files."""
+
+    def __init__(
+        self,
+        state: BaseLoadState,
+        variables: str | list[str] | bool | None,
+        index: int,
+        exout: int,
+    ) -> None:
+        """Initialize the LoadVariables class."""
+        self.state = state
+        self.offsetdata = OffsetData(state)
+        self.load_variables(variables, index, exout)
+
+    def load_variables(
+        self, variables: str | list[str] | bool | None, i: int, exout: int
+    ) -> None:
+        """Load the variables in the class.
+
+        The function checks if the variables to be loaded are valid and then
+        loads them. If the variables are not valid, an error is raised. If the
+        variables are valid, the function loads them in the class through memory
+        mapping. The offset and shape of each variable is computed depenging on
+        the format and typefile characteristics. In case the files are
+        standalone, the relevand time and grid information is loaded.
+
+        Returns
+        -------
+        - None
+
+        Parameters
+        ----------
+        - endian (not optional): bool
+            The endianess of the files. If True the endianess is big, otherwise
+            it is little.
+        - exout (not optional): int
+            The index of the output to be loaded.
+        - i (not optional): int
+            The index of the file to be loaded.
+        - vars (not optional): str | list[str] | bool | None, default True
+            If True all the variables are loaded, otherwise just a selection is
+            loaded.
+
+        ----
+
+        Examples
+        --------
+        - Example #1: Load all the variables
+
+            >>> _load_variables(True, 0, 0, True)
+
+        - Example #2: Load only the selected variables
+
+            >>> _load_variables(["rho", "vx1"], 0, 0, True)
+
+        - Example #3: Load all the variables (little endian)
+
+            >>> _load_variables(True, 0, 0, False)
+
+        - Example #4: Load all the variables from a specific output file
+
+            >>> _load_variables(True, 0, 1, True)
+
+        """
+        # Initialize mm variable to None to avoid uninitialized variable error
+        mm = None
+
+        # Find the class name and find the single_file filepath
+        if self.class_name == "Load":
+            # If the class name is Load (single file), the filepath is data
+            self.filepath = self.pathdir / (
+                "data" + self.d_info["endpath"][exout]
+            )
+        elif self.class_name == "LoadPart":
+            # If the class name is LoadPart, the filepath is particles
+            self.filepath = self.pathdir / (
+                "particles" + self.d_info["endpath"][exout]
+            )
+        else:
+            # If the class name is not recognized, raise an error
+            raise NameError("Invalid class name.")
+
+        # If files in single_file format, inspect the file
+        # or compute the offset and shape
+
+        if self.d_info["typefile"][exout] == "single_file":
+            with open(self.filepath, "rb") as fd:
+                kwargs = {"access": mmap.ACCESS_READ}
+                # if sys.version_info >= (3, 13):
+                #    kwargs["trackfd"] = False
+                mm = mmap.mmap(fd.fileno(), 0, **kwargs)
+            self.offsetdata.compute_offset(i, exout, None, mm)
+            if self.format in ("hdf5", "tab"):
+                return None
+
+        # Check if only specific variables should be loaded
+        if variables is True:
+            # If all the variables are to be loaded, the load_vars
+            # is set to the variables list
+            load_vars = self.d_info["varslist"][exout]
+        elif isinstance(variables, list | str):
+            # If only specific variables are to be loaded, the load_vars
+            # becomes the list of the selected variables
+            load_vars = (
+                variables if isinstance(variables, list) else [variables]
+            )
+        else:
+            # If no variables are to be loaded, return None (WIP)
+            return None
+
+        for j in load_vars:
+            if self.d_info["typefile"][exout] == "multiple_files":
+                self.filepath = self.pathdir / (
+                    j + self.d_info["endpath"][exout]
+                )
+                try:
+                    with open(self.filepath, "rb") as fd:
+                        kwargs = {"access": mmap.ACCESS_READ}
+                        # if sys.version_info >= (3, 13):
+                        #    kwargs["trackfd"] = False
+                        mm = mmap.mmap(fd.fileno(), 0, **kwargs)
+                except (OSError, ValueError) as e:
+                    warnings.warn(
+                        f"AttributeError: Unable to open {self.filepath}.\n"
+                        f"{e}\nSkipping variable {j}.",
+                        UserWarning,
+                        stacklevel=2,
+                    )
+                    continue
+                self.offsetdata.compute_offset(i, exout, j, mm)
+
+            if self.lennout != 1:
+                self.init_vardict(j)
+
+            if mm is None:
+                raise RuntimeError("memmap object not initialized")
+
+            dtype = np.dtype(self.d_info["binformat"][exout])
+            shape = self.varshape[j]
+            offset = self.varoffset[j]
+            scrh = np.ndarray(
+                shape=shape, dtype=dtype, buffer=mm, offset=offset, order="C"
+            ).T
+
+            self.assign_var(exout, j, scrh)
+
+        # ... then after the variable loop ...
+
+    def assign_var(self, time: int, var: str, scrh: np.ndarray) -> None:
+        """Assign the variable data to the class variable.
+
+        The function assigns the data to the class variable. The variable name
+        is constructed based on the variable name and the output index.
+
+        Returns
+        -------
+        - None
+
+        Parameters
+        ----------
+        - exout (not optional): int
+            The index of the output to be loaded.
+        - var (not optional): str
+            The variable name.
+        - data (not optional): np.ndarray
+            The data to be assigned to the variable.
+
+        ----
+
+        Examples
+        --------
+        - Example #1: Assign data to a variable
+
+            >>> _assign_var(0, "rho", data)
+
+        """
+        # Assign the memmap object to the dictionary
+        if self.lennout != 1:
+            # If the number of outputs is not 1, the variable is stored at the
+            # corresponding output
+            self.d_vars[var][time] = scrh
+        else:
+            # If the number of outputs is 1, the variable is stored directly
+            self.d_vars[var] = scrh
+
+        # End of the function
+
+    def init_vardict(self, var: str) -> None:
+        """If not initialized, a new dictionary is created to store the vars.
+
+        The dictionary is stored in the class. The shape of the dictionary is
+        computed depending on the number of outputs and the variable shape.
+
+        Returns
+        -------
+        - None
+
+        Parameters
+        ----------
+        - var (not optional): str
+            The variable to be loaded.
+
+        ----
+
+        Examples
+        --------
+        - Example #1: Initialize the dictionary of a non-initialized variable
+
+            >>> _init_vardict("rho")
+
+        """
+        # If the variable is not initialized, create a new dictionary
+        if var not in self.d_vars:
+            self.d_vars[var] = {}

pyPLUTO/loadfuncs/offsetdata.py

@@ -0,0 +1,87 @@
+"""Docstring for pyPLUTO.loadfuncs.offsetdata."""
+
+import mmap
+
+from pyPLUTO.baseloadmixin import BaseLoadMixin
+from pyPLUTO.baseloadstate import BaseLoadState
+from pyPLUTO.loadfuncs.offsetfluid import OffsetFluid
+from pyPLUTO.loadfuncs.readgridfile import GridFileManager
+from pyPLUTO.loadfuncs.readtab import ReadtabManager
+from pyPLUTO.loadstate import LoadState
+
+
+class OffsetData(BaseLoadMixin[BaseLoadState]):
+    """Class that computes the offset of variables in single_file format."""
+
+    def __init__(self, state: BaseLoadState) -> None:
+        self.state = state
+        if isinstance(state, LoadState):
+            self.GridFileManager = GridFileManager(state)
+            self.Offsetclass = OffsetFluid(state)
+            self.ReadtabManager = ReadtabManager(state)
+
+    def compute_offset(
+        self,
+        i: int,
+        exout: int,
+        varname: str | None,
+        mm: mmap.mmap,
+    ) -> None:
+        """Compute the offset and shape of the variable in single_file format.
+
+        The function computes the offset and shape of the variable in
+        single_file format. The offset is computed based on the variable name,
+        output index, and file index. The shape is computed based on the
+        variable characteristics and the grid information.
+
+        Returns
+        -------
+        - None
+
+        Parameters
+        ----------
+        - i (not optional): int
+            The index of the file to be loaded.
+        - exout (not optional): int
+            The index of the output to be loaded.
+        - varname (not optional): str | None
+            The name of the variable to be loaded. If None, all variables are
+            considered.
+
+        ----
+
+        Examples
+        --------
+        - Example #1: Compute offset for a specific variable
+
+            >>> _compute_offset(0, 0, "rho")
+
+        - Example #2: Compute offset for all variables
+
+            >>> _compute_offset(0, 0, None)
+        """
+        if (
+            self.alone is not True
+            and isinstance(self.state, LoadState)
+            and self.infogrid is True
+        ):
+            self.GridFileManager.read_gridfile()
+            self.infogrid = False
+
+        fmt = "h5" if self.format in {"dbl.h5", "flt.h5"} else self.format
+        fmt = "bin" if fmt in {"dbl", "flt"} else fmt
+
+        if isinstance(self.state, LoadState):
+            handlers = {
+                "tab": self.ReadtabManager.read_tab,
+                "bin": self.Offsetclass.offset_bin,
+                "vtk": self.Offsetclass.offset_vtk,
+                "h5": self.Offsetclass.offset_h5,
+                "hdf5": self.Offsetclass.offset_hdf5,
+            }
+        else:
+            raise TypeError(
+                "OffsetData requires LoadState for now "
+                "(particles still not implemented)."
+            )
+        handlers.get(fmt, self.Offsetclass.offset_bin)(i, varname, exout, mm)