HCTBlog.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8e71c246",
   "metadata": {},
   "outputs": [],
   "source": [
    "#purely to get an input molecule:\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import AllChem\n",
    "\n",
    "# for parameterizing the small molecule with openff:\n",
    "from openff.toolkit import Molecule\n",
    "from openmmforcefields.generators import GAFFTemplateGenerator\n",
    "\n",
    "# OpenMM as a ground truth for the HCT implementation:\n",
    "from openmm.app import ForceField\n",
    "from openmm import app, unit\n",
    "import openmm\n",
    "\n",
    "# general stuff:\n",
    "import requests\n",
    "import numpy as np\n",
    "from scipy.spatial.distance import cdist\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16a7e93a",
   "metadata": {},
   "source": [
    "# Create some input molecule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b5f6edbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "mol = Chem.MolFromSmiles('Cc1c(c2cc(ccc2n1C(=O)c3ccc(cc3)Cl)OC)CC(=O)O')\n",
    "mol = Chem.AddHs(mol)\n",
    "AllChem.EmbedMolecule(mol, randomSeed=1)\n",
    "offmol = Molecule.from_rdkit(mol)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dbc6337",
   "metadata": {},
   "source": [
    "add partial charge data to the offmol object. \n",
    "\n",
    "I'm using Espaloma here, running as a server in the background. For how to do this at home, see: https://ljmartin.github.io/blog/23_flask_4_espaloma.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f84683bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = requests.post('http://localhost:5000/process', json={'input': Chem.MolToMolBlock(mol)})\n",
    "partial_charges = np.array(response.json())\n",
    "offmol.partial_charges = partial_charges*unit.elementary_charge"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1b7b147",
   "metadata": {},
   "source": [
    "with partial charges, we're ready to add the small molecule information to an MD forcefield:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffbf9bb5",
   "metadata": {},
   "source": [
    "# Solvation energy according to the Generalized Born model (HCT variant)\n",
    "\n",
    "First we calculate it with OpenMM. I'll consider this value the ground truth. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4cb3a498",
   "metadata": {},
   "outputs": [],
   "source": [
    "gen = GAFFTemplateGenerator(molecules=offmol)\n",
    "ff = ForceField('implicit/hct.xml')\n",
    "ff.registerTemplateGenerator(gen.generator)\n",
    "\n",
    "top = offmol.to_topology().to_openmm()\n",
    "pos = offmol.conformers[0].to_openmm()\n",
    "\n",
    "system = ff.createSystem(\n",
    "    offmol.to_topology().to_openmm(), \n",
    "    nonbondedMethod=app.NoCutoff,\n",
    "    solventDielectric=78.5, soluteDielectric=1,\n",
    ")\n",
    "\n",
    "# set the force group of the CustomGBForce in order to \n",
    "# evaluate its energy:\n",
    "for force in system.getForces():\n",
    "    if isinstance(force, openmm.CustomGBForce):\n",
    "        gbf = force\n",
    "gbf.setForceGroup(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c507f02a",
   "metadata": {},
   "source": [
    "print the energy of just the HCT force. This is the solvation energy according to the HCT model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4bd078d4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Quantity(value=-125.5433349609375, unit=kilojoule/mole)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "context = openmm.Context(system, openmm.LangevinIntegrator(298, 1, 1))\n",
    "# or: simbo = app.Simulation(top, system, openmm.LangevinIntegrator(298, 1, 1))\n",
    "context.setPositions(pos)\n",
    "context.getState(getEnergy=True, groups={1}).getPotentialEnergy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6af2dc24",
   "metadata": {},
   "source": [
    "# use an isolated CustomGBForce\n",
    "now we've got the HCT energy, but for debugging purposes it's useful to demonstrate how the CustomGBForce object is built. I'm just following the implementation from:\n",
    "\n",
    "- `openmm/wrappers/python/openmm/app/internal/customgbforces.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "022a5fe1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Quantity(value=-125.5433349609375, unit=kilojoule/mole)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from openmm.app.internal.customgbforces import CustomAmberGBForceBase, CustomGBForce\n",
    "from openmm.app.internal import customgbforces\n",
    "from openmm.app import element as E\n",
    "\n",
    "\n",
    "def _screen_parameter(atom):\n",
    "    return _SCREEN_PARAMETERS.get(atom.element, _SCREEN_PARAMETERS[None])\n",
    "\n",
    "\n",
    "# get charges in units of e (but without the openmm.unit wrapper)\n",
    "charges = offmol.partial_charges.to_openmm().value_in_unit(unit.elementary_charge)\n",
    "\n",
    "\n",
    "cabgf = CustomAmberGBForceBase()\n",
    "\n",
    "cabgf.addPerParticleParameter(\"charge\")\n",
    "cabgf.addPerParticleParameter(\"or\") # Offset radius\n",
    "cabgf.addPerParticleParameter(\"sr\") # Scaled offset radius\n",
    "cabgf.addComputedValue(\"I\", \"select(step(r+sr2-or1), 0.5*(1/L-1/U+0.25*(r-sr2^2/r)*(1/(U^2)-1/(L^2))+0.5*log(L/U)/r), 0);\"\n",
    "                           \"U=r+sr2;\"\n",
    "                           \"L=max(or1, D);\"\n",
    "                           \"D=abs(r-sr2)\",\n",
    "                      CustomGBForce.ParticlePairNoExclusions)\n",
    "\n",
    "cabgf.addComputedValue(\"B\", \"1/(1/or-I)\", CustomGBForce.SingleParticle)\n",
    "\n",
    "cabgf.addEnergyTerm(\"-0.5*138.935485*(1/soluteDielectric-1/solventDielectric)*charge^2/B;\"+\\\n",
    "                   \"solventDielectric=78.5; soluteDielectric=1\",\n",
    "                   CustomGBForce.SingleParticle\n",
    "                  )\n",
    "\n",
    "cabgf.addEnergyTerm(\"-138.935485*(1/soluteDielectric-1/solventDielectric)*charge1*charge2/f;\"\n",
    "                        \"f=sqrt(r^2+B1*B2*exp(-r^2/(4*B1*B2)));solventDielectric=78.5; soluteDielectric=1\", CustomGBForce.ParticlePairNoExclusions)\n",
    "\n",
    "# ACE hydrophobic term:\n",
    "cabgf.addEnergyTerm(\"28.3919551*(radius+0.14)^2*(radius/B)^6; radius=or+offset;offset=0.009\", CustomGBForce.SingleParticle)\n",
    "\n",
    "\n",
    "radii = np.empty((top.getNumAtoms(), 2), np.double)\n",
    "radii[:,0] = customgbforces._mbondi_radii(top)/10\n",
    "for rad, atom in zip(radii, top.atoms()):\n",
    "    rad[1] = customgbforces._screen_parameter(atom)[0]\n",
    "params = radii\n",
    "\n",
    "\n",
    "for i, p in enumerate(params):\n",
    "    charge = charges[i]\n",
    "    cabgf.addParticle([charge, p[0], p[1]])\n",
    "\n",
    "cabgf._addParticles()\n",
    "system = ff.createSystem(\n",
    "    offmol.to_topology().to_openmm(), \n",
    "    nonbondedMethod=app.NoCutoff\n",
    ")\n",
    "cabgf.setForceGroup(2)\n",
    "system.addForce(cabgf)\n",
    "\n",
    "simbo = app.Simulation(top, system, openmm.LangevinIntegrator(298, 1, 1))\n",
    "simbo.context.setPositions(pos)\n",
    "state = simbo.context.getState(getEnergy=True, getForces=True, groups={2})\n",
    "state.getPotentialEnergy()\n",
    "\n",
    "#note - output energy is same as before!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "71f550ae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Quantity(value=0.0, unit=kilojoule/mole)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# double check they are same:\n",
    "context.getState(getEnergy=True, groups={1}).getPotentialEnergy() - state.getPotentialEnergy()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebe422c5",
   "metadata": {},
   "source": [
    "# Numpy version\n",
    "\n",
    "Implicit solvent simulations often have large cutoffs or no cutoff at all. I chose zero cutoff to make life easier for the numpy version - we can just do an all-by-all distance calculation, as well as all-by-all parameter comparisons, then sum across one dimension of the the resultant array to get the per-atom computations (like the integral `I` or the `B` values) or energies. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "280d29c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "RADIUS_ARG_POSITION = 0\n",
    "SCREEN_POSITION = 1\n",
    "OFFSET = 0.009\n",
    "params[:,RADIUS_ARG_POSITION] -= OFFSET\n",
    "params[:,SCREEN_POSITION] *= params[:,RADIUS_ARG_POSITION]\n",
    "\n",
    "# the offset radius and the scaled radius:\n",
    "ors, srs = params.T\n",
    "\n",
    "# calculate distance matrix:\n",
    "r = cdist(pos, pos)/10\n",
    "\n",
    "\n",
    "nparticles = system.getNumParticles()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70dac4df",
   "metadata": {},
   "source": [
    "with the parameters and distance matrix at the ready, everything starts to happen: first we do the `addComputedValue` terms from above, which all lead to the 'B' value of each atom, that is, the effective (Born) radius. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4c5c615d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/sh/7bc_43s123v0bfwd325sqs640000gn/T/ipykernel_87186/1359582525.py:23: RuntimeWarning: divide by zero encountered in divide\n",
      "  0.5*(1/L-1/U+0.25*(r-srs[None,:]**2/r)*(1/(U**2)-1/(L**2))+0.5*np.log(L/U)/r)\n",
      "/var/folders/sh/7bc_43s123v0bfwd325sqs640000gn/T/ipykernel_87186/1359582525.py:23: RuntimeWarning: invalid value encountered in add\n",
      "  0.5*(1/L-1/U+0.25*(r-srs[None,:]**2/r)*(1/(U**2)-1/(L**2))+0.5*np.log(L/U)/r)\n"
     ]
    }
   ],
   "source": [
    "# we'll think of the distance matrix as shape m x n where m is considered atom '1'\n",
    "# and n is considered atom '2', in openmm parlance. \n",
    "\n",
    "D = np.abs(r-srs[None,:]) # \"D=abs(r-sr2)\",\n",
    "L = np.maximum(D,ors[:,None]) # \"L=max(or1, D);\"\n",
    "U = r + srs[None,:]# \"U=r+sr2;\"\n",
    "\n",
    "#step(x) = 0 if x is less than 0, 1 \n",
    "#select(x,y,z) = z if x = 0, y otherwise\n",
    "\n",
    "# select(step(r+sr2-or1), 0.5*(1/L-1/U+0.25*(r-sr2^2/r)*(1/(U^2)-1/(L^2))+0.5*log(L/U)/r), 0);\n",
    "step_condition = np.clip(\n",
    "    r + srs[None,:] - ors[:,None],\n",
    "    0,\n",
    "    np.inf\n",
    ")\n",
    "\n",
    "# can safely ignore the 'divided by zero' warnings. \n",
    "# it's from the diagonal where the 'r' value is 0. \n",
    "# but these terms get zeroed out in 7 lines. \n",
    "I = np.where(step_condition==0,\n",
    "         0,\n",
    "         0.5*(1/L-1/U+0.25*(r-srs[None,:]**2/r)*(1/(U**2)-1/(L**2))+0.5*np.log(L/U)/r)\n",
    "        )\n",
    "\n",
    "# zero out the self-pairs:\n",
    "I[range(nparticles), range(nparticles)] = 0\n",
    "I = I.sum(1)\n",
    "\n",
    "\n",
    "B = 1 / (1/ors-I)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44bb0d3e",
   "metadata": {},
   "source": [
    "with `B` in hand we can move to the `addEnergyTerm` lines, sum pairwise interactions into a system-level energy. Keep in mind that each pair is counted _once_, so we'll use the upper triangle of any all-by-all pairwise arrays. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7fed09cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "soluteDielectric = 1\n",
    "solventDielectric = 78.5\n",
    "\n",
    "# \"-0.5*138.935485*(1/soluteDielectric-1/solventDielectric)*charge^2/B\"+params,\n",
    "nrg1 = -0.5*138.935485*(1/soluteDielectric-1/solventDielectric)*charges**2/B\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2796e255",
   "metadata": {},
   "outputs": [],
   "source": [
    "# \"-138.935485*(1/soluteDielectric-1/solventDielectric)*charge1*charge2/f;\"\n",
    "# \"f=sqrt(r^2+B1*B2*exp(-r^2/(4*B1*B2)));solventDielectric=78.5; soluteDielectric=1\"\n",
    "\n",
    "# splitting this into a few lines:\n",
    "\n",
    "bbt = B*B[:,None]\n",
    "t = - (r**2 / (4*bbt))\n",
    "f = r**2 + bbt* np.exp(t)\n",
    "f = np.sqrt(f)\n",
    "\n",
    "a,b = np.triu_indices(B.shape[0],1)\n",
    "\n",
    "nrg2 = -138.935485*(1/soluteDielectric-1/solventDielectric)*charges*charges[:,None]/f\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6efcbd64",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add ACE for the nonpolar solvation energy (i.e. cavitation energy):\n",
    "# \"28.3919551*(radius+0.14)^2*(radius/B)^6; radius=or+offset;offset=0.009\"\n",
    "radius = ors+0.009\n",
    "\n",
    "ace = 28.3919551*(radius+0.14)**2*(radius/B)**6"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5816ae3",
   "metadata": {},
   "source": [
    "note output energy is the same!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e6222923",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-125.54349522167732"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nrg1.sum() + nrg2[a,b].sum() + ace.sum()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}