ruff fixes

mdgo/__init__.py

@@ -1,9 +1,7 @@
 # Copyright (c) Tingzheng Hou.
 # Distributed under the terms of the MIT License.
 
-"""
-This package contains core modules and classes molecular dynamics simulation setup and analysis.
-"""
+"""This package contains core modules and classes molecular dynamics simulation setup and analysis."""
 
 from __future__ import annotations
 

mdgo/core/analysis.py

@@ -201,9 +201,7 @@ def from_lammps(
         )
 
     def get_init_dimension(self) -> np.ndarray:
-        """
-        Returns the initial box dimension.
-        """
+        """Returns the initial box dimension."""
         return self.wrapped_run.trajectory[0].dimensions
 
     def get_equilibrium_dimension(self, npt_range: int, period: int = 200) -> np.ndarray:
@@ -234,9 +232,7 @@ def get_equilibrium_dimension(self, npt_range: int, period: int = 200) -> np.nda
         return np.mean(np.array(d), axis=0)
 
     def get_nvt_dimension(self) -> np.ndarray:
-        """
-        Returns the box dimension at the last frame.
-        """
+        """Returns the box dimension at the last frame."""
         return self.wrapped_run.trajectory[-1].dimensions
 
     def get_cond_array(self) -> np.ndarray:
@@ -249,15 +245,14 @@ def get_cond_array(self) -> np.ndarray:
         nvt_run = self.unwrapped_run
         cations = nvt_run.select_atoms(self.select_dict.get("cation"))
         anions = nvt_run.select_atoms(self.select_dict.get("anion"))
-        cond_array = calc_cond_msd(
+        return calc_cond_msd(
             nvt_run,
             anions,
             cations,
             self.nvt_start,
             self.cation_charge,
             self.anion_charge,
         )
-        return cond_array
 
     def choose_cond_fit_region(self) -> tuple:
         """
@@ -358,10 +353,9 @@ def get_conductivity(self, start: int = -1, end: int = -1) -> float:
         print(f"Start of linear fitting regime: {start} ({self.time_array[start]} {time_units})")
         print(f"End of linear fitting regime: {end} ({self.time_array[end]} {time_units})")
         print(f"Beta value (fit to MSD = t^\u03B2): {beta} (\u03B2 = 1 in the diffusive regime)")
-        cond = conductivity_calculator(
+        return conductivity_calculator(
             self.time_array, self.cond_array, self.nvt_v, self.name, start, end, self.temp, self.units
         )
-        return cond
 
     def coord_num_array_single_species(
         self,
@@ -386,7 +380,7 @@ def coord_num_array_single_species(
         nvt_run = self.wrapped_run
         distance_dict = {species: distance}
         center_atoms = nvt_run.select_atoms(self.select_dict.get(center_atom))
-        num_array = concat_coord_array(
+        return concat_coord_array(
             nvt_run,
             num_of_neighbor,
             center_atoms,
@@ -395,7 +389,6 @@ def coord_num_array_single_species(
             run_start,
             run_end,
         )["total"]
-        return num_array
 
     def coord_num_array_multi_species(
         self,
@@ -418,7 +411,7 @@ def coord_num_array_multi_species(
         """
         nvt_run = self.wrapped_run
         center_atoms = nvt_run.select_atoms(self.select_dict.get(center_atom))
-        num_array_dict = concat_coord_array(
+        return concat_coord_array(
             nvt_run,
             num_of_neighbor,
             center_atoms,
@@ -427,7 +420,6 @@ def coord_num_array_multi_species(
             run_start,
             run_end,
         )
-        return num_array_dict
 
     def coord_num_array_specific(
         self,
@@ -453,7 +445,7 @@ def coord_num_array_specific(
         """
         nvt_run = self.wrapped_run
         center_atoms = nvt_run.select_atoms(self.select_dict.get(center_atom))
-        num_array_dict = concat_coord_array(
+        return concat_coord_array(
             nvt_run,
             num_of_neighbor_specific,
             center_atoms,
@@ -463,7 +455,6 @@ def coord_num_array_specific(
             run_end,
             counter_atom=counter_atom,
         )
-        return num_array_dict
 
     def write_solvation_structure(
         self,
@@ -475,7 +466,7 @@ def write_solvation_structure(
         write_path: str,
         center_atom: str = "cation",
     ):
-        """Writes out a series of desired solvation structures as ``*.xyz`` files
+        """Writes out a series of desired solvation structures as ``*.xyz`` files.
 
         Args:
             distance_dict: A dict of coordination cutoff distance of the neighbor species.
@@ -528,7 +519,7 @@ def coord_type_array(
         nvt_run = self.wrapped_run
         distance_dict = {counter_atom: distance}
         center_atoms = nvt_run.select_atoms(self.select_dict.get(center_atom))
-        num_array = concat_coord_array(
+        return concat_coord_array(
             nvt_run,
             num_of_neighbor_simple,
             center_atoms,
@@ -537,7 +528,6 @@ def coord_type_array(
             run_start,
             run_end,
         )["total"]
-        return num_array
 
     def angle_array(
         self,
@@ -567,8 +557,8 @@ def angle_array(
         nvt_run = self.wrapped_run
         center_atoms = nvt_run.select_atoms(self.select_dict.get(center_atom))
         assert len(distance_dict) == 2, "Only distance a->c, b->c shoud be specified in the distance_dict."
-        distance_dict[center_atom] = list(distance_dict.values())[0]
-        ang_array = concat_coord_array(
+        distance_dict[center_atom] = next(iter(distance_dict.values()))
+        return concat_coord_array(
             nvt_run,
             angular_dist_of_neighbor,
             center_atoms,
@@ -578,7 +568,6 @@ def angle_array(
             run_end,
             cip=cip,
         )["total"]
-        return ang_array
 
     def coordination(
         self,
@@ -612,8 +601,7 @@ def coordination(
             item_list.append(str(int(combined[i, 0])))
             percent_list.append(f"{(combined[i, 1] / combined[:, 1].sum() * 100):.4f}%")
         df_dict = {item_name: item_list, "Percentage": percent_list}
-        df = pd.DataFrame(df_dict)
-        return df
+        return pd.DataFrame(df_dict)
 
     def rdf_integral(
         self,
@@ -644,8 +632,7 @@ def rdf_integral(
                 item_list.append(kw)
                 cn_list.append(cn)
         df_dict = {item_name: item_list, "CN": cn_list}
-        df = pd.DataFrame(df_dict)
-        return df
+        return pd.DataFrame(df_dict)
 
     def coordination_type(
         self,
@@ -655,7 +642,7 @@ def coordination_type(
         center_atom: str = "cation",
         counter_atom: str = "anion",
     ) -> pd.DataFrame:
-        """Tabulates the percentage of each solvation structures (CIP/SSIP/AGG)
+        """Tabulates the percentage of each solvation structures (CIP/SSIP/AGG).
 
         Args:
             distance: The coordination cutoff distance.
@@ -683,8 +670,7 @@ def coordination_type(
             item_list.append(item_dict.get(item))
             percent_list.append(f"{(combined[i, 1] / combined[:, 1].sum() * 100):.4f}%")
         df_dict = {item_name: item_list, "Percentage": percent_list}
-        df = pd.DataFrame(df_dict)
-        return df
+        return pd.DataFrame(df_dict)
 
     def coordination_specific(
         self,
@@ -695,7 +681,7 @@ def coordination_specific(
         counter_atom: str = "anion",
     ) -> pd.DataFrame:
         """Calculates the integral of the coordiantion number of selected species
-        in each type of solvation structures (CIP/SSIP/AGG)
+        in each type of solvation structures (CIP/SSIP/AGG).
 
         Args:
             distance_dict: A dict of coordination cutoff distance of the neighbor species.
@@ -727,8 +713,7 @@ def coordination_specific(
                 else:
                     agg_list.append(cn)
         df_dict = {item_name: item_list, "CN in SSIP": ssip_list, "CN in CIP": cip_list, "CN in AGG": agg_list}
-        df = pd.DataFrame(df_dict)
-        return df
+        return pd.DataFrame(df_dict)
 
     def get_msd_all(
         self,
@@ -758,7 +743,7 @@ def get_msd_all(
         """
         selection = self.select_dict.get(species)
         assert selection is not None
-        msd_array = total_msd(
+        return total_msd(
             self.unwrapped_run,
             start=start,
             end=end,
@@ -768,7 +753,6 @@ def get_msd_all(
             built_in=built_in,
             center_of_mass=center_of_mass,
         )
-        return msd_array
 
     def get_msd_partial(
         self,
@@ -861,7 +845,7 @@ def get_neighbor_corr(
     def get_residence_time(
         self, times: np.ndarray, acf_avg_dict: dict[str, np.ndarray], cutoff_time: int
     ) -> dict[str, np.floating]:
-        """Calculates the residence time of selected species around cation
+        """Calculates the residence time of selected species around cation.
 
         Args:
             times: The time series.
@@ -882,7 +866,7 @@ def get_neighbor_trj(
         center_atom: str = "cation",
         index: int = 0,
     ) -> dict[str, np.ndarray]:
-        """Returns the distance between one center atom and neighbors as a function of time
+        """Returns the distance between one center atom and neighbors as a function of time.
 
         Args:
             run_start: Start frame of analysis.
@@ -1063,12 +1047,12 @@ def get_heat_map(
         hopping_cutoff: float,
         floating_atom: str = "cation",
         cartesian_by_ref: np.ndarray = None,
-        sym_dict: dict[str, list[np.ndarray]] = None,
+        sym_dict: dict[str, list[np.ndarray]] | None = None,
         sample: int | None = None,
         smooth: int = 51,
         dim: str = "xyz",
     ) -> np.ndarray:
-        """Calculates the heatmap matrix of floating ion around a cluster
+        """Calculates the heatmap matrix of floating ion around a cluster.
 
         Args:
             run_start: Start frame of analysis.
@@ -1133,7 +1117,7 @@ def get_heat_map(
     def get_cluster_distance(
         self, run_start: int, run_end: int, neighbor_cutoff: float, cluster_center: str = "center"
     ) -> np.floating:
-        """Calculates the average distance of the center of clusters/molecules
+        """Calculates the average distance of the center of clusters/molecules.
 
         Args:
             run_start: Start frame of analysis.

mdgo/core/run.py

@@ -1,21 +1,15 @@
 # Copyright (c) Tingzheng Hou.
 # Distributed under the terms of the MIT License.
 
-"""
-This module implements a core class MdRun for molecular dynamics job setup.
-"""
+"""This module implements a core class MdRun for molecular dynamics job setup."""
 from __future__ import annotations
 
 
 class MdJob:
-    """
-    A core class for MD results analysis.
-    """
+    """A core class for MD results analysis."""
 
     def __init__(self, name):
-        """
-        Base constructor
-        """
+        """Base constructor."""
         self.name = name
 
     @classmethod

mdgo/forcefield/aqueous.py

@@ -1,9 +1,7 @@
 # Copyright (c) Tingzheng Hou.
 # Distributed under the terms of the MIT License.
 
-"""
-A class for retrieving water and ion force field parameters.
-"""
+"""A class for retrieving water and ion force field parameters."""
 
 from __future__ import annotations
 
@@ -227,10 +225,7 @@ def get_ion(
         parameter_set = alias.get(parameter_set, parameter_set)
 
         # Make the Ion object to get mass and charge
-        if isinstance(ion, Ion):
-            ion_obj = ion
-        else:
-            ion_obj = Ion.from_formula(ion.capitalize())
+        ion_obj = ion if isinstance(ion, Ion) else Ion.from_formula(ion.capitalize())
 
         # load ion data as a list of IonLJData objects
         ion_data = loadfn(os.path.join(DATA_DIR, "ion_lj_params.json"))

mdgo/forcefield/charge.py

@@ -1,9 +1,7 @@
 # Copyright (c) Tingzheng Hou.
 # Distributed under the terms of the MIT License.
 
-"""
-A class for writing, overwriting, scaling charges of a LammpsData object.
-"""
+"""A class for writing, overwriting, scaling charges of a LammpsData object."""
 
 from __future__ import annotations
 
@@ -16,7 +14,7 @@ class ChargeWriter:
     A class for write, overwrite, scale charges of a LammpsData object.
     TODO: Auto determine number of significant figures of charges
     TODO: write to obj or write separate charge file
-    TODO: Read LammpsData or path
+    TODO: Read LammpsData or path.
 
     Args:
         data: The provided LammpsData obj.
@@ -30,7 +28,7 @@ def __init__(self, data: LammpsData, precision: int = 10):
 
     def scale(self, factor: float) -> LammpsData:
         """
-        Scales the charge in of the in self.data and returns a new one. TODO: check if non-destructive
+        Scales the charge in of the in self.data and returns a new one. TODO: check if non-destructive.
 
         Args:
             factor: The charge scaling factor

mdgo/forcefield/crawler.py

@@ -91,10 +91,7 @@ def __init__(
         print("LigParGen server connected.")
 
     def quit(self):
-        """
-        Method for quiting ChromeDriver.
-
-        """
+        """Method for quiting ChromeDriver."""
         self.web.quit()
 
     def data_from_pdb(self, pdb_dir: str):

mdgo/forcefield/maestro.py

@@ -146,8 +146,7 @@ def get_ff(self):
                 FFLD.format(self.mae + ".mae", self.ff),
                 check=True,
                 shell=True,
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
+                capture_output=True,
             )
         except subprocess.CalledProcessError as e:
             raise ValueError(f"Maestro failed with errorcode {e.returncode} and stderr: {e.stderr}") from e

mdgo/forcefield/pubchem.py

@@ -84,10 +84,7 @@ def __init__(
             print("PubChem server connected.")
 
     def quit(self):
-        """
-        Method for quiting ChromeDriver.
-
-        """
+        """Method for quiting ChromeDriver."""
         if not self.api:
             self.web.quit()