Skip to content

Commit 082f884

Browse files
htz1992213htz1992213
committed
ruff fixes
1 parent c81b502 commit 082f884

20 files changed

+94
-147
lines changed

‎mdgo/__init__.py

+1-3
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
11
# Copyright (c) Tingzheng Hou.
22
# Distributed under the terms of the MIT License.
33

4-
"""
5-
This package contains core modules and classes molecular dynamics simulation setup and analysis.
6-
"""
4+
"""This package contains core modules and classes molecular dynamics simulation setup and analysis."""
75

86
from __future__ import annotations
97

‎mdgo/core/analysis.py

+23-39
Original file line numberDiff line numberDiff line change
@@ -201,9 +201,7 @@ def from_lammps(
201201
)
202202

203203
def get_init_dimension(self) -> np.ndarray:
204-
"""
205-
Returns the initial box dimension.
206-
"""
204+
"""Returns the initial box dimension."""
207205
return self.wrapped_run.trajectory[0].dimensions
208206

209207
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
234232
return np.mean(np.array(d), axis=0)
235233

236234
def get_nvt_dimension(self) -> np.ndarray:
237-
"""
238-
Returns the box dimension at the last frame.
239-
"""
235+
"""Returns the box dimension at the last frame."""
240236
return self.wrapped_run.trajectory[-1].dimensions
241237

242238
def get_cond_array(self) -> np.ndarray:
@@ -249,15 +245,14 @@ def get_cond_array(self) -> np.ndarray:
249245
nvt_run = self.unwrapped_run
250246
cations = nvt_run.select_atoms(self.select_dict.get("cation"))
251247
anions = nvt_run.select_atoms(self.select_dict.get("anion"))
252-
cond_array = calc_cond_msd(
248+
return calc_cond_msd(
253249
nvt_run,
254250
anions,
255251
cations,
256252
self.nvt_start,
257253
self.cation_charge,
258254
self.anion_charge,
259255
)
260-
return cond_array
261256

262257
def choose_cond_fit_region(self) -> tuple:
263258
"""
@@ -358,10 +353,9 @@ def get_conductivity(self, start: int = -1, end: int = -1) -> float:
358353
print(f"Start of linear fitting regime: {start} ({self.time_array[start]} {time_units})")
359354
print(f"End of linear fitting regime: {end} ({self.time_array[end]} {time_units})")
360355
print(f"Beta value (fit to MSD = t^\u03B2): {beta} (\u03B2 = 1 in the diffusive regime)")
361-
cond = conductivity_calculator(
356+
return conductivity_calculator(
362357
self.time_array, self.cond_array, self.nvt_v, self.name, start, end, self.temp, self.units
363358
)
364-
return cond
365359

366360
def coord_num_array_single_species(
367361
self,
@@ -386,7 +380,7 @@ def coord_num_array_single_species(
386380
nvt_run = self.wrapped_run
387381
distance_dict = {species: distance}
388382
center_atoms = nvt_run.select_atoms(self.select_dict.get(center_atom))
389-
num_array = concat_coord_array(
383+
return concat_coord_array(
390384
nvt_run,
391385
num_of_neighbor,
392386
center_atoms,
@@ -395,7 +389,6 @@ def coord_num_array_single_species(
395389
run_start,
396390
run_end,
397391
)["total"]
398-
return num_array
399392

400393
def coord_num_array_multi_species(
401394
self,
@@ -418,7 +411,7 @@ def coord_num_array_multi_species(
418411
"""
419412
nvt_run = self.wrapped_run
420413
center_atoms = nvt_run.select_atoms(self.select_dict.get(center_atom))
421-
num_array_dict = concat_coord_array(
414+
return concat_coord_array(
422415
nvt_run,
423416
num_of_neighbor,
424417
center_atoms,
@@ -427,7 +420,6 @@ def coord_num_array_multi_species(
427420
run_start,
428421
run_end,
429422
)
430-
return num_array_dict
431423

432424
def coord_num_array_specific(
433425
self,
@@ -453,7 +445,7 @@ def coord_num_array_specific(
453445
"""
454446
nvt_run = self.wrapped_run
455447
center_atoms = nvt_run.select_atoms(self.select_dict.get(center_atom))
456-
num_array_dict = concat_coord_array(
448+
return concat_coord_array(
457449
nvt_run,
458450
num_of_neighbor_specific,
459451
center_atoms,
@@ -463,7 +455,6 @@ def coord_num_array_specific(
463455
run_end,
464456
counter_atom=counter_atom,
465457
)
466-
return num_array_dict
467458

468459
def write_solvation_structure(
469460
self,
@@ -475,7 +466,7 @@ def write_solvation_structure(
475466
write_path: str,
476467
center_atom: str = "cation",
477468
):
478-
"""Writes out a series of desired solvation structures as ``*.xyz`` files
469+
"""Writes out a series of desired solvation structures as ``*.xyz`` files.
479470
480471
Args:
481472
distance_dict: A dict of coordination cutoff distance of the neighbor species.
@@ -528,7 +519,7 @@ def coord_type_array(
528519
nvt_run = self.wrapped_run
529520
distance_dict = {counter_atom: distance}
530521
center_atoms = nvt_run.select_atoms(self.select_dict.get(center_atom))
531-
num_array = concat_coord_array(
522+
return concat_coord_array(
532523
nvt_run,
533524
num_of_neighbor_simple,
534525
center_atoms,
@@ -537,7 +528,6 @@ def coord_type_array(
537528
run_start,
538529
run_end,
539530
)["total"]
540-
return num_array
541531

542532
def angle_array(
543533
self,
@@ -567,8 +557,8 @@ def angle_array(
567557
nvt_run = self.wrapped_run
568558
center_atoms = nvt_run.select_atoms(self.select_dict.get(center_atom))
569559
assert len(distance_dict) == 2, "Only distance a->c, b->c shoud be specified in the distance_dict."
570-
distance_dict[center_atom] = list(distance_dict.values())[0]
571-
ang_array = concat_coord_array(
560+
distance_dict[center_atom] = next(iter(distance_dict.values()))
561+
return concat_coord_array(
572562
nvt_run,
573563
angular_dist_of_neighbor,
574564
center_atoms,
@@ -578,7 +568,6 @@ def angle_array(
578568
run_end,
579569
cip=cip,
580570
)["total"]
581-
return ang_array
582571

583572
def coordination(
584573
self,
@@ -612,8 +601,7 @@ def coordination(
612601
item_list.append(str(int(combined[i, 0])))
613602
percent_list.append(f"{(combined[i, 1] / combined[:, 1].sum() * 100):.4f}%")
614603
df_dict = {item_name: item_list, "Percentage": percent_list}
615-
df = pd.DataFrame(df_dict)
616-
return df
604+
return pd.DataFrame(df_dict)
617605

618606
def rdf_integral(
619607
self,
@@ -644,8 +632,7 @@ def rdf_integral(
644632
item_list.append(kw)
645633
cn_list.append(cn)
646634
df_dict = {item_name: item_list, "CN": cn_list}
647-
df = pd.DataFrame(df_dict)
648-
return df
635+
return pd.DataFrame(df_dict)
649636

650637
def coordination_type(
651638
self,
@@ -655,7 +642,7 @@ def coordination_type(
655642
center_atom: str = "cation",
656643
counter_atom: str = "anion",
657644
) -> pd.DataFrame:
658-
"""Tabulates the percentage of each solvation structures (CIP/SSIP/AGG)
645+
"""Tabulates the percentage of each solvation structures (CIP/SSIP/AGG).
659646
660647
Args:
661648
distance: The coordination cutoff distance.
@@ -683,8 +670,7 @@ def coordination_type(
683670
item_list.append(item_dict.get(item))
684671
percent_list.append(f"{(combined[i, 1] / combined[:, 1].sum() * 100):.4f}%")
685672
df_dict = {item_name: item_list, "Percentage": percent_list}
686-
df = pd.DataFrame(df_dict)
687-
return df
673+
return pd.DataFrame(df_dict)
688674

689675
def coordination_specific(
690676
self,
@@ -695,7 +681,7 @@ def coordination_specific(
695681
counter_atom: str = "anion",
696682
) -> pd.DataFrame:
697683
"""Calculates the integral of the coordiantion number of selected species
698-
in each type of solvation structures (CIP/SSIP/AGG)
684+
in each type of solvation structures (CIP/SSIP/AGG).
699685
700686
Args:
701687
distance_dict: A dict of coordination cutoff distance of the neighbor species.
@@ -727,8 +713,7 @@ def coordination_specific(
727713
else:
728714
agg_list.append(cn)
729715
df_dict = {item_name: item_list, "CN in SSIP": ssip_list, "CN in CIP": cip_list, "CN in AGG": agg_list}
730-
df = pd.DataFrame(df_dict)
731-
return df
716+
return pd.DataFrame(df_dict)
732717

733718
def get_msd_all(
734719
self,
@@ -758,7 +743,7 @@ def get_msd_all(
758743
"""
759744
selection = self.select_dict.get(species)
760745
assert selection is not None
761-
msd_array = total_msd(
746+
return total_msd(
762747
self.unwrapped_run,
763748
start=start,
764749
end=end,
@@ -768,7 +753,6 @@ def get_msd_all(
768753
built_in=built_in,
769754
center_of_mass=center_of_mass,
770755
)
771-
return msd_array
772756

773757
def get_msd_partial(
774758
self,
@@ -861,7 +845,7 @@ def get_neighbor_corr(
861845
def get_residence_time(
862846
self, times: np.ndarray, acf_avg_dict: dict[str, np.ndarray], cutoff_time: int
863847
) -> dict[str, np.floating]:
864-
"""Calculates the residence time of selected species around cation
848+
"""Calculates the residence time of selected species around cation.
865849
866850
Args:
867851
times: The time series.
@@ -882,7 +866,7 @@ def get_neighbor_trj(
882866
center_atom: str = "cation",
883867
index: int = 0,
884868
) -> dict[str, np.ndarray]:
885-
"""Returns the distance between one center atom and neighbors as a function of time
869+
"""Returns the distance between one center atom and neighbors as a function of time.
886870
887871
Args:
888872
run_start: Start frame of analysis.
@@ -1063,12 +1047,12 @@ def get_heat_map(
10631047
hopping_cutoff: float,
10641048
floating_atom: str = "cation",
10651049
cartesian_by_ref: np.ndarray = None,
1066-
sym_dict: dict[str, list[np.ndarray]] = None,
1050+
sym_dict: dict[str, list[np.ndarray]] | None = None,
10671051
sample: int | None = None,
10681052
smooth: int = 51,
10691053
dim: str = "xyz",
10701054
) -> np.ndarray:
1071-
"""Calculates the heatmap matrix of floating ion around a cluster
1055+
"""Calculates the heatmap matrix of floating ion around a cluster.
10721056
10731057
Args:
10741058
run_start: Start frame of analysis.
@@ -1133,7 +1117,7 @@ def get_heat_map(
11331117
def get_cluster_distance(
11341118
self, run_start: int, run_end: int, neighbor_cutoff: float, cluster_center: str = "center"
11351119
) -> np.floating:
1136-
"""Calculates the average distance of the center of clusters/molecules
1120+
"""Calculates the average distance of the center of clusters/molecules.
11371121
11381122
Args:
11391123
run_start: Start frame of analysis.

‎mdgo/core/run.py

+3-9
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,15 @@
11
# Copyright (c) Tingzheng Hou.
22
# Distributed under the terms of the MIT License.
33

4-
"""
5-
This module implements a core class MdRun for molecular dynamics job setup.
6-
"""
4+
"""This module implements a core class MdRun for molecular dynamics job setup."""
75
from __future__ import annotations
86

97

108
class MdJob:
11-
"""
12-
A core class for MD results analysis.
13-
"""
9+
"""A core class for MD results analysis."""
1410

1511
def __init__(self, name):
16-
"""
17-
Base constructor
18-
"""
12+
"""Base constructor."""
1913
self.name = name
2014

2115
@classmethod

‎mdgo/forcefield/aqueous.py

+2-7
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
11
# Copyright (c) Tingzheng Hou.
22
# Distributed under the terms of the MIT License.
33

4-
"""
5-
A class for retrieving water and ion force field parameters.
6-
"""
4+
"""A class for retrieving water and ion force field parameters."""
75

86
from __future__ import annotations
97

@@ -227,10 +225,7 @@ def get_ion(
227225
parameter_set = alias.get(parameter_set, parameter_set)
228226

229227
# Make the Ion object to get mass and charge
230-
if isinstance(ion, Ion):
231-
ion_obj = ion
232-
else:
233-
ion_obj = Ion.from_formula(ion.capitalize())
228+
ion_obj = ion if isinstance(ion, Ion) else Ion.from_formula(ion.capitalize())
234229

235230
# load ion data as a list of IonLJData objects
236231
ion_data = loadfn(os.path.join(DATA_DIR, "ion_lj_params.json"))

‎mdgo/forcefield/charge.py

+3-5
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
11
# Copyright (c) Tingzheng Hou.
22
# Distributed under the terms of the MIT License.
33

4-
"""
5-
A class for writing, overwriting, scaling charges of a LammpsData object.
6-
"""
4+
"""A class for writing, overwriting, scaling charges of a LammpsData object."""
75

86
from __future__ import annotations
97

@@ -16,7 +14,7 @@ class ChargeWriter:
1614
A class for write, overwrite, scale charges of a LammpsData object.
1715
TODO: Auto determine number of significant figures of charges
1816
TODO: write to obj or write separate charge file
19-
TODO: Read LammpsData or path
17+
TODO: Read LammpsData or path.
2018
2119
Args:
2220
data: The provided LammpsData obj.
@@ -30,7 +28,7 @@ def __init__(self, data: LammpsData, precision: int = 10):
3028

3129
def scale(self, factor: float) -> LammpsData:
3230
"""
33-
Scales the charge in of the in self.data and returns a new one. TODO: check if non-destructive
31+
Scales the charge in of the in self.data and returns a new one. TODO: check if non-destructive.
3432
3533
Args:
3634
factor: The charge scaling factor

‎mdgo/forcefield/crawler.py

+1-4
Original file line numberDiff line numberDiff line change
@@ -91,10 +91,7 @@ def __init__(
9191
print("LigParGen server connected.")
9292

9393
def quit(self):
94-
"""
95-
Method for quiting ChromeDriver.
96-
97-
"""
94+
"""Method for quiting ChromeDriver."""
9895
self.web.quit()
9996

10097
def data_from_pdb(self, pdb_dir: str):

‎mdgo/forcefield/maestro.py

+1-2
Original file line numberDiff line numberDiff line change
@@ -146,8 +146,7 @@ def get_ff(self):
146146
FFLD.format(self.mae + ".mae", self.ff),
147147
check=True,
148148
shell=True,
149-
stdout=subprocess.PIPE,
150-
stderr=subprocess.PIPE,
149+
capture_output=True,
151150
)
152151
except subprocess.CalledProcessError as e:
153152
raise ValueError(f"Maestro failed with errorcode {e.returncode} and stderr: {e.stderr}") from e

‎mdgo/forcefield/pubchem.py

+1-4
Original file line numberDiff line numberDiff line change
@@ -84,10 +84,7 @@ def __init__(
8484
print("PubChem server connected.")
8585

8686
def quit(self):
87-
"""
88-
Method for quiting ChromeDriver.
89-
90-
"""
87+
"""Method for quiting ChromeDriver."""
9188
if not self.api:
9289
self.web.quit()
9390

0 commit comments

Comments
 (0)