Method A and B.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "cdb9e064-ce32-44b2-bce4-ca567bc2a6eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openmm import *\n",
    "from openmm.app import *\n",
    "from openmm.unit import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d54404b1-ac37-4d73-a7a0-a44c3636a50c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Method A: Environment Energy:  -107999.86808286235 kJ/mol\n"
     ]
    }
   ],
   "source": [
    "pdb = PDBFile('villin.pdb')\n",
    "forcefield = ForceField('amber14-all.xml', 'amber14/tip3pfb.xml')\n",
    "\n",
    "env = set([a.index for a in pdb.topology.atoms() if a.residue.name in ('HOH', 'Cl')])\n",
    "protein = set([a.index for a in pdb.topology.atoms() if a.index not in env])\n",
    "\n",
    "\n",
    "env_system = forcefield.createSystem(pdb.topology, nonbondedMethod=PME)\n",
    "\n",
    "for force in env_system.getForces():\n",
    "    if isinstance(force, NonbondedForce):\n",
    "\n",
    "        # zero-out the non-bonded potential energy related to proteins, \n",
    "        # including internal proteins + interactions to the environment\n",
    "        for i in range(force.getNumParticles()):\n",
    "            if i in protein:\n",
    "                charge, sigma, epsilon = force.getParticleParameters(i)\n",
    "                force.setParticleParameters(i, 0.0, 0, 0.0)\n",
    "\n",
    "        # zero-out the 1-4 potential in protein\n",
    "        for i in range(force.getNumExceptions()):\n",
    "            p1, p2, chargeProd, sigma, epsilon = force.getExceptionParameters(i)\n",
    "            if p1 in protein and p2 in protein:\n",
    "                force.setExceptionParameters(i, p1, p2, 0.0, 0, 0.0)\n",
    "\n",
    "\n",
    "# zero-out the bonded potential in protein\n",
    "for force in env_system.getForces():\n",
    "    if isinstance(force, HarmonicBondForce):\n",
    "        for i in range(force.getNumBonds()):\n",
    "            p1, p2, length, k = force.getBondParameters(i)\n",
    "            if p1 in protein and p2 in protein:\n",
    "                force.setBondParameters(i, p1, p2, length, 0.0 * kilojoule_per_mole/nanometer**2)\n",
    "    \n",
    "    elif isinstance(force, HarmonicAngleForce):\n",
    "        for i in range(force.getNumAngles()):\n",
    "            p1, p2, p3, angle, k = force.getAngleParameters(i)\n",
    "            if all(p in protein for p in [p1, p2, p3]):\n",
    "                force.setAngleParameters(i, p1, p2, p3, angle, 0.0 * kilojoule_per_mole/radian**2)\n",
    "    \n",
    "    elif isinstance(force, PeriodicTorsionForce):\n",
    "        for i in range(force.getNumTorsions()):\n",
    "            p1, p2, p3, p4, period, phase, k = force.getTorsionParameters(i)\n",
    "            if all(p in protein for p in [p1, p2, p3, p4]):\n",
    "                force.setTorsionParameters(i, p1, p2, p3, p4, period, phase, 0.0 * kilojoule_per_mole)\n",
    "\n",
    "env_integrator = VerletIntegrator(0.001*picosecond)\n",
    "env_context = Context(env_system, env_integrator)\n",
    "env_context.setPositions(pdb.positions)\n",
    "\n",
    "print('Method A: Environment Energy: ', env_context.getState(getEnergy=True).getPotentialEnergy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "cfbfb5b8-7a6e-4304-b9a3-53dc3970e582",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Method B: Environment Energy:  -108123.90657605253 kJ/mol\n"
     ]
    }
   ],
   "source": [
    "pdb = PDBFile('villin.pdb')\n",
    "forcefield = ForceField('amber14-all.xml', 'amber14/tip3pfb.xml')\n",
    "\n",
    "env = set([a.index for a in pdb.topology.atoms() if a.residue.name in ('HOH', 'Cl')])\n",
    "protein = set([a.index for a in pdb.topology.atoms() if a.index not in env])\n",
    "\n",
    "\n",
    "env_system = forcefield.createSystem(pdb.topology, nonbondedMethod=PME)\n",
    "\n",
    "for force in env_system.getForces():\n",
    "    if isinstance(force, NonbondedForce):\n",
    "        force.addGlobalParameter(\"nb_scale\", 1.0)\n",
    "        # zero-out the non-bonded potential energy related to proteins, \n",
    "        # including internal proteins + interactions to the environment\n",
    "        for i in range(force.getNumParticles()):\n",
    "            if i in protein:\n",
    "                charge, sigma, epsilon = force.getParticleParameters(i)\n",
    "                force.setParticleParameters(i, 0.0, 0.0, 0.0)\n",
    "                force.addParticleParameterOffset(\"nb_scale\", i, charge, sigma, epsilon)\n",
    "\n",
    "        # zero-out the 1-4 potential in protein\n",
    "        for i in range(force.getNumExceptions()):\n",
    "            p1, p2, chargeProd, sigma, epsilon = force.getExceptionParameters(i)\n",
    "            if p1 in protein and p2 in protein:\n",
    "                force.setExceptionParameters(i, p1, p2, 0.0, 0.0, 0.0)\n",
    "                force.addExceptionParameterOffset(\"nb_scale\", i, chargeProd, sigma, epsilon)\n",
    "\n",
    "# zero-out the bonded potential in protein\n",
    "for force in env_system.getForces():\n",
    "    if isinstance(force, HarmonicBondForce):\n",
    "        for i in range(force.getNumBonds()):\n",
    "            p1, p2, length, k = force.getBondParameters(i)\n",
    "            if p1 in protein and p2 in protein:\n",
    "                force.setBondParameters(i, p1, p2, length, 0.0 * kilojoule_per_mole/nanometer**2)\n",
    "    \n",
    "    elif isinstance(force, HarmonicAngleForce):\n",
    "        for i in range(force.getNumAngles()):\n",
    "            p1, p2, p3, angle, k = force.getAngleParameters(i)\n",
    "            if all(p in protein for p in [p1, p2, p3]):\n",
    "                force.setAngleParameters(i, p1, p2, p3, angle, 0.0 * kilojoule_per_mole/radian**2)\n",
    "    \n",
    "    elif isinstance(force, PeriodicTorsionForce):\n",
    "        for i in range(force.getNumTorsions()):\n",
    "            p1, p2, p3, p4, period, phase, k = force.getTorsionParameters(i)\n",
    "            if all(p in protein for p in [p1, p2, p3, p4]):\n",
    "                force.setTorsionParameters(i, p1, p2, p3, p4, period, phase, 0.0 * kilojoule_per_mole)\n",
    "\n",
    "env_integrator = VerletIntegrator(0.001*picosecond)\n",
    "env_context = Context(env_system, env_integrator)\n",
    "env_context.setPositions(pdb.positions)\n",
    "\n",
    "env_context.setParameter(\"nb_scale\", 0)\n",
    "\n",
    "print('Method B: Environment Energy: ', env_context.getState(getEnergy=True).getPotentialEnergy())"
   ]
  }
 ],
 "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.13.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}