py-pluto 1.1.4__py3-none-any.whl

pyPLUTO/baseloadmixin.py

@@ -0,0 +1,258 @@
+"""Mixin class for base load functionality."""
+
+from pathlib import Path
+from typing import Any, TypeVar
+
+from numpy.typing import NDArray
+
+from pyPLUTO.baseloadstate import BaseLoadState
+
+S = TypeVar("S", bound=BaseLoadState)
+
+
+class BaseLoadMixin[S: BaseLoadState]:
+    """Mixin class that provides base functionality for load state."""
+
+    # pylint: disable=too-many-public-methods
+
+    state: S
+
+    @property
+    def alone(self) -> bool:
+        """Get the alone attribute of the load state."""
+        return self.state.alone
+
+    @alone.setter
+    def alone(self, value: bool) -> None:
+        """Set the alone attribute of the load state."""
+        self.state.alone = value
+
+    @property
+    def charsize(self) -> int:
+        """Get the charsize attribute of the load state."""
+        return self.state.charsize
+
+    @charsize.setter
+    def charsize(self, value: int) -> None:
+        """Set the charsize attribute of the load state."""
+        self.state.charsize = value
+
+    @property
+    def class_name(self) -> str:
+        """Get the class_name attribute of the load state."""
+        return self.state.class_name
+
+    @class_name.setter
+    def class_name(self, value: str) -> None:
+        """Set the class_name attribute of the load state."""
+        self.state.class_name = value
+
+    @property
+    def code(self) -> str:
+        """Get the code attribute of the load state."""
+        return self.state.code
+
+    @code.setter
+    def code(self, value: str) -> None:
+        """Set the code attribute of the load state."""
+        self.state.code = value
+
+    @property
+    def d_info(self) -> dict[str, Any]:
+        """Get the d_info attribute of the load state."""
+        return self.state.d_info
+
+    @d_info.setter
+    def d_info(self, value: dict[str, Any]) -> None:
+        """Set the d_info attribute of the load state."""
+        self.state.d_info = value
+
+    @property
+    def d_vars(self) -> dict[str, Any]:
+        """Get the d_vars attribute of the load state."""
+        return self.state.d_vars
+
+    @d_vars.setter
+    def d_vars(self, value: dict[str, Any]) -> None:
+        """Set the d_vars attribute of the load state."""
+        self.state.d_vars = value
+
+    @property
+    def endian(self) -> str | None:
+        """Get the endian attribute of the load state."""
+        return self.state.endian
+
+    @endian.setter
+    def endian(self, value: str | None) -> None:
+        """Set the endian attribute of the load state."""
+        self.state.endian = value
+
+    @property
+    def filepath(self) -> Path:
+        """Get the filepath attribute of the load state."""
+        return self.state.filepath
+
+    @filepath.setter
+    def filepath(self, value: Path) -> None:
+        """Set the filepath attribute of the load state."""
+        self.state.filepath = value
+
+    @property
+    def format(self) -> str:
+        """Get the format attribute of the load state."""
+        return self.state.format
+
+    @format.setter
+    def format(self, value: str) -> None:
+        """Set the format attribute of the load state."""
+        self.state.format = value
+
+    @property
+    def infogrid(self) -> bool:
+        """Get the infogrid attribute of the load state."""
+        return self.state.infogrid
+
+    @infogrid.setter
+    def infogrid(self, value: bool) -> None:
+        """Set the infogrid attribute of the load state."""
+        self.state.infogrid = value
+
+    @property
+    def lennout(self) -> int:
+        """Get the lennout attribute of the load state."""
+        return self.state.lennout
+
+    @lennout.setter
+    def lennout(self, value: int) -> None:
+        """Set the lennout attribute of the load state."""
+        self.state.lennout = value
+
+    @property
+    def lennoutlist(self) -> int:
+        """Get the lenoutlist attribute of the load state."""
+        return self.state.lennoutlist
+
+    @lennoutlist.setter
+    def lennoutlist(self, value: int) -> None:
+        """Set the lennoutlist attribute of the load state."""
+        self.state.lennoutlist = value
+
+    @property
+    def matching_files(self) -> list[str] | None:
+        """Get the matching_files attribute of the load state."""
+        return self.state.matching_files
+
+    @matching_files.setter
+    def matching_files(self, value: list[str] | None) -> None:
+        """Set the matching_files attribute of the load state."""
+        self.state.matching_files = value
+
+    @property
+    def multiple(self) -> bool:
+        """Get the multiple attribute of the load state."""
+        return self.state.multiple
+
+    @multiple.setter
+    def multiple(self, value: bool) -> None:
+        """Set the multiple attribute of the load state."""
+        self.state.multiple = value
+
+    @property
+    def nout(self) -> int | NDArray[Any]:
+        """Get the nout attribute of the load state."""
+        return self.state.nout
+
+    @nout.setter
+    def nout(self, value: int | NDArray[Any]) -> None:
+        """Set the nout attribute of the load state."""
+        self.state.nout = value
+
+    @property
+    def noutlist(self) -> NDArray[Any]:
+        """Get the nout attribute of the load state."""
+        return self.state.noutlist
+
+    @noutlist.setter
+    def noutlist(self, value: NDArray[Any]) -> None:
+        """Set the nout attribute of the load state."""
+        self.state.noutlist = value
+
+    @property
+    def ntime(self) -> int | NDArray[Any]:
+        """Get the ntime attribute of the load state."""
+        return self.state.ntime
+
+    @ntime.setter
+    def ntime(self, value: int | NDArray[Any]) -> None:
+        """Set the ntime attribute of the load state."""
+        self.state.ntime = value
+
+    @property
+    def ntimelist(self) -> NDArray[Any]:
+        """Get the ntime attribute of the load state."""
+        return self.state.ntimelist
+
+    @ntimelist.setter
+    def ntimelist(self, value: NDArray[Any]) -> None:
+        """Set the ntime attribute of the load state."""
+        self.state.ntimelist = value
+
+    @property
+    def outlist(self) -> NDArray[Any]:
+        """Get the outlist attribute of the load state."""
+        return self.state.outlist
+
+    @outlist.setter
+    def outlist(self, value: NDArray[Any]) -> None:
+        """Set the outlist attribute of the load state."""
+        self.state.outlist = value
+
+    @property
+    def pathdir(self) -> str | Path:
+        """Get the pathdir attribute of the load state."""
+        return self.state.pathdir
+
+    @pathdir.setter
+    def pathdir(self, value: str | Path) -> None:
+        """Set the pathdir attribute of the load state."""
+        self.state.pathdir = value
+
+    @property
+    def text(self) -> bool | None:
+        """Get the text attribute of the load state."""
+        return self.state.text
+
+    @text.setter
+    def text(self, value: bool | None) -> None:
+        """Set the text attribute of the load state."""
+        self.state.text = value
+
+    @property
+    def timelist(self) -> NDArray[Any]:
+        """Get the timelist attribute of the load state."""
+        return self.state.timelist
+
+    @timelist.setter
+    def timelist(self, value: NDArray[Any]) -> None:
+        """Set the timelist attribute of the load state."""
+        self.state.timelist = value
+
+    @property
+    def varoffset(self) -> dict[str, Any]:
+        """Get the varoffset attribute of the load state."""
+        return self.state.varoffset
+
+    @varoffset.setter
+    def varoffset(self, value: dict[str, Any]) -> None:
+        """Set the varoffset attribute of the load state."""
+        self.state.varoffset = value
+
+    @property
+    def varshape(self) -> dict[str, Any]:
+        """Get the varshape attribute of the load state."""
+        return self.state.varshape
+
+    @varshape.setter
+    def varshape(self, value: dict[str, Any]) -> None:
+        """Set the varshape attribute of the load state."""
+        self.state.varshape = value

pyPLUTO/baseloadstate.py

@@ -0,0 +1,45 @@
+"""Module that contains the LoadState class."""
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from numpy.typing import NDArray
+
+
+@dataclass(slots=True)
+class BaseLoadState:
+    """Class that stores the state of the Load class.
+
+    Its purpose is to keep track of the current state of the data loading,
+    such as the file paths, data arrays, and other properties and update the
+    key attributes through all the different classes that handle the data
+    loading at runtime.
+    """
+
+    # pylint: disable=too-many-instance-attributes
+
+    alone: bool = field(init=False)
+    charsize: int = field(init=False)
+    class_name: str = field(init=False)
+    code: str = "gPLUTO"
+    d_info: dict[str, Any] = field(init=False)
+    d_vars: dict[str, Any] = field(default_factory=dict)
+    endian: str | None = None
+    filepath: Path = field(init=False)
+    format: str = "Unknown"
+    infogrid: bool = True
+    lennout: int = field(init=False)
+    lennoutlist: int = field(init=False)
+    matching_files: list[str] | None = None
+    multiple: bool = False
+    nout: int | NDArray[Any] = field(init=False)
+    noutlist: NDArray[Any] = field(init=False)
+    ntime: int | NDArray[Any] = field(init=False)
+    ntimelist: NDArray[Any] = field(init=False)
+    outlist: NDArray[Any] = field(init=False)
+    pathdir: str | Path = "./"
+    text: bool | None = None
+    timelist: NDArray[Any] = field(init=False)
+    varoffset: dict[str, Any] = field(init=False)
+    varshape: dict[str, Any] = field(init=False)

pyPLUTO/codes/echo_load.py

@@ -0,0 +1,161 @@
+"""Module to load the data from the output files of the ECHO code."""
+
+import warnings
+from pathlib import Path
+from typing import Any, cast
+
+import h5py
+
+from pyPLUTO.loadmixin import LoadMixin
+from pyPLUTO.loadstate import LoadState
+from pyPLUTO.utils.inspector import track_kwargs
+
+
+@track_kwargs
+class EchoLoadManager(LoadMixin):
+    """Manager for the loading of the ECHO code data."""
+
+    @track_kwargs
+    def __init__(self, state: LoadState) -> None:
+        """Initialize the EchoLoadManager."""
+        self.state = state
+
+    def load_echo(
+        self, nout: int | str | list[int | str] | None, **kwargs: Any
+    ) -> None:
+        """Load the data from the output files of the ECHO code.
+
+        The data are loaded only from h5 files and only a single output is
+        possible. Note that binary files produced by ECHO are not supported
+        by this method. The data are loaded in the PLUTO format, so the
+        variables are renamed to match the PLUTO naming convention.
+
+        Parameters
+        ----------
+        - nout: int | str | list | None, default 0
+            The output number to be loaded.
+        - path: str, default './'
+            The path to the folder containing the data files.
+        - vars: str | list[str] | bool | None, default True
+            The variables to be loaded. If 'True', all the variables are loaded.
+
+        Returns
+        -------
+        - None
+
+        Examples
+        --------
+        Example 1: Load all the variables from the last output in the current
+        folder.
+
+        >>> NOT IMPLEMENTED YET
+
+        """
+        print("load, echo")
+
+        # Geometry is set to CARTESIAN by default
+        self.geom = "CARTESIAN"
+
+        # Dictionary to convert the keys from ECHO to PLUTO
+        conv_dict = {
+            "x": "x1",
+            "y": "x2",
+            "z": "x3",
+            "rh": "rho",
+            "pg": "prs",
+            "se": "ent",
+            "vx": "vx1",
+            "vy": "vx2",
+            "vz": "vx3",
+            "bx": "Bx1",
+            "by": "Bx2",
+            "bz": "Bx3",
+            "ex": "Ex1",
+            "ey": "Ex2",
+            "ez": "Ex3",
+        }
+
+        self.echo_load_grid(conv_dict)
+        self.echo_set_grid_dims()
+
+        if isinstance(nout, str) or nout is None:
+            warnings.warn(
+                "Please specify the output or it will be set to 0.",
+                stacklevel=2,
+            )
+            self.nout = 0
+        elif isinstance(nout, list):
+            if isinstance(nout[0], int):
+                self.nout = nout[0]
+            else:
+                warnings.warn(
+                    "Please specify the output or it will be set to 0.",
+                    stacklevel=2,
+                )
+                self.nout = 0
+        else:
+            self.nout = nout
+
+        file = self.pathdir / Path(f"out{self.nout:03d}.h5")
+
+        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)
+
+        with h5py.File(str(file), "r") as tmp:
+            self.ntime = cast(h5py.Dataset, tmp["time"])[()][0]
+            var = list(tmp.keys()) if loadvars is True else loadvars or []
+            self.echo_load_vars(tmp, conv_dict, var)
+
+    def echo_load_grid(self, conv_dict: dict[str, str]) -> None:
+        """Load grid.h5 and set attributes."""
+        with h5py.File(str(self.pathdir / Path("grid.h5")), "r") as grid:
+            for key, obj in grid.items():
+                if not isinstance(obj, h5py.Dataset):
+                    continue
+                data = obj[()]
+                name = conv_dict.get(key, key)
+                if name is not None:
+                    setattr(self, name, data)
+
+    def echo_set_grid_dims(self) -> None:
+        """Compute nx1, nx2, nx3, dim, gridsize, nshp."""
+        for dim in ["x1", "x2", "x3"]:
+            n = len(getattr(self, dim)) if hasattr(self, dim) else 1
+            setattr(self, f"n{dim}", n)
+
+        self.dim = (self.nx1 > 1) + (self.nx2 > 1) + (self.nx3 > 1)
+        self.gridsize = self.nx1 * self.nx2 * self.nx3
+        dim_dict = {
+            1: self.nx1,
+            2: (self.nx1, self.nx2),
+            3: (self.nx1, self.nx2, self.nx3),
+        }
+        self.nshp = dim_dict[self.dim]
+
+    def echo_load_vars(
+        self, tmp: h5py.File, conv_dict: dict[str, str], var: list[str]
+    ) -> None:
+        """Load variables from output file."""
+        for key in var:
+            if key == "time":
+                continue
+            valkey = next((k for k, v in conv_dict.items() if v == key), key)
+            obj = tmp.get(valkey)
+            if not isinstance(obj, h5py.Dataset):
+                warnings.warn(
+                    f"'{valkey}' not a dataset (found {type(obj)})",
+                    stacklevel=2,
+                )
+                continue
+            loadvar = obj[()]
+            for dim in [self.nx3, self.nx2, self.nx1]:
+                if dim == 1:
+                    loadvar = loadvar[0]
+            setattr(self.state, conv_dict.get(valkey, valkey), loadvar.T)