diff --git a/charge/servers/molecular_property_utils.py b/charge/servers/molecular_property_utils.py index babdba0..dac854a 100644 --- a/charge/servers/molecular_property_utils.py +++ b/charge/servers/molecular_property_utils.py @@ -21,6 +21,9 @@ ) from charge.servers.SMILES_utils import get_synthesizability +from charge.servers.get_chemprop2_preds import predict_with_chemprop +from charge.servers.molecule_pricer import get_chemspace_prices +from charge.servers.nnp_predictor import compute_band_gap import sys import os from typing import Literal, Tuple @@ -220,11 +223,30 @@ def get_molecule_price(smiles): raise ImportError( "Please install the rdkit support packages to use this module." ) - price = get_chemspace_prices([smiles]) return price[0] +def get_gap(smiles: str) -> float: + """ + Retrieve HOMO-LUMO band gap for the molecule specified by the SMILES string, smiles. + + Args: + smiles: A SMILES string for the molecule of interest. + + Returns: + float: Returns float representing the HOMO-LUMO band gap of the given SMILES string. + + Examples: + >>> get_gap("O=CCOC=O") + 0.2217 + """ + + if not HAS_RDKIT: + raise ImportError("Please install the rdkit support packages to use this module.") + return compute_band_gap(smiles) + + def polymerize_monomer(smiles): """ Automatically identify the appropriate polymerization rule for a given monomer and return diff --git a/charge/servers/nnp_predictor.py b/charge/servers/nnp_predictor.py new file mode 100644 index 0000000..73c2519 --- /dev/null +++ b/charge/servers/nnp_predictor.py @@ -0,0 +1,143 @@ +from loguru import logger +import os +try: + from rdkit import Chem + from rdkit.Chem import AllChem + import numpy as np + from ase import Atoms + from ase.optimize.sciopt import SciPyFminCG + from ase.calculators.calculator import Calculator, all_changes + import torch + from torch_geometric.data import Data + from torch_geometric.datasets.qm9 import HAR2EV + from gotennet.models.goten_model import GotenModel + HAS_NNPS = True +except (ImportError, ModuleNotFoundError) as e: + HAS_NNPS = False + logger.warning( + "Please install the nnp support packages to use this module." + "Install it with: pip install charge[nnp]", + ) + +class GotenNetCalculator(Calculator): + def __init__(self, model, device="cuda", weights_only=True, **kwargs): + Calculator.__init__(self, **kwargs) + self.model = model + + self.model.to(device) + self.device = device + self.implemented_properties = [ + "energy", + "forces", + ] + self.results = {} + + def calculate(self, atoms=None, properties=None, system_changes=all_changes): + Calculator.calculate(self, atoms) + z = torch.tensor(atoms.get_atomic_numbers(), device=self.device) + pos = torch.tensor(atoms.get_positions(), dtype=torch.float32, device=self.device, requires_grad=True) + batch = torch.tensor([0] * z.size(0), dtype=torch.int64, device=self.device) + + inp = Data(z=z, pos=pos, batch=batch) + + out = self.model(inp) + + self.results = {} + for property in properties: + if property == "forces": + out["property"].backward() + self.results["forces"] = -1 * pos.grad.cpu().detach().numpy() + # self.results["forces"] = out["forces"].cpu().detach().numpy() + if property == "energy": + self.results["energy"] = float(out["property"].cpu().detach().numpy()[0]) + # self.results["energy"] = float(out["energy"].cpu().detach().numpy()[0]) + +def AtomsFromMol(mol): + """ + Get an ase.Atoms object from an RDKit.Mol, with preexisting conformers. + """ + numbers = [] + + for i, atom in enumerate(mol.GetAtoms()): + numbers.append(atom.GetAtomicNum()) + + positions = mol.GetConformer().GetPositions() + numbers = np.array(numbers) + return Atoms( + numbers=numbers, + positions=positions, + ) + +_calcs = dict() +def initialize_calcs(properties: list[str]) -> None: + """ + Initialize GotenNet calculators for property prediction. If the GOTENNET_BASE_PATH + environment variable is set, models will be loaded from {GOTENNET_BASE_PATH}/{property}.ckpt. + If GOTENNET_BASE_PATH is not set, models will be loaded from the current directory. + + Args: + properties: The properties to load calculators for. + + Raises: + FileNotFoundError: If there isn't a valid model checkpoint at the provided path. + """ + base_path = os.getenv("GOTENNET_BASE_PATH", ".") + for property in properties: + if property not in _calcs or _calcs[property]["base_path"] != base_path: + logger.info(f"Initializing GotenNet {property} calculator") + try: + _calcs[property] = { + "model": GotenNetCalculator(GotenModel.from_pretrained(f"{GOTENNET_BASE_PATH}/{property}.ckpt")), + "base_path": base_path, + } + except FileNotFoundError: + logger.info(f"{property} calculator could not be loaded from {GOTENNET_BASE_PATH}") + raise + + +def compute_band_gap(smiles: str) -> float: + """ + Calculate the HOMO-LUMO gap of a molecule given its SMILES string. + Returns a float of gap in eV. + + Args: + smiles (str): The input SMILES string. + + Returns: + float: The HOMO-LUMO gap of the molecule, returns NaN if there is an error. + """ + + initialize_calcs(["U0", "gap"]) + + mol = Chem.MolFromSmiles(smiles) + if mol is None: + logger.warning("Invalid SMILES string or molecule could not be created.") + return float("nan") + mol = Chem.AddHs(mol) + AllChem.EmbedMolecule(mol, AllChem.ETKDG()) + if mol.GetNumConformers() == 0: + logger.warning("No conformers found for the molecule.") + return float("nan") + + AllChem.UFFOptimizeMolecule(mol, maxIters=490) + atoms = AtomsFromMol(mol) + atoms.calc = _calcs["U0"] + dyn = SciPyFminCG(atoms, logfile="/dev/null") + try: + dyn.run(fmax=0.01, steps=10) + except: + # Sometimes the optimizer will raise an error about precision loss, despite the molecule being optimized successfully. + # pass + raise + + # GotenNet predictions are improperly normalized for some values, see PyG source for QM9 dataset to see which ones + gap = _calcs["gap"].get_potential_energy(atoms) / HAR2EV + logger.info(f"Gap for SMILES {smiles}: {gap}") + + return gap + +def main(smiles: str): + print(f"{smiles} gap: {compute_band_gap(smiles)}") + +if __name__ == "__main__": + main("O=CCOC=O") diff --git a/pyproject.toml b/pyproject.toml index 1855e07..26e532b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ aizynthfinder = ["paretoset", "rdchiral", "wrapt_timeout_decorator", "swifter", chemprop = ["chemprop>=2.2.0", "torch>=2.8.0", "lightning"] chemprice = [] test = ["pytest", "pytest-asyncio", "pytest-mock", "responses", "httpx", "requests-mock"] +nnp = ["rdkit", "ase", "numpy", "torch", "torch_geometric", "gotennet", "omegadict", "pytorch-lightning", "hydra-core"] # Define a set of optional packages for use when deploying persistent data services pds = ["charge[aizynthfinder,autogen,rdkit,chemprice]"]