diff --git a/automol/data/rotor.py b/automol/data/rotor.py index 0768c762..e005adb9 100644 --- a/automol/data/rotor.py +++ b/automol/data/rotor.py @@ -2,6 +2,7 @@ Rotors contain of one or more Torsion data structures """ + import dataclasses import itertools from typing import Dict, List, Optional, Union @@ -210,7 +211,7 @@ def torsion_symmetries(rotor: Rotor) -> List[int]: return tuple(map(tors.symmetry, torsions(rotor))) -def torsion_grids(rotor: Rotor, increment: float=30 * phycon.DEG2RAD) -> List[Grid]: +def torsion_grids(rotor: Rotor, increment: float = 30 * phycon.DEG2RAD) -> List[Grid]: """Get the coordinate grids for the torsions in a rotor :param rotor: A rotor @@ -353,6 +354,7 @@ def rotors_from_zmatrix( """ gra = zmat.graph(zma, stereo=True, dummy=True) if gra is None else gra lin_keys = graph.linear_atom_keys(gra, dummy=True) + znkeys_dct = zmat.neighbor_keys(zma) rot_bkeys = graph.rotational_bond_keys(gra, lin_keys=lin_keys) @@ -369,6 +371,7 @@ def rotors_from_zmatrix( coo=tor_keys, grps=graph.rotational_groups(gra, *tor_axis), symm=graph.rotational_symmetry_number(gra, *tor_axis, lin_keys=lin_keys), + ngrps=list(map(znkeys_dct.get, tor_axis)), ) tor_lst.append(tor) @@ -566,7 +569,7 @@ def rotors_torsion_symmetries( def rotors_torsion_grids( - rotors: List[Rotor], flat: bool = False, increment: float=30 * phycon.DEG2RAD + rotors: List[Rotor], flat: bool = False, increment: float = 30 * phycon.DEG2RAD ) -> Union[List[Grid], List[List[Grid]]]: """Get the torsion coordinate grids from a list of rotors diff --git a/automol/data/tors.py b/automol/data/tors.py index 1983bb66..a4c24aa6 100644 --- a/automol/data/tors.py +++ b/automol/data/tors.py @@ -4,7 +4,7 @@ """ import dataclasses -from typing import Any, List, Optional, Tuple +from typing import Any import numpy import yaml @@ -14,47 +14,43 @@ from .. import zmat from ..util import ZmatConv, zmat_conv -Axis = Tuple[int, int] -DihCoord = Tuple[int, int, int, int] -Groups = Tuple[List[int], List[int]] -Grid = List[float] +Axis = tuple[int, int] +DihCoord = tuple[int, int, int, int] +Group = list[int] +Groups = tuple[Group, Group] +Grid = list[float] @dataclasses.dataclass class Torsion: - """Encodes information for a single torsion, which is one component of a rotor + """Encodes information for a single torsion, which is one component of a rotor. :param name: The z-matrix coordinate name - :type name: str :param coordinate: The z-matrix keys defining the torsion coordinate - :type coordinate: DihKey :param groups: The sets of atoms keys defining the rotational groups - :type groups: Tuple[List[int], List[int]] - :type symmetry: The rotational symmetry number of the torsion - :type symmetry: int + :param symmetry: The rotational symmetry number of the torsion + :param neighbor_groups: The subsets of rotational group atoms that neighbor the axis """ name: str coordinate: DihCoord groups: Groups symmetry: int + neighbor_groups: Groups | None = None # Torsion functions # # Constructors -def from_data(name_: str, coo: DihCoord, grps: Groups, symm: int) -> Torsion: - """Construct a torsion from data +def from_data( + name_: str, coo: DihCoord, grps: Groups, symm: int, ngrps: Groups | None = None +) -> Torsion: + """Construct a torsion from data. :param name_: The z-matrix coordinate name - :type name_: str :param coo: The z-matrix keys defining the torsion coordinate - :type coo: DihKey :param groups: The sets of atoms keys defining the rotational groups - :type groups: Tuple[List[int], List[int]] :type symm: The rotational symmetry number of the torsion - :type symm: int :return: The torsion data structure - :rtype: Torsion """ assert len(coo) == 4, f"Invalid torsion coordinate: {coo}" assert len(grps) == 2, f"Invalid torsion groups: {grps}" @@ -63,54 +59,56 @@ def from_data(name_: str, coo: DihCoord, grps: Groups, symm: int) -> Torsion: coordinate=tuple(coo), groups=tuple(map(tuple, grps)), symmetry=int(symm), + neighbor_groups=None if ngrps is None else tuple(map(tuple, ngrps)), ) # # Getters def name(tor: Torsion) -> str: - """Get the coordinate name of a torsion + """Get the coordinate name of a torsion. :param tor: A torsion - :type tor: Torsion :return: The coordinate name - :rtype: str """ return tor.name def coordinate( - tor: Torsion, key_typ: str = "zmat", zc_: Optional[ZmatConv] = None + tor: Torsion, + key_typ: str = "zmat", + zc_: ZmatConv | None = None, + replace_dummy: bool = True, ) -> DihCoord: - """Get the torsion coordinate keys + """Get the torsion coordinate keys. :param tor: A torsion - :type tor: Torsion :param key_typ: The type of keys to return, "zmat" (default) or "geom" - :type key_typ: str, optional :param zc_: Z-matrix conversion info, to avoid re-calculation, defaults to None - :type zc_: Optional[ZmatConv], optional + :param replace_dummy: Replace dummy atom with other (presumably in-line) neighbor? :return: The torsion rotational keys - :rtype: str """ coo = tor.coordinate if key_typ == "geom": + if replace_dummy: + dkeys = zmat_conv.dummy_keys(zc_) + ngrps = neighbor_groups(tor, dummy=False, zc_=zc_) + end_keys0 = [coo[0], coo[-1]] + end_keys = [ + k if k not in dkeys else next(iter(nks)) + for k, nks in zip(end_keys0, ngrps, strict=True) + ] + coo = (end_keys[0], coo[1], coo[2], end_keys[1]) coo = zmat_conv.relabel_zmatrix_key_sequence(zc_, coo, dummy=True) return coo -def groups( - tor: Torsion, key_typ: str = "zmat", zc_: Optional[ZmatConv] = None -) -> Groups: - """Get the rotational groups of a torsion +def groups(tor: Torsion, key_typ: str = "zmat", zc_: ZmatConv | None = None) -> Groups: + """Get the rotational groups of a torsion. :param tor: A torsion - :type tor: Torsion :param key_typ: The type of keys to return, "zmat" (default) or "geom" - :type key_typ: str, optional :param zc_: Z-matrix conversion info, to avoid re-calculation, defaults to None - :type zc_: Optional[ZmatConv], optional :return: The torsion rotational groups - :rtype: str """ grps = tor.groups if key_typ == "geom": @@ -119,100 +117,145 @@ def groups( def symmetry(tor: Torsion) -> int: - """Get the rotational symmetry of a torsion + """Get the rotational symmetry of a torsion. :param tor: A torsion - :type tor: Torsion :return: The rotational symmetry - :rtype: int """ return tor.symmetry +def neighbor_groups( + tor: Torsion, key_typ: str = "zmat", zc_: ZmatConv | None = None, dummy: bool = True +) -> Groups: + """Get the neighbor groups of a torsion. + + :param tor: A torsion + :param key_typ: The type of keys to return, "zmat" (default) or "geom" + :param zc_: Z-matrix conversion info, to avoid re-calculation, defaults to None + :param dummy: Include dummy atom neighbors? + :return: The neighbor groups + """ + ngrps = tor.neighbor_groups + if not dummy: + dkeys = zmat_conv.dummy_keys(zc_) + ngrps = tuple([k for k in g if k not in dkeys] for g in ngrps) + + if key_typ == "geom": + ngrps = zmat_conv.relabel_zmatrix_key_sequence(zc_, ngrps) + return ngrps + + # setters def set_name(tor: Torsion, name_: str) -> Torsion: - """Set the coordinate name of a torsion + """Set the coordinate name of a torsion. :param tor: A torsion :param name_: The coordinate name :return: The torsion """ return from_data( - name_=name_, coo=coordinate(tor), grps=groups(tor), symm=symmetry(tor) + name_=name_, + coo=coordinate(tor), + grps=groups(tor), + symm=symmetry(tor), + ngrps=neighbor_groups(tor), ) def set_coordinate(tor: Torsion, coo: DihCoord) -> Torsion: - """get the torsion coordinate keys + """Set the torsion coordinate keys. :param tor: A torsion :param coo: The torsion coordinate keys :return: The torsion """ - return from_data(name_=name(tor), coo=coo, grps=groups(tor), symm=symmetry(tor)) + return from_data( + name_=name(tor), + coo=coo, + grps=groups(tor), + symm=symmetry(tor), + ngrps=neighbor_groups(tor), + ) def set_groups(tor: Torsion, grps: Groups) -> Torsion: - """get the rotational groups of a torsion + """Set the rotational groups of a torsion. :param tor: A torsion :param grps: The torsion rotational groups :return: The torsion """ return from_data( - name_=name(tor), coo=coordinate(tor), grps=grps, symm=symmetry(tor) + name_=name(tor), + coo=coordinate(tor), + grps=grps, + symm=symmetry(tor), + ngrps=neighbor_groups(tor), ) def set_symmetry(tor: Torsion, symm: int) -> Torsion: - """get the rotational symmetry of a torsion + """Set the rotational symmetry of a torsion. :param tor: A torsion :param symm: The rotational symmetry :return: The torsion """ - return from_data(name_=name(tor), coo=coordinate(tor), grps=groups(tor), symm=symm) + return from_data( + name_=name(tor), + coo=coordinate(tor), + grps=groups(tor), + symm=symm, + ngrps=neighbor_groups(tor), + ) + + +def set_neighbor_groups(tor: Torsion, ngrps: Groups) -> Torsion: + """Set the neighbor groups of a torsion. + + :param tor: A torsion + :param ngrps: The neighbor groups + :return: The torsion + """ + return from_data( + name_=name(tor), + coo=coordinate(tor), + grps=groups(tor), + symm=symmetry(tor), + ngrps=ngrps, + ) # properties -def axis(tor: Torsion, key_typ: str = "zmat", zc_: Optional[ZmatConv] = None) -> Axis: - """Get the rotational axis of a torsion +def axis(tor: Torsion, key_typ: str = "zmat", zc_: ZmatConv | None = None) -> Axis: + """Get the rotational axis of a torsion. :param tor: A torsion - :type tor: Torsion :param key_typ: The type of keys to return, "zmat" (default) or "geom" - :type key_typ: str, optional :param zc_: Z-matrix conversion info, to avoid re-calculation, defaults to None - :type zc_: Optional[ZmatConv], optional :return: The torsion rotational axis - :rtype: str """ ks_ = coordinate(tor, key_typ=key_typ, zc_=zc_) return ks_[1:3] def span(tor: Torsion) -> float: - """Get the angular span of a torsion, based on the symmetry number + """Get the angular span of a torsion, based on the symmetry number. :param tor: A torsion - :type tor: Torsion :return: The angular span of the torsion, 2 pi / symmetry - :rtype: float """ return 2 * numpy.pi / symmetry(tor) def grid(tor: Torsion, zma, increment: float = 30 * phycon.DEG2RAD) -> Grid: - """Get the coordinate grid for a torsion + """Get the coordinate grid for a torsion. :param tor: A torsion - :type tor: Torsion :param zma: The z-matrix associated with this torsion - :type zma: automol zmat data structure :param increment: The grid increment, in radians - :type increment: float :return: The coordinate grid - :rtype: Grid """ # [0, 30, 60, 90, ...] << in degrees vals = numpy.arange(0, span(tor), increment) @@ -224,27 +267,25 @@ def grid(tor: Torsion, zma, increment: float = 30 * phycon.DEG2RAD) -> Grid: # # Transformations def with_geometry_indices(tor: Torsion, zc_: ZmatConv) -> Torsion: """Given a z-matrix torsion and a z-matrix conversion, return a torsion with - geometry indices + geometry indices. That is, the axis and group indices will skip dummy atoms :param tor: A torsion data structure - :type tor: Torsion :param zc_: A z-matrix conversion data structure - :type zc_: ZmatConv :return: A torsion data structure using geometry indices - :rtype: Torsion """ return from_data( name_=name(tor), coo=coordinate(tor, key_typ="geom", zc_=zc_), grps=groups(tor, key_typ="geom", zc_=zc_), symm=symmetry(tor), + ngrps=neighbor_groups(tor, key_typ="geom", zc_=zc_), ) def update_against_zmatrix(tor: Torsion, zma: Any) -> Torsion: - """Update a torsion object from a z-matrix + """Update a torsion object from a z-matrix. Temporarily needed to make sure the torsion object contains sufficient information @@ -255,49 +296,45 @@ def update_against_zmatrix(tor: Torsion, zma: Any) -> Torsion: coo = list(reversed(zmat.coordinate(zma, name(tor)))) # In case the coordinate was re-ordered, make sure the groups match grps = tuple(sorted(groups(tor), key=lambda g: coo[-1] in g)) - return from_data(name_=name(tor), coo=coo, grps=grps, symm=symmetry(tor)) + # Re-evaluate the neighbor groups, in case they are missing + nkeys_dct = zmat.neighbor_keys(zma) + ngrps = tuple(sorted(nkeys_dct[k]) for k in coo[1:3]) + return from_data( + name_=name(tor), coo=coo, grps=grps, symm=symmetry(tor), ngrps=ngrps + ) -# Torsion List functions -def torsions_string(tor_lst: List[Torsion], one_indexed: bool = True) -> str: - """Write a list of torsions to a string +# Torsion list functions +def torsions_string(tor_lst: list[Torsion], one_indexed: bool = True) -> str: + """Write a list of torsions to a string. :param tor_lst: A list of torsions - :type tor_lst: List[Torsion] :param one_indexed: Is this a one-indexed string? defaults to True - :type one_indexed: bool, optional :returns: A string representations of the torsions, as a flattened list - :rtype: str """ tor_yml_dct = torsions_yaml_data(tor_lst, one_indexed=one_indexed) tor_str = yaml.dump(tor_yml_dct, sort_keys=False) return tor_str -def torsions_from_string(tor_str: str, one_indexed: bool = True) -> List[Torsion]: - """Write a list of torsions to a string +def torsions_from_string(tor_str: str, one_indexed: bool = True) -> list[Torsion]: + """Write a list of torsions to a string. :param tor_str: A string representations of the torsions, as a flattened list - :type tor_str: str :param one_indexed: Is this a one-indexed string? defaults to True - :type one_indexed: bool, optional :returns: A list of torsions - :rtype: List[Torsion] """ tor_yml_dct = yaml.load(tor_str, Loader=yaml.FullLoader) tor_lst = torsions_from_yaml_data(tor_yml_dct, one_indexed=one_indexed) return tor_lst -def torsions_yaml_data(tor_lst: List[Torsion], one_indexed: bool = True) -> dict: - """Write a list of torsions to a yaml-formatted dictionary +def torsions_yaml_data(tor_lst: list[Torsion], one_indexed: bool = True) -> dict: + """Write a list of torsions to a yaml-formatted dictionary. :param tor_lst: A list of torsions - :type tor_lst: List[Torsion] :param one_indexed: Is this a one-indexed string? defaults to True - :type one_indexed: bool, optional :returns: A string representations of the torsions, as a flattened list - :rtype: str """ shift = 1 if one_indexed else 0 @@ -305,40 +342,39 @@ def torsions_yaml_data(tor_lst: List[Torsion], one_indexed: bool = True) -> dict for tor in tor_lst: coo = [k + shift for k in coordinate(tor)] axis1, axis2 = (k + shift for k in axis(tor)) - grps = ([k + shift for k in g] for g in groups(tor)) - grp1, grp2 = ("-".join(map(str, g)) if len(g) > 1 else g[0] for g in grps) coo_str = "-".join("*" if k is None else str(k) for k in coo) - tor_yml_dct[name(tor)] = { + yml_dct = { "axis1": axis1, - "group1": grp1, "axis2": axis2, - "group2": grp2, "symmetry": symmetry(tor), "coordinate": coo_str, } + yml_dct = _write_groups_to_yaml_dict( + yml_dct, groups(tor), shift=shift, key_prefix="group" + ) + yml_dct = _write_groups_to_yaml_dict( + yml_dct, neighbor_groups(tor), shift=shift, key_prefix="neighbor_group" + ) + tor_yml_dct[name(tor)] = yml_dct return tor_yml_dct def torsions_from_yaml_data(tor_yml_dct: dict, one_indexed: bool = True) -> dict: - """Read a list of torsions out of a yaml-formatted torsion dictionary + """Read a list of torsions out of a yaml-formatted torsion dictionary. :param tor_lst: A list of torsions - :type tor_lst: List[Torsion] :param one_indexed: Is this a one-indexed string? defaults to True - :type one_indexed: bool, optional :returns: A string representations of the torsions, as a flattened list - :rtype: str """ shift = -1 if one_indexed else 0 tor_lst = [] - for name_, vals_dct in tor_yml_dct.items(): - ax_ = list(map(vals_dct.__getitem__, ["axis1", "axis2"])) + for name_, yml_dct in tor_yml_dct.items(): + ax_ = list(map(yml_dct.get, ["axis1", "axis2"])) ax_ = [k + shift for k in ax_] - grps = list(map(vals_dct.__getitem__, ["group1", "group2"])) - grps = [[g] if isinstance(g, int) else map(int, g.split("-")) for g in grps] - grps = [[k + shift for k in g] for g in grps] - coo = vals_dct.get("coordinate", None) + grps = _read_groups_from_yaml_dict(yml_dct, shift, "group") + ngrps = _read_groups_from_yaml_dict(yml_dct, shift, "neighbor_group") + coo = yml_dct.get("coordinate", None) if coo is None: coo = [None, *ax_, None] else: @@ -348,9 +384,78 @@ def torsions_from_yaml_data(tor_yml_dct: dict, one_indexed: bool = True) -> dict name_=name_, coo=coo, grps=grps, - symm=vals_dct["symmetry"], + symm=yml_dct["symmetry"], + ngrps=ngrps, ) tor_lst.append(tor) return tuple(tor_lst) + + +# helper functions +def _write_groups_to_yaml_dict( + yml_dct: dict[object, object], + grps: Groups | None, + shift: int, + key_prefix: str, +) -> dict[object, object]: + """Write rotational groups to a YAML dictionary. + + :param yml_dct: YAML dictionary + :param grps: Rotational groups + :param shift: The shift to apply to the keys + :param key_prefix: Prefix for key + :return: Updated YAML dictionary + """ + if grps is None: + return yml_dct + + yml_dct = yml_dct.copy() + grp_vals = tuple(_group_to_yaml_value(g, shift) for g in grps) + yml_dct.update({f"{key_prefix}{i+1}": v for i, v in enumerate(grp_vals)}) + return yml_dct + + +def _read_groups_from_yaml_dict( + yml_dct: dict[object, object], + shift: int, + key_prefix: str, +) -> Groups | None: + """Read rotational groups from a YAML dictionary. + + :param yml_dct: YAML dictionary + :param shift: The shift to apply to the keys + :param key_prefix: Prefix for key + :return: Updated YAML dictionary + """ + grp_vals = [yml_dct.get(f"{key_prefix}{i+1}") for i in range(2)] + any_none = any(v is None for v in grp_vals) + all_none = all(v is None for v in grp_vals) + assert all_none or not any_none, f"In consistent group values: {yml_dct}" + if any_none: + return None + return tuple(_group_from_yaml_value(v, shift) for v in grp_vals) + + +def _group_to_yaml_value(grp: Group, shift: int) -> str | int: + """Get a YAML value to represent a group. + + :param grp: The rotational group + :param shift: The shift to apply to the keys + :return: The YAML value for the group + """ + grp = [k + shift for k in grp] + return "-".join(map(str, grp)) if len(grp) > 1 else grp[0] + + +def _group_from_yaml_value(grp_val: str | int, shift: int) -> Group: + """Get a group from a YAML value. + + :param grp_str: The string representation of the rotational group + :param shift: The shift to apply to the keys + :return: The rotational group + """ + grp = [grp_val] if isinstance(grp_val, int) else map(int, grp_val.split("-")) + grp = [k + shift for k in grp] + return grp diff --git a/automol/geom/_conv.py b/automol/geom/_conv.py index 1f3deea4..a0cf663d 100644 --- a/automol/geom/_conv.py +++ b/automol/geom/_conv.py @@ -624,8 +624,10 @@ def display( # If requested, visualize only a subset of the bonds by removing others if vis_bkeys is not None: - excl_bkeys = graph_base.bond_keys(gra) - set(map(frozenset, vis_bkeys)) - gra = graph_base.remove_bonds(gra, excl_bkeys, stereo=False, check=False) + all_bkeys = graph_base.bond_keys(gra) + vis_bkeys = set(map(frozenset, vis_bkeys)) + gra = graph_base.remove_bonds(gra, all_bkeys - vis_bkeys, stereo=False) + gra = graph_base.add_bonds(gra, vis_bkeys - all_bkeys) view = py3dmol_view(geo, gra=gra, view=view, image_size=image_size, mode=mode) diff --git a/automol/tests/test_rotor.py b/automol/tests/test_rotor.py index 97f99dc8..e4218f2a 100644 --- a/automol/tests/test_rotor.py +++ b/automol/tests/test_rotor.py @@ -393,10 +393,22 @@ def test__consistency(): assert tor_grps == ref_tor_grps, f"{tor_grps} != {ref_tor_grps}" +def test__dummy_replacement(): + """Check that dummy atoms get replaced upon conversion to geometry.""" + zma = C2H5OH_CH3_ZMA + zgra = reac.ts_graph(C2H5OH_CH3_ZRXN) + rotors = rotor.rotors_from_zmatrix(zma, gra=zgra) + + tor_coos = rotor.rotors_torsion_coordinates(rotors, key_typ="geom") + ref_tor_coos = (((2, 0, 1, 3),), ((0, 1, 3, 6),), ((3, 6, 9, 10),)) + assert tor_coos == ref_tor_coos, f"{tor_coos} != {ref_tor_coos}" + + if __name__ == "__main__": # test__rotor() # test__rotor_with_dummy_atoms() # test__rotor_multidimensional() # test__torsion_list_string() # test__rotor_for_ts() - test__consistency() + # test__consistency() + test__dummy_replacement() diff --git a/automol/vmat.py b/automol/vmat.py index f83fb876..0f04ea17 100644 --- a/automol/vmat.py +++ b/automol/vmat.py @@ -1,6 +1,7 @@ """V-Matrix: Variable V-Matrix (V-Matrix without coordinate values).""" import itertools +from collections import defaultdict import more_itertools import numpy @@ -524,6 +525,20 @@ def conversion_info(zma: VMatrix) -> ZmatConv: return zmat_conv.from_zmat_data(zcount, src_zkeys_dct) +def neighbor_keys(vma: VMatrix) -> dict[Key, frozenset[Key]]: + """Identify which atoms are explicit neighbors in the V-Matrix. + + :param vma: V-Matrix + :return: A dictionary mapping atoms onto their neighbors + """ + dist_coos = list(distance_coordinates(vma).values()) + nkeys_dct = defaultdict(set) + for key1, key2 in dist_coos: + nkeys_dct[key1].add(key2) + nkeys_dct[key2].add(key1) + return dict_.transform_values(nkeys_dct, frozenset) + + # # V-Matrix-specific functions # # # setters def set_key_matrix(vma: VMatrix, key_mat: KeyMatrix) -> VMatrix: diff --git a/automol/zmat/__init__.py b/automol/zmat/__init__.py index 54d091b6..9db6a2bc 100644 --- a/automol/zmat/__init__.py +++ b/automol/zmat/__init__.py @@ -39,6 +39,7 @@ from ..vmat import dummy_keys from ..vmat import dummy_source_dict from ..vmat import conversion_info +from ..vmat import neighbor_keys # core functions # # constructors from .base._core import from_data @@ -172,6 +173,7 @@ 'dummy_keys', 'dummy_source_dict', 'conversion_info', + 'neighbor_keys', # extra base functions 'samples', 'constraint_dict', diff --git a/automol/zmat/base/__init__.py b/automol/zmat/base/__init__.py index 06b5aca2..9d17f133 100644 --- a/automol/zmat/base/__init__.py +++ b/automol/zmat/base/__init__.py @@ -32,6 +32,7 @@ from ...vmat import dummy_keys from ...vmat import dummy_source_dict from ...vmat import conversion_info +from ...vmat import neighbor_keys # core functions # # constructors from ._core import from_data @@ -136,6 +137,7 @@ 'dummy_keys', 'dummy_source_dict', 'conversion_info', + 'neighbor_keys', # extra functions 'samples', 'constraint_dict',