flekspy 0.2.0__tar.gz → 0.2.2__tar.gz

{flekspy-0.2.0 → flekspy-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flekspy
-Version: 0.2.0
+Version: 0.2.2
 Summary: Python utilities for processing FLEKS data
 Author: Yuxi Chen
 Author-email: yuxichen@umich.edu>, Hongyang Zhou <hyzhou@umich.edu

{flekspy-0.2.0 → flekspy-0.2.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "flekspy"
-version = "0.2.0"
+version = "0.2.2"
 description = "Python utilities for processing FLEKS data"
 authors = ["Yuxi Chen <yuxichen@umich.edu>, Hongyang Zhou <hyzhou@umich.edu>"]
 readme = "README.md"

{flekspy-0.2.0 → flekspy-0.2.2}/src/flekspy/idl/idl.py

@@ -2,17 +2,22 @@ import glob
 import numpy as np
 import matplotlib.pyplot as plt
 import struct
-import math
-import flekspy.util.utilities as utilities
 import yt
-import flekspy.util.data_container as data_container
+
+from flekspy.util import get_unit
+from flekspy.util import (
+    DataContainer,
+    DataContainer1D,
+    DataContainer2D,
+    DataContainer3D,
+)
 
 x_ = 0
 y_ = 1
 z_ = 2
 
 
-class selector:
+class Selector:
     def __getitem__(self, keys) -> list:
         self.indices = list(keys)
         if len(self.indices) < 3:
@@ -20,12 +25,12 @@ class selector:
         return self.indices
 
 
-class dataframe:
+class Dataframe:
     def __init__(self):
         self.array = None
         self.name = None
         # Selector for the spatial indices
-        self.cut = selector()
+        self.cut = Selector()
         self.cut[:, :, :]
 
     def setData(self, dataIn, nameIn):
@@ -52,6 +57,7 @@ class dataframe:
         else:
             keys = list(keys) + [slice(None, None, None)] * (4 - len(keys))
         ivar = self.name.index(keys[0])
+
         return np.squeeze(self.array[ivar, keys[1], keys[2], keys[3]])
 
 
@@ -62,7 +68,6 @@ class IDLData(object):
     Example:
     >>> ds = IDLData("3d.out")
     >>> dc2d = ds.get_slice("y", 1)
-    >>> dc3d = ds.get_domain()
     """
 
     def __init__(self, filename="none"):
@@ -81,7 +86,7 @@ class IDLData(object):
 
         self.filename = fileList[ifile - 1]
         self.isOuts = self.filename[-4:] == "outs"
-        self.data = dataframe()
+        self.data = Dataframe()
         self.nInstance = None if self.isOuts else 1
         self.npict = 1
         self.fileformat = None
@@ -122,14 +127,12 @@ class IDLData(object):
         )
         return str
 
-    def get_domain(self):
-        r"""
-        Returns a dataContainer3D object that contains all the 3D data.
-        """
+    def get_domain(self) -> DataContainer:
+        """Return data as data container."""
         dataSets = {}
         for varname in self.data.name:
             idx = self.data.name.index(varname)
-            unit = utilities.get_unit(varname, self.unit)
+            unit = get_unit(varname, self.unit)
             dataSets[varname] = yt.YTArray(
                 np.squeeze(self.data.array[idx, :, :, :]), unit, registry=self.registry
             )
@@ -140,15 +143,15 @@ class IDLData(object):
             name = self.variables[idim]
             idx = self.data.name.index(name)
             labels[idim] = name.upper()
-            unit = utilities.get_unit("X", self.unit)
+            unit = get_unit("X", self.unit)
             name = name.upper()
             if name in ["X", "Y", "Z"]:
                 axes[idim] = yt.YTArray(
                     self.data.array[idx, :, :, :], unit, registry=self.registry
                 )
 
-        if self.gencoord:
-            dc = data_container.dataContainer2D(
+        if self.gencoord and self.ndim == 2:
+            dc = DataContainer2D(
                 dataSets,
                 np.squeeze(axes[0]),
                 np.squeeze(axes[1]),
@@ -160,7 +163,7 @@ class IDLData(object):
                 gencoord=True,
             )
         elif self.ndim == 1:
-            dc = data_container.dataContainer1D(
+            dc = DataContainer1D(
                 dataSets,
                 np.squeeze(axes[0]),
                 labels[0],
@@ -169,7 +172,7 @@ class IDLData(object):
                 filename=self.filename,
             )
         elif self.ndim == 2:
-            dc = data_container.dataContainer2D(
+            dc = DataContainer2D(
                 dataSets,
                 np.squeeze(axes[0])[:, 0],
                 np.squeeze(axes[1])[0, :],
@@ -180,7 +183,7 @@ class IDLData(object):
                 filename=self.filename,
             )
         else:
-            dc = data_container.dataContainer3D(
+            dc = DataContainer3D(
                 dataSets,
                 axes[0][:, 0, 0],
                 axes[1][0, :, 0],
@@ -192,6 +195,23 @@ class IDLData(object):
 
         return dc
 
+    def get_slice(self, norm, cut_loc) -> DataContainer2D:
+        """Get a 2D slice from the 3D IDL data.
+
+        Args:
+            norm: str
+                The normal direction of the slice from "x", "y" or "z"
+
+            cur_loc: float
+                The position of slicing.
+
+        Return: DataContainer2D
+        """
+        domain = self.get_domain()
+        ds = domain.get_slice(norm, cut_loc)
+
+        return ds
+
     def save_data(self, saveName, saveFormat="ascii"):
         # Currently only support ascii output.
         self.save_ascii_instance(saveName)
@@ -240,7 +260,7 @@ class IDLData(object):
         elif self.fileformat == "binary":
             self.read_binary()
         else:
-            print("Unknown format = ", self.fileformat)
+            raise ValueError(f"Unknown format = {self.fileformat}")
 
         ndim = self.ndim
         nvar = self.nvar
@@ -295,7 +315,6 @@ class IDLData(object):
         self.npoints = abs(self.grid.prod())
 
         time = self.time
-        npts = self.npoints
         ndim = self.grid.size
         nvar = self.nvar
         npar = self.nparam
@@ -321,9 +340,9 @@ class IDLData(object):
             time % 60.0,
         )
 
-        nRow = ndim + nvar
-        nCol = self.npoints
-        self.data.array = np.zeros((nRow, nCol))
+        nrow = ndim + nvar
+        ncol = self.npoints
+        self.data.array = np.zeros((nrow, ncol))
 
         for i, line in enumerate(infile.readlines()):
             parts = line.split()
@@ -335,7 +354,7 @@ class IDLData(object):
                 self.data.array[j][i] = float(p)
 
         nline += self.npoints
-        shapeNew = np.append([nRow], self.grid)
+        shapeNew = np.append([nrow], self.grid)
         self.data.array = np.reshape(self.data.array, shapeNew, order="F")
 
         return nline
@@ -357,110 +376,123 @@ class IDLData(object):
     def read_binary_instance(self, infile):
         # On the first try, we may fail because of wrong-endianess.
         # If that is the case, swap that endian and try again.
-        EndChar = "<"  # Endian marker (default: little)
+        end_char = "<"  # Endian marker (default: little)
         self.endian = "little"
-        RecLenRaw = infile.read(4)
+        record_len_raw = infile.read(4)
 
-        RecLen = (struct.unpack(EndChar + "l", RecLenRaw))[0]
-        if (RecLen > 10000) or (RecLen < 0):
-            EndChar = ">"
+        record_len = (struct.unpack(end_char + "l", record_len_raw))[0]
+        if (record_len > 10000) or (record_len < 0):
+            end_char = ">"
             self.endian = "big"
-            RecLen = (struct.unpack(EndChar + "l", RecLenRaw))[0]
+            record_len = (struct.unpack(end_char + "l", record_len_raw))[0]
 
         headline = (
-            struct.unpack("{0}{1}s".format(EndChar, RecLen), infile.read(RecLen))
+            struct.unpack(
+                "{0}{1}s".format(end_char, record_len), infile.read(record_len)
+            )
         )[0].strip()
         if str is not bytes:
             headline = headline.decode()
         self.unit = headline.split()[0]
 
-        (OldLen, RecLen) = struct.unpack(EndChar + "2l", infile.read(8))
+        (old_len, record_len) = struct.unpack(end_char + "2l", infile.read(8))
         pformat = "f"
         # Parse rest of header; detect double-precision file
-        if RecLen > 20:
+        if record_len > 20:
             pformat = "d"
         (self.iter, self.time, self.ndim, self.nparam, self.nvar) = struct.unpack(
-            "{0}l{1}3l".format(EndChar, pformat), infile.read(RecLen)
+            "{0}l{1}3l".format(end_char, pformat), infile.read(record_len)
         )
         self.gencoord = self.ndim < 0
         self.ndim = abs(self.ndim)
         # Get gridsize
-        (OldLen, RecLen) = struct.unpack(EndChar + "2l", infile.read(8))
+        (old_len, record_len) = struct.unpack(end_char + "2l", infile.read(8))
 
         self.grid = np.array(
             struct.unpack(
-                "{0}{1}l".format(EndChar, abs(self.ndim)), infile.read(RecLen)
+                "{0}{1}l".format(end_char, abs(self.ndim)), infile.read(record_len)
             )
         )
         self.npoints = abs(self.grid.prod())
 
-        time = self.time
         npts = self.npoints
         ndim = self.grid.size
         nvar = self.nvar
         npar = self.nparam
 
         # Read parameters stored in file
-        self.para = np.zeros(npar)
-        if npar > 0:
-            (OldLen, RecLen) = struct.unpack(EndChar + "2l", infile.read(8))
-            self.para[:] = struct.unpack(
-                "{0}{1}{2}".format(EndChar, npar, pformat), infile.read(RecLen)
-            )
-
-        (OldLen, RecLen) = struct.unpack(EndChar + "2l", infile.read(8))
-        names = (struct.unpack("{0}{1}s".format(EndChar, RecLen), infile.read(RecLen)))[
-            0
-        ]
-        if str is not bytes:
-            names = names.decode()
-
-        names.strip()
-        names = names.split()
-
-        # Save grid names (e.g. "x" or "r") and save associated params
-        self.dims = names[0:ndim]
-        self.variables = np.array(names)
+        self.read_parameters(infile, end_char, pformat, npar)
 
-        self.strtime = "{0:04d}h{1:02d}m{2:06.3f}s".format(
-            int(time // 3600), int(time % 3600 // 60), time % 60
-        )
+        # Read variable names
+        self.read_variable_names(infile, end_char)
 
-        nRow = ndim + nvar
-        nCol = self.npoints
-        self.data.array = np.zeros((nRow, nCol), dtype=np.float32)
+        nrow = ndim + nvar
+        n_col = self.npoints
+        self.data.array = np.zeros((nrow, n_col), dtype=np.float32)
 
         # Get the grid points...
-        (OldLen, RecLen) = struct.unpack(EndChar + "2l", infile.read(8))
+        (old_len, record_len) = struct.unpack(end_char + "2l", infile.read(8))
         for i in range(0, ndim):
             # Read the data into a temporary grid
             tempgrid = np.array(
                 struct.unpack(
-                    "{0}{1}{2}".format(EndChar, npts, pformat),
-                    infile.read(int(RecLen // ndim)),
+                    "{0}{1}{2}".format(end_char, npts, pformat),
+                    infile.read(int(record_len // ndim)),
                 )
             )
             self.data.array[i, :] = tempgrid
 
         # Get the actual data and sort
         for i in range(ndim, nvar + ndim):
-            (OldLen, RecLen) = struct.unpack(EndChar + "2l", infile.read(8))
+            (old_len, record_len) = struct.unpack(end_char + "2l", infile.read(8))
             tmp = np.array(
                 struct.unpack(
-                    "{0}{1}{2}".format(EndChar, npts, pformat), infile.read(RecLen)
+                    "{0}{1}{2}".format(end_char, npts, pformat), infile.read(record_len)
                 )
             )
             self.data.array[i, :] = tmp
         # Consume the last record length
         infile.read(4)
 
-        shapeNew = np.append([nRow], self.grid)
-        self.data.array = np.reshape(self.data.array, shapeNew, order="F")
+        shape_new = np.append([nrow], self.grid)
+        self.data.array = np.reshape(self.data.array, shape_new, order="F")
+
+    def read_parameters(self, infile, end_char, pformat, num_params):
+        """Reads parameters from the binary file."""
+        self.para = np.zeros(num_params)
+        if num_params > 0:
+            (old_len, record_len) = struct.unpack(end_char + "2l", infile.read(8))
+            self.para[:] = struct.unpack(
+                "{0}{1}{2}".format(end_char, num_params, pformat),
+                infile.read(record_len),
+            )
+
+    def read_variable_names(self, infile, end_char):
+        """Reads variable names from the binary file."""
+        (old_len, record_len) = struct.unpack(end_char + "2l", infile.read(8))
+        names = (
+            struct.unpack(
+                "{0}{1}s".format(end_char, record_len), infile.read(record_len)
+            )
+        )[0]
+        if str is not bytes:
+            names = names.decode()
+
+        names.strip()
+        names = names.split()
+
+        # Save grid names (e.g. "x" or "r") and save associated params
+        self.dims = names[0 : self.ndim]
+        self.variables = np.array(names)
+
+        self.strtime = "{0:04d}h{1:02d}m{2:06.3f}s".format(
+            int(self.time // 3600), int(self.time % 3600 // 60), self.time % 60
+        )
 
     def plot(self, *dvname, **kwargs):
         """Plot 1D IDL outputs.
 
-        Returns:
+        Args:
             *dvname (str): variable names
             **kwargs: keyword argument to be passed to `plot`.
         """
@@ -468,7 +500,7 @@ class IDLData(object):
         nvar = len(dvname)
 
         f, axes = plt.subplots(nvar, 1, constrained_layout=True, sharex=True)
-        axes = np.array(axes)  # in case nRow = nCol = 1
+        axes = np.array(axes)  # in case nrows == ncols == 1
         axes = axes.reshape(-1)
         for isub, ax in zip(range(nvar), axes):
             w = self.data[dvname[isub]]

flekspy-0.2.2/src/flekspy/util/__init__.py

@@ -0,0 +1,7 @@
+from .utilities import download_testfile, unit_one, get_unit, get_ticks
+from .data_container import (
+    DataContainer,
+    DataContainer1D,
+    DataContainer2D,
+    DataContainer3D,
+)

{flekspy-0.2.0 → flekspy-0.2.2}/src/flekspy/util/data_container.py

@@ -1,12 +1,13 @@
+from copy import deepcopy
+
 import yt
 import numpy as np
 import matplotlib.pyplot as plt
-import math
-from flekspy.util.utilities import get_unit, get_ticks
-from flekspy.plot.streamplot import streamplot
-from copy import deepcopy
-from scipy.interpolate import griddata
 import matplotlib.tri as tri
+from scipy.interpolate import griddata
+
+from flekspy.util.utilities import get_unit
+from flekspy.plot.streamplot import streamplot
 
 
 def compare(d1, d2):
@@ -27,7 +28,7 @@ def compare(d1, d2):
         print(s.format(*l))
 
 
-class dataContainer(object):
+class DataContainer(object):
     def __init__(
         self,
         dataSets,
@@ -42,12 +43,6 @@ class dataContainer(object):
         gencoord=False,
         filename="",
     ):
-
-        # Re-generate the coordinates to make sure they are equally spaced.
-        # x = np.linspace(x[0], x[-1], len(x))
-        # y = np.linspace(y[0], y[-1], len(y))
-        # z = np.linspace(z[0], z[-1], len(z))
-
         self.data = dataSets
         self.x = x
         self.y = y
@@ -64,22 +59,14 @@ class dataContainer(object):
         self.filename = filename
         self.gencoord = gencoord
 
-    def __repr__(self):
-        print("\n-----------------------------")
-        print("Variables       :", self.vars)
-        print("\nData range      :", self.range)
-        print("\nData dimensions :", self.dimensions)
-        print("\n")
-
-        header = ("var", "min", "max", "mean")
-        s = "{:8}   " + "{:18}" * 3
-        print(s.format(*header))
-        for v in self.vars:
-            d = self.data[v]
-            s = "{:10}" + "{:+.6e},   " * 3
-            print(s.format(v, float(d.min()), float(d.max()), float(d.mean())))
-        print("-----------------------------")
-        return "\n"
+    def __repr__(self) -> str:
+        str = (
+            f"variables : {self.vars}\n"
+            f"data range: {self.range}\n"
+            f"dimension : {self.dimensions}\n"
+        )
+
+        return str
 
     def add_bottom_line(self, f, verbose):
         if verbose > 0:
@@ -94,15 +81,13 @@ class dataContainer(object):
 
     def analyze_variable_string(self, var):
         r"""
-        This method analyzes the plot string and return the plot variable and plot range.
+        Parses the input string and return the plot variable and plot range.
 
-        Parameters
-        ----------------------
-        var: String
-        Example: var = "{bb}<(-10)>(-9.8)"
+        Args:
+            var: str
+            Example: var = "{bb}<(-10)>(-9.8)"
 
-        Return: a tuple contains the variable name, variable min and max.
-        Example: return "{bb}", -10, -9.8
+        Return: a tuple of the variable name, variable min and max.
         """
         vMin = None
         vMax = None
@@ -126,15 +111,13 @@ class dataContainer(object):
 
         return varName, vMin, vMax
 
-    def evaluate_expression(self, expression, unit="planet"):
+    def evaluate_expression(self, expression: str, unit: str = "planet"):
         r"""
-        This method calculates the variable expression and return the result, which is
-        a YTArray.
+        Evaluates the variable expression and return the result of an YTArray.
 
-        Parameters
-        -------------------
-        expression: String
-        Example: expression = "np.log({rhos0}+{rhos1})"
+        Args:
+            expression: str
+            Example: expression = "np.log({rhos0}+{rhos1})"
         """
 
         if expression.find("{") < 0:
@@ -152,16 +135,14 @@ class dataContainer(object):
 
     def add_variable(self, name, val):
         r"""
-        This method adds a variable to the dataset for visualization purpose
-
-        Parameters
-        ------------------
-        name: String
-        The name of the variable to be added into self.data
+        Adds a variable to the dataset.
 
-        val: array-like structure
-        Values of the variable stored in an array
+        Args:
+            name: str
+                The name of the variable to be added into self.data
 
+            val: array-like structure
+                Values of the variable stored in an array
         """
 
         if type(val) != yt.units.yt_array.YTArray:
@@ -173,14 +154,13 @@ class dataContainer(object):
 
     def get_variable(self, var, unit="planet"):
         r"""
-        This method calculate the value of the variable.
+        Return raw variables or calculate derived variables.
 
-        Parameters
-        ------------------
-        var: String
-        Example: var = "pbeta"
+        Args:
+            var: str
+                variable name
 
-        Return: A YTArray
+        Return: YTArray
         """
 
         if var in self.data.keys():
@@ -227,7 +207,7 @@ class dataContainer(object):
         return ytarr if str(ytarr.units) == "dimensionless" else ytarr.in_units(varUnit)
 
 
-class dataContainer3D(dataContainer):
+class DataContainer3D(DataContainer):
     r"""
     A class handles 3D box data sets.
     """
@@ -236,30 +216,27 @@ class dataContainer3D(dataContainer):
         self, dataSets, x, y, z, xlabel="X", ylabel="Y", zlabel="Z", *args, **kwargs
     ):
         r"""
-        Parameters
-        ---------------------
-        dataSets: dictonary
-        The key is the variable name, and the dictionary value is usually a YTArray.
+        Args:
+            dataSets: dictonary
+                The key is the variable name, and the dictionary value is usually a YTArray.
 
-        x/y/z: A 1D YTArray
+            x/y/z: A 1D YTArray
         """
-        super(dataContainer3D, self).__init__(
+        super(DataContainer3D, self).__init__(
             dataSets, x, y, z, xlabel, ylabel, zlabel, *args, **kwargs
         )
 
-    def get_slice(self, norm, cut_loc):
-        r"""
-        Get a 2D slice from the 3D box data.
+    def get_slice(self, norm, cut_loc) -> "DataContainer2D":
+        """Get a 2D slice from the 3D box data.
 
-        Parameters
-        ------------------
-        norm: String
-        The normal direction of the slice. Potential value: 'x', 'y' or 'z'
+        Args:
+            norm: str
+                The normal direction of the slice from "x", "y" or "z"
 
-        cur_loc: Float
-        The position of slicing.
+            cur_loc: float
+                The position of slicing.
 
-        Return: A dataContainer2D object
+        Return: DataContainer2D
         """
 
         axDir = {"X": 0, "Y": 1, "Z": 2}
@@ -286,7 +263,7 @@ class dataContainer3D(dataContainer):
         axLabes = {0: ("Y", "Z"), 1: ("X", "Z"), 2: ("X", "Y")}
         ax = {0: (1, 2), 1: (0, 2), 2: (0, 1)}
 
-        return dataContainer2D(
+        return DataContainer2D(
             dataSets,
             axes[ax[idir][0]],
             axes[ax[idir][1]],
@@ -297,7 +274,7 @@ class dataContainer3D(dataContainer):
         )
 
 
-class dataContainer2D(dataContainer):
+class DataContainer2D(DataContainer):
     r"""
     A class handles 2D Cartesian data.
     """
@@ -312,29 +289,27 @@ class dataContainer2D(dataContainer):
         cut_norm=None,
         cut_loc=None,
         *args,
-        **kwargs
+        **kwargs,
     ):
         r"""
-        Parameters
-        ---------------------
-        dataSets: dictonary
-        The key is the variable name, and the dictionary value is usually a YTArray.
+        Args:
+            dataSets: dictonary
+                The keys are variable names, and the values are YTArrays.
 
-        x/y: A 1D YTArray
+            x/y: 1D YTArray
 
-        xlabel/ylabel: String
+            xlabel/ylabel: str
 
-        cut_norm: String
-        'x', 'y' or 'z'
+            cut_norm: str
+                "x", "y" or "z"
 
-        cut_loc: Float
-        cut_norm and cut_loc are used to record the position of slice if this 2D
-        data set is obtained from a 3D box.
+            cut_loc: float
+                cut_norm and cut_loc are used to record the position of slice if this 2D data set is obtained from a 3D box.
         """
 
         zlabel = None
         z = (0, 0)
-        super(dataContainer2D, self).__init__(
+        super(DataContainer2D, self).__init__(
             dataSets, x, y, z, xlabel, ylabel, zlabel, *args, **kwargs
         )
 
@@ -347,7 +322,7 @@ class dataContainer2D(dataContainer):
             dif.data[var] -= other.data[var]
         return dif
 
-    def contour(
+    def plot(
         self,
         vars,
         xlim=None,
@@ -357,60 +332,45 @@ class dataContainer2D(dataContainer):
         cmap="turbo",
         figsize=(10, 6),
         pcolor=False,
-        log=False,
+        logscale=False,
         addgrid=False,
         bottomline=10,
-        plot=None,
-        cbticks=None,
         showcolorbar=True,
-        createcanvas=False,
         *args,
-        **kwargs
+        **kwargs,
     ):
         r"""
-        Contour plots.
-
-        Parameters
-        ----------------------
-        vars: String
-        Ploting variables and ploting range.
-        Example: vars = "Bx<(50)>(-50) By (np.log(2*{rhos0}))>(-5)"
-
-        xlim/ylim: A list/tuple contains the x- y-axis range
+        2D plots.
 
-        unit: String
-        'planet' or 'si'
+        Args:
+            vars: str
+                Ploting variables and ploting range.
+            Example: vars = "Bx<(50)>(-50) By (np.log(2*{rhos0}))>(-5)"
 
-        nlevels: Integer
-        Number of the countour plot color levels
+            xlim/ylim: A list/tuple contains the x- y-axis range
 
-        cmap: String
-        Color map type
+            unit: str
+                "planet" or "si"
 
-        figsize: A tuple
+            nlevels: int
+                Number of the contour levels
 
-        log: Bool
-        Using log plot or not.
+            cmap: str
+                Color map type
 
-        plot: list or tuple of [f, axes]
-        Used to fine tune the individual subplot
+            figsize: tuple
 
-        createcanvas: boolean
-        If set to True, skip the contour making process to get an empty
-        canvas which can be used to generate subplots in different styles.
-        Set this flag to get the value for 'plot' parameter.
+            logscale: Bool
+                True to scale the variable in log.
 
-        Examples
-        ----------------
-        >>> f,axes = dc.contour("Bx<(50)>(-50) By (np.log(2*{rhos0}))>(-5)",xlim=[-40,-5])
+        Examples:
+            >>> f, axes = dc.contour("Bx<(50)>(-50) By (np.log(2*{rhos0}))>(-5)", xlim=[-40,-5])
         """
 
         if type(vars) == str:
             vars = vars.split()
 
         nvar = len(vars)
-        nRow = int(round(np.sqrt(nvar)))
-        nCol = math.ceil(nvar / nRow)
 
         varNames = []
         varMin = []
@@ -421,17 +381,11 @@ class dataContainer2D(dataContainer):
             varMin.append(vmin)
             varMax.append(vmax)
 
-        if plot:
-            f, axes = plot
-        else:
-            f, axes = plt.subplots(nRow, nCol, figsize=figsize)
-            axes = np.array(axes)  # in case nRow = nCol = 1
-
+        f, axes = plt.subplots(nvar, 1, figsize=figsize, layout="constrained")
+        axes = np.array(axes)  # in case nrows == ncols == 1
         axes = axes.reshape(-1)
 
         for isub, ax in zip(range(nvar), axes):
-            if createcanvas:  # no need to make any contour, just keep
-                break  # the canvas clean
             ytVar = self.evaluate_expression(varNames[isub], unit)
             v = ytVar
             varUnit = "dimensionless"
@@ -442,8 +396,7 @@ class dataContainer2D(dataContainer):
             vmin = v.min() if varMin[isub] == None else varMin[isub]
             vmax = v.max() if varMax[isub] == None else varMax[isub]
 
-            logplot = log and vmin > 0
-            if logplot:
+            if logscale:
                 v = np.log10(v)
 
             levels = np.linspace(vmin, vmax, nlevels)
@@ -461,7 +414,7 @@ class dataContainer2D(dataContainer):
                         cmap=cmap,
                         extend="both",
                         *args,
-                        **kwargs
+                        **kwargs,
                     )
             else:
                 if pcolor or abs(vmin - vmax) < 1e-20 * abs(vmax):
@@ -477,7 +430,7 @@ class dataContainer2D(dataContainer):
                         cmap=cmap,
                         extend="both",
                         *args,
-                        **kwargs
+                        **kwargs,
                     )
             if addgrid:
                 if self.gencoord:
@@ -490,13 +443,8 @@ class dataContainer2D(dataContainer):
 
                 ax.plot(gx, gy, "x")
 
-            if cbticks:
-                ticks = cbticks
-            else:
-                ticks = get_ticks(vmin, vmax)
-
             if showcolorbar:
-                cb = f.colorbar(cs, ax=ax, ticks=ticks)
+                cb = f.colorbar(cs, ax=ax, pad=0.01)
                 cb.formatter.set_powerlimits((0, 0))
 
             ax.set_xlim(xlim)
@@ -505,36 +453,29 @@ class dataContainer2D(dataContainer):
             ax.set_ylabel(self.ylabel)
             title = varNames[isub]
             if varUnit != "dimensionless":
-                title = title + " [" + varUnit + "]"
-            if logplot:
+                title += " [" + varUnit + "]"
+            if logscale:
                 title = "$log_{10}$(" + title + ")"
             ax.set_title(title)
 
-        # Delete axes for empty subplots
-        for ax in axes[nvar : nRow * nCol]:
-            f.delaxes(ax)
-
-        plt.tight_layout()
         if self.cut_norm != None and self.cut_loc != None:
             print("Plots at " + self.cut_norm + " = ", self.cut_loc)
 
         self.add_bottom_line(f, bottomline)
-        return f, axes.reshape(nRow, nCol)
+
+        return f, axes
 
     def add_contour(self, ax, var, unit="planet", rmask=None, *args, **kwargs):
-        r"""
-        Adding contour lines to an axis.
+        r"""Adding contour lines to an axis.
 
-        Parameters
-        --------------------
-        ax: matplotlib axis
-        The aixs to plot contour lines.
+        Args:
+            ax: matplotlib axis
+                The axis to plot contour lines.
 
-        var: String
+            var: string
 
-        Examples
-        ----------------------
-        >>> f,axes = dc.contour("Bx<(50)>(-50) By (np.log(2*{rhos0}))>(-5)",xlim=[-40,-5])
+        Examples:
+        >>> f, axes = dc.plot("Bx<(50)>(-50) By (np.log(2*{rhos0}))>(-5)", xlim=[-40,-5])
         >>> dc.add_contour(axes[0,0], "rhos1>1")
         """
 
@@ -574,24 +515,23 @@ class dataContainer2D(dataContainer):
         ymax=None,
         rmask=None,
         *args,
-        **kwargs
+        **kwargs,
     ):
         r"""
         Adding streamlines to an axis.
 
-        Parameters
-        --------------------
-        ax: matplotlib axis
-        The aixs to add streamlines
+        Args:
+            ax: matplotlib axis
+                The axis to add streamlines
 
-        var1/var2: String
+            var1/var2: str
+                Streamline variable names
 
-        density: Int
-        It controls the number of streamlines.
+            density: int
+                It controls the number of streamlines.
 
-        Examples
-        ----------------------
-        >>> f,axes = dc.contour("Bx<(50)>(-50) By (np.log(2*{rhos0}))>(-5)",xlim=[-40,-5])
+        Examples:
+        >>> f, axes = dc.plot("Bx<(50)>(-50) By (np.log(2*{rhos0}))>(-5)", xlim=[-40,-5])
         >>> dc.add_stream(axes[1,0], "Bx", "Bz", density=2)
         """
 
@@ -618,7 +558,7 @@ class dataContainer2D(dataContainer):
             xy = np.zeros((len(self.x), 2))
             xy[:, 0] = self.x.value
             xy[:, 1] = self.y.value
-            # The first and last row/column may be None. Remove them
+            # Remove the first and last row/column since they may be None
             vect1 = griddata(xy, v1, (gridx, gridy), method="linear")[1:-1, 1:-1]
             vect2 = griddata(xy, v2, (gridx, gridy), method="linear")[1:-1, 1:-1]
             xx = gridx[1:-1, 0]
@@ -638,26 +578,23 @@ class dataContainer2D(dataContainer):
                         vect1[i, j] = np.nan
                         vect2[i, j] = np.nan
 
-        streamplot(
-            ax, xx, yy, vect1.T, vect2.T, density=density, *args, **kwargs
-        )
+        streamplot(ax, xx, yy, vect1.T, vect2.T, density=density, *args, **kwargs)
 
 
-class dataContainer1D(dataContainer):
+class DataContainer1D(DataContainer):
     r"""
     A class handles 1D Cartesian data.
     """
 
     def __init__(self, dataSets, x, xlabel, *args, **kwargs):
         r"""
-        Parameters
-        ---------------------
-        dataSets: dictonary
-        The key is the variable name, and the dictionary value is usually a YTArray.
+        Args:
+            dataSets: dictonary
+                The key is the variable name, and the dictionary value is usually a YTArray.
 
-        x: A 1D YTArray
+            x: 1D YTArray
 
-        xlabel: String
+            xlabel: str
         """
 
         ylabel = None
@@ -666,7 +603,7 @@ class dataContainer1D(dataContainer):
         zlabel = None
         z = (0, 0)
 
-        super(dataContainer1D, self).__init__(
+        super(DataContainer1D, self).__init__(
             dataSets, x, y, z, xlabel, ylabel, zlabel, *args, **kwargs
         )
 
@@ -675,25 +612,21 @@ class dataContainer1D(dataContainer):
         vars,
         xlim=None,
         ylim=None,
-        unit="planet",
+        unit: str = "planet",
         figsize=(12, 8),
-        log=False,
         bottomline=10,
         *args,
-        **kwargs
+        **kwargs,
     ):
         """
-        Examples
-        ----------------
-        >>> f,axes=dc.plot("absdivb bx",xlim=[-5,5],color='k',linestyle='solid',marker='x')
+        Examples:
+        >>> f, axes = dc.plot("absdivb bx", xlim=[-5,5])
         """
 
         if type(vars) == str:
             vars = vars.split()
 
         nvar = len(vars)
-        nRow = int(round(np.sqrt(nvar)))
-        nCol = math.ceil(nvar / nRow)
 
         varNames = []
         varMin = []
@@ -704,18 +637,13 @@ class dataContainer1D(dataContainer):
             varMin.append(vmin)
             varMax.append(vmax)
 
-        f, axes = plt.subplots(nRow, nCol, figsize=figsize)
-        axes = np.array(axes)  # in case nRow = nCol = 1
-
+        f, axes = plt.subplots(nvar, 1, figsize=figsize)
+        axes = np.array(axes)  # in case nrows == ncols == 1
         axes = axes.reshape(-1)
 
         for isub, ax in zip(range(nvar), axes):
             ytVar = self.evaluate_expression(varNames[isub], unit)
             v = ytVar
-            varUnit = "dimensionless"
-            if type(ytVar) == yt.units.yt_array.YTArray:
-                v = ytVar.value
-                varUnit = str(ytVar.units)
 
             vmin = v.min() if varMin[isub] == None else varMin[isub]
             vmax = v.max() if varMax[isub] == None else varMax[isub]
@@ -728,4 +656,5 @@ class dataContainer1D(dataContainer):
             ax.legend()
 
         self.add_bottom_line(f, bottomline)
+
         return f, axes

{flekspy-0.2.0 → flekspy-0.2.2}/src/flekspy/yt/yt.py

@@ -1,8 +1,8 @@
-import yt
 import os
 import glob
 import numpy as np
 
+import yt
 from yt.funcs import setdefaultattr
 from yt.frontends.boxlib.api import BoxlibHierarchy, BoxlibDataset
 from yt.fields.field_info_container import FieldInfoContainer
@@ -11,8 +11,8 @@ from yt.data_objects.selection_objects.data_selection_objects import (
 )
 from yt.visualization.profile_plotter import PhasePlot
 
-import flekspy.util.data_container as data_container
-from flekspy.util.utilities import get_unit
+from flekspy.util import DataContainer2D, DataContainer3D
+from flekspy.util import get_unit
 
 
 class FLEKSFieldInfo(FieldInfoContainer):
@@ -119,16 +119,13 @@ class FLEKSData(BoxlibDataset):
     r"""
     Read and plot the AMReX format output from FLEKS.
 
-    Parameters
-    ----------
-    output_dir : String
-        The path to the data.
-
-    Examples
-    --------
+    Args:
+        output_dir : String
+            The path to the data.
 
-    >>> import flekspy
-    >>> ds = flekspy.FLEKSData("3d_particle*n00004750_amrex")
+    Examples:
+        >>> import flekspy
+        >>> ds = flekspy.FLEKSData("3d_particle*n00004750_amrex")
     """
 
     _index_class = FLEKSHierarchy
@@ -206,15 +203,13 @@ class FLEKSData(BoxlibDataset):
     def pvar(self, var):
         return (self.particle_types[0], var)
 
-    def get_slice(self, norm, cut_loc):
+    def get_slice(self, norm, cut_loc) -> DataContainer2D:
         r"""
-        This method returns a dataContainer2D object that contains a slice along
-        the normal direction.
+        Returns a DataContainer2D object that contains a slice along the normal direction.
 
-        Parameters
-        ---------------------
-        norm (str): slice normal direction in "x", "y" or "z"
-        cut_loc (float): cut location along the normal direction
+        Args:
+            norm (str): slice normal direction in "x", "y" or "z"
+            cut_loc (float): cut location along the normal direction
         """
 
         axDir = {"X": 0, "Y": 1, "Z": 2}
@@ -238,7 +233,8 @@ class FLEKSData(BoxlibDataset):
 
         dataSets = {}
         for var in self.field_list:
-            dataSets[var[1]] = np.squeeze(abArr[var])
+            if abArr[var].size != 0:  # remove empty sliced particle output
+                dataSets[var[1]] = np.squeeze(abArr[var])
 
         axLabes = {0: ("Y", "Z"), 1: ("X", "Z"), 2: ("X", "Y")}
 
@@ -253,7 +249,7 @@ class FLEKSData(BoxlibDataset):
                 )
             )
 
-        return data_container.dataContainer2D(
+        return DataContainer2D(
             dataSets,
             axes[0],
             axes[1],
@@ -263,9 +259,9 @@ class FLEKSData(BoxlibDataset):
             cut_loc,
         )
 
-    def get_domain(self):
+    def get_domain(self) -> DataContainer3D:
         r"""
-        Reading in all the simulation data into a 3D box. It returns a dataContainer3D object.
+        Read all the simulation data into a 3D box.
         """
         domain = self.covering_grid(
             level=0, left_edge=self.domain_left_edge, dims=self.domain_dimensions
@@ -285,25 +281,23 @@ class FLEKSData(BoxlibDataset):
                 )
             )
 
-        return data_container.dataContainer3D(dataSets, axes[0], axes[1], axes[2])
+        return DataContainer3D(dataSets, axes[0], axes[1], axes[2])
 
     def plot_slice(self, norm, cut_loc, vars, unit_type="planet", *args, **kwargs):
         r"""Plot 2D slice
 
-        Parameters
-        ----------
-        norm : string, one of 'x', 'y' or 'z'. Norm direction of the slice.
-
-        cut_loc : float. The location of the slice.
-
-        vars : a list or string of plotting variables.
-            Example: "Bx rhos0" or ["Bx", "rhos0"]
+        Args:
+            norm: str
+                Normal direction of the slice in "x", "y" or "z".
+            cut_loc: float
+                The location of the slice.
 
-        unit_type : The unit system of the plots. "planet" or "si".
+            vars: a list or string of plotting variables.
+                Example: "Bx rhos0" or ["Bx", "rhos0"]
 
-        Examples
-        --------
+        unit_type: The unit system of the plots. "planet" or "si".
 
+        Examples:
         >>> vars = ["rhos0", "uzs0", "Bz", "pxxs1", "Ex"]
         >>> splt = ds.plot_slice("y", 0.0, vars)
         >>> splt.display()
@@ -427,26 +421,23 @@ class FLEKSData(BoxlibDataset):
     ):
         r"""Plot the particle position of particles inside a box.
 
-        Parameters
-        ----------
-        x_field & y_field: string
-            The x- y- axes, from "p_x", "p_y", "p_z".
+        Args:
+            x_field & y_field: str
+                The x- y- axes, from "p_x", "p_y", "p_z".
 
-        z_field: string
-            It is usually the particle weight: "p_w".
+            z_field: str
+                color variable, usually the particle weight "p_w".
 
-        region : YTSelectionContainer
-            Spatial region to be selected, such as all_data, box, region, or sphere.
+            region: YTSelectionContainer
+                Spatial region to be selected, such as all_data, box, region, or sphere.
 
-        unit_type : string
-            The unit system of the plots. "planet" or "si".
+            unit_type: str
+                The unit system of the plots. "planet" or "si".
 
         See more at https://yt-project.org/doc/reference/api/yt.visualization.particle_plots.html#yt.visualization.particle_plots.ParticlePlot
 
-        Examples
-        --------
-        >>> pp = ds.plot_particles([8, -1, -1], [10, 0, 0], "p_x",
-                     "p_y", "p_w", unit_type="planet")
+        Examples:
+        >>> pp = ds.plot_particles([8, -1, -1], [10, 0, 0], "p_x", "p_y", "p_w", unit_type="planet")
         >>> pp.show()
         """
         if region is None:

flekspy-0.2.0/src/flekspy/util/__init__.py

@@ -1 +0,0 @@
-from .utilities import download_testfile, unit_one, get_unit, get_ticks