from typing import List, Optional, Tuple, Dict
import functools
import logging
import numpy as np
from pydantic import Field, validator
from . import base, orutils
from .conformer import Conformer, ConformerGenerationOptions
from .moleculebase import BaseMolecule
logger = logging.getLogger(__name__)
[docs]class Molecule(BaseMolecule):
"""Class to manage a molecule"""
charge: Optional[int] = Field(
default=None,
description=("Charge to apply to the molecule. "
"If `charge=None`, the charge is taken to be "
"the molecular_charge on the QCElemental molecule")
)
multiplicity: Optional[int] = Field(
default=None,
description=("Multiplicity to apply to the molecule. "
"If `multiplicity=None`, the multiplicity is taken to be "
"the molecular_multiplicity on the QCElemental molecule")
)
optimize_geometry: bool = Field(
default=False,
description="Whether to optimize the geometry of conformers",
)
conformers: List[Conformer] = Field(
default_factory=list,
description="List of psiresp.conformer.Conformers of the molecule"
)
conformer_generation_options: ConformerGenerationOptions = Field(
default_factory=ConformerGenerationOptions,
description="Conformer generation options",
)
stage_1_unrestrained_charges: Optional[np.ndarray] = Field(
default=None,
description=("Stage 1 unrestrained charges. "
"These are typically assigned from a "
":class:`psiresp.job.Job`.")
)
stage_1_restrained_charges: Optional[np.ndarray] = Field(
default=None,
description=("Stage 1 restrained charges. "
"These are typically assigned from a "
":class:`psiresp.job.Job`.")
)
stage_2_unrestrained_charges: Optional[np.ndarray] = Field(
default=None,
description=("Stage 2 unrestrained charges. "
"These are typically assigned from a "
":class:`psiresp.job.Job`.")
)
stage_2_restrained_charges: Optional[np.ndarray] = Field(
default=None,
description=("Stage 2 restrained charges. "
"These are typically assigned from a "
":class:`psiresp.job.Job`.")
)
keep_original_orientation: bool = Field(
default=False,
description=("Whether to use the original orientation of the conformer. "
"If `keep_original_orientation=False` but "
"generate_orientations() is called without specifying "
"specific reorientations, rotations, or translations, "
"this is ignored and the original conformer geometry is used.")
)
reorientations: List[Tuple[int, int, int]] = Field(
default_factory=list,
description=("Specific rigid-body reorientations to generate. Each"
"number in the tuple represents an atom. The first atom"
"in the tuple becomes the new origin; the second defines"
"the x-axis; and the third defines the xy plane. This is"
"indexed from 0.")
)
translations: List[Tuple[float, float, float]] = Field(
default_factory=list,
description=("Specific translations to generate. Each item is a tuple of "
"(x, y, z) coordinates. The entire molecule is translated "
"by these coordinates. ")
)
rotations: List[Tuple[int, int, int]] = Field(default_factory=list,
description=("Specific rigid-body rotations to generate. Each"
"number in the tuple represents an atom. The first and"
"second atoms in the tuple define a vector parallel to the"
"x-axis, and the third atom defines a plane parallel to the"
"xy plane. This is indexed from 0.")
)
# _atoms: Tuple["Atom"]
[docs] @validator(
"stage_1_unrestrained_charges",
"stage_1_restrained_charges",
"stage_2_unrestrained_charges",
"stage_2_restrained_charges",
pre=True
)
def validate_charges(v):
if v is not None:
v = np.asarray(v)
return v
[docs] @classmethod
def from_smiles(cls, smiles, order_by_map_number: bool = False, **kwargs):
from . import rdutils
rdmol = rdutils.rdmol_from_smiles(smiles, order_by_map_number=order_by_map_number)
return cls.from_rdkit(rdmol, **kwargs)
[docs] @classmethod
def from_rdkit(cls, molecule, random_seed=-1, **kwargs):
from . import rdutils
return rdutils.molecule_from_rdkit(molecule, cls,
random_seed=random_seed,
**kwargs)
[docs] def to_rdkit(self):
from .rdutils import molecule_to_rdkit
return molecule_to_rdkit(self)
[docs] def to_mdanalysis(self):
from .mdautils import molecule_to_mdanalysis
return molecule_to_mdanalysis(self)
@property
def charges(self):
for charge_prop in [
self.stage_2_restrained_charges,
self.stage_2_unrestrained_charges,
self.stage_1_restrained_charges,
self.stage_1_unrestrained_charges,
]:
if charge_prop is not None:
return charge_prop
def __post_init__(self, **kwargs):
super().__post_init__(**kwargs)
if self.charge is None:
self.charge = self.qcmol.molecular_charge
else:
self.qcmol.__dict__["molecular_charge"] = self.charge
if self.multiplicity is None:
self.multiplicity = self.qcmol.molecular_multiplicity
else:
self.qcmol.__dict__["molecular_multiplicity"] = self.multiplicity
# self._atoms = tuple(Atom.from_molecule(
# self,
# indices=np.arange(len(self.qcmol.symbols)),
# ))
@property
def atoms(self):
return Atom.from_molecule(
self,
indices=np.arange(len(self.qcmol.symbols)),
)
def __repr__(self):
qcmol_repr = self._get_qcmol_repr()
n_confs = len(self.conformers)
return f"{self._clsname}({qcmol_repr}, charge={self.charge}) with {n_confs} conformers"
def __hash__(self):
return hash(self.qcmol.get_hash())
def __eq__(self, other):
return hash(self) == hash(other)
@property
def n_conformers(self):
return len(self.conformers)
@property
def transformations(self):
"""All transformations"""
return [self.reorientations, self.rotations, self.translations]
@property
def n_transformations(self):
"Number of transformations"
return sum(map(len, self.transformations))
@property
def n_orientations(self):
"""Number of orientations in the molecule."""
return sum(conf.n_orientations for conf in self.conformers)
[docs] def generate_orientation_coordinates(
self,
coordinates: Optional[np.ndarray] = None,
) -> np.ndarray:
"""Generate coordinates for orientations in angstrom"""
if coordinates is None:
coordinates = self.coordinates
orientations = []
if self.keep_original_orientation:
orientations.append(coordinates)
for indices in self.reorientations:
orientations.append(orutils.rigid_orient(*indices, coordinates))
for indices in self.rotations:
orientations.append(orutils.rigid_rotate(*indices, coordinates))
for shift in self.translations:
orientations.append(coordinates + shift)
return np.array(orientations)
[docs] def generate_orientations(self, clear_existing_orientations: bool = False):
"""Generate Orientation objects for each conformer."""
if not self.conformers:
self.generate_conformers()
for conf in self.conformers:
if clear_existing_orientations:
conf.orientations = []
coords = self.generate_orientation_coordinates(conf.coordinates)
for coord in coords:
conf.add_orientation_with_coordinates(coord)
if not len(conf.orientations):
conf.add_orientation_with_coordinates(conf.coordinates)
[docs] def get_atoms_from_smarts(self, smiles):
indices = self.get_smarts_matches(smiles)
return [Atom.from_molecule(self, indices=ix) for ix in indices]
[docs] def get_atoms(self, indices: List[int]):
return Atom.from_molecule(self, indices=indices)
[docs] def get_sp3_ch_indices(self) -> Dict[int, List[int]]:
get_connectivity = None
if self._rdmol:
try:
from .rdutils import get_connectivity
except ImportError:
pass
if get_connectivity is None:
try:
from .psi4utils import get_connectivity
except ImportError:
def get_connectivity(x):
return x.qcmol.connectivity
symbols = self.qcmol.symbols
bonds = get_connectivity(self)
if bonds is None:
bonds = np.empty((0, 3))
bonds = np.asarray(bonds)
single_bonds = np.isclose(bonds[:, 2], np.ones_like(bonds[:, 2]))
groups = {}
for i in np.where(symbols == "C")[0]:
contains_index = np.any(bonds[:, :2] == i, axis=1)
c_bonds = bonds[contains_index & single_bonds][:, :2].astype(int)
c_partners = c_bonds[c_bonds != i]
if len(c_partners) == 4:
groups[i] = c_partners[symbols[c_partners] == "H"]
return groups
[docs] def get_symmetric_atom_indices(self) -> List[Tuple[int, ...]]:
"""Get atom indices that are symmetric to each other
Returns
-------
symmetric_atom_indices: List[Tuple[int, ...]]
Each item in the list is a tuple of integers, where
each atom associated with the index is symmetric to the others
"""
from .rdutils import get_symmetric_atom_indices
return get_symmetric_atom_indices(self)
[docs] def get_symmetric_atoms(self) -> List[List["Atom"]]:
"""Get atoms that are symmetric to each other
Returns
-------
symmetric_atoms: List[List[Atom]]
Each item in the list is a list of atoms, where
each atom is symmetric to the others
"""
return [
[self.atoms[i] for i in match]
for match in self.get_symmetric_atom_indices()
]
[docs]@functools.total_ordering
class Atom(base.Model):
"""Class to manage atoms for charge constraints"""
index: int
molecule: Molecule
[docs] @classmethod
def from_molecule(cls, molecule, indices=[]):
if not isinstance(indices, (list, tuple, np.ndarray)):
indices = [indices]
return [cls(molecule=molecule, index=i) for i in indices]
@property
def symbol(self):
return self.molecule.qcmol.symbols[self.index]
@property
def position(self):
return self.molecule.qcmol.geometry[self.index]
@property
def atomic_number(self):
return self.molecule.qcmol.atomic_numbers[self.index]
def __hash__(self):
return hash((self.molecule, self.index))
def __eq__(self, other):
return self.molecule == other.molecule and self.index == other.index
def __lt__(self, other):
return self.index < other.index
Molecule.update_forward_refs()
Atom.update_forward_refs()