From f48e037936ab11a1d5f88ad70946d2338c98be7d Mon Sep 17 00:00:00 2001 From: Sebastian Kloth Date: Fri, 22 Dec 2023 14:22:25 +0100 Subject: [PATCH] Started rewriting, adding type hints --- src/mdevaluate/__init__.py | 42 +++++++---- src/mdevaluate/atoms.py | 69 +++++++++---------- src/mdevaluate/coordinates.py | 9 +-- src/mdevaluate/extra/__init__.py | 0 src/mdevaluate/{ => extra}/chill.py | 0 .../{ => extra}/free_energy_landscape.py | 0 6 files changed, 61 insertions(+), 59 deletions(-) create mode 100644 src/mdevaluate/extra/__init__.py rename src/mdevaluate/{ => extra}/chill.py (100%) rename src/mdevaluate/{ => extra}/free_energy_landscape.py (100%) diff --git a/src/mdevaluate/__init__.py b/src/mdevaluate/__init__.py index dd18ac6..833b520 100644 --- a/src/mdevaluate/__init__.py +++ b/src/mdevaluate/__init__.py @@ -1,5 +1,6 @@ import os from glob import glob +from typing import Optional import pandas as pd @@ -11,20 +12,19 @@ from . import functions from . import pbc from . import autosave from . import reader -from . import chill from . import system -from . import free_energy_landscape +from .extra import free_energy_landscape, chill from .logging import logger def open( - directory="", - topology="*.tpr", - trajectory="*.xtc", - nojump=False, - index_file=None, - charges=None, - masses=None, + directory: str = "", + topology: str = "*.tpr", + trajectory: str = "*.xtc", + nojump: bool = False, + index_file: Optional[str] = None, + charges: Optional[list[float]] = None, + masses: Optional[list[float]] = None, ): """ Open a simulation from a directory. @@ -34,8 +34,17 @@ def open( topology (opt.): Descriptor of the topology file (tpr or gro). By default, a tpr file is used, if there is exactly one in the directoy. - trajectory (opt.): Descriptor of the trajectory (xtc file). - nojump (opt.): If nojump matrixes should be generated. They will alwyas be loaded if present + trajectory (opt.): Descriptor of the trajectory (xtc or trr file). + nojump (opt.): + If nojump matrixes should be generated. They will alwyas be loaded + if present + index_file (opt.): Descriptor of the index file (ndx file). + charges (opt.): + List with charges for each atom. It Has to be the same length as the number + of atoms in the system. Only used if topology file is a gro file. + masses (opt.): + List with masses for each atom. It Has to be the same length as the number + of atoms in the system. Only used if topology file is a gro file. Returns: A Coordinate object of the simulation. @@ -43,13 +52,14 @@ def open( Example: Open a simulation located in '/path/to/sim', where the trajectory is located in a sub-directory '/path/to/sim/out' and named for Example - 'nojump_traj.xtc'. All read frames will be cached in memory. + 'nojump_traj.xtc'. - >>> open('/path/to/sim', trajectory='out/nojump*.xtc', cached=None) + >>> open('/path/to/sim', trajectory='out/nojump*.xtc') The file descriptors can use unix style pathname expansion to define the filenames. The default patterns use the recursive placeholder `**` which matches the base or - any subdirctory, thus files in subdirectories with matching file type will be found too. + any subdirctory, thus files in subdirectories with matching file type will be found + too. For example: 'out/nojump*.xtc' would match xtc files in a subdirectory `out` that start with `nojump` and end with `.xtc`. @@ -91,9 +101,11 @@ def open( return coords -def open_energy(file): +def open_energy(file: str): """Reads a gromacs energy file and output the data in a pandas DataFrame. Args: file: Filename of the energy file + Returns: + A DataFrame with the different energies doe each time. """ return pd.DataFrame(reader.energy_reader(file).data_dict) diff --git a/src/mdevaluate/atoms.py b/src/mdevaluate/atoms.py index 6c0c590..5f0813e 100644 --- a/src/mdevaluate/atoms.py +++ b/src/mdevaluate/atoms.py @@ -1,25 +1,21 @@ import re +from typing import Optional -from .pbc import pbc_diff -from .checksum import checksum import numpy as np +import numpy.typing as npt +from scipy.spatial import KDTree -import scipy - -if scipy.version.version >= "0.17.0": - from scipy.spatial import cKDTree as KDTree -else: - from scipy.spatial import KDTree +from .pbc import pbc_diff, pbc_points +from .checksum import checksum +from .coordinates import CoordinateFrame -def compare_regex(list, exp): +def compare_regex(str_list: list[str], exp: str): """ Compare a list of strings with a regular expression. """ - if not exp.endswith("$"): - exp += "$" regex = re.compile(exp) - return np.array([regex.match(s) is not None for s in list]) + return np.array([regex.match(s) is not None for s in str_list]) class Atoms: @@ -180,34 +176,26 @@ class AtomSubset: return checksum(self.description) -def center_of_mass(position, mass=None): - if mass is not None: - return 1 / mass.sum() * (mass * position).sum(axis=0) - else: - return 1 / len(position) * position.sum(axis=0) - - -def gyration_radius(position): +def gyration_radius(position: CoordinateFrame): r""" - Calculates a list of all radii of gyration of all molecules given in the coordinate frame, - weighted with the masses of the individual atoms. + Calculates a list of all radii of gyration of all molecules given in the coordinate + frame, weighted with the masses of the individual atoms. Args: - position: Coordinate frame object + position: Coordinate frame object ..math:: R_G = \left(\frac{\sum_{i=1}^{n} m_i |\vec{r_i} - \vec{r_{COM}}|^2 }{\sum_{i=1}^{n} m_i } \rigth)^{\frac{1}{2}} """ - gyration_radii = np.array([]) for resid in np.unique(position.residue_ids): pos = position.whole[position.residue_ids == resid] mass = position.masses[position.residue_ids == resid][:, np.newaxis] - COM = center_of_mass(pos, mass) + COM = 1 / mass.sum() * (mass * position).sum(axis=0) r_sq = ((pbc_diff(pos, COM, pos.box.diagonal())) ** 2).sum(1)[:, np.newaxis] - g_radius = ((r_sq * mass).sum() / mass.sum()) ** (0.5) + g_radius = ((r_sq * mass).sum() / mass.sum()) ** 0.5 gyration_radii = np.append(gyration_radii, g_radius) @@ -215,12 +203,16 @@ def gyration_radius(position): def layer_of_atoms( - atoms, thickness, plane_offset=np.array([0, 0, 0]), plane_normal=np.array([1, 0, 0]) + atoms: CoordinateFrame, + thickness: float, + plane_normal: npt.ArrayLike, + plane_offset: Optional[npt.ArrayLike] = np.array([0, 0, 0]), ): - p_ = atoms - plane_offset - distance = np.dot(p_, plane_normal) - - return abs(distance) <= thickness + if plane_offset is None: + np.array([0, 0, 0]) + atoms = atoms - plane_offset + distance = np.dot(atoms, plane_normal) + return np.abs(distance) <= thickness def distance_to_atoms(ref, atoms, box=None): @@ -234,7 +226,7 @@ def distance_to_atoms(ref, atoms, box=None): return out -def distance_to_atoms_cKDtree(ref, atoms, box=None, thickness=None): +def distance_to_atoms_KDtree(ref, atoms, box=None, thickness=None): """ Get the minimal distance from atoms to ref. The result is an array of with length == len(atoms) @@ -251,7 +243,7 @@ def distance_to_atoms_cKDtree(ref, atoms, box=None, thickness=None): start_coords = atoms all_frame_coords = ref - tree = spatial.cKDTree(all_frame_coords) + tree = KDTree(all_frame_coords) return tree.query(start_coords)[0] @@ -266,11 +258,16 @@ def next_neighbors( Find the N next neighbors of a set of atoms. Args: - atoms: The reference atoms and also the atoms which are queried if `query_atoms` is net provided + atoms: + The reference atoms and also the atoms which are queried if `query_atoms` + is net provided query_atoms (opt.): If this is not None, these atoms will be queried number_of_neighbors (int, opt.): Number of neighboring atoms to find - distance_upper_bound (float, opt.): Upper bound of the distance between neighbors - distinct (bool, opt.): If this is true, the atoms and query atoms are taken as distinct sets of atoms + distance_upper_bound (float, opt.): + Upper bound of the distance between neighbors + distinct (bool, opt.): + If this is true, the atoms and query atoms are taken as distinct sets of + atoms """ tree = KDTree(atoms) dnn = 0 diff --git a/src/mdevaluate/coordinates.py b/src/mdevaluate/coordinates.py index bd16541..519c6df 100755 --- a/src/mdevaluate/coordinates.py +++ b/src/mdevaluate/coordinates.py @@ -217,7 +217,7 @@ class Coordinates: """ Coordinates represent trajectory data, which is used for evaluation functions. - Atoms may be selected by specifing a atom_subset or a atom_filter. + Atoms may be selected by specifying an atom_subset or an atom_filter. """ def get_mode(self, mode): @@ -424,13 +424,6 @@ def map_coordinates(func): @map_coordinates def center_of_masses(coordinates, atoms=None, shear: bool = False): - """ - Example: - rd = XTCReader('t.xtc') - coordinates = Coordinates(rd) - com = centers_of_mass(coordinates, (1.0, 2.0, 1.0, 3.0)) - - """ if atoms is None: atoms = list(range(len(coordinates))) res_ids = coordinates.residue_ids[atoms] diff --git a/src/mdevaluate/extra/__init__.py b/src/mdevaluate/extra/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mdevaluate/chill.py b/src/mdevaluate/extra/chill.py similarity index 100% rename from src/mdevaluate/chill.py rename to src/mdevaluate/extra/chill.py diff --git a/src/mdevaluate/free_energy_landscape.py b/src/mdevaluate/extra/free_energy_landscape.py similarity index 100% rename from src/mdevaluate/free_energy_landscape.py rename to src/mdevaluate/extra/free_energy_landscape.py