flekspy 0.2.2__tar.gz → 0.2.4__tar.gz

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

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

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

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

{flekspy-0.2.2 → flekspy-0.2.4}/src/flekspy/tp/test_particles.py

@@ -1,441 +1,466 @@
-from typing import List, Tuple, Dict, Union, Callable
-
-import matplotlib.pyplot as plt
-import os
-import numpy as np
-import glob
-import struct
-
-
-class FLEKSTP(object):
-    r"""
-    A class that is used to read and plot test particles. Each particle ID consists of
-    a CPU index, a particle index on each CPU, and a location index.
-    By default, 7 real numbers saved for each step: time + position + velocity.
-
-    Parameters
-    -----------
-    dirs: String
-        The path to the test particle dataset.
-
-    Examples
-    ----------
-    >>> tp = FLEKSTP("res/run1/PC/test_particles", iSpecies=1)
-    >>> pIDs = list(tp.IDs())
-    >>> tp.plot_trajectory(pIDs[3])
-    >>> tp.save_trajectory_to_csv(pIDs[5])
-    >>> ids, pData = tp.read_particles_at_time(6500.8, doSave=True)
-    >>> f = tp.plot_loc(pData)
-    """
-
-    it_ = 0
-    ix_ = 1
-    iy_ = 2
-    iz_ = 3
-    iu_ = 4
-    iv_ = 5
-    iw_ = 6
-    iBx_ = 7
-    iBy_ = 8
-    iBz_ = 9
-    iEx_ = 10
-    iEy_ = 11
-    iEz_ = 12
-
-    def __init__(
-        self,
-        dirs: Union[str, List[str]],
-        iDomain: int = 0,
-        iSpecies: int = 0,
-        iListStart: int = 0,
-        iListEnd: int = -1,
-        readAllFiles: bool = False,
-    ):
-        if type(dirs) == str:
-            dirs = [dirs]
-
-        header = dirs[0] + "/Header"
-        if os.path.exists(header):
-            with open(header, "r") as f:
-                self.nReal = int(f.readline())
-        else:
-            raise FileNotFoundError(f"Header file not found in {dirs[0]}")
-
-        self.iSpecies = iSpecies
-        self.plistfiles = list()
-        self.pfiles = list()
-
-        for outputDir in dirs:
-            self.plistfiles = self.plistfiles + glob.glob(
-                f"{outputDir}/FLEKS{iDomain}_particle_list_species_{iSpecies}_*"
-            )
-
-            self.pfiles = self.pfiles + glob.glob(
-                f"{outputDir}/FLEKS{iDomain}_particle_species_{iSpecies}_*"
-            )
-
-        self.plistfiles.sort()
-        self.pfiles.sort()
-
-        self.list_index_to_time = []
-        if readAllFiles:
-            for filename in self.pfiles:
-                record = self._read_the_first_record(filename)
-                if record == None:
-                    continue
-                self.list_index_to_time.append(record[FLEKSTP.it_])
-
-        if iListEnd == -1:
-            iListEnd = len(self.plistfiles)
-        self.plistfiles = self.plistfiles[iListStart:iListEnd]
-        self.pfiles = self.pfiles[iListStart:iListEnd]
-
-        self.plists: List[Dict[Tuple[int, int], int]] = []
-        for fileName in self.plistfiles:
-            self.plists.append(self.read_particle_list(fileName))
-
-        self.IDs = set()
-        for plist in self.plists:
-            self.IDs.update(plist.keys())
-
-        self.file_time = []
-        for filename in self.pfiles:
-            record = self._read_the_first_record(filename)
-            if record == None:
-                continue
-            self.file_time.append(record[FLEKSTP.it_])
-
-        print(f"Particles of species {self.iSpecies} are read from {dirs}")
-        print(f"Number of particles: {len(self.IDs)}")
-
-    def getIDs(self):
-        return list(sorted(self.IDs))
-
-    def get_index_to_time(self) -> List:
-        r"""
-        Getter method for accessing list_index_to_time.
-        """
-        if len(self.list_index_to_time) == 0:
-            print("Index to time mapping was not initialized")
-        return self.list_index_to_time
-
-    def read_particle_list(self, fileName: str) -> Dict[Tuple[int, int], int]:
-        r"""
-        Read and return a list of the particle IDs.
-        """
-        # 2 integers + 1 unsigned long long
-        listUnitSize = 2 * 4 + 8
-        nByte = os.path.getsize(fileName)
-        nPart = int(nByte / listUnitSize)
-        plist = {}
-        with open(fileName, "rb") as f:
-            for _ in range(nPart):
-                binaryData = f.read(listUnitSize)
-                (cpu, id, loc) = struct.unpack("iiQ", binaryData)
-                plist.update({(cpu, id): loc})
-        return plist
-
-    def _read_the_first_record(self, fileName: str) -> Union[List[float], None]:
-        r"""
-        Get the first record stored in one file.
-        """
-        dataList = list()
-        with open(fileName, "rb") as f:
-            while True:
-                binaryData = f.read(4 * 4)
-
-                if not binaryData:
-                    break  # EOF
-
-                (cpu, idtmp, nRecord, weight) = struct.unpack("iiif", binaryData)
-                if nRecord > 0:
-                    binaryData = f.read(4 * self.nReal)
-                    dataList = dataList + list(
-                        struct.unpack("f" * self.nReal, binaryData)
-                    )
-                    return dataList
-
-    def read_particles_at_time(
-        self, time: float, doSave: bool = False
-    ) -> Tuple[np.ndarray, np.ndarray]:
-        r"""
-        Get the information of all the particles at a given time, and save to a csv file
-        with the name "particles_t***.csv" in the current directory if doSave is True.
-
-        Returns
-        -----------
-        ids: a numpy array of tuples contains the particle IDs.
-        pData: a numpy real array with the particle weight, location and velocity.
-
-        Examples
-        ----------
-        >>> ids, pData = pt.read_particles_at_time(3700, doSave=True)
-        """
-
-        nFile = len(self.pfiles)
-        for iFile in range(nFile):
-            if iFile == 0 and time < self.file_time[iFile]:
-                raise Exception("Error: There is no particle at the given time")
-
-            if iFile == nFile - 1:
-                break
-            if time >= self.file_time[iFile] and time < self.file_time[iFile + 1]:
-                break
-
-        fileName = self.pfiles[iFile]
-
-        dataList = []
-        idList = []
-        with open(fileName, "rb") as f:
-            while True:
-                binaryData = f.read(4 * 4)
-                if not binaryData:
-                    break  # EOF
-
-                (cpu, idtmp, nRecord, weight) = struct.unpack("iiif", binaryData)
-                binaryData = f.read(4 * self.nReal * nRecord)
-                allRecords = list(struct.unpack("f" * nRecord * self.nReal, binaryData))
-                for i in range(nRecord):
-                    if (
-                        allRecords[self.nReal * i + FLEKSTP.it_] >= time
-                        or i == nRecord - 1
-                    ):
-                        dataList.append(
-                            allRecords[self.nReal * i : self.nReal * (i + 1)]
-                        )
-                        idList.append((cpu, idtmp))
-                        break
-
-        npData = np.array(dataList)
-        idData = np.array(idList, dtype="i,i")
-        if doSave:
-            fileName = f"particles_t{time}.csv"
-            header = "cpu,iid,time,x,y,z,ux,uy,uz"
-            if self.nReal == 10:
-                header += ",bx,by,bz"
-            elif self.nReal == 13:
-                header += ",bx,by,bz,ex,ey,ez"
-
-            with open(fileName, "w") as f:
-                f.write(header + "\n")
-                for id_row, data_row in zip(idData, npData):
-                    f.write(
-                        f"{id_row[0]},{id_row[1]},{','.join(str(x) for x in data_row)}\n"
-                    )
-
-        return idData, npData
-
-    def save_trajectory_to_csv(
-        self,
-        pID: Tuple[int, int],
-        fileName: str = None,
-        shiftTime: bool = False,
-        scaleTime: bool = False,
-    ) -> None:
-        r"""
-        Save the trajectory of a particle to a csv file.
-
-        Parameters
-        ----------
-        pID: particle ID.
-        shiftTime: If set to True, set the initial time to be 0.
-        scaleTime: If set to True, scale the time into [0,1] range, only scale time if
-                    shiftTime = True.
-
-        Example
-        -----------------
-        >>> tp.save_trajectory_to_csv((3,15))
-        """
-        pData = self.read_particle_trajectory(pID)
-        if fileName == None:
-            fileName = "trajectory_" + str(pID[0]) + "_" + str(pID[1]) + ".csv"
-        header = "time [s], X [R], Y [R], Z [R], U_x [km/s], U_y [km/s], U_z [km/s]"
-        if self.nReal == 10:
-            header += ", B_x [nT], B_y [nT], B_z [nT]"
-        if self.nReal == 13:
-            header += (
-                ", B_x [nT], B_y [nT], B_z [nT], E_x [uV/m], E_y [uV/m], E_z [uV/m]"
-            )
-        if shiftTime:
-            pData[:, 0] -= pData[0, 0]
-            if scaleTime:
-                pData[:, 0] /= pData[-1, 0]
-        np.savetxt(fileName, pData, delimiter=",", header=header, comments="")
-
-    def read_particle_trajectory(self, pID: Tuple[int, int]):
-        r"""
-        Read and return the trajectory of a particle.
-
-        Parameters
-        ----------
-        pID: particle ID
-
-        Examples
-        ----------
-        >>> trajectory = tp.read_particle_trajectory((66,888))
-        """
-        dataList = list()
-        for fileName, plist in zip(self.pfiles, self.plists):
-            if pID in plist:
-                ploc = plist[pID]
-                with open(fileName, "rb") as f:
-                    f.seek(ploc)
-                    binaryData = f.read(4 * 4)
-                    (cpu, idtmp, nRecord, weight) = struct.unpack("iiif", binaryData)
-                    binaryData = f.read(4 * self.nReal * nRecord)
-                    dataList = dataList + list(
-                        struct.unpack("f" * nRecord * self.nReal, binaryData)
-                    )
-
-        nRecord = int(len(dataList) / self.nReal)
-        return np.array(dataList).reshape(nRecord, self.nReal)
-
-    def read_initial_location(self, pID):
-        r"""
-        Read and return the initial location of a test particle
-        """
-
-        for fileName, plist in zip(self.pfiles, self.plists):
-            if pID in plist:
-                ploc = plist[pID]
-                with open(fileName, "rb") as f:
-                    f.seek(ploc)
-                    binaryData = f.read(4 * 4)
-                    (cpu, idtmp, nRecord, weight) = struct.unpack("iiif", binaryData)
-                    nRead = 1
-                    binaryData = f.read(4 * self.nReal * nRead)
-                    dataList = list(struct.unpack("f" * nRead * self.nReal, binaryData))
-                return dataList
-
-    def select_particles(self, fSelect: Callable = None) -> List[Tuple[int, int]]:
-        r"""
-        Select and return the particles whose initial condition satisfy the requirement
-        set by the user defined function fSelect. The first argument of fSelect is the
-        particle ID, and the second argument is the initial record (time, location,
-        velocity and weight of a particle) of a particle.
-
-        Examples
-        ----------
-        >>> def fselect(tp, pid):
-        >>>     pdata = tp.read_initial_loc_with_ID(pid)
-        >>>     intime = pdata[FLEKSTP.it_] < 3601
-        >>>     inregion = pdata[FLEKSTP.ix_] > 20
-        >>>     return intime and inregion
-        >>>
-        >>> pselected = tp.select_particles(fselect)
-        >>> tp.plot_trajectory(list(pselected.keys())[1])
-        """
-
-        if fSelect == None:
-
-            def fSelect(tp, pid):
-                return True
-
-        pselected = list(filter(lambda pid: fSelect(self, pid), self.IDs))
-
-        return pselected
-
-    def plot_trajectory(self, pID: Tuple[int, int]):
-        r"""
-        Plots the trajectory and velocities of the particle pID.
-
-        Example
-        -----------------
-        >>> tp.plot_trajectory((3,15))
-        """
-
-        data = self.read_particle_trajectory(pID)
-        t = data[:, FLEKSTP.it_]
-        tNorm = (t - t[0]) / (t[-1] - t[0])
-
-        ncol = 3
-        nrow = 3  # Default for X, V
-        if self.nReal == 10:  # additional B field
-            nrow = 4
-        elif self.nReal == 13:  # additional B and E field
-            nrow = 5
-
-        f, axs = plt.subplots(nrow, ncol, figsize=(12, 6), constrained_layout=True)
-
-        # Plot trajectories
-        for i, ax in enumerate(axs[0, :]):
-            x_id = FLEKSTP.ix_ if i < 2 else FLEKSTP.iy_
-            y_id = FLEKSTP.iy_ if i == 0 else FLEKSTP.iz_
-            ax.plot(data[:, x_id], data[:, y_id], "k")
-            ax.scatter(
-                data[:, x_id],
-                data[:, y_id],
-                c=plt.cm.winter(tNorm),
-                edgecolor="none",
-                marker="o",
-                s=10,
-            )
-            ax.set_xlabel("x" if i < 2 else "y")
-            ax.set_ylabel("y" if i == 0 else "z")
-
-        def plot_data(dd, label, irow, icol):
-            axs[irow, icol].plot(t, dd, label=label)
-            axs[irow, icol].scatter(
-                t, dd, c=plt.cm.winter(tNorm), edgecolor="none", marker="o", s=10
-            )
-            axs[irow, icol].set_xlabel("time")
-            axs[irow, icol].set_ylabel(label)
-
-        def plot_vector(idx, labels, irow):
-            for i, (id, label) in enumerate(zip(idx, labels)):
-                plot_data(data[:, id], label, irow, i)
-
-        plot_vector([FLEKSTP.ix_, FLEKSTP.iy_, FLEKSTP.iz_], ["x", "y", "z"], 1)
-        plot_vector([FLEKSTP.iu_, FLEKSTP.iv_, FLEKSTP.iw_], ["Vx", "Vy", "Vz"], 2)
-
-        if self.nReal > FLEKSTP.iBx_:
-            plot_vector(
-                [FLEKSTP.iBx_, FLEKSTP.iBy_, FLEKSTP.iBz_], ["Bx", "By", "Bz"], 3
-            )
-
-        if self.nReal > FLEKSTP.iEx_:
-            plot_vector(
-                [FLEKSTP.iEx_, FLEKSTP.iEy_, FLEKSTP.iEz_], ["Ex", "Ey", "Ez"], 4
-            )
-
-        return
-
-    def plot_loc(self, pData: np.ndarray):
-        r"""
-        Plot the location of particles pData.
-
-        Examples
-        ----------
-        >>> ids, pData = tp.read_particles_at_time(3700, doSave=True)
-        >>> f = tp.plot_loc(pData)
-        """
-
-        px = pData[:, FLEKSTP.ix_]
-        py = pData[:, FLEKSTP.iy_]
-        pz = pData[:, FLEKSTP.iz_]
-
-        # create subplot mosaic with different keyword arguments
-        skeys = ["A", "B", "C", "D"]
-        f, axs = plt.subplot_mosaic(
-            "AB;CD",
-            per_subplot_kw={("D"): {"projection": "3d"}},
-            gridspec_kw={"width_ratios": [1, 1], "wspace": 0.1, "hspace": 0.1},
-            figsize=(10, 10),
-            constrained_layout=True,
-        )
-
-        # Create 2D scatter plots
-        for i, (x, y, labels) in enumerate(
-            zip([px, px, py], [py, pz, pz], [("x", "y"), ("x", "z"), ("y", "z")])
-        ):
-            axs[skeys[i]].scatter(x, y, s=1)
-            axs[skeys[i]].set_xlabel(labels[0])
-            axs[skeys[i]].set_ylabel(labels[1])
-
-        # Create 3D scatter plot
-        axs[skeys[3]].scatter(px, py, pz, s=1)
-        axs[skeys[3]].set_xlabel("x")
-        axs[skeys[3]].set_ylabel("y")
-        axs[skeys[3]].set_zlabel("z")
-
-        return
+from typing import List, Tuple, Dict, Union, Callable
+
+import matplotlib.pyplot as plt
+import os
+import numpy as np
+import glob
+import struct
+
+
+class FLEKSTP(object):
+    r"""
+    A class that is used to read and plot test particles. Each particle ID consists of
+    a CPU index, a particle index on each CPU, and a location index.
+    By default, 7 real numbers saved for each step: time + position + velocity.
+
+    Args:
+        dirs (str): the path to the test particle dataset.
+
+    Examples:
+    >>> tp = FLEKSTP("res/run1/PC/test_particles", iSpecies=1)
+    >>> pIDs = list(tp.IDs())
+    >>> tp.plot_trajectory(pIDs[3])
+    >>> tp.save_trajectory_to_csv(pIDs[5])
+    >>> ids, pData = tp.read_particles_at_time(6500.8, doSave=True)
+    >>> f = tp.plot_loc(pData)
+    """
+
+    it_ = 0
+    ix_ = 1
+    iy_ = 2
+    iz_ = 3
+    iu_ = 4
+    iv_ = 5
+    iw_ = 6
+    iBx_ = 7
+    iBy_ = 8
+    iBz_ = 9
+    iEx_ = 10
+    iEy_ = 11
+    iEz_ = 12
+
+    def __init__(
+        self,
+        dirs: Union[str, List[str]],
+        iDomain: int = 0,
+        iSpecies: int = 0,
+        iListStart: int = 0,
+        iListEnd: int = -1,
+        readAllFiles: bool = False,
+    ):
+        if type(dirs) == str:
+            dirs = [dirs]
+
+        header = dirs[0] + "/Header"
+        if os.path.exists(header):
+            with open(header, "r") as f:
+                self.nReal = int(f.readline())
+        else:
+            raise FileNotFoundError(f"Header file not found in {dirs[0]}")
+
+        self.iSpecies = iSpecies
+        self.plistfiles = list()
+        self.pfiles = list()
+
+        for outputDir in dirs:
+            self.plistfiles = self.plistfiles + glob.glob(
+                f"{outputDir}/FLEKS{iDomain}_particle_list_species_{iSpecies}_*"
+            )
+
+            self.pfiles = self.pfiles + glob.glob(
+                f"{outputDir}/FLEKS{iDomain}_particle_species_{iSpecies}_*"
+            )
+
+        self.plistfiles.sort()
+        self.pfiles.sort()
+
+        self.indextotime = []
+        if readAllFiles:
+            for filename in self.pfiles:
+                record = self._read_the_first_record(filename)
+                if record == None:
+                    continue
+                self.indextotime.append(record[FLEKSTP.it_])
+
+        if iListEnd == -1:
+            iListEnd = len(self.plistfiles)
+        self.plistfiles = self.plistfiles[iListStart:iListEnd]
+        self.pfiles = self.pfiles[iListStart:iListEnd]
+
+        self.plists: List[Dict[Tuple[int, int], int]] = []
+        for fileName in self.plistfiles:
+            self.plists.append(self.read_particle_list(fileName))
+
+        self.IDs = set()
+        for plist in self.plists:
+            self.IDs.update(plist.keys())
+
+        self.filetime = []
+        for filename in self.pfiles:
+            record = self._read_the_first_record(filename)
+            if record == None:
+                continue
+            self.filetime.append(record[FLEKSTP.it_])
+
+        print(f"Particles of species {self.iSpecies} are read from {dirs}")
+        print(f"Number of particles: {len(self.IDs)}")
+
+    def getIDs(self):
+        return list(sorted(self.IDs))
+
+    def get_index_to_time(self) -> List:
+        """
+        Getter method for accessing indextotime.
+        """
+        if len(self.indextotime) == 0:
+            print("Index to time mapping was not initialized")
+        return self.indextotime
+
+    def read_particle_list(self, fileName: str) -> Dict[Tuple[int, int], int]:
+        """
+        Read and return a list of the particle IDs.
+        """
+        # 2 integers + 1 unsigned long long
+        listUnitSize = 2 * 4 + 8
+        nByte = os.path.getsize(fileName)
+        nPart = int(nByte / listUnitSize)
+        plist = {}
+        with open(fileName, "rb") as f:
+            for _ in range(nPart):
+                binaryData = f.read(listUnitSize)
+                (cpu, id, loc) = struct.unpack("iiQ", binaryData)
+                plist.update({(cpu, id): loc})
+        return plist
+
+    def _read_the_first_record(self, fileName: str) -> Union[List[float], None]:
+        """
+        Get the first record stored in one file.
+        """
+        dataList = list()
+        with open(fileName, "rb") as f:
+            while True:
+                binaryData = f.read(4 * 4)
+
+                if not binaryData:
+                    break  # EOF
+
+                (cpu, idtmp, nRecord, weight) = struct.unpack("iiif", binaryData)
+                if nRecord > 0:
+                    binaryData = f.read(4 * self.nReal)
+                    dataList = dataList + list(
+                        struct.unpack("f" * self.nReal, binaryData)
+                    )
+                    return dataList
+
+    def read_particles_at_time(
+        self, time: float, doSave: bool = False
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        r"""
+        Get the information of all the particles at a given time, and save to a csv file
+        with the name "particles_t***.csv" in the current directory if doSave is True.
+
+        Returns:
+            ids: a numpy array of tuples contains the particle IDs.
+            pData: a numpy real array with the particle weight, location and velocity.
+
+        Examples:
+        >>> ids, pData = pt.read_particles_at_time(3700, doSave=True)
+        """
+
+        nFile = len(self.pfiles)
+        if time < self.filetime[0] or time > self.filetime[-1]:
+            raise Exception(f"There are no particles at time {time}.")
+        iFile = 0
+        while iFile < nFile - 1:
+            if time < self.filetime[iFile + 1]:
+                break
+            iFile += 1
+        
+        fileName = self.pfiles[iFile]
+
+        dataList: list[float] = []
+        idList: list[tuple] = []
+        with open(fileName, "rb") as f:
+            while True:
+                binaryData = f.read(4 * 4)
+                if not binaryData:
+                    break  # EOF
+
+                (cpu, idtmp, nRecord, weight) = struct.unpack("iiif", binaryData)
+                binaryData = f.read(4 * self.nReal * nRecord)
+                allRecords = list(struct.unpack("f" * nRecord * self.nReal, binaryData))
+                for i in range(nRecord):
+                    if (
+                        allRecords[self.nReal * i + FLEKSTP.it_] >= time
+                        or i == nRecord - 1
+                    ):
+                        dataList.append(
+                            allRecords[self.nReal * i : self.nReal * (i + 1)]
+                        )
+                        idList.append((cpu, idtmp))
+                        break
+
+        npData = np.array(dataList)
+        idData = np.array(idList, dtype="i,i")
+        if doSave:
+            fileName = f"particles_t{time}.csv"
+            header = "cpu,iid,time,x,y,z,ux,uy,uz"
+            if self.nReal == 10:
+                header += ",bx,by,bz"
+            elif self.nReal == 13:
+                header += ",bx,by,bz,ex,ey,ez"
+
+            with open(fileName, "w") as f:
+                f.write(header + "\n")
+                for id_row, data_row in zip(idData, npData):
+                    f.write(
+                        f"{id_row[0]},{id_row[1]},{','.join(str(x) for x in data_row)}\n"
+                    )
+
+        return idData, npData
+
+    def save_trajectory_to_csv(
+        self,
+        pID: Tuple[int, int],
+        fileName: str = None,
+        shiftTime: bool = False,
+        scaleTime: bool = False,
+    ) -> None:
+        r"""
+        Save the trajectory of a particle to a csv file.
+
+        Args:
+            pID: particle ID.
+            shiftTime (bool): If set to True, set the initial time to be 0.
+            scaleTime (bool): If set to True, scale the time into [0,1] range, only scale time if shiftTime = True.
+
+        Example:
+        >>> tp.save_trajectory_to_csv((3,15))
+        """
+        pData = self.read_particle_trajectory(pID)
+        if fileName == None:
+            fileName = "trajectory_" + str(pID[0]) + "_" + str(pID[1]) + ".csv"
+        header = "time [s], X [R], Y [R], Z [R], U_x [km/s], U_y [km/s], U_z [km/s]"
+        if self.nReal == 10:
+            header += ", B_x [nT], B_y [nT], B_z [nT]"
+        if self.nReal == 13:
+            header += (
+                ", B_x [nT], B_y [nT], B_z [nT], E_x [uV/m], E_y [uV/m], E_z [uV/m]"
+            )
+        if shiftTime:
+            pData[:, 0] -= pData[0, 0]
+            if scaleTime:
+                pData[:, 0] /= pData[-1, 0]
+        np.savetxt(fileName, pData, delimiter=",", header=header, comments="")
+
+    def read_particle_trajectory(self, pID: Tuple[int, int]):
+        r"""
+        Read and return the trajectory of a particle.
+
+        Args:
+            pID: particle ID
+
+        Examples:
+        >>> trajectory = tp.read_particle_trajectory((66,888))
+        """
+        dataList = list()
+        for fileName, plist in zip(self.pfiles, self.plists):
+            if pID in plist:
+                ploc = plist[pID]
+                with open(fileName, "rb") as f:
+                    f.seek(ploc)
+                    binaryData = f.read(4 * 4)
+                    (cpu, idtmp, nRecord, weight) = struct.unpack("iiif", binaryData)
+                    binaryData = f.read(4 * self.nReal * nRecord)
+                    dataList = dataList + list(
+                        struct.unpack("f" * nRecord * self.nReal, binaryData)
+                    )
+
+        nRecord = int(len(dataList) / self.nReal)
+        return np.array(dataList).reshape(nRecord, self.nReal)
+
+    def read_initial_location(self, pID):
+        r"""
+        Read and return the initial location of a test particle
+        """
+
+        for fileName, plist in zip(self.pfiles, self.plists):
+            if pID in plist:
+                ploc = plist[pID]
+                with open(fileName, "rb") as f:
+                    f.seek(ploc)
+                    binaryData = f.read(4 * 4)
+                    (cpu, idtmp, nRecord, weight) = struct.unpack("iiif", binaryData)
+                    nRead = 1
+                    binaryData = f.read(4 * self.nReal * nRead)
+                    dataList = list(struct.unpack("f" * nRead * self.nReal, binaryData))
+                return dataList
+
+    def select_particles(self, fSelect: Callable = None) -> List[Tuple[int, int]]:
+        r"""
+        Select and return the particles whose initial condition satisfy the requirement
+        set by the user defined function fSelect. The first argument of fSelect is the
+        particle ID, and the second argument is the initial record (time, location,
+        velocity and weight of a particle) of a particle.
+
+        Examples:
+        >>> def fselect(tp, pid):
+        >>>     pdata = tp.read_initial_loc_with_ID(pid)
+        >>>     intime = pdata[FLEKSTP.it_] < 3601
+        >>>     inregion = pdata[FLEKSTP.ix_] > 20
+        >>>     return intime and inregion
+        >>>
+        >>> pselected = tp.select_particles(fselect)
+        >>> tp.plot_trajectory(list(pselected.keys())[1])
+        """
+
+        if fSelect == None:
+
+            def fSelect(tp, pid):
+                return True
+
+        pselected = list(filter(lambda pid: fSelect(self, pid), self.IDs))
+
+        return pselected
+
+    def get_data(self, data, name: str):
+        match name:
+            case "t":
+                x = data[:, FLEKSTP.it_]
+            case "x":
+                x = data[:, FLEKSTP.ix_]
+            case "y":
+                x = data[:, FLEKSTP.iy_]
+            case "z":
+                x = data[:, FLEKSTP.iz_]
+            case "u" | "vx" | "ux":
+                x = data[:, FLEKSTP.iu_]
+            case "v" | "vy" | "uy":
+                x = data[:, FLEKSTP.iv_]
+            case "w" | "vz" | "uz":
+                x = data[:, FLEKSTP.iw_]
+            case _:
+                raise Exception(f"Unknown plot variable {name}")
+
+        return x
+
+    def plot_trajectory(
+        self,
+        pID: Tuple[int, int],
+        *,
+        type="all",
+        xaxis="t",
+        yaxis="x",
+        ax=None,
+        **kwargs,
+    ):
+        r"""
+        Plots the trajectory and velocities of the particle pID.
+
+        Example:
+        >>> tp.plot_trajectory((3,15))
+        """
+
+        def plot_data(dd, label, irow, icol):
+            ax[irow, icol].plot(t, dd, label=label)
+            ax[irow, icol].scatter(
+                t, dd, c=plt.cm.winter(tNorm), edgecolor="none", marker="o", s=10
+            )
+            ax[irow, icol].set_xlabel("time")
+            ax[irow, icol].set_ylabel(label)
+
+        def plot_vector(idx, labels, irow):
+            for i, (id, label) in enumerate(zip(idx, labels)):
+                plot_data(data[:, id], label, irow, i, **kwargs)
+
+        data = self.read_particle_trajectory(pID)
+        t = data[:, FLEKSTP.it_]
+        tNorm = (t - t[0]) / (t[-1] - t[0])
+
+        if type == "single":
+            x = self.get_data(data, xaxis)
+            y = self.get_data(data, yaxis)
+
+            if ax == None:
+                f, ax = plt.subplots(1, 1, figsize=(10, 6), constrained_layout=True)
+
+            ax.plot(x, y, **kwargs)
+        elif type == "all":
+            ncol = 3
+            nrow = 3  # Default for X, V
+            if self.nReal == 10:  # additional B field
+                nrow = 4
+            elif self.nReal == 13:  # additional B and E field
+                nrow = 5
+
+            f, ax = plt.subplots(nrow, ncol, figsize=(12, 6), constrained_layout=True)
+
+            # Plot trajectories
+            for i, a in enumerate(ax[0, :]):
+                x_id = FLEKSTP.ix_ if i < 2 else FLEKSTP.iy_
+                y_id = FLEKSTP.iy_ if i == 0 else FLEKSTP.iz_
+                a.plot(data[:, x_id], data[:, y_id], "k")
+                a.scatter(
+                    data[:, x_id],
+                    data[:, y_id],
+                    c=plt.cm.winter(tNorm),
+                    edgecolor="none",
+                    marker="o",
+                    s=10,
+                )
+                a.set_xlabel("x" if i < 2 else "y")
+                a.set_ylabel("y" if i == 0 else "z")
+
+            plot_vector([FLEKSTP.ix_, FLEKSTP.iy_, FLEKSTP.iz_], ["x", "y", "z"], 1)
+            plot_vector([FLEKSTP.iu_, FLEKSTP.iv_, FLEKSTP.iw_], ["Vx", "Vy", "Vz"], 2)
+
+            if self.nReal > FLEKSTP.iBx_:
+                plot_vector(
+                    [FLEKSTP.iBx_, FLEKSTP.iBy_, FLEKSTP.iBz_], ["Bx", "By", "Bz"], 3
+                )
+
+            if self.nReal > FLEKSTP.iEx_:
+                plot_vector(
+                    [FLEKSTP.iEx_, FLEKSTP.iEy_, FLEKSTP.iEz_], ["Ex", "Ey", "Ez"], 4
+                )
+
+        return ax
+
+    def plot_loc(self, pData: np.ndarray):
+        r"""
+        Plot the location of particles pData.
+
+        Examples:
+        >>> ids, pData = tp.read_particles_at_time(3700, doSave=True)
+        >>> f = tp.plot_loc(pData)
+        """
+
+        px = pData[:, FLEKSTP.ix_]
+        py = pData[:, FLEKSTP.iy_]
+        pz = pData[:, FLEKSTP.iz_]
+
+        # create subplot mosaic with different keyword arguments
+        skeys = ["A", "B", "C", "D"]
+        f, ax = plt.subplot_mosaic(
+            "AB;CD",
+            per_subplot_kw={("D"): {"projection": "3d"}},
+            gridspec_kw={"width_ratios": [1, 1], "wspace": 0.1, "hspace": 0.1},
+            figsize=(10, 10),
+            constrained_layout=True,
+        )
+
+        # Create 2D scatter plots
+        for i, (x, y, labels) in enumerate(
+            zip([px, px, py], [py, pz, pz], [("x", "y"), ("x", "z"), ("y", "z")])
+        ):
+            ax[skeys[i]].scatter(x, y, s=1)
+            ax[skeys[i]].set_xlabel(labels[0])
+            ax[skeys[i]].set_ylabel(labels[1])
+
+        # Create 3D scatter plot
+        ax[skeys[3]].scatter(px, py, pz, s=1)
+        ax[skeys[3]].set_xlabel("x")
+        ax[skeys[3]].set_ylabel("y")
+        ax[skeys[3]].set_zlabel("z")
+
+        return ax

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

@@ -79,12 +79,12 @@ class DataContainer(object):
         if verbose > 0:
             f.text(0.01, 0.01, s)
 
-    def analyze_variable_string(self, var):
+    def analyze_variable_string(self, var: str):
         r"""
         Parses the input string and return the plot variable and plot range.
 
         Args:
-            var: str
+            var (str): string of variables to be parsed
             Example: var = "{bb}<(-10)>(-9.8)"
 
         Return: a tuple of the variable name, variable min and max.
@@ -116,7 +116,7 @@ class DataContainer(object):
         Evaluates the variable expression and return the result of an YTArray.
 
         Args:
-            expression: str
+            expression (str): Python codes to be executed
             Example: expression = "np.log({rhos0}+{rhos1})"
         """
 
@@ -138,8 +138,7 @@ class DataContainer(object):
         Adds a variable to the dataset.
 
         Args:
-            name: str
-                The name of the variable to be added into self.data
+            name (str): The name of the variable to be added.
 
             val: array-like structure
                 Values of the variable stored in an array
@@ -157,8 +156,7 @@ class DataContainer(object):
         Return raw variables or calculate derived variables.
 
         Args:
-            var: str
-                variable name
+            var (str): variable name
 
         Return: YTArray
         """
@@ -217,8 +215,7 @@ class DataContainer3D(DataContainer):
     ):
         r"""
         Args:
-            dataSets: dictonary
-                The key is the variable name, and the dictionary value is usually a YTArray.
+            dataSets (dict): The keys are variable names, and the dictionary values are YTArray.
 
             x/y/z: A 1D YTArray
         """
@@ -230,11 +227,9 @@ class DataContainer3D(DataContainer):
         """Get a 2D slice from the 3D box data.
 
         Args:
-            norm: str
-                The normal direction of the slice from "x", "y" or "z"
+            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: DataContainer2D
         """
@@ -293,18 +288,15 @@ class DataContainer2D(DataContainer):
     ):
         r"""
         Args:
-            dataSets: dictonary
-                The keys are variable names, and the values are YTArrays.
+            dataSets (dict): The keys are variable names, and the values are YTArrays.
 
             x/y: 1D YTArray
 
-            xlabel/ylabel: str
+            xlabel/ylabel (str): x, y dimension labels
 
-            cut_norm: str
-                "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
@@ -327,15 +319,15 @@ class DataContainer2D(DataContainer):
         vars,
         xlim=None,
         ylim=None,
-        unit="planet",
-        nlevels=200,
-        cmap="turbo",
+        unit: str = "planet",
+        nlevels: int = 200,
+        cmap: str = "turbo",
         figsize=(10, 6),
         pcolor=False,
         logscale=False,
         addgrid=False,
         bottomline=10,
-        showcolorbar=True,
+        showcolorbar: bool = True,
         *args,
         **kwargs,
     ):
@@ -343,25 +335,20 @@ class DataContainer2D(DataContainer):
         2D plots.
 
         Args:
-            vars: str
-                Ploting variables and ploting range.
+            vars (str): 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
 
-            unit: str
-                "planet" or "si"
+            unit (str): "planet" or "si"
 
-            nlevels: int
-                Number of the contour levels
+            nlevels (int): Number of the contour levels. Default 200.
 
-            cmap: str
-                Color map type
+            cmap (str): color map type from Matplotlib
 
-            figsize: tuple
+            figsize (tuple): size of figure. Default (10, 6).
 
-            logscale: Bool
-                True to scale the variable in log.
+            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])
@@ -469,10 +456,9 @@ class DataContainer2D(DataContainer):
         r"""Adding contour lines to an axis.
 
         Args:
-            ax: matplotlib axis
-                The axis to plot contour lines.
+            ax (matplotlib axis): the axis to plot contour lines.
 
-            var: string
+            var (str): variable of contours.
 
         Examples:
         >>> f, axes = dc.plot("Bx<(50)>(-50) By (np.log(2*{rhos0}))>(-5)", xlim=[-40,-5])
@@ -521,14 +507,11 @@ class DataContainer2D(DataContainer):
         Adding streamlines to an axis.
 
         Args:
-            ax: matplotlib axis
-                The axis to add streamlines
+            ax (matplotlib axis): the axis to add streamlines.
 
-            var1/var2: str
-                Streamline variable names
+            var1/var2 (str): streamline variable names.
 
-            density: int
-                It controls the number of streamlines.
+            density (float): it controls the number of streamlines.
 
         Examples:
         >>> f, axes = dc.plot("Bx<(50)>(-50) By (np.log(2*{rhos0}))>(-5)", xlim=[-40,-5])
@@ -589,12 +572,11 @@ class DataContainer1D(DataContainer):
     def __init__(self, dataSets, x, xlabel, *args, **kwargs):
         r"""
         Args:
-            dataSets: dictonary
-                The key is the variable name, and the dictionary value is usually a YTArray.
+            dataSets (dict): the keys are variable names, and the values are YTArrays.
 
             x: 1D YTArray
 
-            xlabel: str
+            xlabel (str): x axis label.
         """
 
         ylabel = None

{flekspy-0.2.2 → flekspy-0.2.4}/src/flekspy/util/utilities.py

@@ -76,7 +76,16 @@ plot_unit_si = {
 }
 
 
-def get_unit(var, unit_type="planet"):
+def get_unit(var: str, unit_type="planet") -> str:
+    """Return the unit of variable.
+
+    Args:
+        var (str): variable name.
+        unit_type (str, optional): unit system. Defaults to "planet".
+
+    Returns:
+        str: unit in the specified unit system.
+    """
     if var[-1].isdigit():
         # Example: pxxs0 -> pxx
         var = var[0:-2]
@@ -111,21 +120,24 @@ def get_ticks(vmin, vmax):
     nticks = int((tickMax - tickMin) / dv) + 1
     return np.linspace(tickMin, tickMax, nticks)
 
+
 def unit_one(field, data):
     """Utility function for setting equal weights for macroparticles.
     TBD: add units.
     """
     return np.ones_like(data[("particles", "p_w")])
 
-def download_testfile(url, target_path="."):
+
+def download_testfile(url: str, target_path="."):
     """
     Downloads a tar.gz file from a URL and extracts its contents.
 
     Args:
-      url: The URL of the tar.gz file.
-      target_path: The directory to extract the files to.
-                    Defaults to the current directory.
+      url (str): the URL of the tar.gz file.
+      target_path (str): the directory to extract the files to. Defaults to the current directory.
     """
+    import os
+
     try:
         response = requests.get(url, stream=True)
         response.raise_for_status()  # Raise an exception for bad status codes
@@ -134,6 +146,9 @@ def download_testfile(url, target_path="."):
             for chunk in response.iter_content(chunk_size=8192):
                 f.write(chunk)
 
+        if not os.path.exists(target_path):
+            os.makedirs(target_path)
+
         with tarfile.open("temp.tar.gz", "r:gz") as tar:
             tar.extractall(target_path)
 
@@ -142,7 +157,5 @@ def download_testfile(url, target_path="."):
     except tarfile.TarError as e:
         print(f"Error extracting tar file: {e}")
     finally:
-        import os
-
         if os.path.exists("temp.tar.gz"):
             os.remove("temp.tar.gz")  # Clean up temporary file

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

@@ -120,8 +120,7 @@ class FLEKSData(BoxlibDataset):
     Read and plot the AMReX format output from FLEKS.
 
     Args:
-        output_dir : String
-            The path to the data.
+        output_dir (str): the path to the data.
 
     Examples:
         >>> import flekspy