From 99d31763b8a2d873dd80482e57ff41da7e7a12df Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Sun, 4 Feb 2024 11:08:56 +0800 Subject: [PATCH 01/34] add from __future__ import annotations --- mdgo/conductivity.py | 3 ++- mdgo/coordination.py | 2 +- mdgo/core/analysis.py | 2 +- mdgo/forcefield/aqueous.py | 3 ++- mdgo/forcefield/charge.py | 2 ++ mdgo/forcefield/crawler.py | 2 ++ mdgo/forcefield/maestro.py | 3 ++- mdgo/forcefield/pubchem.py | 2 ++ mdgo/msd.py | 2 +- mdgo/residence_time.py | 4 +++- tests/test_conductivity.py | 1 + tests/test_coordination.py | 1 + tests/test_core.py | 1 + tests/test_forcefield.py | 1 + tests/test_mdgopackmol.py | 1 + 15 files changed, 23 insertions(+), 7 deletions(-) diff --git a/mdgo/conductivity.py b/mdgo/conductivity.py index 7c05999c..191509ee 100644 --- a/mdgo/conductivity.py +++ b/mdgo/conductivity.py @@ -5,7 +5,8 @@ """ This module implements functions to calculate the ionic conductivity. """ -from typing import Union + +from __future__ import annotations import numpy as np from tqdm.auto import tqdm diff --git a/mdgo/coordination.py b/mdgo/coordination.py index 9c3004e8..6d16ac9f 100644 --- a/mdgo/coordination.py +++ b/mdgo/coordination.py @@ -6,7 +6,7 @@ This module implements functions for coordination analysis. """ -from typing import Dict, List, Tuple, Union, Callable, Optional +from __future__ import annotations import numpy as np from tqdm.auto import tqdm diff --git a/mdgo/core/analysis.py b/mdgo/core/analysis.py index ea97b086..33825a47 100644 --- a/mdgo/core/analysis.py +++ b/mdgo/core/analysis.py @@ -8,7 +8,7 @@ """ from __future__ import annotations -from typing import Union, Dict, Tuple, List, Optional + import MDAnalysis import numpy as np import pandas as pd diff --git a/mdgo/forcefield/aqueous.py b/mdgo/forcefield/aqueous.py index 23a1ce5d..1ce3cbdf 100644 --- a/mdgo/forcefield/aqueous.py +++ b/mdgo/forcefield/aqueous.py @@ -6,10 +6,11 @@ A class for retrieving water and ion force field parameters. """ +from __future__ import annotations + import os import re from dataclasses import dataclass -from typing import Optional, Union, Final, Literal from monty.json import MSONable from monty.serialization import loadfn diff --git a/mdgo/forcefield/charge.py b/mdgo/forcefield/charge.py index 3ab0d1b8..925b3787 100644 --- a/mdgo/forcefield/charge.py +++ b/mdgo/forcefield/charge.py @@ -6,6 +6,8 @@ A class for writing, overwriting, scaling charges of a LammpsData object. """ +from __future__ import annotations + import numpy as np from pymatgen.io.lammps.data import LammpsData diff --git a/mdgo/forcefield/crawler.py b/mdgo/forcefield/crawler.py index 6047a768..0d7551e6 100644 --- a/mdgo/forcefield/crawler.py +++ b/mdgo/forcefield/crawler.py @@ -13,6 +13,8 @@ matches your Chrome version via https://chromedriver.chromium.org/downloads """ +from __future__ import annotations + import os import shutil import time diff --git a/mdgo/forcefield/maestro.py b/mdgo/forcefield/maestro.py index 29702cb8..6ef9fd83 100644 --- a/mdgo/forcefield/maestro.py +++ b/mdgo/forcefield/maestro.py @@ -17,12 +17,13 @@ """ +from __future__ import annotations + import os import signal import subprocess import time from string import Template -from typing import Optional, Final from mdgo.util.reformat import ff_parser diff --git a/mdgo/forcefield/pubchem.py b/mdgo/forcefield/pubchem.py index e488aa53..606e1797 100644 --- a/mdgo/forcefield/pubchem.py +++ b/mdgo/forcefield/pubchem.py @@ -7,6 +7,8 @@ can be used to retrieve compound structure and information. """ +from __future__ import annotations + import os import time from typing import Optional, Final diff --git a/mdgo/msd.py b/mdgo/msd.py index 0fd30f71..1bf8ee2e 100644 --- a/mdgo/msd.py +++ b/mdgo/msd.py @@ -10,7 +10,7 @@ http://stackoverflow.com/questions/34222272/computing-mean-square-displacement-using-python-and-fft#34222273 """ -from typing import List, Dict, Tuple, Union, Optional, Literal +from __future__ import annotations try: import MDAnalysis.analysis.msd as mda_msd diff --git a/mdgo/residence_time.py b/mdgo/residence_time.py index 288af091..9966ef6e 100644 --- a/mdgo/residence_time.py +++ b/mdgo/residence_time.py @@ -5,8 +5,10 @@ """ This module calculates species correlation lifetime (residence time). """ + +from __future__ import annotations + import os -from typing import List, Dict, Union, Tuple import numpy as np import matplotlib.pyplot as plt diff --git a/tests/test_conductivity.py b/tests/test_conductivity.py index 9e1d3066..dc5fe230 100644 --- a/tests/test_conductivity.py +++ b/tests/test_conductivity.py @@ -1,3 +1,4 @@ +from __future__ import annotations import os import sys import tempfile diff --git a/tests/test_coordination.py b/tests/test_coordination.py index d6cfbef8..961e79b5 100644 --- a/tests/test_coordination.py +++ b/tests/test_coordination.py @@ -1,3 +1,4 @@ +from __future__ import annotations import os import unittest diff --git a/tests/test_core.py b/tests/test_core.py index d6cfbef8..961e79b5 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -1,3 +1,4 @@ +from __future__ import annotations import os import unittest diff --git a/tests/test_forcefield.py b/tests/test_forcefield.py index 464d7d5f..2ac7f953 100644 --- a/tests/test_forcefield.py +++ b/tests/test_forcefield.py @@ -1,3 +1,4 @@ +from __future__ import annotations import io import sys import tempfile diff --git a/tests/test_mdgopackmol.py b/tests/test_mdgopackmol.py index 9ffa91b3..af02eec3 100644 --- a/tests/test_mdgopackmol.py +++ b/tests/test_mdgopackmol.py @@ -1,3 +1,4 @@ +from __future__ import annotations import os import tempfile from pathlib import Path From 5540f897b7ba89f6301c256437c8c250b624fdfa Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Sun, 4 Feb 2024 11:19:06 +0800 Subject: [PATCH 02/34] Update test_volume.py --- tests/test_volume.py | 36 ++++++++++++++++++------------------ 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/tests/test_volume.py b/tests/test_volume.py index b1bdc356..b3ba5996 100644 --- a/tests/test_volume.py +++ b/tests/test_volume.py @@ -23,23 +23,23 @@ def test_molecular_volume(self) -> None: lipf6_volume_3 = molecular_volume(self.lipf6, radii_type="Lange") lipf6_volume_4 = molecular_volume(self.lipf6, radii_type="pymatgen") lipf6_volume_5 = molecular_volume(self.lipf6, molar_volume=False) - self.assertAlmostEqual(lipf6_volume_1, 47.62, places=2) - self.assertAlmostEqual(lipf6_volume_2, 43.36, places=2) - self.assertAlmostEqual(lipf6_volume_3, 41.49, places=2) - self.assertAlmostEqual(lipf6_volume_4, 51.94, places=2) - self.assertAlmostEqual(lipf6_volume_5, 79.08, places=2) + assert lipf6_volume_1 == 47.62 + assert lipf6_volume_2 == 43.36 + assert lipf6_volume_3 == 41.49 + assert lipf6_volume_4 == 51.94 + assert lipf6_volume_5 == 79.08 ec_volume_1 = molecular_volume(self.ec) ec_volume_2 = molecular_volume(self.ec, exclude_h=False) ec_volume_3 = molecular_volume(self.ec, res=1.0) ec_volume_4 = molecular_volume(self.ec, radii_type="Lange") ec_volume_5 = molecular_volume(self.ec, radii_type="pymatgen") ec_volume_6 = molecular_volume(self.ec, molar_volume=False) - self.assertAlmostEqual(ec_volume_1, 38.44, places=2) - self.assertAlmostEqual(ec_volume_2, 43.17, places=2) - self.assertAlmostEqual(ec_volume_3, 40.95, places=2) - self.assertAlmostEqual(ec_volume_4, 41.07, places=2) - self.assertAlmostEqual(ec_volume_5, 38.44, places=2) - self.assertAlmostEqual(ec_volume_6, 63.83, places=2) + assert ec_volume_1 == 38.44 + assert ec_volume_2 == 43.17 + assert ec_volume_3 == 40.95 + assert ec_volume_4 == 41.07 + assert ec_volume_5 == 38.44 + assert ec_volume_6 == 63.83 litfsi_volume_1 = molecular_volume(self.litfsi) litfsi_volume_2 = molecular_volume(self.litfsi, exclude_h=False) litfsi_volume_3 = molecular_volume(self.litfsi, res=1.0) @@ -47,13 +47,13 @@ def test_molecular_volume(self) -> None: litfsi_volume_5 = molecular_volume(self.litfsi, radii_type="pymatgen") litfsi_volume_6 = molecular_volume(self.litfsi, molar_volume=False) litfsi_volume_7 = molecular_volume(self.litfsi, mode="act", x_size=8, y_size=8, z_size=8) - self.assertAlmostEqual(litfsi_volume_1, 100.16, places=2) - self.assertAlmostEqual(litfsi_volume_2, 100.16, places=2) - self.assertAlmostEqual(litfsi_volume_3, 99.37, places=2) - self.assertAlmostEqual(litfsi_volume_4, 90.78, places=2) - self.assertAlmostEqual(litfsi_volume_5, 105.31, places=2) - self.assertAlmostEqual(litfsi_volume_6, 166.32, places=2) - self.assertAlmostEqual(litfsi_volume_7, 124.66, places=2) + assert litfsi_volume_1 == 100.16 + assert litfsi_volume_2 == 100.16 + assert litfsi_volume_3 == 99.37 + assert litfsi_volume_4 == 90.78 + assert litfsi_volume_5 == 105.31 + assert litfsi_volume_6 == 166.32 + assert litfsi_volume_7 == 124.66 if __name__ == "__main__": From 043fe2269378bc440f6fcb3dde440c9363d3e62d Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 00:36:06 +0800 Subject: [PATCH 03/34] typing to annotations --- mdgo/conductivity.py | 8 +-- mdgo/coordination.py | 105 ++++++++++++++++++----------------- mdgo/core/analysis.py | 56 +++++++++---------- mdgo/forcefield/aqueous.py | 5 +- mdgo/forcefield/crawler.py | 3 +- mdgo/forcefield/maestro.py | 3 +- mdgo/forcefield/pubchem.py | 8 +-- mdgo/msd.py | 15 ++--- mdgo/residence_time.py | 28 +++++----- tests/test_msd.py | 1 + tests/test_residence_time.py | 1 + tests/test_util.py | 1 + tests/test_volume.py | 5 +- 13 files changed, 124 insertions(+), 115 deletions(-) diff --git a/mdgo/conductivity.py b/mdgo/conductivity.py index 191509ee..221b6883 100644 --- a/mdgo/conductivity.py +++ b/mdgo/conductivity.py @@ -27,8 +27,8 @@ def calc_cond_msd( anions: AtomGroup, cations: AtomGroup, run_start: int, - cation_charge: Union[int, float] = 1, - anion_charge: Union[int, float] = -1, + cation_charge: int | float = 1, + anion_charge: int | float = -1, ) -> np.ndarray: """Calculates the conductivity "mean square displacement" over time @@ -128,11 +128,11 @@ def choose_msd_fitting_region( def conductivity_calculator( time_array: np.ndarray, cond_array: np.ndarray, - v: Union[int, float], + v: int | float, name: str, start: int, end: int, - T: Union[int, float], + T: int | float, units: str = "real", ) -> float: """Calculates the overall conductivity of the system diff --git a/mdgo/coordination.py b/mdgo/coordination.py index 6d16ac9f..d39fcfde 100644 --- a/mdgo/coordination.py +++ b/mdgo/coordination.py @@ -7,6 +7,7 @@ """ from __future__ import annotations +from collections.abc import Callable import numpy as np from tqdm.auto import tqdm @@ -30,9 +31,9 @@ def neighbor_distance( run_start: int, run_end: int, species: str, - select_dict: Dict[str, str], + select_dict: dict[str, str], distance: float, -) -> Dict[str, np.ndarray]: +) -> dict[str, np.ndarray]: """ Calculates a dictionary of distances between the ``center_atom`` and neighbor atoms. @@ -74,12 +75,12 @@ def neighbor_distance( def find_nearest( - trj: Dict[str, np.ndarray], + trj: dict[str, np.ndarray], time_step: float, binding_cutoff: float, hopping_cutoff: float, smooth: int = 51, -) -> Tuple[List[int], Union[float, np.floating], List[int]]: +) -> tuple[list[int], float | np.floating, list[int]]: """Using the dictionary of neighbor distance ``trj``, finds the nearest neighbor ``sites`` that the center atom binds to, and calculates the ``frequency`` of hopping between each neighbor, and ``steps`` when each binding site exhibits the closest distance to the center atom. @@ -101,7 +102,7 @@ def find_nearest( for kw in list(trj): trj[kw] = savgol_filter(trj.get(kw), smooth, 2) site_distance = [100 for _ in range(time_span)] - sites: List[Union[int, np.integer]] = [0 for _ in range(time_span)] + sites: list[int | np.integer] = [0 for _ in range(time_span)] start_site = min(trj, key=lambda k: trj[k][0]) kw_start = trj.get(start_site) assert kw_start is not None @@ -133,7 +134,7 @@ def find_nearest( sites = [int(i) for i in sites] sites_and_distance_array = np.array([[sites[i], site_distance[i]] for i in range(len(sites))]) steps = [] - closest_step: Optional[int] = 0 + closest_step: int | None = 0 previous_site = sites_and_distance_array[0][0] if previous_site == 0: closest_step = None @@ -162,12 +163,12 @@ def find_nearest( def find_nearest_free_only( - trj: Dict[str, np.ndarray], + trj: dict[str, np.ndarray], time_step: float, binding_cutoff: float, hopping_cutoff: float, smooth: int = 51, -) -> Tuple[List[int], Union[float, np.floating], List[int]]: +) -> tuple[list[int], float | np.floating, list[int]]: """Using the dictionary of neighbor distance ``trj``, finds the nearest neighbor ``sites`` that the ``center_atom`` binds to, and calculates the ``frequency`` of hopping between each neighbor, and ``steps`` when each binding site exhibits the closest distance to the center atom. @@ -190,7 +191,7 @@ def find_nearest_free_only( for kw in list(trj): trj[kw] = savgol_filter(trj.get(kw), smooth, 2) site_distance = [100 for _ in range(time_span)] - sites: List[Union[int, np.integer]] = [0 for _ in range(time_span)] + sites: list[int | np.integer] = [0 for _ in range(time_span)] start_site = min(trj, key=lambda k: trj[k][0]) kw_start = trj.get(start_site) assert kw_start is not None @@ -222,7 +223,7 @@ def find_nearest_free_only( sites = [int(i) for i in sites] sites_and_distance_array = np.array([[sites[i], site_distance[i]] for i in range(len(sites))]) steps = [] - closest_step: Optional[int] = 0 + closest_step: int | None = 0 previous_site = sites_and_distance_array[0][0] previous_zero = False if previous_site == 0: @@ -257,8 +258,8 @@ def find_nearest_free_only( def find_in_n_out( - trj: Dict[str, np.ndarray], binding_cutoff: float, hopping_cutoff: float, smooth: int = 51, cool: int = 20 -) -> Tuple[List[int], List[int]]: + trj: dict[str, np.ndarray], binding_cutoff: float, hopping_cutoff: float, smooth: int = 51, cool: int = 20 +) -> tuple[list[int], list[int]]: """Finds the frames when the center atom binds with the neighbor (binding) or hopping out (hopping) according to the dictionary of neighbor distance. @@ -309,8 +310,8 @@ def find_in_n_out( sites = [int(i) for i in sites] last = sites[0] - steps_in: List[int] = [] - steps_out: List[int] = [] + steps_in: list[int] = [] + steps_out: list[int] = [] in_cool = cool out_cool = cool for i, s in enumerate(sites): @@ -339,13 +340,13 @@ def find_in_n_out( def check_contiguous_steps( nvt_run: Universe, center_atom: Atom, - distance_dict: Dict[str, float], - select_dict: Dict[str, str], + distance_dict: dict[str, float], + select_dict: dict[str, str], run_start: int, run_end: int, checkpoints: np.ndarray, lag: int = 20, -) -> Dict[str, np.ndarray]: +) -> dict[str, np.ndarray]: """Calculates the distance between the center atom and the neighbor atom in the checkpoint +/- lag time range. @@ -364,7 +365,7 @@ def check_contiguous_steps( An array of distance between the center atom and the neighbor atoms in the checkpoint +/- lag time range. """ - coord_num: Dict[str, Union[List[List[int]], np.ndarray]] = { + coord_num: dict[str, list[list[int]] | np.ndarray] = { x: [[] for _ in range(lag * 2 + 1)] for x in distance_dict } trj_analysis = nvt_run.trajectory[run_start:run_end:] @@ -393,8 +394,8 @@ def check_contiguous_steps( def heat_map( nvt_run: Universe, floating_atom: Atom, - cluster_center_sites: List[int], - cluster_terminal: Union[str, List[str]], + cluster_center_sites: list[int], + cluster_terminal: str | list[str], cartesian_by_ref: np.ndarray, run_start: int, run_end: int, @@ -441,7 +442,7 @@ def heat_map( for species in cluster_terminal ] bind_atoms_xyz = [nvt_run.select_atoms(sel, periodic=True) for sel in selections] - vertex_atoms: List[Atom] = [] + vertex_atoms: list[Atom] = [] for atoms in bind_atoms_xyz: if len(atoms) == 1: vertex_atoms.append(atoms[0]) @@ -501,10 +502,10 @@ def heat_map( def process_evol( nvt_run: Universe, - select_dict: Dict[str, str], - in_list: Dict[str, List[np.ndarray]], - out_list: Dict[str, List[np.ndarray]], - distance_dict: Dict[str, float], + select_dict: dict[str, str], + in_list: dict[str, list[np.ndarray]], + out_list: dict[str, list[np.ndarray]], + distance_dict: dict[str, float], run_start: int, run_end: int, lag: int, @@ -572,10 +573,10 @@ def process_evol( def get_full_coords( coords: np.ndarray, - reflection: Optional[List[np.ndarray]] = None, - rotation: Optional[List[np.ndarray]] = None, - inversion: Optional[List[np.ndarray]] = None, - sample: Optional[int] = None, + reflection: list[np.ndarray] | None = None, + rotation: list[np.ndarray] | None = None, + inversion: list[np.ndarray] | None = None, + sample: int | None = None, dim: str = "xyz", ) -> np.ndarray: """ @@ -640,12 +641,12 @@ def get_full_coords( def cluster_coordinates( # TODO: rewrite the method nvt_run: Universe, - select_dict: Dict[str, str], + select_dict: dict[str, str], run_start: int, run_end: int, - species: List[str], + species: list[str], distance: float, - basis_vectors: Optional[Union[List[np.ndarray], np.ndarray]] = None, + basis_vectors: list[np.ndarray] | np.ndarray | None = None, cluster_center: str = "center", ) -> np.ndarray: """Calculates the average position of a cluster. @@ -708,15 +709,15 @@ def cluster_coordinates( # TODO: rewrite the method def num_of_neighbor( nvt_run: Universe, center_atom: Atom, - distance_dict: Dict[str, float], - select_dict: Dict[str, str], + distance_dict: dict[str, float], + select_dict: dict[str, str], run_start, run_end, write=False, structure_code=None, write_freq=0, write_path=None, -) -> Dict[str, np.ndarray]: +) -> dict[str, np.ndarray]: """Calculates the coordination number of each specified neighbor species and the total coordination number in the specified frame range. @@ -778,11 +779,11 @@ def num_of_neighbor( def num_of_neighbor_simple( nvt_run: Universe, center_atom: Atom, - distance_dict: Dict[str, float], - select_dict: Dict[str, str], + distance_dict: dict[str, float], + select_dict: dict[str, str], run_start: int, run_end: int, -) -> Dict[str, np.ndarray]: +) -> dict[str, np.ndarray]: """Calculates solvation structure type (1 for SSIP, 2 for CIP and 3 for AGG) with respect to the ``enter_atom`` in the specified frame range. @@ -830,12 +831,12 @@ def num_of_neighbor_simple( def angular_dist_of_neighbor( nvt_run: Universe, center_atom: Atom, - distance_dict: Dict[str, float], - select_dict: Dict[str, str], + distance_dict: dict[str, float], + select_dict: dict[str, str], run_start: int, run_end: int, cip: bool = True, -) -> Dict[str, np.ndarray]: +) -> dict[str, np.ndarray]: """ Calculates the angle of a-c-b of center atom c in the specified frames. @@ -893,12 +894,12 @@ def angular_dist_of_neighbor( def num_of_neighbor_specific( nvt_run: Universe, center_atom: Atom, - distance_dict: Dict[str, float], - select_dict: Dict[str, str], + distance_dict: dict[str, float], + select_dict: dict[str, str], run_start: int, run_end: int, counter_atom: str = "anion", -) -> Dict[str, np.ndarray]: +) -> dict[str, np.ndarray]: """ Calculates the coordination number of each specific solvation structure type (SSIP, CIP, AGG). @@ -962,7 +963,7 @@ def full_solvation_structure( # TODO: rewrite the method center_atom: Atom, center_species: str, counter_species: str, - select_dict: Dict[str, str], + select_dict: dict[str, str], distance: float, run_start: int, run_end: int, @@ -1022,8 +1023,8 @@ def counter_shell(this_shell, this_layer, frame): trj_analysis = nvt_run.trajectory[run_start:run_end:] cn_values = np.zeros((int(len(trj_analysis)), depth)) for ts in trj_analysis: - center_ion_list: List[np.int_] = [center_atom.id] - counter_ion_list: List[np.int_] = [] + center_ion_list: list[np.int_] = [center_atom.id] + counter_ion_list: list[np.int_] = [] first_shell = nvt_run.select_atoms( select_counter_ion(counter_selection, distance, center_atom), periodic=True, @@ -1036,12 +1037,12 @@ def concat_coord_array( nvt_run: Universe, func: Callable, center_atoms: AtomGroup, - distance_dict: Dict[str, float], - select_dict: Dict[str, str], + distance_dict: dict[str, float], + select_dict: dict[str, str], run_start: int, run_end: int, - **kwargs: Union[bool, str], -) -> Dict[str, np.ndarray]: + **kwargs: bool | str, +) -> dict[str, np.ndarray]: """ A helper function to analyze the coordination number/structure of every atoms in an ``AtomGroup`` using the specified function. @@ -1104,7 +1105,7 @@ def write_out(center_pos: np.ndarray, center_name: str, neighbors: AtomGroup, pa def select_shell( - select: Union[Dict[str, str], str], distance: Union[Dict[str, float], str], center_atom: Atom, kw: str + select: dict[str, str] | str, distance: dict[str, float] | str, center_atom: Atom, kw: str ) -> str: """ Select a group of atoms that is within a distance of an ``center_atom``. diff --git a/mdgo/core/analysis.py b/mdgo/core/analysis.py index 33825a47..b9440eff 100644 --- a/mdgo/core/analysis.py +++ b/mdgo/core/analysis.py @@ -75,8 +75,8 @@ def __init__( nvt_start: int, time_step: float, name: str, - select_dict: Optional[Dict[str, str]] = None, - res_dict: Optional[Dict[str, str]] = None, + select_dict: dict[str, str] | None = None, + res_dict: dict[str, str] | None = None, cation_name: str = "cation", anion_name: str = "anion", cation_charge: float = 1, @@ -149,8 +149,8 @@ def from_lammps( nvt_start: int, time_step: float, name: str, - select_dict: Optional[Dict[str, str]] = None, - res_dict: Optional[Dict[str, str]] = None, + select_dict: dict[str, str] | None = None, + res_dict: dict[str, str] | None = None, cation_name: str = "cation", anion_name: str = "anion", cation_charge: float = 1, @@ -400,11 +400,11 @@ def coord_num_array_single_species( def coord_num_array_multi_species( self, - distance_dict: Dict[str, float], + distance_dict: dict[str, float], run_start: int, run_end: int, center_atom: str = "cation", - ) -> Dict[str, np.ndarray]: + ) -> dict[str, np.ndarray]: """Calculates the coordination number array of multiple species around the interested ``center_atom``. Args: @@ -432,12 +432,12 @@ def coord_num_array_multi_species( def coord_num_array_specific( self, - distance_dict: Dict[str, float], + distance_dict: dict[str, float], run_start: int, run_end: int, center_atom: str = "cation", counter_atom: str = "anion", - ) -> Dict[str, np.ndarray]: + ) -> dict[str, np.ndarray]: """Calculates the coordination number array of multiple species of specific coordination types (SSIP, CIP, AGG). @@ -468,7 +468,7 @@ def coord_num_array_specific( def write_solvation_structure( self, - distance_dict: Dict[str, float], + distance_dict: dict[str, float], run_start: int, run_end: int, structure_code: int, @@ -542,7 +542,7 @@ def coord_type_array( def angle_array( self, - distance_dict: Dict[str, float], + distance_dict: dict[str, float], run_start: int, run_end: int, center_atom: str = "cation", @@ -618,7 +618,7 @@ def coordination( def rdf_integral( self, - distance_dict: Dict[str, float], + distance_dict: dict[str, float], run_start: int, run_end: int, center_atom: str = "cation", @@ -689,7 +689,7 @@ def coordination_type( def coordination_specific( self, - distance_dict: Dict[str, float], + distance_dict: dict[str, float], run_start: int, run_end: int, center_atom: str = "cation", @@ -779,7 +779,7 @@ def get_msd_partial( largest: int = 1000, center_atom: str = "cation", binding_site: str = "anion", - ) -> Tuple[Optional[List[np.ndarray]], Optional[List[np.ndarray]]]: + ) -> tuple[list[np.ndarray] | None, list[np.ndarray] | None]: """ Calculates the mean square displacement (MSD) of the ``center_atom`` according to coordination states. The returned ``free_array`` include the MSD when ``center_atom`` is not coordinated to ``binding_site``. @@ -832,11 +832,11 @@ def get_d(self, msd_array: np.ndarray, start: int, end: int, percentage: float = def get_neighbor_corr( self, - distance_dict: Dict[str, float], + distance_dict: dict[str, float], run_start: int, run_end: int, center_atom: str = "cation", - ) -> Tuple[np.ndarray, Dict[str, np.ndarray]]: + ) -> tuple[np.ndarray, dict[str, np.ndarray]]: """Calculates the neighbor auto-correlation function (ACF) of selected species around center_atom. @@ -860,8 +860,8 @@ 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]: + 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 Args: @@ -882,7 +882,7 @@ def get_neighbor_trj( neighbor_cutoff: float, center_atom: str = "cation", index: int = 0, - ) -> Dict[str, np.ndarray]: + ) -> dict[str, np.ndarray]: """Returns the distance between one center atom and neighbors as a function of time Args: @@ -917,7 +917,7 @@ def get_hopping_freq_dist( floating_atom: str = "cation", smooth: int = 51, mode: str = "full", - ) -> Tuple[np.floating, np.floating]: + ) -> tuple[np.floating, np.floating]: """Calculates the cation hopping rate and hopping distance. Args: @@ -967,7 +967,7 @@ def get_hopping_freq_dist( def shell_evolution( self, - distance_dict: Dict[str, float], + distance_dict: dict[str, float], run_start: int, run_end: int, lag_step: int, @@ -977,8 +977,8 @@ def shell_evolution( cool: int = 0, binding_site: str = "anion", center_atom: str = "cation", - duplicate_run: Optional[List[MdRun]] = None, - ) -> Dict[str, Dict[str, Union[int, np.ndarray]]]: + duplicate_run: list[MdRun] | None = None, + ) -> dict[str, dict[str, int | np.ndarray]]: """Calculates the coordination number evolution of species around ``center_atom`` as a function of time, the coordination numbers are averaged over all time steps around events when the center_atom hopping to and hopping out from the ``binding_site``. If ``duplicate_run`` is given, it is also averaged over @@ -1001,8 +1001,8 @@ def shell_evolution( A dictionary containing the number of trj logged, the averaged coordination number and standard deviation for each species, and the corresponding time sequence. """ - in_list: Dict[str, List[np.ndarray]] = {} - out_list: Dict[str, List[np.ndarray]] = {} + in_list: dict[str, list[np.ndarray]] = {} + out_list: dict[str, list[np.ndarray]] = {} for k in list(distance_dict): in_list[k] = [] out_list[k] = [] @@ -1059,13 +1059,13 @@ def get_heat_map( run_start: int, run_end: int, cluster_center: str, - cluster_terminal: Union[str, List[str]], + cluster_terminal: str | list[str], binding_cutoff: float, hopping_cutoff: float, floating_atom: str = "cation", cartesian_by_ref: np.ndarray = None, - sym_dict: Dict[str, List[np.ndarray]] = None, - sample: Optional[int] = None, + sym_dict: dict[str, list[np.ndarray]] = None, + sample: int | None = None, smooth: int = 51, dim: str = "xyz", ) -> np.ndarray: @@ -1095,7 +1095,7 @@ def get_heat_map( nvt_run = self.wrapped_run floating_atoms = nvt_run.select_atoms(self.select_dict.get(floating_atom)) if isinstance(cluster_terminal, str): - terminal_atom_type: Union[str, List[str]] = self.select_dict.get(cluster_terminal, "Not defined") + terminal_atom_type: str | list[str] = self.select_dict.get(cluster_terminal, "Not defined") assert terminal_atom_type != "Not defined", f"{cluster_terminal} not defined in select_dict" else: terminal_atom_type = [] diff --git a/mdgo/forcefield/aqueous.py b/mdgo/forcefield/aqueous.py index 1ce3cbdf..a6a68fa6 100644 --- a/mdgo/forcefield/aqueous.py +++ b/mdgo/forcefield/aqueous.py @@ -7,6 +7,7 @@ """ from __future__ import annotations +from typing import Literal, Final import os import re @@ -137,10 +138,10 @@ def get_water(model: str = "spce") -> LammpsData: @staticmethod def get_ion( - ion: Union[Ion, str], + ion: Ion | str, parameter_set: str = "auto", water_model: str = "auto", - mixing_rule: Optional[str] = None, + mixing_rule: str | None = None, ) -> LammpsData: """ Retrieve force field parameters for an ion in water. diff --git a/mdgo/forcefield/crawler.py b/mdgo/forcefield/crawler.py index 0d7551e6..22d92f4e 100644 --- a/mdgo/forcefield/crawler.py +++ b/mdgo/forcefield/crawler.py @@ -18,7 +18,6 @@ import os import shutil import time -from typing import Optional from pymatgen.io.lammps.data import LammpsData @@ -60,7 +59,7 @@ class FFcrawler: def __init__( self, write_dir: str, - chromedriver_dir: Optional[str] = None, + chromedriver_dir: str | None = None, headless: bool = True, xyz: bool = False, gromacs: bool = False, diff --git a/mdgo/forcefield/maestro.py b/mdgo/forcefield/maestro.py index 6ef9fd83..7ad9fe7f 100644 --- a/mdgo/forcefield/maestro.py +++ b/mdgo/forcefield/maestro.py @@ -18,6 +18,7 @@ """ from __future__ import annotations +from typing import Final import os import signal @@ -82,7 +83,7 @@ def __init__( structure_dir: str, working_dir: str, out: str = "lmp", - cmd_template: Optional[str] = None, + cmd_template: str | None = None, assign_bond: bool = False, ): """Base constructor.""" diff --git a/mdgo/forcefield/pubchem.py b/mdgo/forcefield/pubchem.py index 606e1797..2f22d183 100644 --- a/mdgo/forcefield/pubchem.py +++ b/mdgo/forcefield/pubchem.py @@ -11,7 +11,7 @@ import os import time -from typing import Optional, Final +from typing import Final from urllib.parse import quote import pubchempy as pcp @@ -96,7 +96,7 @@ def quit(self): if not self.api: self.web.quit() - def obtain_entry(self, search_text: str, name: str, output_format: str = "sdf") -> Optional[str]: + def obtain_entry(self, search_text: str, name: str, output_format: str = "sdf") -> str | None: """ Search the PubChem database with a text entry and save the structure in desired format. @@ -140,7 +140,7 @@ def smiles_to_pdb(self, smiles: str): print(".", end="") print("\nStructure file saved.") - def _obtain_entry_web(self, search_text: str, name: str, output_format: str) -> Optional[str]: + def _obtain_entry_web(self, search_text: str, name: str, output_format: str) -> str | None: cid = None try: @@ -197,7 +197,7 @@ def _obtain_entry_web(self, search_text: str, name: str, output_format: str) -> self.quit() return cid - def _obtain_entry_api(self, search_text, name, output_format) -> Optional[str]: + def _obtain_entry_api(self, search_text, name, output_format) -> str | None: cid = None cids = pcp.get_cids(search_text, "name", record_type="3d") if len(cids) == 0: diff --git a/mdgo/msd.py b/mdgo/msd.py index 1bf8ee2e..120df2a9 100644 --- a/mdgo/msd.py +++ b/mdgo/msd.py @@ -11,6 +11,7 @@ """ from __future__ import annotations +from typing import Literal try: import MDAnalysis.analysis.msd as mda_msd @@ -259,7 +260,7 @@ def mda_msd_wrapper( return total_array -def parse_msd_type(msd_type: DIM) -> List[int]: +def parse_msd_type(msd_type: DIM) -> list[int]: """ Sets up the desired dimensionality of the MSD. @@ -315,7 +316,7 @@ def _total_msd(nvt_run: Universe, run_start: int, run_end: int, select: str = "a return total_array -def msd_from_frags(coord_list: List[np.ndarray], largest: int) -> np.ndarray: +def msd_from_frags(coord_list: list[np.ndarray], largest: int) -> np.ndarray: """ Calculates the MSD using a list of fragments of trajectory with the conventional algorithm. @@ -326,7 +327,7 @@ def msd_from_frags(coord_list: List[np.ndarray], largest: int) -> np.ndarray: Returns: The MSD series. """ - msd_dict: Dict[Union[int, np.integer], np.ndarray] = {} + msd_dict: dict[int | np.integer, np.ndarray] = {} for state in coord_list: n_frames = state.shape[0] lag_times = np.arange(1, min(n_frames, largest)) @@ -352,12 +353,12 @@ def msd_from_frags(coord_list: List[np.ndarray], largest: int) -> np.ndarray: def states_coord_array( nvt_run: Universe, atom: Atom, - select_dict: Dict[str, str], + select_dict: dict[str, str], distance: float, run_start: int, run_end: int, binding_site: str = "anion", -) -> Tuple[List[np.ndarray], List[np.ndarray]]: +) -> tuple[list[np.ndarray], list[np.ndarray]]: """Cuts the trajectory of an atom into fragments. Each fragment contains consecutive timesteps of coordinates of the atom in either attached or free state. The Attached state is when the atom coordinates with the ``binding_site`` species (distance < ``distance``), and vice versa for the free state. @@ -426,12 +427,12 @@ def partial_msd( nvt_run: Universe, atoms: AtomGroup, largest: int, - select_dict: Dict[str, str], + select_dict: dict[str, str], distance: float, run_start: int, run_end: int, binding_site: str = "anion", -) -> Tuple[Optional[List[np.ndarray]], Optional[List[np.ndarray]]]: +) -> tuple[list[np.ndarray] | None, list[np.ndarray] | None]: """Calculates the mean square displacement (MSD) of the ``atoms`` according to coordination states. The returned ``free_data`` include the MSD when ``atoms`` are not coordinated to ``binding_site``. The ``attach_data`` includes the MSD of ``atoms`` are not coordinated to ``binding_site``. diff --git a/mdgo/residence_time.py b/mdgo/residence_time.py index 9966ef6e..837579d0 100644 --- a/mdgo/residence_time.py +++ b/mdgo/residence_time.py @@ -30,11 +30,11 @@ def neighbors_one_atom( nvt_run: Universe, center_atom: Atom, species: str, - select_dict: Dict[str, str], + select_dict: dict[str, str], distance: float, run_start: int, run_end: int, -) -> Dict[str, np.ndarray]: +) -> dict[str, np.ndarray]: """ Create adjacency matrix for one center atom. @@ -75,7 +75,7 @@ def neighbors_one_atom( return bool_values -def calc_acf(a_values: Dict[str, np.ndarray]) -> List[np.ndarray]: +def calc_acf(a_values: dict[str, np.ndarray]) -> list[np.ndarray]: """ Calculate auto-correlation function (ACF) @@ -94,11 +94,11 @@ def calc_acf(a_values: Dict[str, np.ndarray]) -> List[np.ndarray]: def exponential_func( - x: Union[float, np.floating, np.ndarray], - a: Union[float, np.floating, np.ndarray], - b: Union[float, np.floating, np.ndarray], - c: Union[float, np.floating, np.ndarray], -) -> Union[np.floating, np.ndarray]: + x: float | np.floating | np.ndarray, + a: float | np.floating | np.ndarray, + b: float | np.floating | np.ndarray, + c: float | np.floating | np.ndarray, +) -> np.floating | np.ndarray: """ An exponential decay function @@ -116,13 +116,13 @@ def exponential_func( def calc_neigh_corr( nvt_run: Universe, - distance_dict: Dict[str, float], - select_dict: Dict[str, str], + distance_dict: dict[str, float], + select_dict: dict[str, str], time_step: float, run_start: int, run_end: int, center_atom: str = "cation", -) -> Tuple[np.ndarray, Dict[str, np.ndarray]]: +) -> tuple[np.ndarray, dict[str, np.ndarray]]: """Calculates the neighbor auto-correlation function (ACF) of selected species around center atom. @@ -171,11 +171,11 @@ def calc_neigh_corr( def fit_residence_time( times: np.ndarray, - acf_avg_dict: Dict[str, np.ndarray], + acf_avg_dict: dict[str, np.ndarray], cutoff_time: int, time_step: float, - save_curve: Union[str, bool] = False, -) -> Dict[str, np.floating]: + save_curve: str | bool = False, +) -> dict[str, np.floating]: """ Use the ACF to fit the residence time (Exponential decay constant). TODO: allow defining the residence time according to a threshold value of the decay diff --git a/tests/test_msd.py b/tests/test_msd.py index e3b70506..3f5ca573 100644 --- a/tests/test_msd.py +++ b/tests/test_msd.py @@ -1,3 +1,4 @@ +from __future__ import annotations import os import unittest diff --git a/tests/test_residence_time.py b/tests/test_residence_time.py index d6cfbef8..961e79b5 100644 --- a/tests/test_residence_time.py +++ b/tests/test_residence_time.py @@ -1,3 +1,4 @@ +from __future__ import annotations import os import unittest diff --git a/tests/test_util.py b/tests/test_util.py index d6cfbef8..961e79b5 100644 --- a/tests/test_util.py +++ b/tests/test_util.py @@ -1,3 +1,4 @@ +from __future__ import annotations import os import unittest diff --git a/tests/test_volume.py b/tests/test_volume.py index b3ba5996..2e60001f 100644 --- a/tests/test_volume.py +++ b/tests/test_volume.py @@ -1,3 +1,4 @@ +from __future__ import annotations import os import unittest from pymatgen.core import Molecule @@ -46,7 +47,9 @@ def test_molecular_volume(self) -> None: litfsi_volume_4 = molecular_volume(self.litfsi, radii_type="Lange") litfsi_volume_5 = molecular_volume(self.litfsi, radii_type="pymatgen") litfsi_volume_6 = molecular_volume(self.litfsi, molar_volume=False) - litfsi_volume_7 = molecular_volume(self.litfsi, mode="act", x_size=8, y_size=8, z_size=8) + litfsi_volume_7 = molecular_volume( + self.litfsi, mode="act", x_size=8, y_size=8, z_size=8 + ) assert litfsi_volume_1 == 100.16 assert litfsi_volume_2 == 100.16 assert litfsi_volume_3 == 99.37 From 5d949fe612ca0fe08eb71c72f139bc644e0efc20 Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 00:44:20 +0800 Subject: [PATCH 04/34] typing to future annotations 2 --- mdgo/util/coord.py | 10 +++++----- mdgo/util/dict_utils.py | 20 ++++++++++---------- mdgo/util/num.py | 4 ++-- mdgo/util/packmol.py | 12 ++++++------ mdgo/util/reformat.py | 7 ++++--- mdgo/util/volume.py | 37 +++++++++++++++++++------------------ 6 files changed, 46 insertions(+), 44 deletions(-) diff --git a/mdgo/util/coord.py b/mdgo/util/coord.py index 2d536c7b..3f6df5c4 100644 --- a/mdgo/util/coord.py +++ b/mdgo/util/coord.py @@ -4,7 +4,7 @@ """Utilities for manipulating coordinates under periodic boundary conditions.""" -from typing import List, Union +from __future__ import annotations import numpy as np from MDAnalysis.core.groups import Atom @@ -35,9 +35,9 @@ def atom_vec(atom1: Atom, atom2: Atom, dimension: np.ndarray) -> np.ndarray: def position_vec( - pos1: Union[List[float], np.ndarray], - pos2: Union[List[float], np.ndarray], - dimension: Union[List[float], np.ndarray], + pos1: list[float] | np.ndarray, + pos2: list[float] | np.ndarray, + dimension: list[float] | np.ndarray, ) -> np.ndarray: """ Calculate the vector from pos2 to pos2. @@ -50,7 +50,7 @@ def position_vec( Return: The obtained vector. """ - vec: List[Union[int, float, np.floating]] = [0, 0, 0] + vec: list[int | float | np.floating] = [0, 0, 0] for i in range(3): diff = pos1[i] - pos2[i] if diff > dimension[i] / 2: diff --git a/mdgo/util/dict_utils.py b/mdgo/util/dict_utils.py index 945dbe5f..7a5d384c 100644 --- a/mdgo/util/dict_utils.py +++ b/mdgo/util/dict_utils.py @@ -7,7 +7,7 @@ import string import re import math -from typing import Dict, Union +from __future__ import annotations import numpy as np import pandas as pd @@ -38,7 +38,7 @@ def mass_to_name(masses: np.ndarray) -> np.ndarray: return np.array(names) -def lmp_mass_to_name(df: pd.DataFrame) -> Dict[int, str]: +def lmp_mass_to_name(df: pd.DataFrame) -> dict[int, str]: """ Create a dict for mapping atom type id to element from the mass information. Args: @@ -65,7 +65,7 @@ def assign_name(u: Universe, names: np.ndarray): u.add_TopologyAttr("name", values=names) -def assign_resname(u: Universe, res_dict: Dict[str, str]): +def assign_resname(u: Universe, res_dict: dict[str, str]): """ Assign resnames to residues in a MDAnalysis.universe object. The function will not overwrite existing resnames. @@ -82,7 +82,7 @@ def assign_resname(u: Universe, res_dict: Dict[str, str]): res_group.residues.resnames = res_names -def res_dict_from_select_dict(u: Universe, select_dict: Dict[str, str]) -> Dict[str, str]: +def res_dict_from_select_dict(u: Universe, select_dict: dict[str, str]) -> dict[str, str]: """ Infer res_dict (residue selection) from select_dict (atom selection) in a MDAnalysis.universe object. @@ -112,7 +112,7 @@ def res_dict_from_select_dict(u: Universe, select_dict: Dict[str, str]) -> Dict[ return res_dict -def res_dict_from_datafile(filename: str) -> Dict[str, str]: +def res_dict_from_datafile(filename: str) -> dict[str, str]: """ Infer res_dict (residue selection) from a LAMMPS data file. @@ -148,7 +148,7 @@ def res_dict_from_datafile(filename: str) -> Dict[str, str]: raise ValueError("The LAMMPS data file should be generated by pymatgen.io.lammps.data.") -def res_dict_from_lammpsdata(lammps_data: CombinedData) -> Dict[str, str]: +def res_dict_from_lammpsdata(lammps_data: CombinedData) -> dict[str, str]: """ Infer res_dict (residue selection) from a LAMMPS data file. @@ -184,7 +184,7 @@ def res_dict_from_lammpsdata(lammps_data: CombinedData) -> Dict[str, str]: return res_dict -def select_dict_from_resname(u: Universe) -> Dict[str, str]: +def select_dict_from_resname(u: Universe) -> dict[str, str]: """ Infer select_dict (possibly interested atom species selection) from resnames in a MDAnalysis.universe object. The resname must be pre-assigned already. @@ -195,7 +195,7 @@ def select_dict_from_resname(u: Universe) -> Dict[str, str]: return: A dictionary of atom species. """ - select_dict: Dict[str, str] = {} + select_dict: dict[str, str] = {} resnames = np.unique(u.residues.resnames) for resname in resnames: if resname == "": @@ -235,7 +235,7 @@ def select_dict_from_resname(u: Universe) -> Dict[str, str]: return select_dict -def extract_atom_from_ion(positive: bool, ion: Union[Residue, AtomGroup], select_dict: Dict[str, str], number: int = 0): +def extract_atom_from_ion(positive: bool, ion: Residue | AtomGroup, select_dict: dict[str, str], number: int = 0): """ Assign the most most charged atom and/or one unique atom in the ion into select_dict. @@ -285,7 +285,7 @@ def extract_atom_from_ion(positive: bool, ion: Union[Residue, AtomGroup], select def extract_atom_from_molecule( - resname: str, molecule: Union[Residue, AtomGroup], select_dict: Dict[str, str], number: int = 0 + resname: str, molecule: Residue | AtomGroup, select_dict: dict[str, str], number: int = 0 ): """ Assign the most negatively charged atom in the molecule into select_dict diff --git a/mdgo/util/num.py b/mdgo/util/num.py index ef68b57d..1358ff60 100644 --- a/mdgo/util/num.py +++ b/mdgo/util/num.py @@ -4,10 +4,10 @@ """Utilities for manipulating numbers in data structures.""" -from typing import List, Union, Optional +from __future__ import annotations -def strip_zeros(items: Union[List[Union[str, float, int]], str]) -> Optional[List[int]]: +def strip_zeros(items: list[str | float | int] | str) -> list[int] | None: """ Strip the trailing zeros of a sequence. diff --git a/mdgo/util/packmol.py b/mdgo/util/packmol.py index da5cabb1..76cba654 100644 --- a/mdgo/util/packmol.py +++ b/mdgo/util/packmol.py @@ -16,7 +16,7 @@ import os import subprocess from pathlib import Path -from typing import Dict, List, Optional, Union +from __future__ import annotations from shutil import which from pymatgen.core import Molecule @@ -54,13 +54,13 @@ class PackmolWrapper: def __init__( self, path: str, - molecules: List[Dict], - box: Optional[List[float]] = None, + molecules: list[dict], + box: list[float] | None = None, tolerance: float = 2.0, seed: int = 1, - control_params: Optional[Dict] = None, - inputfile: Union[str, Path] = "packmol.inp", - outputfile: Union[str, Path] = "packmol_out.xyz", + control_params: dict | None = None, + inputfile: str | Path = "packmol.inp", + outputfile: str | Path = "packmol_out.xyz", ): """ Args: diff --git a/mdgo/util/reformat.py b/mdgo/util/reformat.py index 4b967967..a9ee4c6f 100644 --- a/mdgo/util/reformat.py +++ b/mdgo/util/reformat.py @@ -8,14 +8,15 @@ from io import StringIO import re -from typing import List, Dict, Any, Final +from __future__ import annotations +from typing import Any, Final import pandas as pd from mdgo.util.dict_utils import MM_of_Elements from . import __author__ -SECTION_SORTER: Final[Dict[str, Dict[str, Any]]] = { +SECTION_SORTER: Final[dict[str, dict[str, Any]]] = { "atoms": { "in_kw": None, "in_header": ["atom", "charge", "sigma", "epsilon"], @@ -212,7 +213,7 @@ def sdf_to_pdb( title = "cid_" else: title = "" - pdb_atoms: List[Dict[str, Any]] = [] + pdb_atoms: list[dict[str, Any]] = [] # create pdb list of dictionaries atoms = 0 bonds = 0 diff --git a/mdgo/util/volume.py b/mdgo/util/volume.py index 2a77bba7..b9a052d4 100644 --- a/mdgo/util/volume.py +++ b/mdgo/util/volume.py @@ -13,11 +13,12 @@ the cube is defined by the -xsize, -ysize and -zsize options. """ +from __future__ import annotations import sys import os import argparse -from typing import Optional, List, Dict, Union, Tuple, Final +from typing import Final import numpy as np from pymatgen.core import Molecule, Element @@ -25,9 +26,9 @@ DEFAULT_VDW = 1.5 # See Ev:130902 -MOLAR_VOLUME: Final[Dict[str, float]] = {"lipf6": 18, "litfsi": 100} # empirical value +MOLAR_VOLUME: Final[dict[str, float]] = {"lipf6": 18, "litfsi": 100} # empirical value -ALIAS: Final[Dict[str, str]] = { +ALIAS: Final[dict[str, str]] = { "ethylene carbonate": "ec", "ec": "ec", "propylene carbonate": "pc", @@ -68,7 +69,7 @@ } # From PubChem -MOLAR_MASS: Final[Dict[str, float]] = { +MOLAR_MASS: Final[dict[str, float]] = { "ec": 88.06, "pc": 102.09, "dec": 118.13, @@ -88,7 +89,7 @@ } # from Sigma-Aldrich -DENSITY: Final[Dict[str, float]] = { +DENSITY: Final[dict[str, float]] = { "ec": 1.321, "pc": 1.204, "dec": 0.975, @@ -284,7 +285,7 @@ def parse_command_line(): return args -def get_max_dimensions(mol: Molecule) -> Tuple[float, float, float, float, float, float]: +def get_max_dimensions(mol: Molecule) -> tuple[float, float, float, float, float, float]: """ Calculates the dimension of a Molecule @@ -319,7 +320,7 @@ def get_max_dimensions(mol: Molecule) -> Tuple[float, float, float, float, float def set_max_dimensions( x: float = 0.0, y: float = 0.0, z: float = 0.0, x_size: float = 10.0, y_size: float = 10.0, z_size: float = 10.0 -) -> Tuple[float, float, float, float, float, float]: +) -> tuple[float, float, float, float, float, float]: """ Set the max dimensions for calculating active site volume. @@ -345,7 +346,7 @@ def set_max_dimensions( def round_dimensions( x_min: float, x_max: float, y_min: float, y_max: float, z_min: float, z_max: float, mode: str = "lig" -) -> Tuple[float, float, float, float, float, float]: +) -> tuple[float, float, float, float, float, float]: """ Round dimensions to a larger box size (+ buffer). @@ -394,7 +395,7 @@ def dsq(a1: float, a2: float, a3: float, b1: float, b2: float, b3: float) -> flo def get_dimensions( x0: float, x1: float, y0: float, y1: float, z0: float, z1: float, res: float = 0.1 -) -> Tuple[int, int, int]: +) -> tuple[int, int, int]: """ Mesh dimensions in unit of res. @@ -438,7 +439,7 @@ def make_matrix(x_num: int, y_num: int, z_num: int) -> np.ndarray: return matrix -def get_radii(radii_type: str = "Bondi") -> Dict[str, float]: +def get_radii(radii_type: str = "Bondi") -> dict[str, float]: """ Get a radii dict by type. @@ -560,7 +561,7 @@ def fill_volume_matrix( return matrix -def get_occupied_volume(matrix: np.ndarray, res: float, name: Optional[str] = None, molar_volume=True) -> float: +def get_occupied_volume(matrix: np.ndarray, res: float, name: str | None = None, molar_volume=True) -> float: """ Get the occupied volume of the molecule in the box. @@ -581,7 +582,7 @@ def get_occupied_volume(matrix: np.ndarray, res: float, name: Optional[str] = No return v # Å^3 -def get_unoccupied_volume(matrix: np.ndarray, res: float, name: Optional[str] = None, molar_volume=True) -> float: +def get_unoccupied_volume(matrix: np.ndarray, res: float, name: str | None = None, molar_volume=True) -> float: """ Get the unoccupied volume of the molecule in the box. @@ -603,8 +604,8 @@ def get_unoccupied_volume(matrix: np.ndarray, res: float, name: Optional[str] = def molecular_volume( - mol: Union[str, Molecule], - name: Optional[str] = None, + mol: str | Molecule, + name: str | None = None, res: float = 0.1, radii_type: str = "Bondi", molar_volume: bool = True, @@ -674,13 +675,13 @@ def molecular_volume( def concentration_matcher( concentration: float, - salt: Union[float, int, str, Molecule], - solvents: List[Union[str, Dict[str, float]]], - solv_ratio: List[float], + salt: float | int | str | Molecule, + solvents: list[str | dict[str, float]], + solv_ratio: list[float], num_salt: int = 100, mode: str = "v", radii_type: str = "Bondi", -) -> Tuple[List, float]: +) -> tuple[list, float]: """ Estimate the number of molecules of each species in a box, given the salt concentration, salt type, solvent molecular weight, From 9cd855e34a7bdb58068387d3a70bec3b1bb9c02a Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 00:48:59 +0800 Subject: [PATCH 05/34] drop UTF-8 encoding declaration --- mdgo/__init__.py | 1 - mdgo/conductivity.py | 9 ++++----- mdgo/coordination.py | 1 - mdgo/forcefield/__init__.py | 1 - mdgo/forcefield/aqueous.py | 1 - mdgo/forcefield/charge.py | 1 - mdgo/forcefield/crawler.py | 1 - mdgo/forcefield/maestro.py | 1 - mdgo/forcefield/pubchem.py | 1 - mdgo/msd.py | 1 - mdgo/residence_time.py | 1 - mdgo/util/__init__.py | 1 - mdgo/util/coord.py | 1 - mdgo/util/dict_utils.py | 1 - mdgo/util/num.py | 1 - mdgo/util/packmol.py | 1 - mdgo/util/reformat.py | 1 - mdgo/util/volume.py | 1 - 18 files changed, 4 insertions(+), 22 deletions(-) diff --git a/mdgo/__init__.py b/mdgo/__init__.py index fa4ec20d..27e88577 100644 --- a/mdgo/__init__.py +++ b/mdgo/__init__.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. diff --git a/mdgo/conductivity.py b/mdgo/conductivity.py index 221b6883..2ef3651a 100644 --- a/mdgo/conductivity.py +++ b/mdgo/conductivity.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. @@ -27,8 +26,8 @@ def calc_cond_msd( anions: AtomGroup, cations: AtomGroup, run_start: int, - cation_charge: int | float = 1, - anion_charge: int | float = -1, + cation_charge: float = 1, + anion_charge: float = -1, ) -> np.ndarray: """Calculates the conductivity "mean square displacement" over time @@ -128,11 +127,11 @@ def choose_msd_fitting_region( def conductivity_calculator( time_array: np.ndarray, cond_array: np.ndarray, - v: int | float, + v: float, name: str, start: int, end: int, - T: int | float, + T: float, units: str = "real", ) -> float: """Calculates the overall conductivity of the system diff --git a/mdgo/coordination.py b/mdgo/coordination.py index d39fcfde..35f39a3b 100644 --- a/mdgo/coordination.py +++ b/mdgo/coordination.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. diff --git a/mdgo/forcefield/__init__.py b/mdgo/forcefield/__init__.py index 69f445c9..0a5532f6 100644 --- a/mdgo/forcefield/__init__.py +++ b/mdgo/forcefield/__init__.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. diff --git a/mdgo/forcefield/aqueous.py b/mdgo/forcefield/aqueous.py index a6a68fa6..238bbc53 100644 --- a/mdgo/forcefield/aqueous.py +++ b/mdgo/forcefield/aqueous.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. diff --git a/mdgo/forcefield/charge.py b/mdgo/forcefield/charge.py index 925b3787..cc26a250 100644 --- a/mdgo/forcefield/charge.py +++ b/mdgo/forcefield/charge.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. diff --git a/mdgo/forcefield/crawler.py b/mdgo/forcefield/crawler.py index 22d92f4e..488e7750 100644 --- a/mdgo/forcefield/crawler.py +++ b/mdgo/forcefield/crawler.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. diff --git a/mdgo/forcefield/maestro.py b/mdgo/forcefield/maestro.py index 7ad9fe7f..74f85f3e 100644 --- a/mdgo/forcefield/maestro.py +++ b/mdgo/forcefield/maestro.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. diff --git a/mdgo/forcefield/pubchem.py b/mdgo/forcefield/pubchem.py index 2f22d183..b75863ba 100644 --- a/mdgo/forcefield/pubchem.py +++ b/mdgo/forcefield/pubchem.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. diff --git a/mdgo/msd.py b/mdgo/msd.py index 120df2a9..e9d560f9 100644 --- a/mdgo/msd.py +++ b/mdgo/msd.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. diff --git a/mdgo/residence_time.py b/mdgo/residence_time.py index 837579d0..02924ac2 100644 --- a/mdgo/residence_time.py +++ b/mdgo/residence_time.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. diff --git a/mdgo/util/__init__.py b/mdgo/util/__init__.py index eedbf9ee..ab7c3598 100644 --- a/mdgo/util/__init__.py +++ b/mdgo/util/__init__.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. diff --git a/mdgo/util/coord.py b/mdgo/util/coord.py index 3f6df5c4..3d6542fe 100644 --- a/mdgo/util/coord.py +++ b/mdgo/util/coord.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. diff --git a/mdgo/util/dict_utils.py b/mdgo/util/dict_utils.py index 7a5d384c..b04f0623 100644 --- a/mdgo/util/dict_utils.py +++ b/mdgo/util/dict_utils.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. diff --git a/mdgo/util/num.py b/mdgo/util/num.py index 1358ff60..12089111 100644 --- a/mdgo/util/num.py +++ b/mdgo/util/num.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. diff --git a/mdgo/util/packmol.py b/mdgo/util/packmol.py index 76cba654..6969d2ab 100644 --- a/mdgo/util/packmol.py +++ b/mdgo/util/packmol.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. diff --git a/mdgo/util/reformat.py b/mdgo/util/reformat.py index a9ee4c6f..efb32847 100644 --- a/mdgo/util/reformat.py +++ b/mdgo/util/reformat.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. diff --git a/mdgo/util/volume.py b/mdgo/util/volume.py index b9a052d4..3dc4fdbe 100644 --- a/mdgo/util/volume.py +++ b/mdgo/util/volume.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. From 38b34489f69dbbcad5fd3c13a64027a9c86f1d21 Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 00:50:15 +0800 Subject: [PATCH 06/34] add more annotations imports --- mdgo/__init__.py | 3 +++ mdgo/core/__init__.py | 2 ++ mdgo/forcefield/__init__.py | 2 ++ mdgo/util/__init__.py | 2 ++ 4 files changed, 9 insertions(+) diff --git a/mdgo/__init__.py b/mdgo/__init__.py index 27e88577..ba1bf5d2 100644 --- a/mdgo/__init__.py +++ b/mdgo/__init__.py @@ -4,6 +4,9 @@ """ This package contains core modules and classes molecular dynamics simulation setup and analysis. """ + +from __future__ import annotations + __author__ = "Mdgo Development Team" __email__ = "tingzheng_hou@berkeley.edu" __maintainer__ = "Tingzheng Hou" diff --git a/mdgo/core/__init__.py b/mdgo/core/__init__.py index 1f80549c..b4e1559d 100644 --- a/mdgo/core/__init__.py +++ b/mdgo/core/__init__.py @@ -7,6 +7,8 @@ setup, run and analysis. """ +from __future__ import annotations + __author__ = "Tingzheng Hou" __version__ = "0.3.1" __maintainer__ = "Tingzheng Hou" diff --git a/mdgo/forcefield/__init__.py b/mdgo/forcefield/__init__.py index 0a5532f6..df9c3de0 100644 --- a/mdgo/forcefield/__init__.py +++ b/mdgo/forcefield/__init__.py @@ -6,6 +6,8 @@ modifying MD force filed data. """ +from __future__ import annotations + __author__ = "Tingzheng Hou, Ryan Kingsbury" __version__ = "0.3.1" __maintainer__ = "Tingzheng Hou, Ryan Kingsbury" diff --git a/mdgo/util/__init__.py b/mdgo/util/__init__.py index ab7c3598..f08a9560 100644 --- a/mdgo/util/__init__.py +++ b/mdgo/util/__init__.py @@ -3,6 +3,8 @@ """The util package implements various utilities that are commonly used by various packages.""" +from __future__ import annotations + __author__ = "Tingzheng Hou" __version__ = "0.3.0" __maintainer__ = "Tingzheng Hou" From 27511ef1d426a98f6fac97e20955eb3038d4a265 Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 01:04:40 +0800 Subject: [PATCH 07/34] lint fixes 1 --- mdgo/conductivity.py | 7 +++---- mdgo/coordination.py | 6 +++--- setup.py | 4 +++- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/mdgo/conductivity.py b/mdgo/conductivity.py index 2ef3651a..01cc5c36 100644 --- a/mdgo/conductivity.py +++ b/mdgo/conductivity.py @@ -29,7 +29,7 @@ def calc_cond_msd( cation_charge: float = 1, anion_charge: float = -1, ) -> np.ndarray: - """Calculates the conductivity "mean square displacement" over time + """Calculates the conductivity "mean square displacement" over time. Note: Coordinates must be unwrapped (in dcd file when creating MDAnalysis Universe) @@ -58,8 +58,7 @@ def calc_cond_msd( for cation in cation_list: qr_temp += cation.center_of_mass() * cation_charge qr.append(qr_temp) - msd = msd_fft(np.array(qr)) - return msd + return msd_fft(np.array(qr)) def get_beta( @@ -134,7 +133,7 @@ def conductivity_calculator( T: float, units: str = "real", ) -> float: - """Calculates the overall conductivity of the system + """Calculates the overall conductivity of the system. Args: time_array: times at which position data was collected in the simulation diff --git a/mdgo/coordination.py b/mdgo/coordination.py index 35f39a3b..d14b6242 100644 --- a/mdgo/coordination.py +++ b/mdgo/coordination.py @@ -96,7 +96,7 @@ def find_nearest( the ``frequency`` of hopping between sites, and ``steps`` when each binding site exhibits the closest distance to the center atom. """ - time_span = len(list(trj.values())[0]) + time_span = len(next(iter(trj.values()))) if smooth > 0: for kw in list(trj): trj[kw] = savgol_filter(trj.get(kw), smooth, 2) @@ -185,7 +185,7 @@ def find_nearest_free_only( the ``frequency`` of hopping between sites, and ``steps`` when each binding site exhibits the closest distance to the center atom. """ - time_span = len(list(trj.values())[0]) + time_span = len(next(iter(trj.values()))) if smooth > 0: for kw in list(trj): trj[kw] = savgol_filter(trj.get(kw), smooth, 2) @@ -272,7 +272,7 @@ def find_in_n_out( Returns: Two arrays of numberings of frames with hopping in and hopping out event, respectively. """ - time_span = len(list(trj.values())[0]) + time_span = len(next(iter(trj.values()))) if smooth > 0: for kw in list(trj): trj[kw] = savgol_filter(trj.get(kw), smooth, 2) diff --git a/setup.py b/setup.py index d2dc1d78..2ccc1745 100644 --- a/setup.py +++ b/setup.py @@ -1,9 +1,11 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. """Setup.py for MDGO.""" +from __future__ import annotations + + import os from setuptools import setup, find_packages From 4e6e6d8c37044f8df05295b9049a4c689df9b381 Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 01:15:03 +0800 Subject: [PATCH 08/34] ruff 1 --- mdgo/conductivity.py | 4 ++-- mdgo/coordination.py | 14 ++++++------- mdgo/core/__init__.py | 1 - mdgo/core/analysis.py | 39 ++++++++++++++++++------------------- mdgo/core/run.py | 2 +- mdgo/forcefield/__init__.py | 2 +- mdgo/forcefield/aqueous.py | 4 ++-- mdgo/forcefield/crawler.py | 7 +------ mdgo/forcefield/maestro.py | 7 +++---- mdgo/forcefield/pubchem.py | 6 +----- mdgo/msd.py | 8 ++++---- mdgo/residence_time.py | 13 ++++++------- mdgo/util/__init__.py | 2 +- mdgo/util/coord.py | 2 +- mdgo/util/dict_utils.py | 25 ++++++++++++------------ mdgo/util/packmol.py | 20 ++++++++++--------- mdgo/util/reformat.py | 29 ++++++++++++++------------- mdgo/util/volume.py | 17 +++++++--------- tests/test_msd.py | 2 +- 19 files changed, 96 insertions(+), 108 deletions(-) diff --git a/mdgo/conductivity.py b/mdgo/conductivity.py index 01cc5c36..ed23b8db 100644 --- a/mdgo/conductivity.py +++ b/mdgo/conductivity.py @@ -8,9 +8,9 @@ from __future__ import annotations import numpy as np -from tqdm.auto import tqdm +from MDAnalysis import AtomGroup, Universe from scipy import stats -from MDAnalysis import Universe, AtomGroup +from tqdm.auto import tqdm from mdgo.msd import msd_fft diff --git a/mdgo/coordination.py b/mdgo/coordination.py index d14b6242..0079489e 100644 --- a/mdgo/coordination.py +++ b/mdgo/coordination.py @@ -6,16 +6,17 @@ """ from __future__ import annotations + from collections.abc import Callable import numpy as np -from tqdm.auto import tqdm -from MDAnalysis import Universe, AtomGroup -from MDAnalysis.core.groups import Atom +from MDAnalysis import AtomGroup, Universe from MDAnalysis.analysis.distances import distance_array +from MDAnalysis.core.groups import Atom from scipy.signal import savgol_filter -from mdgo.util.coord import atom_vec, angle +from tqdm.auto import tqdm +from mdgo.util.coord import angle, atom_vec __author__ = "Tingzheng Hou" __version__ = "0.3.0" @@ -451,8 +452,8 @@ def heat_map( vertex_atoms.append(atoms[idx[0]]) else: raise ValueError( - f"There should be at least 1 cluster_terminal atom in the {str(dim[i])} dimension." - f"Try broadening the selection at index {str(i + 1)} of the cluster_terminal " + f"There should be at least 1 cluster_terminal atom in the {dim[i]!s} dimension." + f"Try broadening the selection at index {i + 1!s} of the cluster_terminal " ) else: assert isinstance(cluster_terminal, str) @@ -799,7 +800,6 @@ def num_of_neighbor_simple( A dict with "total" as the key and an array of the solvation structure type in the specified frame range as the value. """ - time_count = 0 trj_analysis = nvt_run.trajectory[run_start:run_end:] center_selection = "same type as index " + str(center_atom.index) diff --git a/mdgo/core/__init__.py b/mdgo/core/__init__.py index b4e1559d..a817d7b1 100644 --- a/mdgo/core/__init__.py +++ b/mdgo/core/__init__.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. diff --git a/mdgo/core/analysis.py b/mdgo/core/analysis.py index b9440eff..fbc4db3a 100644 --- a/mdgo/core/analysis.py +++ b/mdgo/core/analysis.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. @@ -9,38 +8,39 @@ from __future__ import annotations +import matplotlib.pyplot as plt import MDAnalysis import numpy as np import pandas as pd -import matplotlib.pyplot as plt from MDAnalysis import Universe from MDAnalysis.analysis.distances import distance_array from MDAnalysis.lib.distances import capped_distance from tqdm.auto import tqdm -from mdgo.util.dict_utils import ( - mass_to_name, - assign_name, - assign_resname, - res_dict_from_select_dict, - res_dict_from_datafile, - select_dict_from_resname, -) -from mdgo.conductivity import calc_cond_msd, conductivity_calculator, choose_msd_fitting_region, get_beta + +from mdgo.conductivity import calc_cond_msd, choose_msd_fitting_region, conductivity_calculator, get_beta from mdgo.coordination import ( + angular_dist_of_neighbor, concat_coord_array, + find_nearest, + find_nearest_free_only, + get_full_coords, + heat_map, + neighbor_distance, num_of_neighbor, num_of_neighbor_simple, num_of_neighbor_specific, - angular_dist_of_neighbor, - neighbor_distance, - find_nearest, - find_nearest_free_only, process_evol, - heat_map, - get_full_coords, ) -from mdgo.msd import total_msd, partial_msd, DIM +from mdgo.msd import DIM, partial_msd, total_msd from mdgo.residence_time import calc_neigh_corr, fit_residence_time +from mdgo.util.dict_utils import ( + assign_name, + assign_resname, + mass_to_name, + res_dict_from_datafile, + res_dict_from_select_dict, + select_dict_from_resname, +) class MdRun: @@ -90,7 +90,6 @@ def __init__( parsed data (``Universe``) or other bridging objects (``CombinedData``). Not recommended to use directly. """ - self.wrapped_run = wrapped_run self.unwrapped_run = unwrapped_run self.nvt_start = nvt_start @@ -1124,7 +1123,7 @@ def get_heat_map( run_start, run_end, ) - if not coords.size == 0: + if coords.size != 0: coord_list = np.concatenate((coord_list, coords), axis=0) coord_list = coord_list[1:] if sym_dict: diff --git a/mdgo/core/run.py b/mdgo/core/run.py index ffe60cba..cc30ae6c 100644 --- a/mdgo/core/run.py +++ b/mdgo/core/run.py @@ -1,10 +1,10 @@ -# coding: utf-8 # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. """ This module implements a core class MdRun for molecular dynamics job setup. """ +from __future__ import annotations class MdJob: diff --git a/mdgo/forcefield/__init__.py b/mdgo/forcefield/__init__.py index df9c3de0..64225456 100644 --- a/mdgo/forcefield/__init__.py +++ b/mdgo/forcefield/__init__.py @@ -15,7 +15,7 @@ __date__ = "Dec 19, 2023" -from .aqueous import IonLJData, Aqueous +from .aqueous import Aqueous, IonLJData from .charge import ChargeWriter from .crawler import FFcrawler from .maestro import MaestroRunner diff --git a/mdgo/forcefield/aqueous.py b/mdgo/forcefield/aqueous.py index 238bbc53..c071d4db 100644 --- a/mdgo/forcefield/aqueous.py +++ b/mdgo/forcefield/aqueous.py @@ -6,11 +6,11 @@ """ from __future__ import annotations -from typing import Literal, Final import os import re from dataclasses import dataclass +from typing import Final, Literal from monty.json import MSONable from monty.serialization import loadfn @@ -18,7 +18,6 @@ from pymatgen.core.ion import Ion from pymatgen.io.lammps.data import ForceField, LammpsData, Topology, lattice_2_lmpbox - MODULE_DIR: Final[str] = os.path.dirname(os.path.abspath(__file__)) DATA_DIR: Final[str] = os.path.join(MODULE_DIR, "data") DATA_MODELS: Final[dict] = { @@ -126,6 +125,7 @@ def get_water(model: str = "spce") -> LammpsData: model: Water model to use. Valid choices are "spc", "spce", "opc3", "tip3pew", "tip3pfb", "tip4p2005", "tip4pew", "tip4pfb", and "opc". (Default: "spce") + Returns: LammpsData: Force field parameters for the chosen water model. If you specify an invalid water model, None is returned. diff --git a/mdgo/forcefield/crawler.py b/mdgo/forcefield/crawler.py index 488e7750..27596f7d 100644 --- a/mdgo/forcefield/crawler.py +++ b/mdgo/forcefield/crawler.py @@ -18,13 +18,9 @@ import shutil import time - from pymatgen.io.lammps.data import LammpsData from selenium import webdriver -from selenium.common.exceptions import ( - TimeoutException, - WebDriverException, -) +from selenium.common.exceptions import TimeoutException, WebDriverException from selenium.webdriver.common.by import By from selenium.webdriver.support import expected_conditions as EC from selenium.webdriver.support.ui import WebDriverWait @@ -50,7 +46,6 @@ class FFcrawler: Default to False. Examples: - >>> lpg = FFcrawler('/path/to/work/dir', '/path/to/chromedriver') >>> lpg.data_from_pdb("/path/to/pdb") """ diff --git a/mdgo/forcefield/maestro.py b/mdgo/forcefield/maestro.py index 74f85f3e..c1a61d09 100644 --- a/mdgo/forcefield/maestro.py +++ b/mdgo/forcefield/maestro.py @@ -17,13 +17,13 @@ """ from __future__ import annotations -from typing import Final import os import signal import subprocess import time from string import Template +from typing import Final from mdgo.util.reformat import ff_parser @@ -67,7 +67,6 @@ class MaestroRunner: $SCHRODINGER/mmshare-vversion/data/f14/ Examples: - >>> mr = MaestroRunner('/path/to/structure', '/path/to/working/dir') >>> mr.get_mae() >>> mr.get_ff() @@ -100,11 +99,11 @@ def __init__( self.cmd_template = cmd_template else: if assign_bond: - with open(self.template_assignbond, "r") as f: + with open(self.template_assignbond) as f: cmd_template = f.read() self.cmd_template = cmd_template else: - with open(self.template_noassignbond, "r") as f: + with open(self.template_noassignbond) as f: cmd_template = f.read() self.cmd_template = cmd_template diff --git a/mdgo/forcefield/pubchem.py b/mdgo/forcefield/pubchem.py index b75863ba..da8617f5 100644 --- a/mdgo/forcefield/pubchem.py +++ b/mdgo/forcefield/pubchem.py @@ -15,17 +15,13 @@ import pubchempy as pcp from selenium import webdriver -from selenium.common.exceptions import ( - NoSuchElementException, - TimeoutException, -) +from selenium.common.exceptions import NoSuchElementException, TimeoutException from selenium.webdriver.common.by import By from selenium.webdriver.support import expected_conditions as EC from selenium.webdriver.support.ui import WebDriverWait from mdgo.util.reformat import sdf_to_pdb - MAESTRO: Final[str] = "$SCHRODINGER/maestro -console -nosplash" FFLD: Final[str] = "$SCHRODINGER/utilities/ffld_server -imae {} -version 14 -print_parameters -out_file {}" MolecularWeight: Final[str] = "https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{}/property/MolecularWeight/txt" diff --git a/mdgo/msd.py b/mdgo/msd.py index e9d560f9..035ac157 100644 --- a/mdgo/msd.py +++ b/mdgo/msd.py @@ -10,6 +10,7 @@ """ from __future__ import annotations + from typing import Literal try: @@ -22,10 +23,9 @@ td = None import numpy as np -from tqdm.auto import trange - -from MDAnalysis import Universe, AtomGroup +from MDAnalysis import AtomGroup, Universe from MDAnalysis.core.groups import Atom +from tqdm.auto import trange __author__ = "Tingzheng Hou" __version__ = "0.3.0" @@ -333,7 +333,7 @@ def msd_from_frags(coord_list: list[np.ndarray], largest: int) -> np.ndarray: for lag in lag_times: disp = state[:-lag, :] - state[lag:, :] sqdist = np.square(disp).sum(axis=-1) - if lag in msd_dict.keys(): + if lag in msd_dict: msd_dict[lag] = np.concatenate((msd_dict[lag], sqdist), axis=0) else: msd_dict[lag] = sqdist diff --git a/mdgo/residence_time.py b/mdgo/residence_time.py index 02924ac2..ce26678e 100644 --- a/mdgo/residence_time.py +++ b/mdgo/residence_time.py @@ -9,14 +9,13 @@ import os -import numpy as np import matplotlib.pyplot as plt -from statsmodels.tsa.stattools import acovf -from scipy.optimize import curve_fit -from tqdm.auto import tqdm - +import numpy as np from MDAnalysis import Universe from MDAnalysis.core.groups import Atom +from scipy.optimize import curve_fit +from statsmodels.tsa.stattools import acovf +from tqdm.auto import tqdm __author__ = "Kara Fong, Tingzheng Hou" __version__ = "0.3.0" @@ -53,7 +52,7 @@ def neighbors_one_atom( bool_values = {} time_count = 0 for ts in nvt_run.trajectory[run_start:run_end:]: - if species in select_dict.keys(): + if species in select_dict: selection = ( "(" + select_dict[species] @@ -147,7 +146,7 @@ def calc_neigh_corr( times = np.array(times) acf_avg = {} - for kw in distance_dict.keys(): + for kw in distance_dict: acf_all = [] for atom in tqdm(center_atoms[::]): distance = distance_dict.get(kw) diff --git a/mdgo/util/__init__.py b/mdgo/util/__init__.py index f08a9560..46ca2dd3 100644 --- a/mdgo/util/__init__.py +++ b/mdgo/util/__init__.py @@ -11,7 +11,7 @@ __email__ = "tingzheng_hou@berkeley.edu" __date__ = "Jul 19, 2021" -from typing import Final, Dict +from typing import Dict, Final MM_of_Elements: Final[Dict[str, float]] = { "H": 1.00794, diff --git a/mdgo/util/coord.py b/mdgo/util/coord.py index 3d6542fe..1db0b6a8 100644 --- a/mdgo/util/coord.py +++ b/mdgo/util/coord.py @@ -4,8 +4,8 @@ """Utilities for manipulating coordinates under periodic boundary conditions.""" from __future__ import annotations -import numpy as np +import numpy as np from MDAnalysis.core.groups import Atom diff --git a/mdgo/util/dict_utils.py b/mdgo/util/dict_utils.py index b04f0623..d5eb8f79 100644 --- a/mdgo/util/dict_utils.py +++ b/mdgo/util/dict_utils.py @@ -3,17 +3,17 @@ """Utilities for manipulating dictionaries.""" -import string -import re -import math from __future__ import annotations -import numpy as np -import pandas as pd -from pymatgen.io.lammps.data import CombinedData +import math +import re +import string +import numpy as np +import pandas as pd from MDAnalysis import Universe -from MDAnalysis.core.groups import Residue, AtomGroup +from MDAnalysis.core.groups import AtomGroup, Residue +from pymatgen.io.lammps.data import CombinedData from . import MM_of_Elements @@ -40,6 +40,7 @@ def mass_to_name(masses: np.ndarray) -> np.ndarray: def lmp_mass_to_name(df: pd.DataFrame) -> dict[int, str]: """ Create a dict for mapping atom type id to element from the mass information. + Args: df: The masses attribute from LammpsData object Return: @@ -90,7 +91,7 @@ def res_dict_from_select_dict(u: Universe, select_dict: dict[str, str]) -> dict[ select_dict: A dictionary of atom species, where each atom species name is a key and the corresponding values are the selection language. - return: + Return: A dictionary of resnames. """ saved_select = [] @@ -118,11 +119,11 @@ def res_dict_from_datafile(filename: str) -> dict[str, str]: Args: filename: Path to the data file. The data file must be generated by a CombinedData object. - return: + Return: A dictionary of resnames. """ res_dict = {} - with open(filename, "r") as f: + with open(filename) as f: lines = f.readlines() if lines[0] == "Generated by pymatgen.io.lammps.data.LammpsData\n" and lines[1].startswith("#"): elyte_info = re.findall(r"\w+", lines[1]) @@ -154,7 +155,7 @@ def res_dict_from_lammpsdata(lammps_data: CombinedData) -> dict[str, str]: Args: lammps_data: A CombinedData object. - return: + Return: A dictionary of resnames. """ assert isinstance(lammps_data, CombinedData) @@ -191,7 +192,7 @@ def select_dict_from_resname(u: Universe) -> dict[str, str]: Args: u: The universe object to work with. - return: + Return: A dictionary of atom species. """ select_dict: dict[str, str] = {} diff --git a/mdgo/util/packmol.py b/mdgo/util/packmol.py index 6969d2ab..d0f8fb36 100644 --- a/mdgo/util/packmol.py +++ b/mdgo/util/packmol.py @@ -12,15 +12,16 @@ set the folder of the packmol executable to the PATH environment variable. """ +from __future__ import annotations + import os import subprocess from pathlib import Path -from __future__ import annotations from shutil import which + from pymatgen.core import Molecule # from pymatgen.io.core import InputFile, InputSet, InputGenerator - from mdgo.util.volume import molecular_volume __author__ = "Tingzheng Hou, Ryan Kingsbury" @@ -38,7 +39,6 @@ class PackmolWrapper: molecules into a one single unit. Examples: - >>> molecules = [{"name": "EMC", "number": 2, "coords": "/Users/th/Downloads/test_selenium/EMC.lmp.xyz"}] @@ -71,7 +71,8 @@ def __init__( 2. "number" - the number of that molecule to pack into the box 3. "coords" - Coordinates in the form of either a Molecule object or a path to a file. - Example: + + Example: {"name": "water", "number": 500, "coords": "/path/to/input/file.xyz"} @@ -96,8 +97,10 @@ def __init__( def run_packmol(self, timeout=30): """Run packmol and write out the packed structure. + Args: timeout: Timeout in seconds. + Raises: ValueError if packmol does not succeed in packing the box. TimeoutExpiredError if packmold does not finish within the timeout. @@ -142,7 +145,6 @@ def run_packmol(self, timeout=30): # InputSet def make_packmol_input(self): """Make a Packmol usable input file.""" - if self.box: box_list = " ".join(str(i) for i in self.box) else: @@ -171,7 +173,7 @@ def make_packmol_input(self): if isinstance(v, list): out.write(f'{k} {" ".join(str(x) for x in v)}\n') else: - out.write(f"{k} {str(v)}\n") + out.write(f"{k} {v!s}\n") out.write(f"seed {self.seed}\n") out.write(f"tolerance {self.tolerance}\n\n") @@ -189,9 +191,9 @@ def make_packmol_input(self): out.write(f'structure {d["coords"]}\n') elif isinstance(d["coords"], Path): if " " in str(d["coords"]): - out.write(f'structure "{str(d["coords"])}"\n') + out.write(f'structure "{d["coords"]!s}"\n') else: - out.write(f'structure {str(d["coords"])}\n') + out.write(f'structure {d["coords"]!s}\n') elif isinstance(d["coords"], Molecule): fname = os.path.join(self.path, f'packmol_{d["name"]}.xyz') d["coords"].to(filename=fname) @@ -199,7 +201,7 @@ def make_packmol_input(self): out.write(f'structure "{fname}"\n') else: out.write(f"structure {fname}\n") - out.write(f' number {str(d["number"])}\n') + out.write(f' number {d["number"]!s}\n') out.write(f" inside box {box_list}\n") out.write("end structure\n\n") diff --git a/mdgo/util/reformat.py b/mdgo/util/reformat.py index efb32847..2a78f03c 100644 --- a/mdgo/util/reformat.py +++ b/mdgo/util/reformat.py @@ -5,15 +5,17 @@ Utilities for converting data file formats. """ -from io import StringIO -import re from __future__ import annotations + +import re +from io import StringIO from typing import Any, Final + import pandas as pd from mdgo.util.dict_utils import MM_of_Elements -from . import __author__ +from . import __author__ SECTION_SORTER: Final[dict[str, dict[str, Any]]] = { "atoms": { @@ -87,12 +89,12 @@ def ff_parser(ff_dir: str, xyz_dir: str) -> str: Return: The output LAMMPS data string. """ - with open(xyz_dir, "r") as f_xyz: + with open(xyz_dir) as f_xyz: molecule = pd.read_table(f_xyz, skiprows=2, delim_whitespace=True, names=["atom", "x", "y", "z"]) coordinates = molecule[["x", "y", "z"]] lo = coordinates.min().min() - 0.5 hi = coordinates.max().max() + 0.5 - with open(ff_dir, "r") as f: + with open(ff_dir) as f: lines_org = f.read() lines = lines_org.split("\n\n") atoms = "\n".join(lines[4].split("\n", 4)[4].split("\n")[:-1]) @@ -203,9 +205,8 @@ def sdf_to_pdb( credit: Whether to credit line (remark 888) in the pdb file. Default to True. pubchem: Whether the sdf is downloaded from PubChem. Default to True. """ - # parse sdf file file - with open(sdf_file, "r") as inp: + with open(sdf_file) as inp: sdf_lines = inp.readlines() sdf = list(map(str.strip, sdf_lines)) if pubchem: @@ -243,12 +244,12 @@ def sdf_to_pdb( "z": float(line_split[2]), "occupancy": 1.00, "tempFactor": 0.00, - "altLoc": str(""), - "chainID": str(""), - "iCode": str(""), + "altLoc": "", + "chainID": "", + "iCode": "", "element": str(line_split[3]), - "charge": str(""), - "segment": str(""), + "charge": "", + "segment": "", } pdb_atoms.append(newline) elif i in list(range(4 + atoms, 4 + atoms + bonds)): @@ -266,13 +267,13 @@ def sdf_to_pdb( pass # write pdb file - with open(pdb_file, "wt") as outp: + with open(pdb_file, "w") as outp: if write_title: outp.write(f"TITLE {title:70s}\n") if version: outp.write("REMARK 4 COMPLIES WITH FORMAT V. 3.3, 21-NOV-2012\n") if credit: - outp.write("REMARK 888\n" "REMARK 888 WRITTEN BY MDGO (CREATED BY TINGZHENG HOU)\n") + outp.write("REMARK 888\nREMARK 888 WRITTEN BY MDGO (CREATED BY TINGZHENG HOU)\n") for n in range(atoms): line_dict = pdb_atoms[n].copy() if len(line_dict["name"]) == 3: diff --git a/mdgo/util/volume.py b/mdgo/util/volume.py index 3dc4fdbe..d6cef08f 100644 --- a/mdgo/util/volume.py +++ b/mdgo/util/volume.py @@ -14,14 +14,13 @@ """ from __future__ import annotations -import sys -import os import argparse +import os +import sys from typing import Final import numpy as np -from pymatgen.core import Molecule, Element - +from pymatgen.core import Element, Molecule DEFAULT_VDW = 1.5 # See Ev:130902 @@ -294,7 +293,6 @@ def get_max_dimensions(mol: Molecule) -> tuple[float, float, float, float, float Returns: xmin, xmax, ymin, ymax, zmin, zmax """ - xmin = 9999 ymin = 9999 zmin = 9999 @@ -433,7 +431,6 @@ def make_matrix(x_num: int, y_num: int, z_num: int) -> np.ndarray: Returns: matrix """ - matrix = np.array([[[None for _ in range(z_num)] for _ in range(y_num)] for _ in range(x_num)]) return matrix @@ -537,11 +534,11 @@ def fill_volume_matrix( if element == "H": continue radius = radii.get(element, DEFAULT_VDW) - for i in range(0, xsteps): + for i in range(xsteps): if abs(a.x - (x0 + 0.5 * res + i * res)) < radius: - for j in range(0, ysteps): + for j in range(ysteps): if abs(a.y - (y0 + 0.5 * res + j * res)) < radius: - for k in range(0, zsteps): + for k in range(zsteps): if matrix[i][j][k] != 1: if abs(a.z - (z0 + 0.5 * res + k * res)) < radius: if dsq( @@ -745,7 +742,7 @@ def concentration_matcher( sys.exit(1) name = os.path.splitext(os.path.split(salt)[-1])[0] ext = os.path.splitext(os.path.split(salt)[-1])[1] - if not ext == ".xyz": + if ext != ".xyz": print("Error: Wrong file format, please use a .xyz file.\n") sys.exit(1) salt_molar_volume = molecular_volume(salt, name, radii_type=radii_type) diff --git a/tests/test_msd.py b/tests/test_msd.py index 3f5ca573..46c03c92 100644 --- a/tests/test_msd.py +++ b/tests/test_msd.py @@ -56,7 +56,7 @@ def test_create_position_arrays(self): def test_parse_msd_type(self): xyz = parse_msd_type("xyz") - self.assertEqual(["x", "y", "z"], self.dims[xyz[0]:xyz[1]:xyz[2]]) + assert ["x", "y", "z"] == self.dims[xyz[0]:xyz[1]:xyz[2]] xy = parse_msd_type("xy") self.assertEqual(["x", "y"], self.dims[xy[0]:xy[1]:xy[2]]) yz = parse_msd_type("yz") From e1ba1e341d48f26e5647c1a27f82c80656c70c09 Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 01:18:23 +0800 Subject: [PATCH 09/34] ruff 2 --- docs/source/conf.py | 2 ++ setup.py | 5 ++--- tasks.py | 11 ++++++----- tests/test_conductivity.py | 6 +++--- tests/test_coordination.py | 2 +- tests/test_core.py | 2 +- tests/test_forcefield.py | 4 ++-- tests/test_mdgopackmol.py | 11 ++++++----- tests/test_msd.py | 2 +- tests/test_residence_time.py | 2 +- tests/test_util.py | 2 +- tests/test_volume.py | 3 +++ 12 files changed, 29 insertions(+), 23 deletions(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index 950b8eb3..b61f8c72 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -10,6 +10,8 @@ # add these directories to sys.path here. If the directory is relative to the # documentation root, use os.path.abspath to make it absolute, like shown here. # +from __future__ import annotations + import os import sys import sphinx_rtd_theme, sphinx_autodoc_typehints diff --git a/setup.py b/setup.py index 2ccc1745..d08a09bc 100644 --- a/setup.py +++ b/setup.py @@ -5,14 +5,13 @@ from __future__ import annotations - import os -from setuptools import setup, find_packages +from setuptools import find_packages, setup module_dir = os.path.dirname(os.path.abspath(__file__)) -with open(os.path.join(module_dir, "README.md"), "r") as f: +with open(os.path.join(module_dir, "README.md")) as f: readme = f.read() INSTALL_REQUIRES = [ diff --git a/tasks.py b/tasks.py index 825cde6a..c0d226e2 100644 --- a/tasks.py +++ b/tasks.py @@ -2,6 +2,7 @@ Pyinvoke tasks.py file for automating releases and admin stuff. Author: Tingzheng Hou """ +from __future__ import annotations import glob import json @@ -109,35 +110,35 @@ def set_ver(ctx, version): contents = f.read() contents = re.sub(r"__version__ = .*\n", '__version__ = "%s"\n' % version, contents) - with open("mdgo/__init__.py", "wt") as f: + with open("mdgo/__init__.py", "w") as f: f.write(contents) with open("mdgo/core/__init__.py") as f: contents = f.read() contents = re.sub(r"__version__ = .*\n", '__version__ = "%s"\n' % version, contents) - with open("mdgo/core/__init__.py", "wt") as f: + with open("mdgo/core/__init__.py", "w") as f: f.write(contents) with open("mdgo/forcefield/__init__.py") as f: contents = f.read() contents = re.sub(r"__version__ = .*\n", '__version__ = "%s"\n' % version, contents) - with open("mdgo/forcefield/__init__.py", "wt") as f: + with open("mdgo/forcefield/__init__.py", "w") as f: f.write(contents) with open("mdgo/util/__init__.py") as f: contents = f.read() contents = re.sub(r"__version__ = .*\n", '__version__ = "%s"\n' % version, contents) - with open("mdgo/util/__init__.py", "wt") as f: + with open("mdgo/util/__init__.py", "w") as f: f.write(contents) with open("setup.py") as f: contents = f.read() contents = re.sub(r"version=([^,]+),", 'version="%s",' % version, contents) - with open("setup.py", "wt") as f: + with open("setup.py", "w") as f: f.write(contents) diff --git a/tests/test_conductivity.py b/tests/test_conductivity.py index dc5fe230..dcfa89fb 100644 --- a/tests/test_conductivity.py +++ b/tests/test_conductivity.py @@ -1,10 +1,10 @@ from __future__ import annotations + import os -import sys -import tempfile -from io import StringIO import unittest + import MDAnalysis + from mdgo.conductivity import * test_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_files") diff --git a/tests/test_coordination.py b/tests/test_coordination.py index 961e79b5..20c90463 100644 --- a/tests/test_coordination.py +++ b/tests/test_coordination.py @@ -1,5 +1,5 @@ from __future__ import annotations -import os + import unittest diff --git a/tests/test_core.py b/tests/test_core.py index 961e79b5..20c90463 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -1,5 +1,5 @@ from __future__ import annotations -import os + import unittest diff --git a/tests/test_forcefield.py b/tests/test_forcefield.py index 2ac7f953..8b454fc6 100644 --- a/tests/test_forcefield.py +++ b/tests/test_forcefield.py @@ -1,5 +1,5 @@ from __future__ import annotations -import io + import sys import tempfile import unittest @@ -8,8 +8,8 @@ import numpy as np import pytest -from mdgo.forcefield.crawler import * from mdgo.forcefield.aqueous import * +from mdgo.forcefield.crawler import * test_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_files") diff --git a/tests/test_mdgopackmol.py b/tests/test_mdgopackmol.py index af02eec3..78bbe699 100644 --- a/tests/test_mdgopackmol.py +++ b/tests/test_mdgopackmol.py @@ -1,11 +1,12 @@ from __future__ import annotations + import os import tempfile from pathlib import Path from subprocess import TimeoutExpired -import pytest import numpy as np +import pytest from pymatgen.core import Molecule from mdgo.util.packmol import PackmolWrapper @@ -13,7 +14,7 @@ test_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_files") -@pytest.fixture +@pytest.fixture() def ethanol(): """ Returns a Molecule of ethanol @@ -34,7 +35,7 @@ def ethanol(): return Molecule(ethanol_atoms, ethanol_coords) -@pytest.fixture +@pytest.fixture() def water(): """ Returns a Molecule of water @@ -121,7 +122,7 @@ def test_control_params(self, water, ethanol): control_params={"maxit": 0, "nloop": 0}, ) pw.make_packmol_input() - with open(os.path.join(scratch_dir, "packmol.inp"), "r") as f: + with open(os.path.join(scratch_dir, "packmol.inp")) as f: input_string = f.read() assert "maxit 0" in input_string assert "nloop 0" in input_string @@ -159,7 +160,7 @@ def test_no_return_and_box(self, water, ethanol): box=[0, 0, 0, 2, 2, 2], ) pw.make_packmol_input() - with open(os.path.join(scratch_dir, "packmol.inp"), "r") as f: + with open(os.path.join(scratch_dir, "packmol.inp")) as f: input_string = f.read() assert "inside box 0 0 0 2 2 2" in input_string with pytest.raises(ValueError): diff --git a/tests/test_msd.py b/tests/test_msd.py index 46c03c92..7d4b5577 100644 --- a/tests/test_msd.py +++ b/tests/test_msd.py @@ -1,4 +1,5 @@ from __future__ import annotations + import os import unittest @@ -11,7 +12,6 @@ from mdgo.msd import * - test_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_files") diff --git a/tests/test_residence_time.py b/tests/test_residence_time.py index 961e79b5..20c90463 100644 --- a/tests/test_residence_time.py +++ b/tests/test_residence_time.py @@ -1,5 +1,5 @@ from __future__ import annotations -import os + import unittest diff --git a/tests/test_util.py b/tests/test_util.py index 961e79b5..20c90463 100644 --- a/tests/test_util.py +++ b/tests/test_util.py @@ -1,5 +1,5 @@ from __future__ import annotations -import os + import unittest diff --git a/tests/test_volume.py b/tests/test_volume.py index 2e60001f..86b10c54 100644 --- a/tests/test_volume.py +++ b/tests/test_volume.py @@ -1,7 +1,10 @@ from __future__ import annotations + import os import unittest + from pymatgen.core import Molecule + from mdgo.util.volume import molecular_volume test_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_files") From 2f886d919a613b6d36fb99280036c0a3ab026515 Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 01:25:43 +0800 Subject: [PATCH 10/34] self.assertAlmostEqual to assert_allclose --- tests/test_msd.py | 35 ++++++++++++++++++----------------- 1 file changed, 18 insertions(+), 17 deletions(-) diff --git a/tests/test_msd.py b/tests/test_msd.py index 7d4b5577..851d5e92 100644 --- a/tests/test_msd.py +++ b/tests/test_msd.py @@ -4,6 +4,7 @@ import unittest import MDAnalysis +from numpy.testing import assert_allclose try: import tidynamics as td @@ -75,25 +76,25 @@ def test_onsager_ii_self(self): onsager_ii_self_nocom = onsager_ii_self(self.gen2, 0, 100, select="type 3", center_of_mass=False) onsager_ii_self_nofft = onsager_ii_self(self.gen2, 0, 100, select="type 3", fft=False) - self.assertAlmostEqual(32.14254152556588, onsager_ii_self_fft[50]) - self.assertAlmostEqual(63.62190983, onsager_ii_self_fft[98]) - self.assertAlmostEqual(67.29990019, onsager_ii_self_fft[99]) - self.assertAlmostEqual(32.14254152556588, onsager_ii_self_nofft[50]) - self.assertAlmostEqual(63.62190983, onsager_ii_self_nofft[98]) - self.assertAlmostEqual(67.29990019, onsager_ii_self_nofft[99]) - self.assertAlmostEqual(32.338364098424634, onsager_ii_self_nocom[50]) - self.assertAlmostEqual(63.52915984813752, onsager_ii_self_nocom[98]) - self.assertAlmostEqual(67.29599346166411, onsager_ii_self_nocom[99]) + assert_allclose(32.14254152556588, onsager_ii_self_fft[50]) + assert_allclose(63.62190983, onsager_ii_self_fft[98]) + assert_allclose(67.29990019, onsager_ii_self_fft[99]) + assert_allclose(32.14254152556588, onsager_ii_self_nofft[50]) + assert_allclose(63.62190983, onsager_ii_self_nofft[98]) + assert_allclose(67.29990019, onsager_ii_self_nofft[99]) + assert_allclose(32.338364098424634, onsager_ii_self_nocom[50]) + assert_allclose(63.52915984813752, onsager_ii_self_nocom[98]) + assert_allclose(67.29599346166411, onsager_ii_self_nocom[99]) def test_mda_msd_wrapper(self): mda_msd_cation = mda_msd_wrapper(self.gen2, 0, 100, select="type 3", fft=False) mda_msd_anion = mda_msd_wrapper(self.gen2, 0, 100, select="type 1", fft=False) - self.assertAlmostEqual(32.338364098424634, mda_msd_cation[50]) - self.assertAlmostEqual(63.52915984813752, mda_msd_cation[98]) - self.assertAlmostEqual(67.29599346166411, mda_msd_cation[99]) - self.assertAlmostEqual(42.69200176568008, mda_msd_anion[50]) - self.assertAlmostEqual(86.9209518, mda_msd_anion[98]) - self.assertAlmostEqual(89.84668178, mda_msd_anion[99]) + assert_allclose(32.338364098424634, mda_msd_cation[50]) + assert_allclose(63.52915984813752, mda_msd_cation[98]) + assert_allclose(67.29599346166411, mda_msd_cation[99]) + assert_allclose(42.69200176568008, mda_msd_anion[50]) + assert_allclose(86.9209518, mda_msd_anion[98]) + assert_allclose(89.84668178, mda_msd_anion[99]) assert np.allclose( onsager_ii_self(self.gen2, 0, 10, select="type 3", msd_type="x", center_of_mass=False), mda_msd_wrapper(self.gen2, 0, 10, select="type 3", msd_type="x", fft=False), @@ -126,8 +127,8 @@ def test_total_msd(self): total_mda_cation = total_msd( self.gen2, 0, 100, select="type 3", fft=False, built_in=False, center_of_mass=False ) - self.assertAlmostEqual(32.14254152556588, total_builtin_cation[50]) - self.assertAlmostEqual(32.338364098424634, total_mda_cation[50]) + assert_allclose(32.14254152556588, total_builtin_cation[50]) + assert_allclose(32.338364098424634, total_mda_cation[50]) with self.assertRaises(ValueError): total_msd(self.gen2, 0, 100, select="type 3", fft=True, built_in=False, center_of_mass=True) From 8b6cf53b77b79176205bf7ace28d15a47dcee7ea Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 08:16:54 +0800 Subject: [PATCH 11/34] try import numpy twice? --- tests/test_msd.py | 1 + tests/test_residence_time.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_msd.py b/tests/test_msd.py index 851d5e92..e7906854 100644 --- a/tests/test_msd.py +++ b/tests/test_msd.py @@ -4,6 +4,7 @@ import unittest import MDAnalysis +import numpy as np from numpy.testing import assert_allclose try: diff --git a/tests/test_residence_time.py b/tests/test_residence_time.py index 20c90463..cf1cc95f 100644 --- a/tests/test_residence_time.py +++ b/tests/test_residence_time.py @@ -5,7 +5,7 @@ class MyTestCase(unittest.TestCase): def test_something(self): - self.assertEqual(True, True) + assert True is True if __name__ == "__main__": From ec8ebc3f6297e058847511dec2a17eadaa2dd27d Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 08:33:42 +0800 Subject: [PATCH 12/34] update test_msd forms --- tests/test_msd.py | 54 ++++++++++++++++++++++++++++------------------- 1 file changed, 32 insertions(+), 22 deletions(-) diff --git a/tests/test_msd.py b/tests/test_msd.py index e7906854..20287d56 100644 --- a/tests/test_msd.py +++ b/tests/test_msd.py @@ -12,7 +12,17 @@ except ImportError: td = None -from mdgo.msd import * +import pytest + +from mdgo.msd import ( + create_position_arrays, + mda_msd_wrapper, + msd_fft, + msd_straight_forward, + onsager_ii_self, + parse_msd_type, + total_msd, +) test_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_files") @@ -32,26 +42,26 @@ def setUpClass(cls) -> None: cls.dims = ["x", "y", "z"] def test_msd_straight_forward(self): - assert np.allclose(self.fft, msd_straight_forward(self.arr1)) + assert_allclose(self.fft, msd_straight_forward(self.arr1)) def test_msd_fft(self): - assert np.allclose(self.fft, msd_fft(self.arr1)) - assert np.allclose(msd_straight_forward(self.arr2), msd_fft(self.arr2)) + assert_allclose(self.fft, msd_fft(self.arr1)) + assert_allclose(msd_straight_forward(self.arr2), msd_fft(self.arr2)) def test_create_position_arrays(self): - assert np.allclose( + assert_allclose( np.array([21.53381769, 14.97501839, -3.87998785]), create_position_arrays(self.gen2, 0, 100, select="type 3")[50][2] ) - assert np.allclose( + assert_allclose( np.array([-2.78550047, -11.85487624, -17.1221954]), create_position_arrays(self.gen2, 0, 100, select="type 3")[99][10] ) - assert np.allclose( + assert_allclose( np.array([41.1079216 , 34.95127106, 18.00482368]), create_position_arrays(self.gen2, 0, 100, select="type 3", center_of_mass=False)[50][2] ) - assert np.allclose( + assert_allclose( np.array([16.98478317, 8.27190208, 5.07116079]), create_position_arrays(self.gen2, 0, 100, select="type 3", center_of_mass=False)[99][10] ) @@ -60,17 +70,17 @@ def test_parse_msd_type(self): xyz = parse_msd_type("xyz") assert ["x", "y", "z"] == self.dims[xyz[0]:xyz[1]:xyz[2]] xy = parse_msd_type("xy") - self.assertEqual(["x", "y"], self.dims[xy[0]:xy[1]:xy[2]]) + assert ["x", "y"] == self.dims[xy[0]:xy[1]:xy[2]] yz = parse_msd_type("yz") - self.assertEqual(["y", "z"], self.dims[yz[0]:yz[1]:yz[2]]) + assert ["y", "z"] == self.dims[yz[0]:yz[1]:yz[2]] xz = parse_msd_type("xz") - self.assertEqual(["x", "z"], self.dims[xz[0]:xz[1]:xz[2]]) + assert ["x", "z"] == self.dims[xz[0]:xz[1]:xz[2]] x = parse_msd_type("x") - self.assertEqual(["x"], self.dims[x[0]:x[1]:x[2]]) + assert ["x"] == self.dims[x[0]:x[1]:x[2]] y = parse_msd_type("y") - self.assertEqual(["y"], self.dims[y[0]:y[1]:y[2]]) + assert ["y"] == self.dims[y[0]:y[1]:y[2]] z = parse_msd_type("z") - self.assertEqual(["z"], self.dims[z[0]:z[1]:z[2]]) + assert ["z"] == self.dims[z[0]:z[1]:z[2]] def test_onsager_ii_self(self): onsager_ii_self_fft = onsager_ii_self(self.gen2, 0, 100, select="type 3") @@ -96,32 +106,32 @@ def test_mda_msd_wrapper(self): assert_allclose(42.69200176568008, mda_msd_anion[50]) assert_allclose(86.9209518, mda_msd_anion[98]) assert_allclose(89.84668178, mda_msd_anion[99]) - assert np.allclose( + assert_allclose( onsager_ii_self(self.gen2, 0, 10, select="type 3", msd_type="x", center_of_mass=False), mda_msd_wrapper(self.gen2, 0, 10, select="type 3", msd_type="x", fft=False), ) - assert np.allclose( + assert_allclose( onsager_ii_self(self.gen2, 0, 10, select="type 3", msd_type="y", center_of_mass=False), mda_msd_wrapper(self.gen2, 0, 10, select="type 3", msd_type="y", fft=False), ) - assert np.allclose( + assert_allclose( onsager_ii_self(self.gen2, 0, 10, select="type 3", msd_type="z", center_of_mass=False), mda_msd_wrapper(self.gen2, 0, 10, select="type 3", msd_type="z", fft=False), ) - assert np.allclose( + assert_allclose( onsager_ii_self(self.gen2, 0, 100, select="type 3", msd_type="xy", center_of_mass=False), mda_msd_wrapper(self.gen2, 0, 100, select="type 3", msd_type="xy", fft=False), ) - assert np.allclose( + assert_allclose( onsager_ii_self(self.gen2, 0, 100, select="type 3", msd_type="yz", center_of_mass=False), mda_msd_wrapper(self.gen2, 0, 100, select="type 3", msd_type="yz", fft=False), ) - assert np.allclose( + assert_allclose( onsager_ii_self(self.gen2, 0, 100, select="type 3", msd_type="xz", center_of_mass=False), mda_msd_wrapper(self.gen2, 0, 100, select="type 3", msd_type="xz", fft=False), ) if td is not None: - assert np.allclose(mda_msd_cation, mda_msd_wrapper(self.gen2, 0, 100, select="type 3")) + assert_allclose(mda_msd_cation, mda_msd_wrapper(self.gen2, 0, 100, select="type 3")) def test_total_msd(self): total_builtin_cation = total_msd(self.gen2, 0, 100, select="type 3", fft=True, built_in=True) @@ -130,7 +140,7 @@ def test_total_msd(self): ) assert_allclose(32.14254152556588, total_builtin_cation[50]) assert_allclose(32.338364098424634, total_mda_cation[50]) - with self.assertRaises(ValueError): + with pytest.raises(ValueError): total_msd(self.gen2, 0, 100, select="type 3", fft=True, built_in=False, center_of_mass=True) From 6d2e8b053f51886143286ae0a873a392cd815ec2 Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 08:43:54 +0800 Subject: [PATCH 13/34] update test_forcefield forms --- tests/test_forcefield.py | 49 ++++++++++++++++++---------------------- 1 file changed, 22 insertions(+), 27 deletions(-) diff --git a/tests/test_forcefield.py b/tests/test_forcefield.py index 8b454fc6..8a0814b2 100644 --- a/tests/test_forcefield.py +++ b/tests/test_forcefield.py @@ -1,5 +1,7 @@ from __future__ import annotations +import os +import shutil import sys import tempfile import unittest @@ -7,9 +9,11 @@ import numpy as np import pytest +from numpy.testing import assert_allclose +from pymatgen.io.lammps.data import LammpsData -from mdgo.forcefield.aqueous import * -from mdgo.forcefield.crawler import * +from mdgo.forcefield.aqueous import Aqueous, Ion +from mdgo.forcefield.crawler import FFcrawler test_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_files") @@ -35,43 +39,34 @@ def test_chrome(self) -> None: lpg = FFcrawler(download_dir, xyz=True, gromacs=True) lpg.data_from_pdb(os.path.join(test_dir, "EMC.pdb")) - self.assertIn( - "LigParGen server connected.\n" - "Structure info uploaded. Rendering force field...\n", - out.getvalue(), - ) - self.assertIn( - "Force field file downloaded.\n" - ".xyz file saved.\n" - "Force field file saved.\n", - out.getvalue(), - ) - self.assertTrue(os.path.exists(os.path.join(download_dir, "EMC.lmp"))) - self.assertTrue(os.path.exists(os.path.join(download_dir, "EMC.lmp.xyz"))) - self.assertTrue(os.path.exists(os.path.join(download_dir, "EMC.gro"))) - self.assertTrue(os.path.exists(os.path.join(download_dir, "EMC.itp"))) + assert "LigParGen server connected.\nStructure info uploaded. Rendering force field...\n" in out.getvalue() + assert "Force field file downloaded.\n.xyz file saved.\nForce field file saved.\n" in out.getvalue() + assert os.path.exists(os.path.join(download_dir, "EMC.lmp")) + assert os.path.exists(os.path.join(download_dir, "EMC.lmp.xyz")) + assert os.path.exists(os.path.join(download_dir, "EMC.gro")) + assert os.path.exists(os.path.join(download_dir, "EMC.itp")) with open(os.path.join(download_dir, "EMC.lmp")) as f: pdf_actual = f.readlines() - self.assertListEqual(pdf, pdf_actual) + assert_allclose(pdf, pdf_actual) with open(os.path.join(download_dir, "EMC.lmp.xyz")) as f: xyz_actual = f.readlines() - self.assertListEqual(xyz, xyz_actual) + assert_allclose(xyz, xyz_actual) with open(os.path.join(download_dir, "EMC.gro")) as f: gro_actual = f.readlines() - self.assertListEqual(gro, gro_actual) + assert_allclose(gro, gro_actual) with open(os.path.join(download_dir, "EMC.itp")) as f: itp_actual = f.readlines() - self.assertListEqual(itp, itp_actual) + assert_allclose(itp, itp_actual) lpg = FFcrawler(download_dir) lpg.data_from_smiles("CCOC(=O)OC") with open(os.path.join(download_dir, "CCOC(=O)OC.lmp")) as f: smiles_actual = f.readlines() - self.assertListEqual(smiles[:13], smiles_actual[:13]) - self.assertListEqual(smiles[18:131], smiles_actual[18:131]) - self.assertEqual(" 1 1 1 -0.28", smiles_actual[131][:26]) - self.assertEqual(" 2 1 2 0.01", smiles_actual[132][:25]) - self.assertEqual(" 15 1 15 0.10", smiles_actual[145][:25]) - self.assertListEqual(smiles_actual[146:], smiles[146:]) + assert_allclose(smiles[:13], smiles_actual[:13]) + assert_allclose(smiles[18:131], smiles_actual[18:131]) + assert smiles_actual[131][:26] == " 1 1 1 -0.28" + assert smiles_actual[132][:25] == " 2 1 2 0.01" + assert smiles_actual[145][:25] == " 15 1 15 0.10" + assert_allclose(smiles_actual[146:], smiles[146:]) finally: sys.stdout = saved_stdout shutil.rmtree(download_dir) From a9dc5d9174108a52f38864fbb1c847d04f33598d Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 09:03:56 +0800 Subject: [PATCH 14/34] update test_conductivity forms --- tests/test_conductivity.py | 29 +++++++++-------------------- 1 file changed, 9 insertions(+), 20 deletions(-) diff --git a/tests/test_conductivity.py b/tests/test_conductivity.py index dcfa89fb..b5f864da 100644 --- a/tests/test_conductivity.py +++ b/tests/test_conductivity.py @@ -4,8 +4,9 @@ import unittest import MDAnalysis +import numpy as np -from mdgo.conductivity import * +from mdgo.conductivity import calc_cond_msd, get_beta test_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_files") @@ -24,27 +25,15 @@ def setUpClass(cls) -> None: cls.time_array = np.array([i * 10 for i in range(cls.gen2.trajectory.n_frames - 100)]) def test_calc_cond_msd(self): - self.assertEqual(-2.9103830456733704e-11, self.cond_array[0]) - self.assertEqual(112.66080481783138, self.cond_array[1]) - self.assertEqual(236007.76624833583, self.cond_array[-1]) + assert self.cond_array[0] == -2.9103830456733704e-11 + assert self.cond_array[1] == 112.66080481783138 + assert self.cond_array[-1] == 236007.76624833583 def test_get_beta(self): - self.assertEqual( - (0.8188201425517928, 0.2535110576154693), - get_beta(self.cond_array, self.time_array, 10, 100), - ) - self.assertEqual( - (1.2525648107674503, 1.0120346984003845), - get_beta(self.cond_array, self.time_array, 1000, 2000), - ) - self.assertEqual( - (1.4075552564189142, 1.3748981878979976), - get_beta(self.cond_array, self.time_array, 1500, 2500), - ) - self.assertEqual( - (1.5021915651236932, 51.79451695748163), - get_beta(self.cond_array, self.time_array, 2000, 4000), - ) + assert get_beta(self.cond_array, self.time_array, 10, 100) == (0.8188201425517928, 0.2535110576154693) + assert get_beta(self.cond_array, self.time_array, 1000, 2000) == (1.2525648107674503, 1.0120346984003845) + assert get_beta(self.cond_array, self.time_array, 1500, 2500) == (1.4075552564189142, 1.3748981878979976) + assert get_beta(self.cond_array, self.time_array, 2000, 4000) == (1.5021915651236932, 51.79451695748163) if __name__ == "__main__": From 362fcb0aad6b1f1f373f461e4c14f0c305b7ee4e Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 09:34:00 +0800 Subject: [PATCH 15/34] ruff fix conductivity --- mdgo/conductivity.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/mdgo/conductivity.py b/mdgo/conductivity.py index ed23b8db..0c76bdcc 100644 --- a/mdgo/conductivity.py +++ b/mdgo/conductivity.py @@ -1,19 +1,21 @@ # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. -""" -This module implements functions to calculate the ionic conductivity. -""" +"""This module implements functions to calculate the ionic conductivity.""" from __future__ import annotations +from typing import TYPE_CHECKING + import numpy as np -from MDAnalysis import AtomGroup, Universe from scipy import stats from tqdm.auto import tqdm from mdgo.msd import msd_fft +if TYPE_CHECKING: + from MDAnalysis import AtomGroup, Universe + __author__ = "Kara Fong, Tingzheng Hou" __version__ = "0.3.0" __maintainer__ = "Tingzheng Hou" @@ -51,7 +53,7 @@ def calc_cond_msd( anion_list = anions.split("residue") # compute sum over all charges and positions qr = [] - for ts in tqdm(u.trajectory[run_start:]): + for _ts in tqdm(u.trajectory[run_start:]): qr_temp = np.zeros(3) for anion in anion_list: qr_temp += anion.center_of_mass() * anion_charge From c81b502f9d1ffa5b4d65c5419c605e7bbd5539f4 Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 09:53:32 +0800 Subject: [PATCH 16/34] ruff fix coordination --- mdgo/coordination.py | 61 +++++++++++++++++--------------------------- 1 file changed, 24 insertions(+), 37 deletions(-) diff --git a/mdgo/coordination.py b/mdgo/coordination.py index 0079489e..dddb7b37 100644 --- a/mdgo/coordination.py +++ b/mdgo/coordination.py @@ -1,23 +1,25 @@ # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. -""" -This module implements functions for coordination analysis. -""" +"""This module implements functions for coordination analysis.""" from __future__ import annotations -from collections.abc import Callable +from typing import TYPE_CHECKING import numpy as np -from MDAnalysis import AtomGroup, Universe from MDAnalysis.analysis.distances import distance_array -from MDAnalysis.core.groups import Atom from scipy.signal import savgol_filter from tqdm.auto import tqdm from mdgo.util.coord import angle, atom_vec +if TYPE_CHECKING: + from collections.abc import Callable + + from MDAnalysis import AtomGroup, Universe + from MDAnalysis.core.groups import Atom + __author__ = "Tingzheng Hou" __version__ = "0.3.0" __maintainer__ = "Tingzheng Hou" @@ -51,12 +53,11 @@ def neighbor_distance( A dictionary of distance of neighbor atoms to the ``center_atom``. The keys are atom indexes in string type . """ dist_dict = {} - time_count = 0 trj_analysis = nvt_run.trajectory[run_start:run_end:] species_selection = select_dict.get(species) if species_selection is None: raise ValueError("Invalid species selection") - for ts in trj_analysis: + for _ts in trj_analysis: selection = ( "(" + species_selection + ") and (around " + str(distance) + " index " + str(center_atom.index) + ")" ) @@ -64,13 +65,10 @@ def neighbor_distance( for atom in shell.atoms: if str(atom.index) not in dist_dict: dist_dict[str(atom.index)] = np.full(run_end - run_start, 100.0) - time_count += 1 - time_count = 0 - for ts in trj_analysis: + for time_count, ts in enumerate(trj_analysis): for atom_index, val in dist_dict.items(): dist = distance_array(ts[center_atom.index], ts[int(atom_index)], ts.dimensions) val[time_count] = dist - time_count += 1 return dist_dict @@ -370,7 +368,7 @@ def check_contiguous_steps( } trj_analysis = nvt_run.trajectory[run_start:run_end:] has = False - for i, ts in enumerate(trj_analysis): + for i, _ts in enumerate(trj_analysis): log = False checkpoint = -1 for j in checkpoints: @@ -680,7 +678,7 @@ def cluster_coordinates( # TODO: rewrite the method cluster = [] for atom in shell: coord_list = [] - for ts in trj_analysis: + for _ts in trj_analysis: coord_list.append(atom.position) cluster.append(np.mean(np.array(coord_list), axis=0)) cluster_array = np.array(cluster) @@ -701,8 +699,7 @@ def cluster_coordinates( # TODO: rewrite the method vec3 = vec3 / np.linalg.norm(vec3) basis_xyz = np.transpose([vec1, vec2, vec3]) cluster_norm = np.linalg.solve(basis_xyz, cluster_array.T).T - cluster_norm = cluster_norm - np.mean(cluster_norm, axis=0) - return cluster_norm + return cluster_norm - np.mean(cluster_norm, axis=0) return cluster_array @@ -739,14 +736,13 @@ def num_of_neighbor( A diction containing the coordination number sequence of each specified neighbor species and the total coordination number sequence in the specified frame range . """ - time_count = 0 trj_analysis = nvt_run.trajectory[run_start:run_end:] cn_values = {} species = list(distance_dict.keys()) for kw in species: cn_values[kw] = np.zeros(int(len(trj_analysis))) cn_values["total"] = np.zeros(int(len(trj_analysis))) - for ts in trj_analysis: + for time_count, ts in enumerate(trj_analysis): digit_of_species = len(species) - 1 for kw in species: selection = select_shell(select_dict, distance_dict, center_atom, kw) @@ -772,7 +768,6 @@ def num_of_neighbor( center_name = center_atom.name path = write_path + str(center_atom.id) + "_" + str(int(ts.time)) + "_" + str(structure_code) + ".xyz" write_out(center_pos, center_name, structure, path) - time_count += 1 return cn_values @@ -800,13 +795,12 @@ def num_of_neighbor_simple( A dict with "total" as the key and an array of the solvation structure type in the specified frame range as the value. """ - time_count = 0 trj_analysis = nvt_run.trajectory[run_start:run_end:] center_selection = "same type as index " + str(center_atom.index) assert len(distance_dict) == 1, "Please only specify the counter-ion species in the distance_dict" - species = list(distance_dict.keys())[0] + species = next(iter(distance_dict.keys())) cn_values = np.zeros(int(len(trj_analysis))) - for ts in trj_analysis: + for time_count, _ts in enumerate(trj_analysis): selection = select_shell(select_dict, distance_dict, center_atom, species) shell = nvt_run.select_atoms(selection, periodic=True) shell_len = len(shell) @@ -822,9 +816,7 @@ def num_of_neighbor_simple( cn_values[time_count] = 3 else: cn_values[time_count] = 3 - time_count += 1 - cn_values = {"total": cn_values} - return cn_values + return {"total": cn_values} def angular_dist_of_neighbor( @@ -860,7 +852,7 @@ def angular_dist_of_neighbor( neighbor_a, neighbor_b, center_c = tuple(names) acb_angle = [] trj_analysis = nvt_run.trajectory[run_start:run_end:] - for ts in trj_analysis: + for _ts in trj_analysis: a_selection = select_shell(select_dict, distance_dict, center_atom, neighbor_a) a_group = nvt_run.select_atoms(a_selection, periodic=True) a_num = len(a_group) @@ -870,10 +862,7 @@ def angular_dist_of_neighbor( c_selection = select_shell(select_dict, distance_dict, a_group.atoms[0], center_c) c_atoms = nvt_run.select_atoms(c_selection, periodic=True) shell_species_len = len(c_atoms) - 1 - if shell_species_len == 0: - shell_type = "cip" - else: - shell_type = "agg" + shell_type = "cip" if shell_species_len == 0 else "agg" else: shell_type = "agg" if shell_type == "agg" and cip: @@ -916,7 +905,6 @@ def num_of_neighbor_specific( A tuple containing three dictionary of the coordination number of each neighbor species and total coordination number for the three solvation structure type, respectively. """ - time_count = 0 trj_analysis = nvt_run.trajectory[run_start:run_end:] cip_step = [] ssip_step = [] @@ -925,7 +913,7 @@ def num_of_neighbor_specific( for kw in distance_dict: cn_values[kw] = np.zeros(int(len(trj_analysis))) cn_values["total"] = np.zeros(int(len(trj_analysis))) - for ts in trj_analysis: + for time_count, _ts in enumerate(trj_analysis): for kw in distance_dict: kw_selection = select_shell(select_dict, distance_dict, center_atom, kw) kw_shell = nvt_run.select_atoms(kw_selection, periodic=True) @@ -948,7 +936,6 @@ def num_of_neighbor_specific( agg_step.append(time_count) else: agg_step.append(time_count) - time_count += 1 cn_dict = {} for kw in distance_dict: cn_dict["ssip_" + kw] = cn_values[kw][ssip_step] @@ -989,7 +976,8 @@ def full_solvation_structure( # TODO: rewrite the method """ center_selection = select_dict.get(center_species) counter_selection = select_dict.get(counter_species) - assert (center_selection is not None) and (counter_selection is not None) + assert center_selection is not None + assert counter_selection is not None def select_counter_ion(selection, dist, atom): return "(" + selection + " and around " + str(dist) + " same fragment as index " + str(atom.index) + ")" @@ -1021,7 +1009,7 @@ def counter_shell(this_shell, this_layer, frame): time_count = 0 trj_analysis = nvt_run.trajectory[run_start:run_end:] cn_values = np.zeros((int(len(trj_analysis)), depth)) - for ts in trj_analysis: + for _ts in trj_analysis: center_ion_list: list[np.int_] = [center_atom.id] counter_ion_list: list[np.int_] = [] first_shell = nvt_run.select_atoms( @@ -1132,5 +1120,4 @@ def select_shell( distance_str = str(distance_value) else: distance_str = distance - selection = "(" + species_selection + ") and (around " + distance_str + " index " + str(center_atom.index) + ")" - return selection + return "(" + species_selection + ") and (around " + distance_str + " index " + str(center_atom.index) + ")" From 082f88424d91d9147d82f5ab37e020fcd6f0829e Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 10:02:44 +0800 Subject: [PATCH 17/34] ruff fixes --- mdgo/__init__.py | 4 +-- mdgo/core/analysis.py | 62 ++++++++++++++------------------------ mdgo/core/run.py | 12 ++------ mdgo/forcefield/aqueous.py | 9 ++---- mdgo/forcefield/charge.py | 8 ++--- mdgo/forcefield/crawler.py | 5 +-- mdgo/forcefield/maestro.py | 3 +- mdgo/forcefield/pubchem.py | 5 +-- mdgo/msd.py | 23 +++++++------- mdgo/residence_time.py | 21 +++++++------ mdgo/util/__init__.py | 2 +- mdgo/util/coord.py | 6 +++- mdgo/util/dict_utils.py | 21 ++++++------- mdgo/util/packmol.py | 12 +++----- mdgo/util/reformat.py | 16 +++------- mdgo/util/volume.py | 24 ++++++--------- tasks.py | 2 +- tests/test_coordination.py | 2 +- tests/test_core.py | 2 +- tests/test_util.py | 2 +- 20 files changed, 94 insertions(+), 147 deletions(-) diff --git a/mdgo/__init__.py b/mdgo/__init__.py index ba1bf5d2..689a18ae 100644 --- a/mdgo/__init__.py +++ b/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 diff --git a/mdgo/core/analysis.py b/mdgo/core/analysis.py index fbc4db3a..3c64012f 100644 --- a/mdgo/core/analysis.py +++ b/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,7 +245,7 @@ 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, @@ -257,7 +253,6 @@ def get_cond_array(self) -> np.ndarray: 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. diff --git a/mdgo/core/run.py b/mdgo/core/run.py index cc30ae6c..9de5e172 100644 --- a/mdgo/core/run.py +++ b/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 diff --git a/mdgo/forcefield/aqueous.py b/mdgo/forcefield/aqueous.py index c071d4db..44f34615 100644 --- a/mdgo/forcefield/aqueous.py +++ b/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")) diff --git a/mdgo/forcefield/charge.py b/mdgo/forcefield/charge.py index cc26a250..ae3b30d9 100644 --- a/mdgo/forcefield/charge.py +++ b/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 diff --git a/mdgo/forcefield/crawler.py b/mdgo/forcefield/crawler.py index 27596f7d..a50b66ca 100644 --- a/mdgo/forcefield/crawler.py +++ b/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): diff --git a/mdgo/forcefield/maestro.py b/mdgo/forcefield/maestro.py index c1a61d09..8b35eb6a 100644 --- a/mdgo/forcefield/maestro.py +++ b/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 diff --git a/mdgo/forcefield/pubchem.py b/mdgo/forcefield/pubchem.py index da8617f5..e35fcc74 100644 --- a/mdgo/forcefield/pubchem.py +++ b/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() diff --git a/mdgo/msd.py b/mdgo/msd.py index 035ac157..04e3c442 100644 --- a/mdgo/msd.py +++ b/mdgo/msd.py @@ -11,7 +11,7 @@ from __future__ import annotations -from typing import Literal +from typing import TYPE_CHECKING, Literal try: import MDAnalysis.analysis.msd as mda_msd @@ -23,10 +23,12 @@ td = None import numpy as np -from MDAnalysis import AtomGroup, Universe -from MDAnalysis.core.groups import Atom from tqdm.auto import trange +if TYPE_CHECKING: + from MDAnalysis import AtomGroup, Universe + from MDAnalysis.core.groups import Atom + __author__ = "Tingzheng Hou" __version__ = "0.3.0" __maintainer__ = "Tingzheng Hou" @@ -169,11 +171,11 @@ def create_position_arrays( atom_group = nvt_run.select_atoms(select) atom_positions = np.zeros((end - start, len(atom_group), 3)) if center_of_mass: - for ts in nvt_run.trajectory[start:end]: + for _ts in nvt_run.trajectory[start:end]: atom_positions[time, :, :] = atom_group.positions - nvt_run.atoms.center_of_mass() time += 1 else: - for ts in nvt_run.trajectory[start:end]: + for _ts in nvt_run.trajectory[start:end]: atom_positions[time, :, :] = atom_group.positions time += 1 return atom_positions @@ -223,8 +225,7 @@ def onsager_ii_self( r = atom_positions[:, atom_num, dim[0] : dim[1] : dim[2]] msd_temp = msd_straight_forward(np.array(r)) # [start:end] ii_self += msd_temp - msd = np.array(ii_self) / n_atoms - return msd + return np.array(ii_self) / n_atoms def mda_msd_wrapper( @@ -311,8 +312,7 @@ def _total_msd(nvt_run: Universe, run_start: int, run_end: int, select: str = "a current_coord = ts[li_atom.id - 1] coords.append(current_coord) all_list.append(np.array(coords)) - total_array = msd_from_frags(all_list, run_end - run_start - 1) - return total_array + return msd_from_frags(all_list, run_end - run_start - 1) def msd_from_frags(coord_list: list[np.ndarray], largest: int) -> np.ndarray: @@ -345,8 +345,7 @@ def msd_from_frags(coord_list: list[np.ndarray], largest: int) -> np.ndarray: assert msds is not None msds_by_state[kw] = msds.mean() timeseries.append(msds_by_state[kw]) - timeseries = np.array(timeseries) - return timeseries + return np.array(timeseries) def states_coord_array( @@ -361,7 +360,7 @@ def states_coord_array( """Cuts the trajectory of an atom into fragments. Each fragment contains consecutive timesteps of coordinates of the atom in either attached or free state. The Attached state is when the atom coordinates with the ``binding_site`` species (distance < ``distance``), and vice versa for the free state. - TODO: check if need wrapped trj + TODO: check if need wrapped trj. Args: nvt_run: An MDAnalysis ``Universe`` containing unwrapped trajectory. diff --git a/mdgo/residence_time.py b/mdgo/residence_time.py index ce26678e..c3a6a843 100644 --- a/mdgo/residence_time.py +++ b/mdgo/residence_time.py @@ -1,22 +1,23 @@ # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. -""" -This module calculates species correlation lifetime (residence time). -""" +"""This module calculates species correlation lifetime (residence time).""" from __future__ import annotations import os +from typing import TYPE_CHECKING import matplotlib.pyplot as plt import numpy as np -from MDAnalysis import Universe -from MDAnalysis.core.groups import Atom from scipy.optimize import curve_fit from statsmodels.tsa.stattools import acovf from tqdm.auto import tqdm +if TYPE_CHECKING: + from MDAnalysis import Universe + from MDAnalysis.core.groups import Atom + __author__ = "Kara Fong, Tingzheng Hou" __version__ = "0.3.0" __maintainer__ = "Tingzheng Hou" @@ -51,7 +52,7 @@ def neighbors_one_atom( """ bool_values = {} time_count = 0 - for ts in nvt_run.trajectory[run_start:run_end:]: + for _ts in nvt_run.trajectory[run_start:run_end:]: if species in select_dict: selection = ( "(" @@ -75,7 +76,7 @@ def neighbors_one_atom( def calc_acf(a_values: dict[str, np.ndarray]) -> list[np.ndarray]: """ - Calculate auto-correlation function (ACF) + Calculate auto-correlation function (ACF). Args: a_values: A dict of adjacency matrix with neighbor atom id as keys and arrays @@ -98,7 +99,7 @@ def exponential_func( c: float | np.floating | np.ndarray, ) -> np.floating | np.ndarray: """ - An exponential decay function + An exponential decay function. Args: x: Independent variable. @@ -140,7 +141,7 @@ def calc_neigh_corr( times = [] step = 0 center_atoms = nvt_run.select_atoms(select_dict[center_atom]) - for ts in nvt_run.trajectory[run_start:run_end]: + for _ts in nvt_run.trajectory[run_start:run_end]: times.append(step * time_step) step += 1 times = np.array(times) @@ -176,7 +177,7 @@ def fit_residence_time( ) -> dict[str, np.floating]: """ Use the ACF to fit the residence time (Exponential decay constant). - TODO: allow defining the residence time according to a threshold value of the decay + TODO: allow defining the residence time according to a threshold value of the decay. Args: times: A time series. diff --git a/mdgo/util/__init__.py b/mdgo/util/__init__.py index 46ca2dd3..05c04b5a 100644 --- a/mdgo/util/__init__.py +++ b/mdgo/util/__init__.py @@ -13,7 +13,7 @@ from typing import Dict, Final -MM_of_Elements: Final[Dict[str, float]] = { +MM_of_Elements: Final[dict[str, float]] = { "H": 1.00794, "He": 4.002602, "Li": 6.941, diff --git a/mdgo/util/coord.py b/mdgo/util/coord.py index 1db0b6a8..32570f0b 100644 --- a/mdgo/util/coord.py +++ b/mdgo/util/coord.py @@ -5,8 +5,12 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import numpy as np -from MDAnalysis.core.groups import Atom + +if TYPE_CHECKING: + from MDAnalysis.core.groups import Atom def atom_vec(atom1: Atom, atom2: Atom, dimension: np.ndarray) -> np.ndarray: diff --git a/mdgo/util/dict_utils.py b/mdgo/util/dict_utils.py index d5eb8f79..bcc80195 100644 --- a/mdgo/util/dict_utils.py +++ b/mdgo/util/dict_utils.py @@ -8,15 +8,18 @@ import math import re import string +from typing import TYPE_CHECKING import numpy as np -import pandas as pd -from MDAnalysis import Universe -from MDAnalysis.core.groups import AtomGroup, Residue from pymatgen.io.lammps.data import CombinedData from . import MM_of_Elements +if TYPE_CHECKING: + import pandas as pd + from MDAnalysis import Universe + from MDAnalysis.core.groups import AtomGroup, Residue + def mass_to_name(masses: np.ndarray) -> np.ndarray: """ @@ -247,10 +250,7 @@ def extract_atom_from_ion(positive: bool, ion: Residue | AtomGroup, select_dict: number: The serial number of the ion. """ if positive: - if number == 0: - cation_name = "cation" - else: - cation_name = "cation_" + str(number) + cation_name = "cation" if number == 0 else "cation_" + str(number) if len(ion.atoms.types) == 1: select_dict[cation_name] = "type " + ion.atoms.types[0] else: @@ -265,10 +265,7 @@ def extract_atom_from_ion(positive: bool, ion: Residue | AtomGroup, select_dict: select_dict[cation_name + "_" + pos_center.name + pos_center.type] = "type " + pos_center.type select_dict[cation_name] = "type " + uni_center else: - if number == 0: - anion_name = "anion" - else: - anion_name = "anion_" + str(number) + anion_name = "anion" if number == 0 else "anion_" + str(number) if len(ion.atoms.types) == 1: select_dict[anion_name] = "type " + ion.atoms.types[0] else: @@ -288,7 +285,7 @@ def extract_atom_from_molecule( resname: str, molecule: Residue | AtomGroup, select_dict: dict[str, str], number: int = 0 ): """ - Assign the most negatively charged atom in the molecule into select_dict + Assign the most negatively charged atom in the molecule into select_dict. Args: resname: The name of the molecule diff --git a/mdgo/util/packmol.py b/mdgo/util/packmol.py index d0f8fb36..9d7fa175 100644 --- a/mdgo/util/packmol.py +++ b/mdgo/util/packmol.py @@ -119,8 +119,7 @@ def run_packmol(self, timeout=30): check=True, shell=True, timeout=timeout, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, + capture_output=True, ) # this workaround is needed because packmol can fail to find # a solution but still return a zero exit code @@ -150,11 +149,8 @@ def make_packmol_input(self): else: # estimate the total volume of all molecules net_volume = 0.0 - for idx, d in enumerate(self.molecules): - if not isinstance(d["coords"], Molecule): - mol = Molecule.from_file(d["coords"]) - else: - mol = d["coords"] + for _idx, d in enumerate(self.molecules): + mol = Molecule.from_file(d["coords"]) if not isinstance(d["coords"], Molecule) else d["coords"] # molecular volume in cubic Å vol = molecular_volume(mol, radii_type="pymatgen", molar_volume=False) # pad the calculated length by an amount related to the tolerance parameter @@ -183,7 +179,7 @@ def make_packmol_input(self): else: out.write(f"output {self.output}\n\n") - for i, d in enumerate(self.molecules): + for _i, d in enumerate(self.molecules): if isinstance(d["coords"], str): if " " in d["coords"]: out.write(f'structure "{d["coords"]}"\n') diff --git a/mdgo/util/reformat.py b/mdgo/util/reformat.py index 2a78f03c..5583306b 100644 --- a/mdgo/util/reformat.py +++ b/mdgo/util/reformat.py @@ -1,9 +1,7 @@ # Copyright (c) Tingzheng Hou. # Distributed under the terms of the MIT License. -""" -Utilities for converting data file formats. -""" +"""Utilities for converting data file formats.""" from __future__ import annotations @@ -120,7 +118,7 @@ def ff_parser(ff_dir: str, xyz_dir: str) -> str: counts = {} counts["atoms"] = len(dfs["atoms"].index) mass_list = [] - for index, row in dfs["atoms"].iterrows(): + for _index, row in dfs["atoms"].iterrows(): mass_list.append(MM_of_Elements.get(re.split(r"(\d+)", row["atom"])[0])) mass_df = pd.DataFrame(mass_list) mass_df.index += 1 @@ -175,15 +173,14 @@ def ff_parser(ff_dir: str, xyz_dir: str) -> str: stats_template = "{:>" + str(max_stats) + "} {}" count_lines = [stats_template.format(v, k) for k, v in counts.items()] type_lines = [stats_template.format(v, k[:-1] + " types") for k, v in counts.items()] - stats = "\n".join(count_lines + [""] + type_lines) + stats = "\n".join([*count_lines, "", *type_lines]) header = [ f"LAMMPS data file created by mdgo (by {__author__})\n" "# OPLS force field: harmonic, harmonic, opls, cvff", stats, BOX.format(lo, hi), ] - data_string = "\n\n".join(header + masses + ff + topo) + "\n" - return data_string + return "\n\n".join(header + masses + ff + topo) + "\n" def sdf_to_pdb( @@ -209,10 +206,7 @@ def sdf_to_pdb( with open(sdf_file) as inp: sdf_lines = inp.readlines() sdf = list(map(str.strip, sdf_lines)) - if pubchem: - title = "cid_" - else: - title = "" + title = "cid_" if pubchem else "" pdb_atoms: list[dict[str, Any]] = [] # create pdb list of dictionaries atoms = 0 diff --git a/mdgo/util/volume.py b/mdgo/util/volume.py index d6cef08f..645dce3c 100644 --- a/mdgo/util/volume.py +++ b/mdgo/util/volume.py @@ -285,7 +285,7 @@ def parse_command_line(): def get_max_dimensions(mol: Molecule) -> tuple[float, float, float, float, float, float]: """ - Calculates the dimension of a Molecule + Calculates the dimension of a Molecule. Args: mol: A Molecule object. @@ -373,7 +373,7 @@ def round_dimensions( def dsq(a1: float, a2: float, a3: float, b1: float, b2: float, b3: float) -> float: """ - Squared distance between a and b + Squared distance between a and b. Args: a1: x coordinate of a @@ -386,8 +386,7 @@ def dsq(a1: float, a2: float, a3: float, b1: float, b2: float, b3: float) -> flo Returns: squared distance """ - d2 = (b1 - a1) ** 2 + (b2 - a2) ** 2 + (b3 - a3) ** 2 - return d2 + return (b1 - a1) ** 2 + (b2 - a2) ** 2 + (b3 - a3) ** 2 def get_dimensions( @@ -431,8 +430,7 @@ def make_matrix(x_num: int, y_num: int, z_num: int) -> np.ndarray: Returns: matrix """ - matrix = np.array([[[None for _ in range(z_num)] for _ in range(y_num)] for _ in range(x_num)]) - return matrix + return np.array([[[None for _ in range(z_num)] for _ in range(y_num)] for _ in range(x_num)]) def get_radii(radii_type: str = "Bondi") -> dict[str, float]: @@ -530,9 +528,8 @@ def fill_volume_matrix( for a in mol.sites: element = str(a.species.elements[0]) - if exclude_h: - if element == "H": - continue + if exclude_h and element == "H": + continue radius = radii.get(element, DEFAULT_VDW) for i in range(xsteps): if abs(a.x - (x0 + 0.5 * res + i * res)) < radius: @@ -615,7 +612,7 @@ def molecular_volume( z_size: float = 10.0, ) -> float: """ - Estimate the molar volume in cm^3/mol or volume in Å^3 + Estimate the molar volume in cm^3/mol or volume in Å^3. Args: mol: Molecule object or path to .xyz or other file that can be read @@ -644,10 +641,7 @@ def molecular_volume( Returns: The molar volume in cm^3/mol or volume in Å^3. """ - if isinstance(mol, str): - molecule = Molecule.from_file(mol) - else: - molecule = mol + molecule = Molecule.from_file(mol) if isinstance(mol, str) else mol if mode == "lig": print("Calculating occupied volume...") x_min, x_max, y_min, y_max, z_min, z_max = get_max_dimensions(molecule) @@ -682,7 +676,7 @@ def concentration_matcher( Estimate the number of molecules of each species in a box, given the salt concentration, salt type, solvent molecular weight, solvent density, solvent ratio and total number of salt. - TODO: Auto box size according to Debye screening length + TODO: Auto box size according to Debye screening length. Args: concentration: Salt concentration in mol/L. diff --git a/tasks.py b/tasks.py index c0d226e2..d3f563d8 100644 --- a/tasks.py +++ b/tasks.py @@ -191,7 +191,7 @@ def update_changelog(ctx, version, sim=False): ll = ll.strip() if ll in ["", "## Summary"]: continue - elif ll.startswith("## Checklist") or ll.startswith("## TODO"): + elif ll.startswith(("## Checklist", "## TODO")): break lines.append(f" {ll}") misc.append(l) diff --git a/tests/test_coordination.py b/tests/test_coordination.py index 20c90463..cf1cc95f 100644 --- a/tests/test_coordination.py +++ b/tests/test_coordination.py @@ -5,7 +5,7 @@ class MyTestCase(unittest.TestCase): def test_something(self): - self.assertEqual(True, True) + assert True is True if __name__ == "__main__": diff --git a/tests/test_core.py b/tests/test_core.py index 20c90463..cf1cc95f 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -5,7 +5,7 @@ class MyTestCase(unittest.TestCase): def test_something(self): - self.assertEqual(True, True) + assert True is True if __name__ == "__main__": diff --git a/tests/test_util.py b/tests/test_util.py index 20c90463..cf1cc95f 100644 --- a/tests/test_util.py +++ b/tests/test_util.py @@ -5,7 +5,7 @@ class MyTestCase(unittest.TestCase): def test_something(self): - self.assertEqual(True, True) + assert True is True if __name__ == "__main__": From 73d8a331870c69aada293f98fd766dcf212ca495 Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 10:28:19 +0800 Subject: [PATCH 18/34] ruff fixes 2 --- mdgo/core/analysis.py | 2 +- mdgo/core/run.py | 3 +- mdgo/forcefield/aqueous.py | 2 +- mdgo/msd.py | 6 +-- mdgo/residence_time.py | 12 ++---- mdgo/util/__init__.py | 2 +- mdgo/util/dict_utils.py | 84 +++++++++++++++++++++----------------- mdgo/util/reformat.py | 2 +- mdgo/util/volume.py | 2 +- 9 files changed, 61 insertions(+), 54 deletions(-) diff --git a/mdgo/core/analysis.py b/mdgo/core/analysis.py index 3c64012f..ed24532d 100644 --- a/mdgo/core/analysis.py +++ b/mdgo/core/analysis.py @@ -730,7 +730,7 @@ def get_msd_all( Args: start: Start time step. end: End time step. - msd_type: Desired dimensions to be included in the MSD. Defaults to ‘xyz’. + msd_type: Desired dimensions to be included in the MSD. Defaults to "xyz". fft: Whether to use FFT to accelerate the calculation. Default to True. built_in: Whether to use built in method to calculate msd or use function from mds. Default to True. diff --git a/mdgo/core/run.py b/mdgo/core/run.py index 9de5e172..ed1abec3 100644 --- a/mdgo/core/run.py +++ b/mdgo/core/run.py @@ -18,6 +18,7 @@ def from_dict(cls): Constructor. Returns: + name: The name of the class """ return cls("name") @@ -28,6 +29,6 @@ def from_recipe(cls): Constructor. Returns: - + name: The name of the class """ return cls("name") diff --git a/mdgo/forcefield/aqueous.py b/mdgo/forcefield/aqueous.py index 44f34615..27e766ee 100644 --- a/mdgo/forcefield/aqueous.py +++ b/mdgo/forcefield/aqueous.py @@ -177,7 +177,7 @@ def get_ion( Sachini et al., Systematic Comparison of the Structural and Dynamic Properties of Commonly Used Water Models for Molecular Dynamics Simulations. J. Chem. Inf. Model. - 2021, 61, 9, 4521–4536. https://doi.org/10.1021/acs.jcim.1c00794 + 2021, 61, 9, 4521-4536. https://doi.org/10.1021/acs.jcim.1c00794 mixing_rule: The mixing rule to use for the ion parameter. Default to None, which does not change the original mixing rule of the parameter set. Available choices are 'LB' diff --git a/mdgo/msd.py b/mdgo/msd.py index 04e3c442..b0c4e97a 100644 --- a/mdgo/msd.py +++ b/mdgo/msd.py @@ -57,7 +57,7 @@ def total_msd( start: Start frame of analysis. end: End frame of analysis. select: A selection string. Defaults to “all” in which case all atoms are selected. - msd_type: Desired dimensions to be included in the MSD. Defaults to ‘xyz’. + msd_type: Desired dimensions to be included in the MSD. Defaults to "xyz". fft: Whether to use FFT to accelerate the calculation. Default to True. built_in: Whether to use built in method to calculate msd or use function from mds. Default to True. center_of_mass: Whether to subtract center of mass at each step for atom coordinates. Default to True. @@ -199,7 +199,7 @@ def onsager_ii_self( start: Start frame of analysis. end: End frame of analysis. select: A selection string. Defaults to “all” in which case all atoms are selected. - msd_type: Desired dimensions to be included in the MSD. Defaults to ‘xyz’. + msd_type: Desired dimensions to be included in the MSD. Defaults to "xyz". center_of_mass: Whether to subtract center of mass at each step for atom coordinates. Default to True. fft: Whether to use FFT to accelerate the calculation. Default to True. @@ -239,7 +239,7 @@ def mda_msd_wrapper( start: Start frame of analysis. end: End frame of analysis. select: A selection string. Defaults to “all” in which case all atoms are selected. - msd_type: Desired dimensions to be included in the MSD. Defaults to ‘xyz’. + msd_type: Desired dimensions to be included in the MSD. Defaults to "xyz". fft: Whether to use FFT to accelerate the calculation. Default to True. Warning: diff --git a/mdgo/residence_time.py b/mdgo/residence_time.py index c3a6a843..cd1814d6 100644 --- a/mdgo/residence_time.py +++ b/mdgo/residence_time.py @@ -51,8 +51,7 @@ def neighbors_one_atom( A neighbor dict with neighbor atom id as keys and arrays of adjacent boolean (0/1) as values. """ bool_values = {} - time_count = 0 - for _ts in nvt_run.trajectory[run_start:run_end:]: + for time_count, _ts in enumerate(nvt_run.trajectory[run_start:run_end:]): if species in select_dict: selection = ( "(" @@ -70,7 +69,6 @@ def neighbors_one_atom( if str(atom.id) not in bool_values: bool_values[str(atom.id)] = np.zeros(int((run_end - run_start) / 1)) bool_values[str(atom.id)][time_count] = 1 - time_count += 1 return bool_values @@ -86,8 +84,8 @@ def calc_acf(a_values: dict[str, np.ndarray]) -> list[np.ndarray]: A list of auto-correlation functions for each neighbor species. """ acfs = [] - for atom_id, neighbors in a_values.items(): - # atom_id_numeric = int(re.search(r"\d+", atom_id).group()) + for _atom_id, neighbors in a_values.items(): + # atom_id_numeric = int(re.search(r"\d+", _atom_id).group()) acfs.append(acovf(neighbors, demean=False, unbiased=True, fft=True)) return acfs @@ -139,11 +137,9 @@ def calc_neigh_corr( """ # Set up times array times = [] - step = 0 center_atoms = nvt_run.select_atoms(select_dict[center_atom]) - for _ts in nvt_run.trajectory[run_start:run_end]: + for step, _ts in enumerate(nvt_run.trajectory[run_start:run_end]): times.append(step * time_step) - step += 1 times = np.array(times) acf_avg = {} diff --git a/mdgo/util/__init__.py b/mdgo/util/__init__.py index 05c04b5a..88100714 100644 --- a/mdgo/util/__init__.py +++ b/mdgo/util/__init__.py @@ -11,7 +11,7 @@ __email__ = "tingzheng_hou@berkeley.edu" __date__ = "Jul 19, 2021" -from typing import Dict, Final +from typing import Final MM_of_Elements: Final[dict[str, float]] = { "H": 1.00794, diff --git a/mdgo/util/dict_utils.py b/mdgo/util/dict_utils.py index bcc80195..2ff75dd0 100644 --- a/mdgo/util/dict_utils.py +++ b/mdgo/util/dict_utils.py @@ -209,9 +209,9 @@ def select_dict_from_resname(u: Universe) -> dict[str, str]: for i, frag in enumerate(residue.atoms.fragments): charge = np.sum(frag.charges) if charge > 0.001: - extract_atom_from_ion(True, frag, select_dict) + extract_atom_from_cation(frag, select_dict) elif charge < -0.001: - extract_atom_from_ion(False, frag, select_dict) + extract_atom_from_anion(frag, select_dict) else: extract_atom_from_molecule(resname, frag, select_dict, number=i + 1) elif len(residue.atoms.fragments) >= 2: @@ -221,10 +221,10 @@ def select_dict_from_resname(u: Universe) -> dict[str, str]: for frag in residue.atoms.fragments: charge = np.sum(frag.charges) if charge > 0.001: - extract_atom_from_ion(True, frag, select_dict, cation_number) + extract_atom_from_cation(frag, select_dict, cation_number) cation_number += 1 elif charge < -0.001: - extract_atom_from_ion(False, frag, select_dict, anion_number) + extract_atom_from_anion(frag, select_dict, anion_number) anion_number += 1 else: extract_atom_from_molecule(resname, frag, select_dict, molecule_number) @@ -232,53 +232,63 @@ def select_dict_from_resname(u: Universe) -> dict[str, str]: else: extract_atom_from_molecule(resname, residue, select_dict) elif residue.charge > 0: - extract_atom_from_ion(True, residue, select_dict) + extract_atom_from_cation(residue, select_dict) else: - extract_atom_from_ion(False, residue, select_dict) + extract_atom_from_anion(residue, select_dict) return select_dict -def extract_atom_from_ion(positive: bool, ion: Residue | AtomGroup, select_dict: dict[str, str], number: int = 0): +def extract_atom_from_cation(ion: Residue | AtomGroup, select_dict: dict[str, str], number: int = 0): """ - Assign the most most charged atom and/or one unique atom in the ion into select_dict. + Assign the most charged atom and/or one unique atom in the cation into select_dict. Args: - positive: Whether the charge of ion is positive. Otherwise negative. Default to True. ion: Residue or AtomGroup select_dict: A dictionary of atom species, where each atom species name is a key and the corresponding values are the selection language. - number: The serial number of the ion. + number: The serial number of the cation. """ - if positive: - cation_name = "cation" if number == 0 else "cation_" + str(number) - if len(ion.atoms.types) == 1: - select_dict[cation_name] = "type " + ion.atoms.types[0] + cation_name = "cation" if number == 0 else "cation_" + str(number) + if len(ion.atoms.types) == 1: + select_dict[cation_name] = "type " + ion.atoms.types[0] + else: + # The most positively charged atom in the cation + pos_center = ion.atoms[np.argmax(ion.atoms.charges)] + unique_types = np.unique(ion.atoms.types, return_counts=True) + # One unique atom in the cation + uni_center = unique_types[0][np.argmin(unique_types[1])] + if pos_center.type == uni_center: + select_dict[cation_name] = "type " + uni_center else: - # The most positively charged atom in the cation - pos_center = ion.atoms[np.argmax(ion.atoms.charges)] - unique_types = np.unique(ion.atoms.types, return_counts=True) - # One unique atom in the cation - uni_center = unique_types[0][np.argmin(unique_types[1])] - if pos_center.type == uni_center: - select_dict[cation_name] = "type " + uni_center - else: - select_dict[cation_name + "_" + pos_center.name + pos_center.type] = "type " + pos_center.type - select_dict[cation_name] = "type " + uni_center + select_dict[cation_name + "_" + pos_center.name + pos_center.type] = "type " + pos_center.type + select_dict[cation_name] = "type " + uni_center + + +def extract_atom_from_anion(ion: Residue | AtomGroup, select_dict: dict[str, str], number: int = 0): + """ + Assign the most charged atom and/or one unique atom in the anion into select_dict. + + Args: + ion: Residue or AtomGroup + select_dict: A dictionary of atom species, where each atom species name is a key + and the corresponding values are the selection language. + number: The serial number of the anion. + """ + + anion_name = "anion" if number == 0 else "anion_" + str(number) + if len(ion.atoms.types) == 1: + select_dict[anion_name] = "type " + ion.atoms.types[0] else: - anion_name = "anion" if number == 0 else "anion_" + str(number) - if len(ion.atoms.types) == 1: - select_dict[anion_name] = "type " + ion.atoms.types[0] + # The most negatively charged atom in the anion + neg_center = ion.atoms[np.argmin(ion.atoms.charges)] + unique_types = np.unique(ion.atoms.types, return_counts=True) + # One unique atom in the anion + uni_center = unique_types[0][np.argmin(unique_types[1])] + if neg_center.type == uni_center: + select_dict[anion_name] = "type " + uni_center else: - # The most negatively charged atom in the anion - neg_center = ion.atoms[np.argmin(ion.atoms.charges)] - unique_types = np.unique(ion.atoms.types, return_counts=True) - # One unique atom in the anion - uni_center = unique_types[0][np.argmin(unique_types[1])] - if neg_center.type == uni_center: - select_dict[anion_name] = "type " + uni_center - else: - select_dict[anion_name + "_" + neg_center.name + neg_center.type] = "type " + neg_center.type - select_dict[anion_name] = "type " + uni_center + select_dict[anion_name + "_" + neg_center.name + neg_center.type] = "type " + neg_center.type + select_dict[anion_name] = "type " + uni_center def extract_atom_from_molecule( diff --git a/mdgo/util/reformat.py b/mdgo/util/reformat.py index 5583306b..616390d0 100644 --- a/mdgo/util/reformat.py +++ b/mdgo/util/reformat.py @@ -303,7 +303,7 @@ def sdf_to_pdb( bond_lines[atom].append(atom2s[i]) for i, atom in enumerate(atom2s): bond_lines[atom].append(atom1s[i]) - for i, odr in enumerate(orders): + for _i, odr in enumerate(orders): for j, ln in enumerate(bond_lines): if ln[0] == odr[0]: bond_lines.insert(j + 1, odr) diff --git a/mdgo/util/volume.py b/mdgo/util/volume.py index 645dce3c..923fd073 100644 --- a/mdgo/util/volume.py +++ b/mdgo/util/volume.py @@ -665,7 +665,7 @@ def molecular_volume( def concentration_matcher( concentration: float, - salt: float | int | str | Molecule, + salt: float | str | Molecule, solvents: list[str | dict[str, float]], solv_ratio: list[float], num_salt: int = 100, From 80325cfd3278deb035b409d2bdde324ca806df11 Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 10:29:16 +0800 Subject: [PATCH 19/34] ruff doc conf --- docs/source/conf.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index b61f8c72..44522a11 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -1,4 +1,4 @@ -# Configuration file for the Sphinx documentation builder. +"""Configuration file for the Sphinx documentation builder.""" # # This file only contains a selection of the most common options. For a full # list see the documentation: @@ -14,8 +14,6 @@ import os import sys -import sphinx_rtd_theme, sphinx_autodoc_typehints -from typing import List sys.path.insert(0, os.path.abspath("../../mdgo")) @@ -57,7 +55,7 @@ # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. # This pattern also affects html_static_path and html_extra_path. -exclude_patterns: List[str] = [] +exclude_patterns: list[str] = [] # -- Options for HTML output ------------------------------------------------- From a6cc3fd766af8ed9c78e3637d03007d1aa267f4a Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 10:40:11 +0800 Subject: [PATCH 20/34] more ruff fixes --- mdgo/forcefield/pubchem.py | 8 +++----- mdgo/residence_time.py | 2 +- mdgo/util/dict_utils.py | 1 - mdgo/util/packmol.py | 10 +++++----- 4 files changed, 9 insertions(+), 12 deletions(-) diff --git a/mdgo/forcefield/pubchem.py b/mdgo/forcefield/pubchem.py index e35fcc74..9344331b 100644 --- a/mdgo/forcefield/pubchem.py +++ b/mdgo/forcefield/pubchem.py @@ -10,7 +10,7 @@ import os import time -from typing import Final +from typing import Final, NoReturn from urllib.parse import quote import pubchempy as pcp @@ -103,15 +103,13 @@ def obtain_entry(self, search_text: str, name: str, output_format: str = "sdf") return self._obtain_entry_api(search_text, name, output_format=output_format) return self._obtain_entry_web(search_text, name, output_format=output_format) - def smiles_to_pdb(self, smiles: str): + def smiles_to_pdb(self, smiles: str) -> NoReturn: """ Obtain pdf file based on SMILES code. Args: smiles: SMILES code. - Returns: - """ convertor_url = "https://cactus.nci.nih.gov/translate/" input_xpath = "/html/body/div/div[2]/div[1]/form/table[1]/tbody/tr[2]/td[1]/input[1]" @@ -141,7 +139,7 @@ def _obtain_entry_web(self, search_text: str, name: str, output_format: str) -> self.web.get(url) loaded_element_path = '//*[@id="main-results"]/div[1]/div/ul' self.wait.until(EC.presence_of_element_located((By.XPATH, loaded_element_path))) - best_xpath = '//*[@id="featured-results"]/div/div[2]' "/div/div[1]/div[2]/div[1]/a/span/span" + best_xpath = '//*[@id="featured-results"]/div/div[2]/div/div[1]/div[2]/div[1]/a/span/span' relevant_xpath = ( '//*[@id="collection-results-container"]' "/div/div/div[2]/ul/li[1]/div/div/div[1]" diff --git a/mdgo/residence_time.py b/mdgo/residence_time.py index cd1814d6..876928ba 100644 --- a/mdgo/residence_time.py +++ b/mdgo/residence_time.py @@ -84,7 +84,7 @@ def calc_acf(a_values: dict[str, np.ndarray]) -> list[np.ndarray]: A list of auto-correlation functions for each neighbor species. """ acfs = [] - for _atom_id, neighbors in a_values.items(): + for neighbors in a_values.values(): # for _atom_id, neighbors in a_values.items(): # atom_id_numeric = int(re.search(r"\d+", _atom_id).group()) acfs.append(acovf(neighbors, demean=False, unbiased=True, fft=True)) return acfs diff --git a/mdgo/util/dict_utils.py b/mdgo/util/dict_utils.py index 2ff75dd0..2dec52c1 100644 --- a/mdgo/util/dict_utils.py +++ b/mdgo/util/dict_utils.py @@ -274,7 +274,6 @@ def extract_atom_from_anion(ion: Residue | AtomGroup, select_dict: dict[str, str and the corresponding values are the selection language. number: The serial number of the anion. """ - anion_name = "anion" if number == 0 else "anion_" + str(number) if len(ion.atoms.types) == 1: select_dict[anion_name] = "type " + ion.atoms.types[0] diff --git a/mdgo/util/packmol.py b/mdgo/util/packmol.py index 9d7fa175..67a38d6c 100644 --- a/mdgo/util/packmol.py +++ b/mdgo/util/packmol.py @@ -71,17 +71,17 @@ def __init__( 2. "number" - the number of that molecule to pack into the box 3. "coords" - Coordinates in the form of either a Molecule object or a path to a file. - - Example: - {"name": "water", - "number": 500, - "coords": "/path/to/input/file.xyz"} + For Example, + {"name": "water", + "number": 500, + "coords": "/path/to/input/file.xyz"} box: A list of box dimensions xlo, ylo, zlo, xhi, yhi, zhi, in Å. If set to None (default), mdgo will estimate the required box size based on the volumes of the provided molecules using mdgo.volume.molecular_volume() tolerance: Tolerance for packmol, in Å. seed: Random seed for packmol. Use a value of 1 (default) for deterministic output, or -1 to generate a new random seed from the current time. + control_params: Specify custom control parapeters, e,g, "maxit" and "nloop", in a dict inputfile: Path to the input file. Default to 'packmol.inp'. outputfile: Path to the output file. Default to 'output.xyz'. """ From 4ab775d66ec78ddfd2e6c6ce38b5cc1712650581 Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 10:56:59 +0800 Subject: [PATCH 21/34] more ruff fixes 2 --- mdgo/coordination.py | 1 + mdgo/msd.py | 2 +- mdgo/residence_time.py | 9 ++++----- tasks.py | 19 +++++++++++-------- 4 files changed, 17 insertions(+), 14 deletions(-) diff --git a/mdgo/coordination.py b/mdgo/coordination.py index dddb7b37..c142ae2e 100644 --- a/mdgo/coordination.py +++ b/mdgo/coordination.py @@ -1043,6 +1043,7 @@ def concat_coord_array( and the corresponding values are the selection language. run_start: Start frame of analysis. run_end: End frame of analysis. + kwargs: Keyword arguments in the func. Returns: A diction containing the coordination number sequence of each specified neighbor species diff --git a/mdgo/msd.py b/mdgo/msd.py index b0c4e97a..b3a50412 100644 --- a/mdgo/msd.py +++ b/mdgo/msd.py @@ -285,7 +285,7 @@ def parse_msd_type(msd_type: DIM) -> list[int]: dim = keys[msd_type_str] except KeyError: raise ValueError( - f"invalid msd_type: {msd_type_str} specified, please specify one of xyz, " "xy, xz, yz, x, y, z" + f"invalid msd_type: {msd_type_str} specified, please specify one of xyz, xy, xz, yz, x, y, z" ) return dim diff --git a/mdgo/residence_time.py b/mdgo/residence_time.py index 876928ba..5808ee6a 100644 --- a/mdgo/residence_time.py +++ b/mdgo/residence_time.py @@ -125,9 +125,9 @@ def calc_neigh_corr( Args: nvt_run: An MDAnalysis ``Universe``. - distance_dict: - select_dict: - time_step: + distance_dict: A dict of coordination cutoff distance of the neighbor species. + select_dict: A dictionary of atom species selection. + time_step: Timestep between each frame, in ps. run_start: Start frame of analysis. run_end: End frame of analysis. center_atom: The center atom to calculate the ACF for. Default to "cation". @@ -158,8 +158,7 @@ def calc_neigh_corr( run_end, ) acfs = calc_acf(adjacency_matrix) - for acf in acfs: - acf_all.append(acf) + acf_all.extend(list(acfs)) acf_avg[kw] = np.mean(acf_all, axis=0) return times, acf_avg diff --git a/tasks.py b/tasks.py index d3f563d8..0f46e223 100644 --- a/tasks.py +++ b/tasks.py @@ -1,5 +1,8 @@ """ Pyinvoke tasks.py file for automating releases and admin stuff. + +To cut a new mdgo release, use `invoke update-changelog` followed by `invoke release`. + Author: Tingzheng Hou """ from __future__ import annotations @@ -179,8 +182,8 @@ def update_changelog(ctx, version, sim=False): output = subprocess.check_output(["git", "log", "--pretty=format:%s", "v%s..HEAD" % CURRENT_VER]) lines = [] misc = [] - for l in output.decode("utf-8").strip().split("\n"): - m = re.match(r"Merge pull request \#(\d+) from (.*)", l) + for line in output.decode("utf-8").strip().split("\n"): + m = re.match(r"Merge pull request \#(\d+) from (.*)", line) if m: pr_number = m.group(1) contrib, pr_name = m.group(2).split("/", 1) @@ -191,22 +194,22 @@ def update_changelog(ctx, version, sim=False): ll = ll.strip() if ll in ["", "## Summary"]: continue - elif ll.startswith(("## Checklist", "## TODO")): + if ll.startswith(("## Checklist", "## TODO")): break lines.append(f" {ll}") - misc.append(l) + misc.append(line) with open("CHANGES.rst") as f: contents = f.read() - l = "==========" - toks = contents.split(l) + line = "==========" + toks = contents.split(line) head = "\n\nv%s\n" % version + "-" * (len(version) + 1) + "\n" toks.insert(-1, head + "\n".join(lines)) if not sim: with open("CHANGES.rst", "w") as f: - f.write(toks[0] + l + "".join(toks[1:])) + f.write(toks[0] + line + "".join(toks[1:])) ctx.run("open CHANGES.rst") else: - print(toks[0] + l + "".join(toks[1:])) + print(toks[0] + line + "".join(toks[1:])) print("The following commit messages were not included...") print("\n".join(misc)) From 54953f2685f5b8d81cc17ce193232d31cec146db Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 11:09:10 +0800 Subject: [PATCH 22/34] update forcefield test forms --- tests/test_forcefield.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/tests/test_forcefield.py b/tests/test_forcefield.py index 8a0814b2..2a5703e0 100644 --- a/tests/test_forcefield.py +++ b/tests/test_forcefield.py @@ -9,7 +9,6 @@ import numpy as np import pytest -from numpy.testing import assert_allclose from pymatgen.io.lammps.data import LammpsData from mdgo.forcefield.aqueous import Aqueous, Ion @@ -47,26 +46,26 @@ def test_chrome(self) -> None: assert os.path.exists(os.path.join(download_dir, "EMC.itp")) with open(os.path.join(download_dir, "EMC.lmp")) as f: pdf_actual = f.readlines() - assert_allclose(pdf, pdf_actual) + assert pdf == pdf_actual with open(os.path.join(download_dir, "EMC.lmp.xyz")) as f: xyz_actual = f.readlines() - assert_allclose(xyz, xyz_actual) + assert xyz == xyz_actual with open(os.path.join(download_dir, "EMC.gro")) as f: gro_actual = f.readlines() - assert_allclose(gro, gro_actual) + assert gro == gro_actual with open(os.path.join(download_dir, "EMC.itp")) as f: itp_actual = f.readlines() - assert_allclose(itp, itp_actual) + assert itp == itp_actual lpg = FFcrawler(download_dir) lpg.data_from_smiles("CCOC(=O)OC") with open(os.path.join(download_dir, "CCOC(=O)OC.lmp")) as f: smiles_actual = f.readlines() - assert_allclose(smiles[:13], smiles_actual[:13]) - assert_allclose(smiles[18:131], smiles_actual[18:131]) + assert smiles_actual[:13] == smiles[:13] + assert smiles_actual[18:131] == smiles[18:131] assert smiles_actual[131][:26] == " 1 1 1 -0.28" assert smiles_actual[132][:25] == " 2 1 2 0.01" assert smiles_actual[145][:25] == " 15 1 15 0.10" - assert_allclose(smiles_actual[146:], smiles[146:]) + assert smiles_actual[146:] == smiles[146:] finally: sys.stdout = saved_stdout shutil.rmtree(download_dir) From f9eda4cffaeff32adb759e03a710a1615c4e484d Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 11:22:33 +0800 Subject: [PATCH 23/34] add absolute error to msd test --- tests/test_msd.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/tests/test_msd.py b/tests/test_msd.py index 20287d56..0a4e9135 100644 --- a/tests/test_msd.py +++ b/tests/test_msd.py @@ -42,7 +42,11 @@ def setUpClass(cls) -> None: cls.dims = ["x", "y", "z"] def test_msd_straight_forward(self): - assert_allclose(self.fft, msd_straight_forward(self.arr1)) + assert_allclose( + self.fft, + msd_straight_forward(self.arr1), + atol=1e-12, + ) def test_msd_fft(self): assert_allclose(self.fft, msd_fft(self.arr1)) @@ -109,29 +113,39 @@ def test_mda_msd_wrapper(self): assert_allclose( onsager_ii_self(self.gen2, 0, 10, select="type 3", msd_type="x", center_of_mass=False), mda_msd_wrapper(self.gen2, 0, 10, select="type 3", msd_type="x", fft=False), + atol=1e-12, ) assert_allclose( onsager_ii_self(self.gen2, 0, 10, select="type 3", msd_type="y", center_of_mass=False), mda_msd_wrapper(self.gen2, 0, 10, select="type 3", msd_type="y", fft=False), + atol=1e-12, ) assert_allclose( onsager_ii_self(self.gen2, 0, 10, select="type 3", msd_type="z", center_of_mass=False), mda_msd_wrapper(self.gen2, 0, 10, select="type 3", msd_type="z", fft=False), + atol=1e-12, ) assert_allclose( onsager_ii_self(self.gen2, 0, 100, select="type 3", msd_type="xy", center_of_mass=False), mda_msd_wrapper(self.gen2, 0, 100, select="type 3", msd_type="xy", fft=False), + atol=1e-12, ) assert_allclose( onsager_ii_self(self.gen2, 0, 100, select="type 3", msd_type="yz", center_of_mass=False), mda_msd_wrapper(self.gen2, 0, 100, select="type 3", msd_type="yz", fft=False), + atol=1e-12, ) assert_allclose( onsager_ii_self(self.gen2, 0, 100, select="type 3", msd_type="xz", center_of_mass=False), mda_msd_wrapper(self.gen2, 0, 100, select="type 3", msd_type="xz", fft=False), + atol=1e-12, ) if td is not None: - assert_allclose(mda_msd_cation, mda_msd_wrapper(self.gen2, 0, 100, select="type 3")) + assert_allclose( + mda_msd_cation, + mda_msd_wrapper(self.gen2, 0, 100, select="type 3"), + atol=1e-12, + ) def test_total_msd(self): total_builtin_cation = total_msd(self.gen2, 0, 100, select="type 3", fft=True, built_in=True) From f5eb2dbfa41d5f4688dc6609c000c25d4e0cd98c Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 11:28:31 +0800 Subject: [PATCH 24/34] ruff fixes 3 --- mdgo/forcefield/maestro.py | 2 +- tests/test_msd.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/mdgo/forcefield/maestro.py b/mdgo/forcefield/maestro.py index 8b35eb6a..4ea465f2 100644 --- a/mdgo/forcefield/maestro.py +++ b/mdgo/forcefield/maestro.py @@ -124,7 +124,7 @@ def get_mae(self, wait: float = 30): shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, - preexec_fn=os.setsid, + start_new_session=True, ) except subprocess.CalledProcessError as e: raise ValueError(f"Maestro failed with errorcode {e.returncode} and stderr: {e.stderr}") from e diff --git a/tests/test_msd.py b/tests/test_msd.py index 0a4e9135..11898846 100644 --- a/tests/test_msd.py +++ b/tests/test_msd.py @@ -152,8 +152,8 @@ def test_total_msd(self): total_mda_cation = total_msd( self.gen2, 0, 100, select="type 3", fft=False, built_in=False, center_of_mass=False ) - assert_allclose(32.14254152556588, total_builtin_cation[50]) - assert_allclose(32.338364098424634, total_mda_cation[50]) + assert total_builtin_cation[50] == 32.14254152556588 + assert total_mda_cation[50] == 32.338364098424634 with pytest.raises(ValueError): total_msd(self.gen2, 0, 100, select="type 3", fft=True, built_in=False, center_of_mass=True) From 7bd1704e737f9c7f0163e5f186e466876a8e6658 Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 11:35:30 +0800 Subject: [PATCH 25/34] revert assertion change --- tests/test_msd.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_msd.py b/tests/test_msd.py index 11898846..6f80b511 100644 --- a/tests/test_msd.py +++ b/tests/test_msd.py @@ -152,8 +152,8 @@ def test_total_msd(self): total_mda_cation = total_msd( self.gen2, 0, 100, select="type 3", fft=False, built_in=False, center_of_mass=False ) - assert total_builtin_cation[50] == 32.14254152556588 - assert total_mda_cation[50] == 32.338364098424634 + assert_allclose(total_builtin_cation[50], 32.14254152556588) + assert_allclose(total_mda_cation[50], 32.338364098424634) with pytest.raises(ValueError): total_msd(self.gen2, 0, 100, select="type 3", fft=True, built_in=False, center_of_mass=True) From dd99bc4b31945b61baa5723f7176972aa791037c Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 11:40:14 +0800 Subject: [PATCH 26/34] finish ruff fixes --- tests/test_msd.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_msd.py b/tests/test_msd.py index 6f80b511..b7e9d2df 100644 --- a/tests/test_msd.py +++ b/tests/test_msd.py @@ -154,7 +154,7 @@ def test_total_msd(self): ) assert_allclose(total_builtin_cation[50], 32.14254152556588) assert_allclose(total_mda_cation[50], 32.338364098424634) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Packmol failed with"): total_msd(self.gen2, 0, 100, select="type 3", fft=True, built_in=False, center_of_mass=True) From e8a094c69c3b23f74099ba5a2c27a141f10d443a Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 11:40:50 +0800 Subject: [PATCH 27/34] final ruff fixes 2 --- tests/test_mdgopackmol.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_mdgopackmol.py b/tests/test_mdgopackmol.py index 78bbe699..2e20e49e 100644 --- a/tests/test_mdgopackmol.py +++ b/tests/test_mdgopackmol.py @@ -126,7 +126,7 @@ def test_control_params(self, water, ethanol): input_string = f.read() assert "maxit 0" in input_string assert "nloop 0" in input_string - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="MDAnalysis"): pw.run_packmol() def test_timeout(self, water, ethanol): @@ -163,7 +163,7 @@ def test_no_return_and_box(self, water, ethanol): with open(os.path.join(scratch_dir, "packmol.inp")) as f: input_string = f.read() assert "inside box 0 0 0 2 2 2" in input_string - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Packmol failed with"): pw.run_packmol() def test_random_seed(self, water, ethanol): From fe7c547699b1aca64623c3316df18366ecbdbf1b Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 11:42:27 +0800 Subject: [PATCH 28/34] set atol for more assert_allclose --- tests/test_msd.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_msd.py b/tests/test_msd.py index b7e9d2df..d999c4b0 100644 --- a/tests/test_msd.py +++ b/tests/test_msd.py @@ -49,8 +49,8 @@ def test_msd_straight_forward(self): ) def test_msd_fft(self): - assert_allclose(self.fft, msd_fft(self.arr1)) - assert_allclose(msd_straight_forward(self.arr2), msd_fft(self.arr2)) + assert_allclose(self.fft, msd_fft(self.arr1), atol=1e-12) + assert_allclose(msd_straight_forward(self.arr2), msd_fft(self.arr2), atol=1e-12) def test_create_position_arrays(self): assert_allclose( From d878a1cdf5dcfab523dc08d919c285db9d7110fc Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 11:49:13 +0800 Subject: [PATCH 29/34] fix test raises --- tests/test_mdgopackmol.py | 2 +- tests/test_msd.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/test_mdgopackmol.py b/tests/test_mdgopackmol.py index 2e20e49e..9e21d7e9 100644 --- a/tests/test_mdgopackmol.py +++ b/tests/test_mdgopackmol.py @@ -126,7 +126,7 @@ def test_control_params(self, water, ethanol): input_string = f.read() assert "maxit 0" in input_string assert "nloop 0" in input_string - with pytest.raises(ValueError, match="MDAnalysis"): + with pytest.raises(ValueError, match="Packmol failed with 1"): pw.run_packmol() def test_timeout(self, water, ethanol): diff --git a/tests/test_msd.py b/tests/test_msd.py index d999c4b0..aa226b6d 100644 --- a/tests/test_msd.py +++ b/tests/test_msd.py @@ -154,7 +154,11 @@ def test_total_msd(self): ) assert_allclose(total_builtin_cation[50], 32.14254152556588) assert_allclose(total_mda_cation[50], 32.338364098424634) - with pytest.raises(ValueError, match="Packmol failed with"): + with pytest.raises( + ValueError, + match="Warning! MDAnalysis does not support subtracting center " + "of mass. Calculating without subtracting...", + ): total_msd(self.gen2, 0, 100, select="type 3", fft=True, built_in=False, center_of_mass=True) From 867001b123a4d24e30e0ef4510a52a36a9a41d20 Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 11:56:25 +0800 Subject: [PATCH 30/34] lint --- mdgo/coordination.py | 8 ++------ mdgo/msd.py | 4 +--- mdgo/util/reformat.py | 4 +--- setup.py | 2 +- tests/test_msd.py | 32 ++++++++++++++++---------------- tests/test_volume.py | 4 +--- 6 files changed, 22 insertions(+), 32 deletions(-) diff --git a/mdgo/coordination.py b/mdgo/coordination.py index c142ae2e..f7c350eb 100644 --- a/mdgo/coordination.py +++ b/mdgo/coordination.py @@ -363,9 +363,7 @@ def check_contiguous_steps( An array of distance between the center atom and the neighbor atoms in the checkpoint +/- lag time range. """ - coord_num: dict[str, list[list[int]] | np.ndarray] = { - x: [[] for _ in range(lag * 2 + 1)] for x in distance_dict - } + coord_num: dict[str, list[list[int]] | np.ndarray] = {x: [[] for _ in range(lag * 2 + 1)] for x in distance_dict} trj_analysis = nvt_run.trajectory[run_start:run_end:] has = False for i, _ts in enumerate(trj_analysis): @@ -1092,9 +1090,7 @@ def write_out(center_pos: np.ndarray, center_name: str, neighbors: AtomGroup, pa xyz_file.write("\n".join(lines)) -def select_shell( - select: dict[str, str] | str, distance: dict[str, float] | str, center_atom: Atom, kw: str -) -> str: +def select_shell(select: dict[str, str] | str, distance: dict[str, float] | str, center_atom: Atom, kw: str) -> str: """ Select a group of atoms that is within a distance of an ``center_atom``. diff --git a/mdgo/msd.py b/mdgo/msd.py index b3a50412..f8c6d6c9 100644 --- a/mdgo/msd.py +++ b/mdgo/msd.py @@ -284,9 +284,7 @@ def parse_msd_type(msd_type: DIM) -> list[int]: try: dim = keys[msd_type_str] except KeyError: - raise ValueError( - f"invalid msd_type: {msd_type_str} specified, please specify one of xyz, xy, xz, yz, x, y, z" - ) + raise ValueError(f"invalid msd_type: {msd_type_str} specified, please specify one of xyz, xy, xz, yz, x, y, z") return dim diff --git a/mdgo/util/reformat.py b/mdgo/util/reformat.py index 616390d0..ec5d4a68 100644 --- a/mdgo/util/reformat.py +++ b/mdgo/util/reformat.py @@ -68,9 +68,7 @@ }, } -BOX: Final[ - str -] = """{0:6f} {1:6f} xlo xhi +BOX: Final[str] = """{0:6f} {1:6f} xlo xhi {0:6f} {1:6f} ylo yhi {0:6f} {1:6f} zlo zhi""" diff --git a/setup.py b/setup.py index d08a09bc..fa75a4ea 100644 --- a/setup.py +++ b/setup.py @@ -71,7 +71,7 @@ "Topic :: Scientific/Engineering :: Physics", "Topic :: Scientific/Engineering :: Chemistry", "Topic :: Software Development :: Libraries :: Python Modules", - ], + ], packages=find_packages(), install_requires=INSTALL_REQUIRES, extras_require={ diff --git a/tests/test_msd.py b/tests/test_msd.py index aa226b6d..bc7dab97 100644 --- a/tests/test_msd.py +++ b/tests/test_msd.py @@ -55,36 +55,36 @@ def test_msd_fft(self): def test_create_position_arrays(self): assert_allclose( np.array([21.53381769, 14.97501839, -3.87998785]), - create_position_arrays(self.gen2, 0, 100, select="type 3")[50][2] + create_position_arrays(self.gen2, 0, 100, select="type 3")[50][2], ) assert_allclose( np.array([-2.78550047, -11.85487624, -17.1221954]), - create_position_arrays(self.gen2, 0, 100, select="type 3")[99][10] + create_position_arrays(self.gen2, 0, 100, select="type 3")[99][10], ) assert_allclose( - np.array([41.1079216 , 34.95127106, 18.00482368]), - create_position_arrays(self.gen2, 0, 100, select="type 3", center_of_mass=False)[50][2] + np.array([41.1079216, 34.95127106, 18.00482368]), + create_position_arrays(self.gen2, 0, 100, select="type 3", center_of_mass=False)[50][2], ) assert_allclose( - np.array([16.98478317, 8.27190208, 5.07116079]), - create_position_arrays(self.gen2, 0, 100, select="type 3", center_of_mass=False)[99][10] + np.array([16.98478317, 8.27190208, 5.07116079]), + create_position_arrays(self.gen2, 0, 100, select="type 3", center_of_mass=False)[99][10], ) def test_parse_msd_type(self): xyz = parse_msd_type("xyz") - assert ["x", "y", "z"] == self.dims[xyz[0]:xyz[1]:xyz[2]] + assert ["x", "y", "z"] == self.dims[xyz[0] : xyz[1] : xyz[2]] xy = parse_msd_type("xy") - assert ["x", "y"] == self.dims[xy[0]:xy[1]:xy[2]] + assert ["x", "y"] == self.dims[xy[0] : xy[1] : xy[2]] yz = parse_msd_type("yz") - assert ["y", "z"] == self.dims[yz[0]:yz[1]:yz[2]] + assert ["y", "z"] == self.dims[yz[0] : yz[1] : yz[2]] xz = parse_msd_type("xz") - assert ["x", "z"] == self.dims[xz[0]:xz[1]:xz[2]] + assert ["x", "z"] == self.dims[xz[0] : xz[1] : xz[2]] x = parse_msd_type("x") - assert ["x"] == self.dims[x[0]:x[1]:x[2]] + assert ["x"] == self.dims[x[0] : x[1] : x[2]] y = parse_msd_type("y") - assert ["y"] == self.dims[y[0]:y[1]:y[2]] + assert ["y"] == self.dims[y[0] : y[1] : y[2]] z = parse_msd_type("z") - assert ["z"] == self.dims[z[0]:z[1]:z[2]] + assert ["z"] == self.dims[z[0] : z[1] : z[2]] def test_onsager_ii_self(self): onsager_ii_self_fft = onsager_ii_self(self.gen2, 0, 100, select="type 3") @@ -155,9 +155,9 @@ def test_total_msd(self): assert_allclose(total_builtin_cation[50], 32.14254152556588) assert_allclose(total_mda_cation[50], 32.338364098424634) with pytest.raises( - ValueError, - match="Warning! MDAnalysis does not support subtracting center " - "of mass. Calculating without subtracting...", + ValueError, + match="Warning! MDAnalysis does not support subtracting center " + "of mass. Calculating without subtracting...", ): total_msd(self.gen2, 0, 100, select="type 3", fft=True, built_in=False, center_of_mass=True) diff --git a/tests/test_volume.py b/tests/test_volume.py index 86b10c54..2f2efcba 100644 --- a/tests/test_volume.py +++ b/tests/test_volume.py @@ -50,9 +50,7 @@ def test_molecular_volume(self) -> None: litfsi_volume_4 = molecular_volume(self.litfsi, radii_type="Lange") litfsi_volume_5 = molecular_volume(self.litfsi, radii_type="pymatgen") litfsi_volume_6 = molecular_volume(self.litfsi, molar_volume=False) - litfsi_volume_7 = molecular_volume( - self.litfsi, mode="act", x_size=8, y_size=8, z_size=8 - ) + litfsi_volume_7 = molecular_volume(self.litfsi, mode="act", x_size=8, y_size=8, z_size=8) assert litfsi_volume_1 == 100.16 assert litfsi_volume_2 == 100.16 assert litfsi_volume_3 == 99.37 From 6f9d3d504fc03053b8892b647f685b716d97417f Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 11:59:42 +0800 Subject: [PATCH 31/34] drop noreturn --- mdgo/forcefield/pubchem.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mdgo/forcefield/pubchem.py b/mdgo/forcefield/pubchem.py index 9344331b..56cab4a6 100644 --- a/mdgo/forcefield/pubchem.py +++ b/mdgo/forcefield/pubchem.py @@ -10,7 +10,7 @@ import os import time -from typing import Final, NoReturn +from typing import Final from urllib.parse import quote import pubchempy as pcp @@ -103,7 +103,7 @@ def obtain_entry(self, search_text: str, name: str, output_format: str = "sdf") return self._obtain_entry_api(search_text, name, output_format=output_format) return self._obtain_entry_web(search_text, name, output_format=output_format) - def smiles_to_pdb(self, smiles: str) -> NoReturn: + def smiles_to_pdb(self, smiles: str): """ Obtain pdf file based on SMILES code. From 5da174a60e9056a2ebb6a16b0b1f62e5eceacbf7 Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 12:05:17 +0800 Subject: [PATCH 32/34] Update pyproject.toml --- pyproject.toml | 45 +++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 43 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e1480f13..348fcae9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,6 @@ [build-system] requires = [ - "setuptools>=42", - "wheel" + "setuptools>=65.0.0", ] build-backend = "setuptools.build_meta" @@ -75,3 +74,45 @@ lint.isort.split-on-trailing-comma = false "pymatgen/vis/*" = ["D"] "pymatgen/io/*" = ["D"] "dev_scripts/*" = ["D"] + +[tool.pytest.ini_options] +addopts = "--durations=30 --quiet -r xXs --color=yes -p no:warnings --import-mode=importlib" + +[tool.coverage.run] +parallel = true + +[tool.coverage.report] +exclude_also = [ + "@deprecated", + "@np.deprecate", + "def __repr__", + "except ImportError:", + "if 0:", + "if TYPE_CHECKING:", + "if __name__ == .__main__.:", + "if self.debug:", + "if settings.DEBUG", + "if typing.TYPE_CHECKING:", + "pragma: no cover", + "raise AssertionError", + "raise NotImplementedError", + "show_plot", +] + +[tool.mypy] +ignore_missing_imports = true +namespace_packages = true +explicit_package_bases = true +no_implicit_optional = false +disable_error_code = "annotation-unchecked" + +[[tool.mypy.overrides]] +module = ["requests.*", "tabulate.*"] +ignore_missing_imports = true + +[tool.codespell] +ignore-words-list = """ +titel,alls,ans,nd,mater,nwo,te,hart,ontop,ist,ot,fo,nax,coo,coul,ser,leary,thre,fase, +rute,reson,titels,ges,scalr,strat,struc,hda,nin,ons,pres,kno,loos,lamda,lew,atomate +""" +check-filenames = true \ No newline at end of file From e7929de762919d0014a7d7098c7a27e4e29fd2a5 Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 12:10:06 +0800 Subject: [PATCH 33/34] change assertion format --- tests/test_volume.py | 37 +++++++++++++++++++------------------ 1 file changed, 19 insertions(+), 18 deletions(-) diff --git a/tests/test_volume.py b/tests/test_volume.py index 2f2efcba..f0018e9c 100644 --- a/tests/test_volume.py +++ b/tests/test_volume.py @@ -3,6 +3,7 @@ import os import unittest +from numpy.testing import assert_allclose from pymatgen.core import Molecule from mdgo.util.volume import molecular_volume @@ -27,23 +28,23 @@ def test_molecular_volume(self) -> None: lipf6_volume_3 = molecular_volume(self.lipf6, radii_type="Lange") lipf6_volume_4 = molecular_volume(self.lipf6, radii_type="pymatgen") lipf6_volume_5 = molecular_volume(self.lipf6, molar_volume=False) - assert lipf6_volume_1 == 47.62 - assert lipf6_volume_2 == 43.36 - assert lipf6_volume_3 == 41.49 - assert lipf6_volume_4 == 51.94 - assert lipf6_volume_5 == 79.08 + assert_allclose(lipf6_volume_1, 47.62) + assert_allclose(lipf6_volume_2, 43.36) + assert_allclose(lipf6_volume_3, 41.49) + assert_allclose(lipf6_volume_4, 51.94) + assert_allclose(lipf6_volume_5, 79.08) ec_volume_1 = molecular_volume(self.ec) ec_volume_2 = molecular_volume(self.ec, exclude_h=False) ec_volume_3 = molecular_volume(self.ec, res=1.0) ec_volume_4 = molecular_volume(self.ec, radii_type="Lange") ec_volume_5 = molecular_volume(self.ec, radii_type="pymatgen") ec_volume_6 = molecular_volume(self.ec, molar_volume=False) - assert ec_volume_1 == 38.44 - assert ec_volume_2 == 43.17 - assert ec_volume_3 == 40.95 - assert ec_volume_4 == 41.07 - assert ec_volume_5 == 38.44 - assert ec_volume_6 == 63.83 + assert_allclose(ec_volume_1, 38.44) + assert_allclose(ec_volume_2, 43.17) + assert_allclose(ec_volume_3, 40.95) + assert_allclose(ec_volume_4, 41.07) + assert_allclose(ec_volume_5, 38.44) + assert_allclose(ec_volume_6, 63.83) litfsi_volume_1 = molecular_volume(self.litfsi) litfsi_volume_2 = molecular_volume(self.litfsi, exclude_h=False) litfsi_volume_3 = molecular_volume(self.litfsi, res=1.0) @@ -51,13 +52,13 @@ def test_molecular_volume(self) -> None: litfsi_volume_5 = molecular_volume(self.litfsi, radii_type="pymatgen") litfsi_volume_6 = molecular_volume(self.litfsi, molar_volume=False) litfsi_volume_7 = molecular_volume(self.litfsi, mode="act", x_size=8, y_size=8, z_size=8) - assert litfsi_volume_1 == 100.16 - assert litfsi_volume_2 == 100.16 - assert litfsi_volume_3 == 99.37 - assert litfsi_volume_4 == 90.78 - assert litfsi_volume_5 == 105.31 - assert litfsi_volume_6 == 166.32 - assert litfsi_volume_7 == 124.66 + assert_allclose(litfsi_volume_1, 100.16) + assert_allclose(litfsi_volume_2, 100.16) + assert_allclose(litfsi_volume_3, 99.37) + assert_allclose(litfsi_volume_4, 90.78) + assert_allclose(litfsi_volume_5, 105.31) + assert_allclose(litfsi_volume_6, 166.32) + assert_allclose(litfsi_volume_7, 124.66) if __name__ == "__main__": From bcc3572ae754e5b6d2d891b190697b7095f66559 Mon Sep 17 00:00:00 2001 From: Tingzheng Hou <25351437+htz1992213@users.noreply.github.com> Date: Mon, 5 Feb 2024 12:15:45 +0800 Subject: [PATCH 34/34] add additional tol for volume --- tests/test_volume.py | 36 ++++++++++++++++++------------------ 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/tests/test_volume.py b/tests/test_volume.py index f0018e9c..9a2f893e 100644 --- a/tests/test_volume.py +++ b/tests/test_volume.py @@ -28,23 +28,23 @@ def test_molecular_volume(self) -> None: lipf6_volume_3 = molecular_volume(self.lipf6, radii_type="Lange") lipf6_volume_4 = molecular_volume(self.lipf6, radii_type="pymatgen") lipf6_volume_5 = molecular_volume(self.lipf6, molar_volume=False) - assert_allclose(lipf6_volume_1, 47.62) - assert_allclose(lipf6_volume_2, 43.36) - assert_allclose(lipf6_volume_3, 41.49) - assert_allclose(lipf6_volume_4, 51.94) - assert_allclose(lipf6_volume_5, 79.08) + assert_allclose(lipf6_volume_1, 47.62, atol=0.01) + assert_allclose(lipf6_volume_2, 43.36, atol=0.01) + assert_allclose(lipf6_volume_3, 41.49, atol=0.01) + assert_allclose(lipf6_volume_4, 51.94, atol=0.01) + assert_allclose(lipf6_volume_5, 79.08, atol=0.01) ec_volume_1 = molecular_volume(self.ec) ec_volume_2 = molecular_volume(self.ec, exclude_h=False) ec_volume_3 = molecular_volume(self.ec, res=1.0) ec_volume_4 = molecular_volume(self.ec, radii_type="Lange") ec_volume_5 = molecular_volume(self.ec, radii_type="pymatgen") ec_volume_6 = molecular_volume(self.ec, molar_volume=False) - assert_allclose(ec_volume_1, 38.44) - assert_allclose(ec_volume_2, 43.17) - assert_allclose(ec_volume_3, 40.95) - assert_allclose(ec_volume_4, 41.07) - assert_allclose(ec_volume_5, 38.44) - assert_allclose(ec_volume_6, 63.83) + assert_allclose(ec_volume_1, 38.44, atol=0.01) + assert_allclose(ec_volume_2, 43.17, atol=0.01) + assert_allclose(ec_volume_3, 40.95, atol=0.01) + assert_allclose(ec_volume_4, 41.07, atol=0.01) + assert_allclose(ec_volume_5, 38.44, atol=0.01) + assert_allclose(ec_volume_6, 63.83, atol=0.01) litfsi_volume_1 = molecular_volume(self.litfsi) litfsi_volume_2 = molecular_volume(self.litfsi, exclude_h=False) litfsi_volume_3 = molecular_volume(self.litfsi, res=1.0) @@ -52,13 +52,13 @@ def test_molecular_volume(self) -> None: litfsi_volume_5 = molecular_volume(self.litfsi, radii_type="pymatgen") litfsi_volume_6 = molecular_volume(self.litfsi, molar_volume=False) litfsi_volume_7 = molecular_volume(self.litfsi, mode="act", x_size=8, y_size=8, z_size=8) - assert_allclose(litfsi_volume_1, 100.16) - assert_allclose(litfsi_volume_2, 100.16) - assert_allclose(litfsi_volume_3, 99.37) - assert_allclose(litfsi_volume_4, 90.78) - assert_allclose(litfsi_volume_5, 105.31) - assert_allclose(litfsi_volume_6, 166.32) - assert_allclose(litfsi_volume_7, 124.66) + assert_allclose(litfsi_volume_1, 100.16, atol=0.01) + assert_allclose(litfsi_volume_2, 100.16, atol=0.01) + assert_allclose(litfsi_volume_3, 99.37, atol=0.01) + assert_allclose(litfsi_volume_4, 90.78, atol=0.01) + assert_allclose(litfsi_volume_5, 105.31, atol=0.01) + assert_allclose(litfsi_volume_6, 166.32, atol=0.01) + assert_allclose(litfsi_volume_7, 124.66, atol=0.01) if __name__ == "__main__":