from typing import List, Optional, Tuple
import functools
import logging
import numpy as np
from pydantic import Field
from . import base, orutils, rdutils
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.")
)
[docs] @classmethod
def from_smiles(cls, smiles, **kwargs):
rdmol = rdutils.rdmol_from_smiles(smiles)
return cls.from_rdkit(rdmol, **kwargs)
[docs] @classmethod
def from_rdkit(cls, molecule, random_seed=-1, **kwargs):
qcmol = rdutils.rdmol_to_qcelemental(molecule, random_seed=random_seed)
kwargs = dict(**kwargs)
kwargs["qcmol"] = qcmol
obj = cls(**kwargs)
obj._rdmol = molecule
return obj
def __init__(self, *args, **kwargs):
super().__init__(*args, **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
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"
@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]@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