GBN2Blog.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8e71c246",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit import Chem\n",
    "from rdkit.Chem import AllChem\n",
    "import numpy as np\n",
    "from scipy.spatial.distance import cdist\n",
    "\n",
    "# for parameterizing a small molecule with openff:\n",
    "from openff.toolkit import Molecule\n",
    "from openmmforcefields.generators import GAFFTemplateGenerator\n",
    "\n",
    "# OpenMM stuff to sanity check the GBN2 implementation.\n",
    "from openmm.app import ForceField\n",
    "from openmm import app, unit\n",
    "import openmm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16a7e93a",
   "metadata": {},
   "source": [
    "# Create some input molecule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b5f6edbd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5de2ddba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5fa3c4de32824c78908a108ffb5e06e6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "062edda6efac47e3b8ec5811d275859f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "NGLWidget()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "offmol = Molecule.from_rdkit(mol)\n",
    "offmol"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dbc6337",
   "metadata": {},
   "source": [
    "add partial charge data to the offmol object. \n",
    "\n",
    "for how to do this, see: https://ljmartin.github.io/blog/23_flask_4_espaloma.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f84683bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "from openmm import unit\n",
    "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": [
    "# Calculate solvation energy according to Generalized Born (GBN2 variant)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4cb3a498",
   "metadata": {},
   "outputs": [],
   "source": [
    "gen = GAFFTemplateGenerator(molecules=offmol)\n",
    "ff = ForceField('implicit/gbn2.xml')\n",
    "ff.registerTemplateGenerator(gen.generator)\n",
    "\n",
    "top = offmol.to_topology().to_openmm()\n",
    "pos = offmol.conformers[0].to_openmm()\n",
    "system = ff.createSystem(\n",
    "    offmol.to_topology().to_openmm(), \n",
    "    nonbondedMethod=app.NoCutoff,\n",
    "     solventDielectric=78.5, soluteDielectric=1,\n",
    ")\n",
    "\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 GBN2 force - this is the solvation free energy and we'll use it as the reference value to validate the other calculations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d07117d2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Quantity(value=-125.0518798828125, unit=kilojoule/mole)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "simbo = app.Simulation(top, system, openmm.LangevinIntegrator(298, 1, 1))\n",
    "simbo.context.setPositions(pos)\n",
    "simbo.context.getState(getEnergy=True, groups={1}).getPotentialEnergy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ab9d4cde",
   "metadata": {},
   "outputs": [],
   "source": [
    "gbfParams = [gbf.getParticleParameters(i) for i in range(system.getNumParticles())]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6af2dc24",
   "metadata": {},
   "source": [
    "# repeat but with new CustomGBForce\n",
    "while that's fine, I'd like to build it up bit-by-bit for debugging purposes. This rebuilds the CustomGBForce in the GBN2 style. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "822da34c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-0.10519525,  0.23538291, -0.10908066, -0.12013286, -0.20175458,\n",
       "        0.18494274, -0.4304733 , -0.04817048,  0.12787934, -0.34457088,\n",
       "        0.82396853, -0.67302334, -0.41611397, -0.05485828, -0.1991355 ,\n",
       "       -0.09165348, -0.1991355 , -0.05485828, -0.10338877, -0.43119672,\n",
       "        0.06052691, -0.18721738,  0.85249525, -0.53630644, -0.6577695 ,\n",
       "        0.0913897 ,  0.0913897 ,  0.0913897 ,  0.20572913,  0.21527055,\n",
       "        0.19669569,  0.20759341,  0.21494271,  0.21494271,  0.20759341,\n",
       "        0.07578675,  0.07578675,  0.07578675,  0.14993969,  0.14993969,\n",
       "        0.4146632 ])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "charges = offmol.partial_charges.to_openmm().value_in_unit(unit.elementary_charge)\n",
    "charges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ac4ef545",
   "metadata": {},
   "outputs": [],
   "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",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# radii = np.empty((top.getNumAtoms(), 2), np.double)\n",
    "# radii[:,0] = customgbforces._mbondi3_radii(top)/10\n",
    "\n",
    "# for rad, atom in zip(radii, top.atoms()):\n",
    "#     rad[1] = customgbforces._screen_parameter(atom)[1]\n",
    "# params = radii\n",
    "\n",
    "\n",
    "natoms = top.getNumAtoms()\n",
    "radii = np.empty([natoms,5], np.double)\n",
    "radii[:,0] = customgbforces._mbondi3_radii(top)/10\n",
    "\n",
    "for atom, rad in zip(top.atoms(), radii):\n",
    "    rad[1] = customgbforces._screen_parameter(atom)[2]\n",
    "    rad[2:] = customgbforces.GBSAGBn2Force._atom_params.get(atom.element, customgbforces.GBSAGBn2Force._default_atom_params)\n",
    "params = radii\n",
    "\n",
    "cabgf = CustomAmberGBForceBase()\n",
    "\n",
    "cabgf.addPerParticleParameter(\"charge\")\n",
    "cabgf.addPerParticleParameter(\"or\") # Offset radius\n",
    "cabgf.addPerParticleParameter(\"sr\") # Scaled offset radius\n",
    "cabgf.addPerParticleParameter(\"alpha\")\n",
    "cabgf.addPerParticleParameter(\"beta\")\n",
    "cabgf.addPerParticleParameter(\"gamma\")\n",
    "cabgf.addPerParticleParameter(\"radindex\")\n",
    "\n",
    "for i, p in enumerate(params):\n",
    "    #charge, sigma, epsilon = nonbonded.getParticleParameters(i)\n",
    "    charge = charges[i]\n",
    "    cabgf.addParticle([charge, p[0], p[1], p[2], p[3], p[4]])\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "30b01540",
   "metadata": {},
   "outputs": [],
   "source": [
    "radii = [p[cabgf.RADIUS_ARG_POSITION] for p in cabgf.parameters]\n",
    "cabgf._uniqueRadii = list(sorted(set(radii)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "00d7479b",
   "metadata": {},
   "outputs": [],
   "source": [
    "cabgf._radiusToIndex = {r: i for i, r in enumerate(cabgf._uniqueRadii)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e15e44fb",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def createUniqueTable(cabgf, fullTable):\n",
    "    \n",
    "    OFFSET = 0.009\n",
    "    \n",
    "    import math\n",
    "    n = len(cabgf._uniqueRadii)\n",
    "    tablePositions = [(r+OFFSET-0.1)*200 for r in cabgf._uniqueRadii]\n",
    "    numRadii = len(cabgf._uniqueRadii)\n",
    "    index1 = [0]*numRadii\n",
    "    index2 = [0]*numRadii\n",
    "    weight1 = [0]*numRadii\n",
    "    weight2 = [0]*numRadii\n",
    "    \n",
    "    for i,p in enumerate(tablePositions):\n",
    "        if p <= 0:\n",
    "            weight1[i] = 1.0\n",
    "        elif p >= 20:\n",
    "            index1[i] = 20\n",
    "            weight1[i] = 1.0\n",
    "        else:\n",
    "            index1[i] = int(math.floor(p))\n",
    "            index2[i] = index1[i]+1\n",
    "            weight1[i] = index2[i]-p\n",
    "            weight2[i] = 1.0-weight1[i]\n",
    "\n",
    "    table = []\n",
    "    for i in range(numRadii):\n",
    "        for j in range(numRadii):\n",
    "            table.append(weight1[i]*weight1[j]*fullTable[index1[i]*21+index1[j]] +\n",
    "                          weight1[i]*weight2[j]*fullTable[index1[i]*21+index2[j]] +\n",
    "                          weight2[i]*weight1[j]*fullTable[index2[i]*21+index1[j]] +\n",
    "                          weight2[i]*weight2[j]*fullTable[index2[i]*21+index2[j]])\n",
    "    return table\n",
    "\n",
    "m0Table = createUniqueTable(cabgf, customgbforces.m0)\n",
    "d0Table = createUniqueTable(cabgf, customgbforces.d0)\n",
    "n = len(cabgf._uniqueRadii)\n",
    "cabgf.addTabulatedFunction(\"getd0\", openmm.Discrete2DFunction(n, n, d0Table))\n",
    "cabgf.addTabulatedFunction(\"getm0\", openmm.Discrete2DFunction(n, n, m0Table))\n",
    "\n",
    "cabgf.addComputedValue(\"I\", \"Ivdw+neckScale*Ineck;\"\n",
    "                           \"Ineck=step(radius1+radius2+neckCut-r)*getm0(radindex1,radindex2)/(1+100*(r-getd0(radindex1,radindex2))^2+\"\n",
    "                           \"0.3*1000000*(r-getd0(radindex1,radindex2))^6);\"\n",
    "                           \"Ivdw=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",
    "                           \"radius1=or1+offset; radius2=or2+offset;\"\n",
    "                           \"neckScale=0.826836; neckCut=0.68; offset=0.0195141\", CustomGBForce.ParticlePairNoExclusions)\n",
    "\n",
    "# real one:\n",
    "cabgf.addComputedValue(\"B\", \"1/(1/or-tanh(alpha*psi-beta*psi^2+gamma*psi^3)/radius);\"\n",
    "                             \"psi=I*or; radius=or+offset; offset=0.0195141\", 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",
    "                  ) # correct up to here. \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.0195141\", CustomGBForce.SingleParticle)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "60ab4cb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "cabgf._addParticles()\n",
    "\n",
    "radindices = [cabgf._radiusToIndex[p[cabgf.RADIUS_ARG_POSITION]] for p in cabgf.parameters]\n",
    "\n",
    "# re-set the per-atom parameters using a different offset. \n",
    "# why? I don't know. but the d0 and m0 table creation uses offset 0.009, \n",
    "# and then now we use a different one. Maybe it's a GBn vs GBn2 thing.  \n",
    "\n",
    "cabgf.parameters = []\n",
    "cabgf.OFFSET = 0.0195141\n",
    "for i, p in enumerate(params):\n",
    "    #charge, sigma, epsilon = nonbonded.getParticleParameters(i)\n",
    "    charge = charges[i]\n",
    "    cabgf.addParticle([charge, p[0], p[1], p[2], p[3], p[4]])\n",
    "    \n",
    "for i in range(natoms):\n",
    "    p = cabgf.parameters[i]\n",
    "    radindex = radindices[i]\n",
    "    cabgf.setParticleParameters(i, p+[radindex] )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "9c9a0405",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Quantity(value=-125.0518798828125, unit=kilojoule/mole)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "system = ff.createSystem(\n",
    "    offmol.to_topology().to_openmm(), \n",
    "    nonbondedMethod=app.NoCutoff\n",
    ")\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()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fa9b0ab",
   "metadata": {},
   "source": [
    "# now do it in numpy:\n",
    "here we step through OpenMM's implementation, but using numpy functions instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a0278450",
   "metadata": {},
   "outputs": [],
   "source": [
    "charges, ors, srs, alphas, betas, gammas, radindexs = np.array(gbfParams).T\n",
    "radindexs = radindexs.astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "dee8a17c",
   "metadata": {},
   "outputs": [],
   "source": [
    "nparticles = system.getNumParticles()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac3b0bcd",
   "metadata": {},
   "source": [
    "create tables of m0 and d0 - which correspond to the height (m0) and distance (d0) of the maximum of the integrated neck volume: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "1a2e7fd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "m0_square = np.array(m0Table).reshape(n, -1)\n",
    "d0_square = np.array(d0Table).reshape(n, -1)\n",
    "\n",
    "a,b= np.where(\n",
    "    np.ones([system.getNumParticles(), system.getNumParticles()]).astype(bool)\n",
    ")\n",
    "m0_pairs = m0_square[\n",
    "    radindexs[a], radindexs[b]\n",
    "].reshape(system.getNumParticles(), -1)\n",
    "d0_pairs = d0_square[\n",
    "    radindexs[a], radindexs[b]\n",
    "].reshape(system.getNumParticles(), -1)\n",
    "\n",
    "m0_pairs[range(nparticles), range(nparticles)] = 0\n",
    "d0_pairs[range(nparticles), range(nparticles)] = 0\n",
    "\n",
    "r = cdist(pos, pos)/10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7201f864",
   "metadata": {},
   "source": [
    "calculate Ivdw:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "5344546e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/sh/7bc_43s123v0bfwd325sqs640000gn/T/ipykernel_86697/2152491378.py:19: 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_86697/2152491378.py:19: RuntimeWarning: invalid value encountered in multiply\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_86697/2152491378.py:19: RuntimeWarning: invalid value 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_86697/2152491378.py:19: 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": [
    "neckScale = 0.826836\n",
    "neckCut=0.68\n",
    "offset=0.0195141\n",
    "\n",
    "\n",
    "\n",
    "radius1 = ors+offset\n",
    "\n",
    "\n",
    "D = np.abs(r-srs[None,:]) \n",
    "L = np.maximum(D,ors[:,None])\n",
    "U = r + srs[None,:]\n",
    "\n",
    "step_condition = np.clip(\n",
    "    r + srs[None,:] - ors[:,None],\n",
    "    0,\n",
    "    np.inf\n",
    ")\n",
    "Ivdw = 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",
    "Ivdw[range(nparticles), range(nparticles)] = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ecdea8e",
   "metadata": {},
   "source": [
    "and then `I`, the coulomb field approximation Integral:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "eda697ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "t1 = np.where(\n",
    "    (radius1[:,None] + radius1 + neckCut - r) <= 0,\n",
    "    0, \n",
    "    1\n",
    ")\n",
    "t1[range(nparticles), range(nparticles)]=0\n",
    "t2 = m0_pairs.T\n",
    "t3 = 1+100*(r-d0_pairs.T)**2\n",
    "t4 = 0.3*1000000*(r-d0_pairs.T)**6\n",
    "\n",
    "Ineck = t1 * t2 / (t3+t4)\n",
    "Ineck[range(nparticles), range(nparticles)] = 0\n",
    "\n",
    "I = (Ivdw+neckScale*Ineck).sum(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e97fa629",
   "metadata": {},
   "source": [
    "which leads to `B`, the Born radius:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "297476f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "psi=I*ors\n",
    "offset=0.0195141\n",
    "radius=ors+offset\n",
    "\n",
    "B = 1/(1/ors-np.tanh(alphas*psi-betas*psi**2+gammas*psi**3)/radius)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b81ecb51",
   "metadata": {},
   "source": [
    "with `B` in hand, sum up all solvation energies using the Generalized Born equation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "ef9848f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "soluteDielectric = 1\n",
    "solventDielectric = 78.5\n",
    "\n",
    "nrg1 = -0.5*138.935485*(1/soluteDielectric-1/solventDielectric)*charges**2/B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "c16ca157",
   "metadata": {},
   "outputs": [],
   "source": [
    "bbt = B*B[:,None]\n",
    "t = - (r**2 / (4*bbt))\n",
    "f = r**2 + bbt* np.exp(t)\n",
    "f = np.sqrt(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "a2caf61e",
   "metadata": {},
   "outputs": [],
   "source": [
    "nrg2 = -138.935485*(1/soluteDielectric-1/solventDielectric)*charges*charges[:,None]/f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "b55c1a22",
   "metadata": {},
   "outputs": [],
   "source": [
    "radius = ors+offset\n",
    "ace = 28.3919551*(radius+0.14)**2*(radius/B)**6\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "773de77a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-125.05181054538855"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a,b = np.triu_indices(B.shape[0],1)\n",
    "nrg1.sum() + nrg2[a,b].sum() + ace.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa655417",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}