example/using_custom_NNP.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using custom NNPs with Auto3D\n",
    "\n",
    "Auto3D (>= 2.3.0) is compatible with any jitable NNPs. This notebook demonstrates how to wrapper and jit a custom NNP to a specific format that Auto3D can use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.3.0\n"
     ]
    }
   ],
   "source": [
    "import os, sys\n",
    "root = os.path.dirname(os.path.dirname(os.path.abspath(\"__file__\")))\n",
    "\n",
    "import torch\n",
    "import torchani\n",
    "import Auto3D\n",
    "from Auto3D.auto3D import options, main\n",
    "\n",
    "print(Auto3D.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Below is the template of the wrapper that you need to implement with your custom NNP.\n",
    "class userNNP(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(userNNP, self).__init__()\n",
    "        \"\"\"This is an example NNP model that can be used with Auto3D.\n",
    "        You can initialize an NNP model however you want,\n",
    "        just make sure that:\n",
    "            - It contains the coord_pad and species_pad attributes \n",
    "              (These values will be used when processing the molecules in batch.)\n",
    "            - The signature of the forward method is the same as below.\n",
    "        \"\"\"\n",
    "        # Here I constructed an example NNP using ANI2x.\n",
    "        # In your case, you can replace this with your own NNP model.\n",
    "        self.model = torchani.models.ANI2x(periodic_table_index=True)\n",
    "\n",
    "        self.coord_pad = 0  # int, the padding value for coordinates\n",
    "        self.species_pad = -1  # int, the padding value for species.\n",
    "\n",
    "    def forward(self,\n",
    "                species: torch.Tensor,\n",
    "                coords: torch.Tensor,\n",
    "                charges: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Your NNP should take species, coords, and charges as input\n",
    "        and return the energies of the molecules.\n",
    "\n",
    "        species contains the atomic numbers of the atoms in the molecule: [B, N]\n",
    "        where B is the batch size, N is the number of atoms in the largest molecule.\n",
    "        \n",
    "        coords contains the coordinates of the atoms in the molecule: [B, N, 3]\n",
    "        where B is the batch size, N is the number of atoms in the largest molecule,\n",
    "        and 3 represents the x, y, z coordinates.\n",
    "        \n",
    "        charges contains the molecular charges: [B]\n",
    "        \n",
    "        The forward function returns the energies of the molecules: [B],\n",
    "        output energy unit: eV\"\"\"\n",
    "\n",
    "        # an example for computing molecular energy, replace with your NNP model\n",
    "        energies = self.model((species, coords)).energies * 27.211386245988\n",
    "        return energies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/jack/miniconda3/envs/py39/lib/python3.9/site-packages/torchani/resources/\n"
     ]
    }
   ],
   "source": [
    "# initialize and jit the wrapper with your NNP model\n",
    "myNNP = userNNP()\n",
    "myNNP_jit = torch.jit.script(myNNP)\n",
    "\n",
    "# save the model to a file for later use\n",
    "model_path = os.path.join(root, 'myNNP.pt')\n",
    "torch.jit.save(myNNP_jit, model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Checking input file...\n",
      "\tThere are 4 SMILES in the input file /home/jack/Auto3D_pkg/example/files/smiles.smi. \n",
      "\tAll SMILES and IDs are valid.\n",
      "Suggestions for choosing isomer_engine and optimizing_engine: \n",
      "\tIsomer engine options: RDKit and Omega.\n",
      "\tOptimizing engine options: ANI2x, ANI2xt, AIMNET or your own NNP.\n",
      "The available memory is 24 GB.\n",
      "The task will be divided into 1 jobs.\n",
      "Job1, number of inputs: 4\n",
      "\n",
      "\n",
      "Isomer generation for job1\n",
      "Enumerating cis/tran isomers for unspecified double bonds...\n",
      "Enumerating R/S isomers for unspecified atomic centers...\n",
      "Removing enantiomers...\n",
      "Enumerating conformers/rotamers, removing duplicates...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4/4 [00:00<00:00, 13.09it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Optimizing on job1\n",
      "Loading model from /home/jack/Auto3D_pkg/myNNP.pt\n",
      "Preparing for parallel optimizing... (Max optimization steps: 5000)\n",
      "Total 3D conformers: 45\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 10%|█         | 500/5000 [00:13<02:14, 33.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total 3D structures: 45  Converged: 18   Dropped(Oscillating): 0    Active: 27\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|█▉        | 999/5000 [00:25<01:27, 45.60it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total 3D structures: 45  Converged: 36   Dropped(Oscillating): 0    Active: 9\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 30%|██▉       | 1499/5000 [00:37<01:16, 45.68it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total 3D structures: 45  Converged: 41   Dropped(Oscillating): 0    Active: 4\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|████      | 2000/5000 [00:47<01:03, 46.89it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total 3D structures: 45  Converged: 44   Dropped(Oscillating): 0    Active: 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 2126/5000 [00:49<01:07, 42.59it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Optimization finished at step 2127:   Total 3D structures: 45  Converged: 45   Dropped(Oscillating): 0    Active: 0\n",
      "Begin to select structures that satisfy the requirements...\n",
      "Energy unit: Hartree if implicit.\n",
      "Program running time: 1 minute(s)\n",
      "Output path: /home/jack/Auto3D_pkg/example/files/smiles_20240802-014545-537526/smiles_out.sdf\n",
      "/home/jack/Auto3D_pkg/example/files/smiles_20240802-014545-537526/smiles_out.sdf\n"
     ]
    }
   ],
   "source": [
    "# Now you can run Auto3D with your custom NNP model.\n",
    "# Simply parse the model_path to the optimizing_engine arguement\n",
    "\n",
    "smi_path = os.path.join(root, \"example/files/smiles.smi\")  # You can specify the path to your file here\n",
    "args = options(smi_path, k=1, optimizing_engine=model_path, use_gpu=True, gpu_idx=0)\n",
    "out = main(args)\n",
    "print(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py39",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}