Skip to content

Instantly share code, notes, and snippets.

@ljmartin

ljmartin/HCTBlog.ipynb

Last active June 18, 2025 23:39
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save ljmartin/7c1d8386ca84972e8bdc317cd7c09d11 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save ljmartin/7c1d8386ca84972e8bdc317cd7c09d11 to your computer and use it in GitHub Desktop.
Download ZIP
Display the source blob
Display the rendered blob
Raw
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
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment