Skip to content

Commit

Permalink
ruff fixes 2
Browse files Browse the repository at this point in the history
  • Loading branch information
htz1992213 committed Feb 5, 2024
1 parent 082f884 commit 73d8a33
Show file tree
Hide file tree
Showing 9 changed files with 61 additions and 54 deletions.
2 changes: 1 addition & 1 deletion mdgo/core/analysis.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
3 changes: 2 additions & 1 deletion mdgo/core/run.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ def from_dict(cls):
Constructor.
Returns:
name: The name of the class
"""
return cls("name")
Expand All @@ -28,6 +29,6 @@ def from_recipe(cls):
Constructor.
Returns:
name: The name of the class
"""
return cls("name")
2 changes: 1 addition & 1 deletion mdgo/forcefield/aqueous.py
Original file line number Diff line number Diff line change
Expand Up @@ -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, 45214536. 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'
Expand Down
6 changes: 3 additions & 3 deletions mdgo/msd.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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:
Expand Down
12 changes: 4 additions & 8 deletions mdgo/residence_time.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = (
"("
Expand All @@ -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


Expand All @@ -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

Expand Down Expand Up @@ -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 = {}
Expand Down
2 changes: 1 addition & 1 deletion mdgo/util/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
__email__ = "[email protected]"
__date__ = "Jul 19, 2021"

from typing import Dict, Final
from typing import Final

MM_of_Elements: Final[dict[str, float]] = {
"H": 1.00794,
Expand Down
84 changes: 47 additions & 37 deletions mdgo/util/dict_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -221,64 +221,74 @@ 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)
molecule_number += 1
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(
Expand Down
2 changes: 1 addition & 1 deletion mdgo/util/reformat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion mdgo/util/volume.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit 73d8a33

Please sign in to comment.