<commit_before><commit_msg>add notebook for generating example<commit_after>{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:numexpr.utils:Note: detected 72 virtual cores but NumExpr set to maximum of 64, check \"NUMEXPR_MAX_THREADS\" environment variable.\n",
      "INFO:numexpr.utils:Note: NumExpr detected 72 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
      "INFO:numexpr.utils:NumExpr defaulting to 8 threads.\n",
      "RDKit WARNING: [17:44:28] Enabling RDKit 2019.09.3 jupyter extensions\n"
     ]
    }
   ],
   "source": [
    "from perses.utils.openeye import *\n",
    "from perses.annihilation.relative import HybridTopologyFactory\n",
    "from perses.rjmc.topology_proposal import PointMutationEngine\n",
    "from perses.rjmc.geometry import FFAllAngleGeometryEngine\n",
    "\n",
    "import simtk.openmm as openmm\n",
    "import simtk.openmm.app as app\n",
    "import simtk.unit as unit\n",
    "import numpy as np\n",
    "from openmoltools import forcefield_generators\n",
    "import mdtraj as md\n",
    "from openmmtools.constants import kB\n",
    "from perses.tests.utils import validate_endstate_energies\n",
    "from openforcefield.topology import Molecule\n",
    "from openmmforcefields.generators import SystemGenerator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "ENERGY_THRESHOLD = 1e-2\n",
    "temperature = 300 * unit.kelvin\n",
    "kT = kB * temperature\n",
    "beta = 1.0/kT\n",
    "ring_amino_acids = ['TYR', 'PHE', 'TRP', 'PRO', 'HIS']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First thing to do is load the apo protein to mutate...\n",
    "protein_filename = '../input/ala_vacuum.pdb'\n",
    "protein_pdbfile = open(protein_filename, 'r')\n",
    "protein_pdb = app.PDBFile(protein_pdbfile)\n",
    "protein_pdbfile.close()\n",
    "protein_positions, protein_topology, protein_md_topology = protein_pdb.positions, protein_pdb.topology, md.Topology.from_openmm(protein_pdb.topology)\n",
    "protein_topology = protein_md_topology.to_openmm()\n",
    "protein_n_atoms = protein_md_topology.n_atoms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "DEBUG:perses.forcefields.system_generators:Trying GAFFTemplateGenerator to load openff-1.0.0\n",
      "DEBUG:perses.forcefields.system_generators:  GAFFTemplateGenerator cannot load openff-1.0.0\n",
      "DEBUG:perses.forcefields.system_generators:Specified 'forcefield' (openff-1.0.0) must be one of ['gaff-1.4', 'gaff-1.8', 'gaff-1.81', 'gaff-2.1', 'gaff-2.11']\n",
      "DEBUG:perses.forcefields.system_generators:Trying SMIRNOFFTemplateGenerator to load openff-1.0.0\n"
     ]
    }
   ],
   "source": [
    "forcefield_files = ['amber14/protein.ff14SB.xml', 'amber14/tip3p.xml']\n",
    "barostat = None\n",
    "forcefield_kwargs = {'removeCMMotion': False, 'ewaldErrorTolerance': 1e-4, 'constraints' : app.HBonds, 'hydrogenMass' : 4 * unit.amus}\n",
    "periodic_forcefield_kwargs = None\n",
    "nonperiodic_forcefield_kwargs = nonperiodic_forcefield_kwargs={'nonbondedMethod': app.NoCutoff}\n",
    "molecules = []\n",
    "system_generator = SystemGenerator(forcefields = forcefield_files,\n",
    "                                                barostat=barostat,\n",
    "                                                forcefield_kwargs=forcefield_kwargs,\n",
    "                                                periodic_forcefield_kwargs=periodic_forcefield_kwargs,\n",
    "                                   nonperiodic_forcefield_kwargs=nonperiodic_forcefield_kwargs,\n",
    "                                                molecules=molecules,\n",
    "                                                cache=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _solvate(topology,\n",
    "                    positions,\n",
    "                    model,\n",
    "                    phase):\n",
    "    \"\"\"\n",
    "    Generate a solvated topology, positions, and system for a given input topology and positions.\n",
    "    For generating the system, the forcefield files provided in the constructor will be used.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    topology : app.Topology\n",
    "        Topology of the system to solvate\n",
    "    positions : [n, 3] ndarray of Quantity nm\n",
    "        the positions of the unsolvated system\n",
    "    forcefield : SystemGenerator.forcefield\n",
    "        forcefield file of solvent to add\n",
    "    model : str, default 'tip3p'\n",
    "        solvent model to use for solvation\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    solvated_topology : app.Topology\n",
    "        Topology of the system with added waters\n",
    "    solvated_positions : [n + 3(n_waters), 3] ndarray of Quantity nm\n",
    "        Solvated positions\n",
    "    solvated_system : openmm.System\n",
    "        The parameterized system, containing a barostat if one was specified.\n",
    "    \"\"\"\n",
    "    modeller = app.Modeller(topology, positions)\n",
    "\n",
    "\n",
    "    # Now we have to add missing atoms\n",
    "    if phase != 'vacuum':\n",
    "        modeller.addSolvent(system_generator.forcefield, model=model, padding=1.0 * unit.nanometers, ionicStrength=0.05*unit.molar)\n",
    "    else:\n",
    "        pass\n",
    "\n",
    "    solvated_topology = modeller.getTopology()\n",
    "    solvated_positions = modeller.getPositions()\n",
    "\n",
    "    # Canonicalize the solvated positions: turn tuples into np.array\n",
    "    solvated_positions = unit.quantity.Quantity(value=np.array([list(atom_pos) for atom_pos in solvated_positions.value_in_unit_system(unit.md_unit_system)]), unit=unit.nanometers)\n",
    "    solvated_system = system_generator.create_system(solvated_topology)\n",
    "\n",
    "    return solvated_topology, solvated_positions, solvated_system\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "top, pos, sys = list(_solvate(protein_topology, protein_positions, 'tip3p', phase='vacuum'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "outdir = \"/home/zhangi/choderalab/perses/examples/protein-atom-mapping\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "with open(os.path.join(outdir, f\"ala_ser_top.pickle\"), 'wb') as f:\n",
    "    pickle.dump(top, f)\n",
    "with open(os.path.join(outdir, f\"ala_ser_pos.pickle\"), 'wb') as f:\n",
    "    pickle.dump(pos, f)\n",
    "with open(os.path.join(outdir, f\"ala_ser_sys.pickle\"), 'wb') as f:\n",
    "    pickle.dump(sys, f)\n",
    "with open(os.path.join(outdir, f\"system_generator.pickle\"), 'wb') as f:\n",
    "    pickle.dump(system_generator, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
{'fraction_non_alphanumeric': 0.16398390342052313, 'fraction_numerical': 0.011569416498993963, 'mean_word_length': 2.6331658291457285, 'pattern_counts': {'":': 69, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 2}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 11, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3634837', 'n_tokens_mistral': 2904, 'n_tokens_neox': 2761, 'n_words': 521}