Added support for NNP property (gap) prediction

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

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")

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]"]