From ddfa738ef999083fe4320d77f5198c0359473235 Mon Sep 17 00:00:00 2001 From: Obie Smolenski Date: Thu, 6 Nov 2025 21:22:48 -0800 Subject: [PATCH 1/4] Added support for NNP property (gap) prediction --- charge/servers/molecular_property_utils.py | 25 +++++ charge/servers/nnp_predictor.py | 125 +++++++++++++++++++++ pyproject.toml | 1 + 3 files changed, 151 insertions(+) create mode 100644 charge/servers/nnp_predictor.py diff --git a/charge/servers/molecular_property_utils.py b/charge/servers/molecular_property_utils.py index c2244c3b..90451ccf 100644 --- a/charge/servers/molecular_property_utils.py +++ b/charge/servers/molecular_property_utils.py @@ -21,6 +21,7 @@ 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 @@ -175,3 +176,27 @@ 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. + + Parameters + ---------- + smiles : str + 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) diff --git a/charge/servers/nnp_predictor.py b/charge/servers/nnp_predictor.py new file mode 100644 index 00000000..069227a2 --- /dev/null +++ b/charge/servers/nnp_predictor.py @@ -0,0 +1,125 @@ +from loguru import logger +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 = { + "U0": None, + "homo": None, + "lumo": None, +} + +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. + """ + + if _calcs["U0"] is None: + logger.info("Initializing GotenNet U0 calculator") + _calcs["U0"] = GotenNetCalculator(GotenModel.from_pretrained("QM9_small_U0")) + logger.info("Initializing GotenNet homo calculator") + _calcs["homo"] = GotenNetCalculator(GotenModel.from_pretrained("QM9_small_homo")) + logger.info("Initializing GotenNet lumo calculator") + _calcs["lumo"] = GotenNetCalculator(GotenModel.from_pretrained("QM9_small_lumo")) + + 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 + + gap = _calcs["lumo"].get_potential_energy(atoms) - _calcs["homo"].get_potential_energy(atoms) + logger.info(f"Gap for SMILES {smiles}: {gap}") + + # GotenNet predictions are improperly normalized for some values, see PyG source for QM9 dataset to see which ones + return gap / HAR2EV + +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 cfba1560..e8a9f28d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,6 +42,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]"] From 8d2fd55bf163d86e36fa6ff71a3f61ad65892096 Mon Sep 17 00:00:00 2001 From: Obie Smolenski Date: Mon, 15 Dec 2025 22:47:31 +0000 Subject: [PATCH 2/4] Fixed docstring style --- charge/servers/molecular_property_utils.py | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/charge/servers/molecular_property_utils.py b/charge/servers/molecular_property_utils.py index f651c9e1..0366d3dc 100644 --- a/charge/servers/molecular_property_utils.py +++ b/charge/servers/molecular_property_utils.py @@ -231,20 +231,15 @@ def get_gap(smiles: str) -> float: """ Retrieve HOMO-LUMO band gap for the molecule specified by the SMILES string, smiles. - Parameters - ---------- - smiles : str - A SMILES string for the molecule of interest. + 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. + Returns: + float: Returns float representing the HOMO-LUMO band gap of the given SMILES string. - Examples - -------- - >>> get_gap("O=CCOC=O") - 0.2217 + Examples: + >>> get_gap("O=CCOC=O") + 0.2217 """ if not HAS_RDKIT: From a34f702953ae76bcb948df3a766d1bcbee259d00 Mon Sep 17 00:00:00 2001 From: Obie Smolenski Date: Mon, 15 Dec 2025 17:48:37 -0500 Subject: [PATCH 3/4] Fix newline spacing Co-authored-by: Shehtab Zaman --- charge/servers/molecular_property_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/charge/servers/molecular_property_utils.py b/charge/servers/molecular_property_utils.py index 0366d3dc..dac854a2 100644 --- a/charge/servers/molecular_property_utils.py +++ b/charge/servers/molecular_property_utils.py @@ -275,4 +275,4 @@ def polymerize_monomer(smiles): logger.warning("Unable to find polymerizer tool. Returning input smiles:") return smiles PSMILES = polymerize_auto(smiles) - return PSMILES \ No newline at end of file + return PSMILES From 8acbbb91b0509e30d87ed05f801b83328f802839 Mon Sep 17 00:00:00 2001 From: Obie Smolenski Date: Tue, 16 Dec 2025 14:56:25 -0500 Subject: [PATCH 4/4] Changed model loading to be local --- charge/servers/nnp_predictor.py | 48 ++++++++++++++++++++++----------- 1 file changed, 33 insertions(+), 15 deletions(-) diff --git a/charge/servers/nnp_predictor.py b/charge/servers/nnp_predictor.py index 069227a2..73c25197 100644 --- a/charge/servers/nnp_predictor.py +++ b/charge/servers/nnp_predictor.py @@ -1,4 +1,5 @@ from loguru import logger +import os try: from rdkit import Chem from rdkit.Chem import AllChem @@ -67,29 +68,46 @@ def AtomsFromMol(mol): positions=positions, ) -_calcs = { - "U0": None, - "homo": None, - "lumo": None, -} +_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. """ - if _calcs["U0"] is None: - logger.info("Initializing GotenNet U0 calculator") - _calcs["U0"] = GotenNetCalculator(GotenModel.from_pretrained("QM9_small_U0")) - logger.info("Initializing GotenNet homo calculator") - _calcs["homo"] = GotenNetCalculator(GotenModel.from_pretrained("QM9_small_homo")) - logger.info("Initializing GotenNet lumo calculator") - _calcs["lumo"] = GotenNetCalculator(GotenModel.from_pretrained("QM9_small_lumo")) + initialize_calcs(["U0", "gap"]) mol = Chem.MolFromSmiles(smiles) if mol is None: @@ -112,11 +130,11 @@ def compute_band_gap(smiles: str) -> float: # pass raise - gap = _calcs["lumo"].get_potential_energy(atoms) - _calcs["homo"].get_potential_energy(atoms) + # 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}") - # GotenNet predictions are improperly normalized for some values, see PyG source for QM9 dataset to see which ones - return gap / HAR2EV + return gap def main(smiles: str): print(f"{smiles} gap: {compute_band_gap(smiles)}")