From 9f3ab15ce473ff383d7004af4d3493b1d4e93b66 Mon Sep 17 00:00:00 2001 From: Rosalba Date: Wed, 14 Aug 2024 16:22:55 -0400 Subject: [PATCH] vmat pt 2 --- automol/vmat.py | 152 +++++++++++++++++++----------------------------- 1 file changed, 61 insertions(+), 91 deletions(-) diff --git a/automol/vmat.py b/automol/vmat.py index 6feea6b0..b8356076 100644 --- a/automol/vmat.py +++ b/automol/vmat.py @@ -1,7 +1,6 @@ """V-Matrix: Variable V-Matrix (V-Matrix without coordinate values).""" import itertools -from typing import Any import more_itertools import numpy @@ -19,7 +18,8 @@ NameRow = tuple[Name, Name, Name] KeyMatrix = tuple[KeyRow, ...] NameMatrix = tuple[NameRow, ...] -VMatrix = tuple[tuple[Symbol, KeyRow, NameRow], ...] +VMatrixRow = tuple[Symbol, KeyRow, NameRow] | tuple[Symbol, KeyRow, NameRow, object] +VMatrix = tuple[VMatrixRow, ...] # Build the v-xmatrix parser CHAR = pp.Char(pp.alphas) @@ -57,11 +57,11 @@ def from_data( ) -> VMatrix: """V-Matrix constructor (V-Matrix without numerical coordinate values). - :param symbs: atomic symbols - :param key_mat: key/index columns of the v-matrix, zero-indexed - :param name_mat: coordinate name columns of the v-matrix - :param one_indexed: parameter to store keys in one-indexing - :return: automol V-Matrix data structure + :param symbs: Atomic symbols + :param key_mat: Key/index columns of the v-matrix, zero-indexed + :param name_mat: Coordinate name columns of the v-matrix + :param one_indexed: Parameter to store keys in one-indexing + :return: Automol V-Matrix data structure """ symbs = list(map(ptab.to_symbol, symbs)) natms = len(symbs) @@ -80,7 +80,7 @@ def symbols(vma: VMatrix, idxs: list[int] | None = None) -> list[str]: """Obtain the atomic symbols for all atoms defined in the V-Matrix. :param vma: V-Matrix - :param idxs: indices of atoms to obtain information for + :param idxs: Indices of atoms to obtain information for :return: The list of atomic symbols """ if vma: @@ -96,9 +96,8 @@ def key_matrix(vma: VMatrix, shift: int = 0) -> KeyMatrix: coordinate atom keys by row and column. :param vma: V-Matrix - :type vma: automol V-Matrix data structure - :param shift: value to shift the keys by when obtaining the key matrix - :type shift: int + :type vma: Automol V-Matrix data structure + :param shift: Value to shift the keys by when obtaining the key matrix :return: Key matrix """ if vma: @@ -121,7 +120,7 @@ def name_matrix(vma: VMatrix) -> NameMatrix: coordinate names by row and column. :param vma: V-Matrix - :type vma: automol V-Matrix data structure + :type vma: Automol V-Matrix data structure :return: Name matrix """ if vma: @@ -141,7 +140,7 @@ def count(vma: VMatrix) -> int: real and dummy atoms. :param vma: V-Matrix - :type vma: automol V-Matrix data structure + :type vma: Automol V-Matrix data structure :return: Number of rows """ return len(symbols(vma)) @@ -151,9 +150,9 @@ def atom_indices(vma: VMatrix, symb: str, match: bool = True) -> tuple[int]: """Obtain the indices of a atoms of a particular type in the v-matrix. :param vma: V-Matrix - :param match: grab idxs that match given atom type - :param symb: atomic symbol - :param match: obtain indices of symbols that match the type? + :param match: Grab idxs that match given atom type + :param symb: Atomic symbol + :param match: Obtain indices of symbols that match the type? :return: Indices """ symbs = symbols(vma) @@ -189,18 +188,17 @@ def coordinate_key_matrix(vma: VMatrix, shift: int = 0) -> key_matrix: return tuple(map(tuple, coo_key_mat)) -def coordinates(vma: VMatrix, shift=0, multi=True): +def coordinates( + vma: VMatrix, shift: int = 0, multi: bool = True +) -> dict[str : tuple(int)]: """Obtain the coordinate keys associated with each coordinate name, as a dictionary. Values are sequences of coordinate keys, since there may be multiple. :param vma: V-Matrix - :type vma: automol V-Matrix data structure - :param shift: value to shift the keys by when obtaining the keys - :type shift: int - :param multi: parameter to grab multiple coordinate keys - :type multi: bool - :rtype: dict[str: tuple(int)] + :param shift: Value to shift the keys by when obtaining the keys + :param multi: Parameter to grab multiple coordinate keys + :return: Dictionary of coordinate keys """ _names = numpy.ravel(name_matrix(vma)) coo_keys = numpy.ravel(numpy.array(coordinate_key_matrix(vma, shift), dtype=object)) @@ -217,59 +215,50 @@ def coordinates(vma: VMatrix, shift=0, multi=True): return coo_dct -def distance_coordinates(vma: Any) -> dict[str, tuple[int, ...]]: +def distance_coordinates(vma: VMatrix) -> dict[str, tuple[int, ...]]: """Get the distance coordinates by coordinate name. :param vma: V-Matrix - :type vma: automol V-Matrix data structure :return: The coordinates, by coordinate name """ return dict_.by_key(coordinates(vma, multi=False), distance_names(vma)) -def central_angle_coordinates(vma: Any) -> dict[str, tuple[int, ...]]: +def central_angle_coordinates(vma: VMatrix) -> dict[str, tuple[int, ...]]: """Get the central angle coordinates by coordinate name. :param vma: V-Matrix - :type vma: automol V-Matrix data structure :return: The coordinates, by coordinate name """ return dict_.by_key(coordinates(vma, multi=False), central_angle_names(vma)) -def dihedral_angle_coordinates(vma: Any) -> dict[str, tuple[int, ...]]: +def dihedral_angle_coordinates(vma: VMatrix) -> dict[str, tuple[int, ...]]: """Get the dihedral angle coordinates by coordinate name. :param vma: V-Matrix - :type vma: automol V-Matrix data structure :return: The coordinates, by coordinate name """ return dict_.by_key(coordinates(vma, multi=False), dihedral_angle_names(vma)) -def coordinate(vma, name: str) -> list[int]: +def coordinate(vma: VMatrix, name: str) -> list[int]: """Get the atom keys defining a coordinate by name. :param vma: A v-matrix or z-matrix - :type vma: automol vmat or zmat data structure :param name: The name of the coordinate, e.g. "R5" - :type name: str :return: The atom keys defining the coordinate - :rtype: List[int] """ coo, *_ = coordinates(vma)[name] return coo -def torsion_axis(vma, dih_name: str) -> tuple[int, int]: +def torsion_axis(vma: VMatrix, dih_name: str) -> tuple[int, int]: """Get the rotational axis of a torsion from the dihedral angle name. :param vma: A v-matrix or z-matrix - :type vma: automol vmat or zmat data structure :param dih_name: The dihedral angle name of the torsion - :type dih_name: str :return: The axis, i.e. the central two atoms in the coordinate - :rtype: Tuple[int, int] """ dih_coo = coordinate(vma, dih_name) assert len(dih_coo) == 4, f"{dih_name} is not a dihedral angle:\n{vma}" @@ -278,12 +267,11 @@ def torsion_axis(vma, dih_name: str) -> tuple[int, int]: # # # names and standard naming -def names(vma): +def names(vma: VMatrix) -> tuple[str]: """Obtain names of all coordinates defined in the V-Matrix. :param vma: V-Matrix - :type vma: automol V-Matrix data structure - :rtype: tuple(str) + :return: Names of coordinates """ name_mat = name_matrix(vma) _names = filter(lambda x: x is not None, numpy.ravel(numpy.transpose(name_mat))) @@ -291,53 +279,49 @@ def names(vma): return tuple(more_itertools.unique_everseen(_names)) -def distance_names(vma): +def distance_names(vma: VMatrix) -> tuple[str]: """Obtain names of all distance coordinates defined in the V-Matrix. :param vma: V-Matrix - :type vma: automol V-Matrix data structure - :rtype: tuple(str) + :return: Names of defined coordinates """ name_mat = numpy.array(name_matrix(vma)) return tuple(more_itertools.unique_everseen(name_mat[1:, 0])) -def central_angle_names(vma): +def central_angle_names(vma: VMatrix) -> tuple[str]: """Obtain names of all central-angle coordinates defined in the V-Matrix. :param vma: V-Matrix - :type vma: automol V-Matrix data structure - :rtype: tuple(str) + :return: Names of central angle coordinates """ name_mat = numpy.array(name_matrix(vma)) return tuple(more_itertools.unique_everseen(name_mat[2:, 1])) -def dihedral_angle_names(vma): +def dihedral_angle_names(vma: VMatrix) -> tuple[str]: """Obtain names of all dihedral angle coordinates defined in the V-Matrix. :param vma: V-Matrix - :type vma: automol V-Matrix data structure - :rtype: tuple(str) + :return: Names of dihedral angles """ name_mat = numpy.array(name_matrix(vma)) return tuple(more_itertools.unique_everseen(name_mat[3:, 2])) -def angle_names(vma): +def angle_names(vma: VMatrix) -> tuple[str]: """Obtain names of all angle coordinates defined in the V-Matrix. :param vma: V-Matrix - :type vma: automol V-Matrix data structure - :rtype: tuple(str) + :return: Name of angle coordinates """ return tuple(itertools.chain(central_angle_names(vma), dihedral_angle_names(vma))) -def standard_names(vma, shift=0): +def standard_names(vma: VMatrix, shift: int = 0) -> dict[str, str]: """Build a dictionary that can mas the coordinate names of the input V-Matrix to their name in a standard-form V-Matrix: RN: (1<=N<=Ncoords) @@ -345,10 +329,8 @@ def standard_names(vma, shift=0): DN: (1<=N<=Ncoords). :param vma: V-Matrix - :type vma: automol V-Matrix data structure - :param shift: value to shift the keys by when obtaining the keys - :type shift: int - :rtype: dict[str: str] + :param shift: Value to shift the keys by when obtaining the keys + :return: Dictionary with names """ dist_names = distance_names(vma) cent_ang_names = central_angle_names(vma) @@ -373,7 +355,7 @@ def standard_names(vma, shift=0): return name_dct -def standard_name_matrix(vma, shift=0): +def standard_name_matrix(vma: VMatrix, shift: int = 0) -> tuple[tuple[str]]: """Build a name matrix of the V-Matrix where all of the coordinate names have been standardized: RN: (1<=N<=Ncoords) @@ -398,16 +380,13 @@ def standard_name_matrix(vma, shift=0): return name_mat -def distance_coordinate_name(zma, key1, key2): +def distance_coordinate_name(zma: VMatrix, key1: int, key2: int) -> str: """Get the name of a distance coordinate for a given bond. :param zma: the z-matrix - :type zma: automol Z-Matrix data structure :param key1: the first key - :type key1: int :param key2: the second key - :type key2: int - :rtype: str + :return: Name of coordinate """ key1, key2 = sorted([key1, key2]) name_mat = name_matrix(zma) @@ -420,18 +399,14 @@ def distance_coordinate_name(zma, key1, key2): return name -def central_angle_coordinate_name(zma, key1, key2, key3): +def central_angle_coordinate_name(zma: VMatrix, key1: int, key2: int, key3: int) -> str: """Get the name of angle coordinate for a set of 3 atoms. - :param zma: the z-matrix - :type zma: automol Z-Matrix data structure - :param key1: the first key - :type key1: int - :param key2: the second key (central atom) - :type key2: int - :param key3: the third key - :type key3: int - :rtype: str + :param zma: The z-matrix + :param key1: The first key + :param key2: The second key (central atom) + :param key3: The third key + :return: Name of angle coordinates """ key1, key3 = sorted([key1, key3]) name_mat = name_matrix(zma) @@ -444,20 +419,17 @@ def central_angle_coordinate_name(zma, key1, key2, key3): return name -def dihedral_angle_coordinate_name(zma, key1, key2, key3, key4): +def dihedral_angle_coordinate_name( + zma: VMatrix, key1: int, key2: int, key3: int, key4: int +): """Get the name of dihedral coordinate for a set of 4 atoms. - :param zma: the z-matrix - :type zma: automol Z-Matrix data structure - :param key1: the first key - :type key1: int - :param key2: the second key - :type key2: int - :param key3: the third key - :type key3: int - :param key4: the fourth key - :type key4: int - :rtype: str + :param zma: The z-matrix + :param key1: The first key + :param key2: The second key + :param key3: The third key + :param key4: The fourth key + :return:Name for dihedral coordinate """ if key1 > key4: key1, key2, key3, key4 = key4, key3, key2, key1 @@ -476,24 +448,22 @@ def dihedral_angle_coordinate_name(zma, key1, key2, key3, key4): # # dummy atom functions -def dummy_keys(zma): +def dummy_keys(zma: VMatrix) -> tuple[int]: """Obtain keys to dummy atoms in the Z-Matrix. :param zma: Z-Matrix - :type zma: automol Z-Matrix data structure - :rtype: tuple[int] + :return: Key to dummy atoms """ keys = tuple(key for key, sym in enumerate(symbols(zma)) if sym == "X") return keys -def dummy_coordinate_names(vma): +def dummy_coordinate_names(vma: VMatrix) -> tuple[str]: """Obtain names of all coordinates associated with dummy atoms defined in the V-Matrix. :param vma: V-Matrix - :type vma: automol V-Matrix data structure - :rtype: tuple(str) + :return: Name of coordinates """ symbs = symbols(vma) name_mat = numpy.array(name_matrix(vma))