add userNNP option to allow the use of other NNP

example/tautomer_with_userNNP.ipynb

@@ -0,0 +1,40 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, sys\n",
+    "root = os.path.dirname(os.path.dirname(os.path.abspath(\"__file__\")))\n",
+    "sys.path.append(os.path.dirname(__file__))\n",
+    "\n",
+    "import Auto3D\n",
+    "from Auto3D.auto3D import options\n",
+    "from Auto3D.tautomer import get_stable_tautomers\n",
+    "\n",
+    "input_path = os.path.join(root, \"example\", \"files\", \"sildnafil.smi\")\n",
+    "\n",
+    "if __name__ == \"__main__\":\n",
+    "    args = options(input_path, k=1, enumerate_tautomer=True, tauto_engine=\"rdkit\",\n",
+    "                   optimizing_engine=\"userNNP\",  #Use userNNP for tautomers\n",
+    "                   max_confs=10, patience=200, use_gpu=False)\n",
+    "    tautomer_out = get_stable_tautomers(args, tauto_k=3)\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.8.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

example/userModel.py

@@ -0,0 +1,29 @@
+import sys
+import torch
+import torchani
+from torch import Tensor
+from typing import NamedTuple
+
+class SpeciesEnergies(NamedTuple):
+    species: Tensor
+    energies: Tensor
+
+class userNNP(torch.nn.Module):
+    def __init__(self, periodic_table_index=True, model_choice=0, device=None):
+        super().__init__()
+        try:
+            from ani2x_ext.custom_emsemble_ani2x_ext import CustomEnsemble
+        except:
+            print("ani2x_ext is not installed, please check out https://github.com/plin1112/ani_ext.")
+            sys.exit("userNNP is used, but NNP model is not available.")
+
+        self.model = CustomEnsemble(model_choice=model_choice, 
+                                    periodic_table_index=periodic_table_index,
+                                    device=device)
+        self.ase = self.model.ase
+
+    def forward(self, species_coords):
+        species, energy = self.model(species_coords)
+        return SpeciesEnergies(species, energy)
+
+

src/Auto3D/ASE/geometry.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 """
-Geometry optimization with ANI2xt, AIMNET or ANI2x
+Geometry optimization with ANI2xt, AIMNET, userNNP or ANI2x
 """
 import sys
 import os
@@ -22,7 +22,7 @@ def opt_geometry(path: str, model_name:str, gpu_idx=0, opt_tol=0.003, opt_steps=
 
     :param path: Input sdf file
     :type path: str
-    :param model_name: ANI2x, ANI2xt or AIMNET
+    :param model_name: ANI2x, ANI2xt, userNNP or AIMNET
     :type model_name: str
     :param gpu_idx: GPU cuda index, defaults to 0
     :type gpu_idx: int, optional
@@ -61,6 +61,8 @@ def opt_geometry(path: str, model_name:str, gpu_idx=0, opt_tol=0.003, opt_steps=
     path = '/home/jack/Auto3D_pkg/tests/files/DA.sdf'
     out = opt_geometry(path, 'ANI2x', gpu_idx=0, opt_tol=0.003, opt_steps=5000)
     print(out)
+    out = opt_geometry(path, 'userNNP', gpu_idx=0, opt_tol=0.003, opt_steps=5000)
+    print(out)
     out = opt_geometry(path, 'AIMNET', gpu_idx=0, opt_tol=0.003, opt_steps=5000)
     print(out)
     out = opt_geometry(path, 'ANI2xt', gpu_idx=0, opt_tol=0.003, opt_steps=5000)

src/Auto3D/ASE/thermo.py

@@ -23,6 +23,10 @@
 import torchani
 from Auto3D.batch_opt.batchopt import EnForce_ANI
 from Auto3D.batch_opt.ANI2xt_no_rep import ANI2xt
+try:
+    from userModel import userNNP
+except:
+    pass
 from Auto3D.utils import hartree2ev
 
 torch.backends.cuda.matmul.allow_tf32 = False
@@ -103,8 +107,13 @@ def model_name2model_calculator(model_name: str, device=torch.device('cpu'), cha
         ani2x = torchani.models.ANI2x(periodic_table_index=True).to(device).double()
         model = EnForce_ANI(ani2x, model_name)
         calculator = ani2x.ase()
+    elif model_name == "userNNP":
+        user_nnp = userNNP().to(device).double()
+        model = EnForce_ANI(user_nnp, model_name)
+        calculator = user_nnp.ase()
+
     else:
-        raise ValueError("model has to be 'ANI2x', 'ANI2xt' or 'AIMNET'")
+        raise ValueError("model has to be 'ANI2x', 'ANI2xt', 'userNNP' or 'AIMNET'")
     return model, calculator
 
 def mol2atoms(mol: Chem.Mol):
@@ -117,7 +126,7 @@ def mol2atoms(mol: Chem.Mol):
 def vib_hessian(mol: Chem.Mol, ase_calculator, model,
                 device=torch.device('cpu'), model_name='AIMNET'):
     '''return a VibrationsData object
-    model: ANI2xt or AIMNet2 or ANI2x that can be used to calculate Hessian'''
+    model: ANI2xt or AIMNet2 or ANI2x or userNNP that can be used to calculate Hessian'''
     # get the ASE atoms object
     coord = mol.GetConformer().GetPositions()
     species = [a.GetSymbol() for a in mol.GetAtoms()]
@@ -163,7 +172,7 @@ def do_mol_thermo(mol: Chem.Mol,
                   device=torch.device('cpu'),
                   T=298.0, model_name='AIMNET'):
     """For a RDKit mol object, calculate its thermochemistry properties.
-    model: ANI2xt or AIMNet2 or ANI2x that can be used to calculate Hessian"""
+    model: ANI2xt or AIMNet2 or ANI2x or userNNP that can be used to calculate Hessian"""
     vib = vib_hessian(mol, atoms.get_calculator(), model, device, model_name=model_name)
     vib_e = vib.get_energies()
     e = atoms.get_potential_energy()
@@ -205,7 +214,7 @@ def aimnet_hessian_helper(coord:torch.tensor,
         numbers2 = torch.tensor([periodict2idx[num.item()] for num in numbers.squeeze()], device=device).unsqueeze(0)
         e = model(numbers2, coord)
         return e  # energy unit: eV
-    elif model_name == 'ANI2x':
+    elif model_name == 'ANI2x' or model_name == 'userNNP':
         e = model((numbers, coord)).energies * hartree2ev
         return e  # energy unit: eV
 
@@ -244,6 +253,8 @@ def calc_thermo(path: str, model_name: str, get_mol_idx_t=None, gpu_idx=0, opt_t
         hessian_model = ANI2xt(device).double()
     elif model_name == 'ANI2x':
         hessian_model = torchani.models.ANI2x(periodic_table_index=True).to(device).double()
+    elif model_name == 'userNNP':
+        hessian_model = userNNP().to(device).double()
     model, calculator = model_name2model_calculator(model_name, device)
 
     mols = list(Chem.SDMolSupplier(path, removeHs=False))

src/Auto3D/SPE.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-"""Calculating single point energy using ANI2xt, ANI2x or AIMNET"""
+"""Calculating single point energy using ANI2xt, ANI2x, 'userNNP' or AIMNET"""
 import sys
 import os
 root = os.path.dirname(os.path.abspath(__file__))
@@ -33,7 +33,7 @@ def calc_spe(path: str, model_name: str, gpu_idx=0):
 
     :param path: Input sdf file
     :type path: str
-    :param model_name: AIMNET, ANI2x or ANI2xt
+    :param model_name: AIMNET, ANI2x, userNNP, or ANI2xt
     :type model_name: str
     :param gpu_idx: GPU cuda index, defaults to 0
     :type gpu_idx: int, optional
@@ -58,8 +58,12 @@ def calc_spe(path: str, model_name: str, gpu_idx=0):
     elif model_name == "ANI2x":
         calculator = torchani.models.ANI2x(periodic_table_index=True).to(device)
         model = EnForce_ANI(calculator, model_name)
+    elif model_name == "userNNP":
+        from userModel import userNNP
+        calculator = userNNP().to(device)
+        model = EnForce_ANI(calculator, model_name)
     else:
-        raise ValueError("model has to be 'ANI2x', 'ANI2xt' or 'AIMNET'")
+        raise ValueError("model has to be 'ANI2x', 'ANI2xt' 'userNNP' or 'AIMNET'")
 
     mols = list(Chem.SDMolSupplier(path, removeHs=False))
     coord, numbers, charges = mols2lists(mols, model_name)

src/Auto3D/auto3D.py

@@ -206,7 +206,7 @@ def options(path: Optional[str]=None, k=False, window=False, verbose=False, job_
     :type gpu_idx: int or list of int, optional
     :param capacity: Number of SMILES that the model will handle for 1 G memory, defaults to 42
     :type capacity: int, optional
-    :param optimizing_engine: Choose either 'ANI2x', 'ANI2xt', or 'AIMNET' for energy calculation and geometry optimization, defaults to "AIMNET"
+    :param optimizing_engine: Choose either 'ANI2x', 'ANI2xt', 'userNNP', or 'AIMNET' for energy calculation and geometry optimization, defaults to "AIMNET"
     :type optimizing_engine: str, optional
     :param patience: If the force does not decrease for a continuous patience steps, the conformer will drop out of the optimization loop, defaults to 1000
     :type patience: int, optional

src/Auto3D/auto3Dcli.py

@@ -85,7 +85,7 @@ def cli():
         parser.add_argument('--mpi_np', type=int, default=4,
                             help="Number of CPU cores for the isomer generation step.")
         parser.add_argument('--optimizing_engine', type=str, default='AIMNET',
-                            help=("Choose either 'ANI2x', 'ANI2xt', or 'AIMNET' for energy "
+                            help=("Choose either 'ANI2x', 'ANI2xt', 'userNNP', or 'AIMNET' for energy "
                                 "calculation and geometry optimization."))
         parser.add_argument('--use_gpu', default=True, type=lambda x: (str(x).lower() == 'true'),
                             help="If True, the program will use GPU.")

src/Auto3D/batch_opt/batchopt.py

@@ -16,6 +16,12 @@
     from .ANI2xt_no_rep import ANI2xt
 except:
     pass
+
+try:
+    from userModel import userNNP
+except:
+    pass
+
 from tqdm import tqdm
 from Auto3D.utils import hartree2ev
 
@@ -161,6 +167,14 @@ def forward(self, coord, numbers, charges):
             # ANI ASE interface unit is eV
             g = torch.autograd.grad([e.sum()], [coord])[0]
             f = -g
+        elif self.name == "userNNP":
+            # use userNNP as genetic name for a ANI type NNP
+            e = self.ani((numbers, coord)).energies
+            e = e * hartree2ev  # ANI type NNP output energy unit is Hatree;
+            # ANI ASE interface unit is eV
+            g = torch.autograd.grad([e.sum()], [coord])[0]
+            f = -g
+
         return e, f
 
     #    @torch.jit.script_method
@@ -298,7 +312,7 @@ def ensemble_opt(net, coord, numbers, charges, param, model, device):
     numbers: atomic numbers in the molecule (include H). (N, m)
     charges: (N,)
     param: a dictionary containing parameters
-    model: "AIMNET", "ANI2xt" or "ANI2x"
+    model: "AIMNET", "ANI2xt", "ANI2x" or "userNNP"
     device
     """
     coord = torch.tensor(coord, dtype=torch.float, device=device)
@@ -399,8 +413,10 @@ def __init__(self, in_f, out_f, model, device, config):
             self.ani = ANI2xt(device)
         elif model == "ANI2x":
             self.ani = torchani.models.ANI2x(periodic_table_index=True).to(device)
+        elif model == "userNNP":
+            self.ani = userNNP().to(device)
         else:
-            raise ValueError("Model has to be ANI2x, ANI2xt or AIMNET.")
+            raise ValueError("Model has to be ANI2x, ANI2xt, userNNP or AIMNET.")
 
     def run(self):
         print("Preparing for parallel optimizing... (Max optimization steps: %i)" % self.config[

src/Auto3D/utils.py

@@ -97,6 +97,11 @@ def check_input(args):
             import torchani
         except:
             sys.exit("ANI2x is used as optimizing engine, but TorchANI is not installed.")
+    if args.optimizing_engine == "userNNP":
+        try:
+            from userModel import userNNP
+        except:
+            sys.exit("userNNP is used as optimizing engine, but userNNP is not installed.")
     if int(args.opt_steps) < 10:
         sys.exit(f"Number of optimization steps cannot be smaller than 10, but received {args.opt_steps}")
 
@@ -110,9 +115,9 @@ def check_input(args):
     logger.info(f"Suggestions for choosing isomer_engine and optimizing_engine: ")
     if ANI:
         print("\tIsomer engine options: RDKit and Omega.\n"
-              "\tOptimizing engine options: ANI2x, ANI2xt and AIMNET.", flush=True)
+              "\tOptimizing engine options: ANI2x, ANI2xt, userNNP and AIMNET.", flush=True)
         logger.info("\tIsomer engine options: RDKit and Omega.")
-        logger.info("\tOptimizing engine options: ANI2x, ANI2xt and AIMNET.")
+        logger.info("\tOptimizing engine options: ANI2x, ANI2xt, userNNP and AIMNET.")
     else:
         print("\tIsomer engine options: RDKit and Omega.\n"
               "\tOptimizing engine options: AIMNET.", flush=True)