Source code for morfeus.cone_angle

"""Cone angle code."""

from __future__ import annotations

from collections.abc import Iterable
import functools
import itertools
import math
import typing
from typing import Any
import warnings

import numpy as np

from morfeus.data import atomic_symbols, jmol_colors
from morfeus.geometry import Atom, Cone
from morfeus.io import read_geometry
from morfeus.plotting import get_drawing_cone
from morfeus.typing import Array1DFloat, Array2DFloat, ArrayLike1D, ArrayLike2D
from morfeus.utils import (
    check_distances,
    convert_elements,
    get_radii,
    Import,
    requires_dependency,
)

if typing.TYPE_CHECKING:
    from matplotlib.colors import hex2color
    import pyvista as pv
    from pyvistaqt import BackgroundPlotter


[docs] class ConeAngle: """Calculates and stores the results of exact cone angle calculation. As described in J. Comput. Chem. 2013, 34, 1189. Args: elements: Elements as atomic symbols or numbers coordinates: Coordinates (Å) atom_1: Index of central atom (1-inexed) radii: vdW radii (Å) radii_type: Type of vdW radii: 'alvarez', 'bondi', 'crc' or 'truhlar' method: Method of calculation: 'internal' or 'libconeangle' (default) Attributes: cone_angle: Exact cone angle (degrees) tangent_atoms: Atoms tangent to cone (1-indexed) Raises: RunTimeError: If cone angle could not be found by internal algorithm ValueError: If atoms within vdW radius of central atom or if exception happened with libconeangle or if wrong method chosen """ cone_angle: float tangent_atoms: list[int] _atoms: list[Atom] _max_2_cone: Cone def __init__( # noqa: C901 self, elements: Iterable[int] | Iterable[str], coordinates: ArrayLike2D, atom_1: int, radii: ArrayLike1D | None = None, radii_type: str = "crc", method: str = "libconeangle", ) -> None: # Convert elements to atomic numbers if they are symbols elements = convert_elements(elements, output="numbers") coordinates: Array2DFloat = np.array(coordinates) # Get radii if they are not supplied if radii is None: radii = get_radii(elements, radii_type=radii_type) radii: Array1DFloat = np.array(radii) # Check so that no atom is within vdW distance of atom 1 within = check_distances(elements, coordinates, atom_1, radii=radii) if len(within) > 0: atom_string = " ".join([str(i) for i in within]) raise ValueError("Atoms within vdW radius of central atom:", atom_string) # Set up coordinate array and translate coordinates coordinates -= coordinates[atom_1 - 1] # Get list of atoms as Atom objects atoms: list[Atom] = [] for i, (element, coord, radius) in enumerate( zip(elements, coordinates, radii), start=1 ): if i != atom_1: atom = Atom(element, coord, radius, i) atom.get_cone() atoms.append(atom) self._atoms = atoms # Calculate cone angle if method == "libconeangle": try: from libconeangle import cone_angle angle, axis, tangent_atoms = cone_angle(coordinates, radii, atom_1 - 1) self.cone_angle = angle self.tangent_atoms = [i + 1 for i in tangent_atoms] atoms = [ atom for atom in self._atoms if atom.index in self.tangent_atoms ] self._cone = Cone(self.cone_angle, atoms, axis) except ImportError: warnings.warn( "Failed to import libconeangle. Defaulting to method='internal'", stacklevel=2, ) self._cone_angle_internal() elif method == "internal": self._cone_angle_internal() else: raise ValueError( "Method not implemented. Choose between 'libconeangle' and 'internal'" )
[docs] def print_report(self) -> None: """Prints report of results.""" tangent_atoms = [ atom for atom in self._atoms if atom.index in self.tangent_atoms ] tangent_labels = [ f"{atomic_symbols[atom.element]}{atom.index}" for atom in tangent_atoms ] tangent_string = " ".join(tangent_labels) print(f"Cone angle: {self.cone_angle:.1f}") print(f"No. tangent atoms: {len(tangent_atoms)}") print(f"Tangent to: {tangent_string}")
def _cone_angle_internal(self) -> None: """Calculates cone angle with internal algorithm. Raises: RuntimeError: If cone cannot be found. """ # Search for cone over single atoms cone = self._search_one_cones() # Prune out atoms that lie in the shadow of another atom's cone if cone is None: loop_atoms = list(self._atoms) remove_atoms: set[Atom] = set() for cone_atom in loop_atoms: for test_atom in loop_atoms: if cone_atom is not test_atom: if cone_atom.cone.is_inside(test_atom): remove_atoms.add(test_atom) for atom in remove_atoms: loop_atoms.remove(atom) self._loop_atoms = loop_atoms # Search for cone over pairs of atoms if cone is None: cone = self._search_two_cones() # Search for cones over triples of atoms if cone is None: cone = self._search_three_cones() # Check if no cone was found if cone is None: raise RuntimeError("Cone not found") # Set attributes self._cone = cone self.cone_angle = math.degrees(cone.angle * 2) self.tangent_atoms = [atom.index for atom in cone.atoms] def _get_upper_bound(self) -> float: """Calculates upper bound for apex angle. Returns: upper_bound: Upper bound to apex angle (radians) """ # Calculate unit vector to centroid coordinates: Array2DFloat = np.array([atom.coordinates for atom in self._atoms]) centroid_vector = np.mean(coordinates, axis=0) centroid_unit_vector = centroid_vector / np.linalg.norm(centroid_vector) # Getting sums of angle to centroid and vertex angle. angle_sums = [] for atom in self._atoms: cone = atom.cone cos_angle = np.dot(centroid_unit_vector, cone.normal) vertex_angle = math.acos(cos_angle) angle_sum = cone.angle + vertex_angle angle_sums.append(angle_sum) # Select upper bound as the maximum angle upper_bound = max(angle_sums) return upper_bound def _search_one_cones(self) -> Cone | None: """Searches over cones tangent to one atom. Returns: max_1_cone: Largest cone tangent to one atom """ # Get the largest cone atoms = self._atoms alphas: list[float] = [] for atom in atoms: alphas.append(atom.cone.angle) idx = int(np.argmax(alphas)) max_1_cone = atoms[idx].cone # Check if all atoms are contained in cone. If yes, return cone, # otherwise, return None. in_atoms = [] test_atoms = [atom for atom in atoms if atom not in max_1_cone.atoms] for atom in test_atoms: in_atoms.append(max_1_cone.is_inside(atom)) if all(in_atoms): return max_1_cone else: return None def _search_two_cones(self) -> Cone | None: """Search over cones tangent to two atoms. Returns: max_2_cone: Largest cone tangent to two atoms """ # Create two-atom cones loop_atoms = self._loop_atoms cones = [] for atom_i, atom_j in itertools.combinations(loop_atoms, r=2): cone = _get_two_atom_cone(atom_i, atom_j) cones.append(cone) # Select largest two-atom cone angles = [cone.angle for cone in cones] idx = int(np.argmax(angles)) max_2_cone = cones[idx] self._max_2_cone = max_2_cone # Check if all atoms are contained in cone. If yes, return cone, # otherwise, return None in_atoms = [] for atom in loop_atoms: in_atoms.append(max_2_cone.is_inside(atom)) if all(in_atoms): return max_2_cone else: return None def _search_three_cones(self) -> Cone: """Search over cones tangent to three atoms. Returns: min_3_cone: Smallest cone tangent to three atoms """ # Create three-atom cones loop_atoms = self._loop_atoms cones = [] for atom_i, atom_j, atom_k in itertools.combinations(loop_atoms, r=3): three_cones = _get_three_atom_cones(atom_i, atom_j, atom_k) cones.extend(three_cones) # Get upper and lower bound to apex angle upper_bound = self._get_upper_bound() lower_bound = self._max_2_cone.angle # Remove cones from consideration which are outside the bounds remove_cones = [] for cone in cones: if cone.angle - lower_bound < -1e-5 or upper_bound - cone.angle < -1e-5: remove_cones.append(cone) for cone in reversed(remove_cones): cones.remove(cone) # Keep only cones that encompass all atoms keep_cones: list[Cone] = [] for cone in cones: in_atoms = [] for atom in loop_atoms: in_atoms.append(cone.is_inside(atom)) if all(in_atoms): keep_cones.append(cone) # Take the smallest cone that encompasses all atoms angles = [cone.angle for cone in keep_cones] idx = int(np.argmin(angles)) min_3_cone = keep_cones[idx] return min_3_cone @requires_dependency( [ Import(module="matplotlib.colors", item="hex2color"), Import(module="pyvista", alias="pv"), Import(module="pyvistaqt", item="BackgroundPlotter"), ], globals(), ) def draw_3D( self, atom_scale: float = 1, background_color: str = "white", cone_color: str = "steelblue", cone_opacity: float = 0.75, ) -> None: """Draw a 3D representation of the molecule with the cone. Args: atom_scale: Scaling factor for atom size background_color: Background color for plot cone_color: Cone color cone_opacity: Cone opacity """ # Set up plotter p = BackgroundPlotter() p.set_background(background_color) # Draw molecule for atom in self._atoms: color = hex2color(jmol_colors[atom.element]) radius = atom.radius * atom_scale sphere = pv.Sphere(center=list(atom.coordinates), radius=radius) p.add_mesh(sphere, color=color, opacity=1, name=str(atom.index)) # Determine direction and extension of cone angle = math.degrees(self._cone.angle) coordinates: Array2DFloat = np.array([atom.coordinates for atom in self._atoms]) radii: Array1DFloat = np.array([atom.radius for atom in self._atoms]) if angle > 180: normal = -self._cone.normal else: normal = self._cone.normal projected = np.dot(normal, coordinates.T) + np.array(radii) max_extension = np.max(projected) if angle > 180: max_extension += 1 # Make the cone cone = get_drawing_cone( center=[0, 0, 0] + (max_extension * normal) / 2, direction=-normal, angle=angle, height=max_extension, capping=False, resolution=100, ) p.add_mesh(cone, opacity=cone_opacity, color=cone_color) def __repr__(self) -> str: return f"{self.__class__.__name__}({len(self._atoms)!r} atoms)"
def _get_two_atom_cone(atom_i: Atom, atom_j: Atom) -> Cone: """Creates a cone tangent to two atoms. Args: atom_i: First tangent atom atom_j: Second tangent atom Returns: cone: Cone tangent to the two atoms """ # Get the cone angle cone_i = atom_i.cone cone_j = atom_j.cone beta_i = cone_i.angle beta_j = cone_j.angle beta_ij = math.acos(np.dot(atom_i.cone.normal, atom_j.cone.normal)) alpha_ij = (beta_ij + beta_i + beta_j) / 2 # Get the cone normal a_ij = (1 / math.sin(beta_ij)) * math.sin(0.5 * (beta_ij + beta_i - beta_j)) b_ij = (1 / math.sin(beta_ij)) * math.sin(0.5 * (beta_ij - beta_i + beta_j)) c_ij = 0 n = a_ij * cone_i.normal + b_ij * cone_j.normal + c_ij n = n / np.linalg.norm(n) # Create cone angle = alpha_ij normal = n cone = Cone(angle, [atom_i, atom_j], normal) return cone def _get_three_atom_cones(atom_i: Atom, atom_j: Atom, atom_k: Atom) -> list[Cone]: """Creates cones tangent to three atoms. Args: atom_i: First tangent atom atom_j: Second tangent atom atom_k: Third tangent atom Returns: cones: Cones tangent to the three atoms """ # Set up vertex angles beta_i = atom_i.cone.angle beta_j = atom_j.cone.angle beta_k = atom_k.cone.angle # Set up angles between atom vectors beta_ij = math.acos(np.dot(atom_i.cone.normal, atom_j.cone.normal)) # Set up normal vectors to atoms m_i = atom_i.cone.normal m_j = atom_j.cone.normal m_k = atom_k.cone.normal # Setup matrices u: Array1DFloat = np.array([math.cos(beta_i), math.cos(beta_j), math.cos(beta_k)]) v: Array1DFloat = np.array([math.sin(beta_i), math.sin(beta_j), math.sin(beta_k)]) N: Array2DFloat = np.array( [np.cross(m_j, m_k), np.cross(m_k, m_i), np.cross(m_i, m_j)] ).T P: Array2DFloat = N.T @ N gamma = np.dot(m_i, np.cross(m_j, m_k)) # Set up coefficients of quadratic equation A = u @ P @ u B = v.T @ P @ v C = u.T @ P @ v D = gamma**2 # Solve quadratic equation p2 = (A - B) ** 2 + 4 * C**2 p1 = 2 * (A - B) * (A + B - 2 * D) p0 = (A + B - 2 * D) ** 2 - 4 * C**2 roots = np.roots([p2, p1, p0]) roots = np.real_if_close(roots, tol=1e10) roots[np.isclose(roots, 1, rtol=1e-9, atol=0.0)] = 1 roots[np.isclose(roots, -1, rtol=1e-9, atol=0.0)] = -1 cos_roots = [ math.acos(roots[0]), 2 * np.pi - math.acos(roots[0]), math.acos(roots[1]), 2 * np.pi - math.acos(roots[1]), ] # Test roots and keep only those that are physical angles = [] D_tests = [] for root in cos_roots: alpha = root / 2 test = ( A * math.cos(alpha) ** 2 + B * math.sin(alpha) ** 2 + 2 * C * math.sin(alpha) * math.cos(alpha) ) D_test = abs(test - D) angles.append(alpha) D_tests.append(D_test) angles = np.array(angles) D_tests = np.array(D_tests) physical_angles = angles[np.argsort(D_tests)][:2] # Create cones for physical angles cones = [] for alpha in physical_angles: # Calculate normal vector a_ij = ( math.cos(alpha - beta_i) - math.cos(alpha - beta_j) * math.cos(beta_ij) ) / math.sin(beta_ij) ** 2 b_ij = ( math.cos(alpha - beta_j) - math.cos(alpha - beta_i) * math.cos(beta_ij) ) / math.sin(beta_ij) ** 2 c_ij_squared = 1 - a_ij**2 - b_ij**2 - 2 * a_ij * b_ij * math.cos(beta_ij) # Set c_ij_squared to 0 if negative due to numerical precision. if c_ij_squared < 0: c_ij_squared = 0 c_ij = math.sqrt(c_ij_squared) p = N @ (u * math.cos(alpha) + v * math.sin(alpha)).reshape(-1) sign = np.sign(gamma) * np.sign(np.dot(p, np.cross(m_i, m_j))) if np.sign(c_ij) != sign: c_ij = -c_ij n = a_ij * m_i + b_ij * m_j + c_ij * 1 / math.sin(beta_ij) * np.cross(m_i, m_j) # Create cone cone = Cone(alpha, [atom_i, atom_j, atom_k], n) cones.append(cone) return cones
[docs] def cli(file: str) -> Any: """CLI for cone angle. Args: file: Geometry file Returns: Partially instantiated class """ elements, coordinates = read_geometry(file) return functools.partial(ConeAngle, elements, coordinates)