MD_in_Colab(KOR).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "collapsed_sections": [
        "xoA5kbSMqkv5",
        "g0Ou4jzVqqi5",
        "7Xf2GSRAqxDg",
        "ktFrmmAnqyp1",
        "1EGMmauoHQmT",
        "B8BgKwrGEKsT"
      ],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyMV6HidJH/OVf4ePCJ6pk1/",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/kimjc95/computational-chemistry/blob/main/MD_in_Colab(KOR).ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Colab에서 돌리는 분자동역학 (MD) simulation"
      ],
      "metadata": {
        "id": "80CnOAJSxXyW"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "2024-08-21 by Joo-Chan Kim at [MSBL](https://msbl.kaist.ac.kr), KAIST\n",
        "\n",
        "연구에 이 노트북을 사용하셨다면 인용 부탁드립니다. DOI:[10.5281/zenodo.13133762](https://doi.org/10.5281/zenodo.13685874)\n",
        "\n",
        "아래 코드는 출처 표기 하에 배포, 수정, 상업적 이용이 가능합니다. (BSD-3 라이센스)\n",
        "\n",
        "사용 중 버그나 문제를 발견하셨다면 제 [GitHub](https://github.com/kimjc95/computational-chemistry/issues)에 보고주시거나 이메일 (kimjoochan@kaist.ac.kr)로 연락 주세요.\n",
        "\n",
        "-------------------------------------\n",
        "\n",
        "**준비물 :**\n",
        "\n",
        "1. PDB 파일 혹은 PDB ID\n",
        "\n",
        "2. (선택사항) custom 잔기/분자의 SDF 파일\n",
        "\n",
        "3. 충분한 여유 공간의 구글 드라이브, 인터넷 연결\n",
        "\n",
        "-------------------------------------\n",
        "\n",
        "**이 노트북으로 할 수 있는 일:**\n",
        "1. 비정규 residue나 small molecule ligand를 포함할 수 있는 단백질/핵산 시스템의 all-atom MD 시뮬레이션 (용매 유무 선택 가능)\n",
        "\n",
        "2. 기계학습 (Machine-learning) 역장 또는 양자역학 (XTB) 역장을 도입한 QM/MM\n",
        "\n",
        "3. MD 결과 분석: clustering, RMSD, RMSF, RoG, FEL, H-bonds, SASA, DSSP, DCCM, PCA, 그리고 MM/GBSA\n",
        "\n",
        "**이 노트북으로 못하는 일:**\n",
        "1. 물이 아닌 용매 조건이나 인지질 이중층 같은 막단백질의 MD 분석\n",
        "\n",
        "2. steered MD, replica exchange sampling, simulated tempering, metadynamics 등의 고급 MD 기술들\n",
        "\n",
        "3. Umbrella sampling, alchemical transformation 등을 통한 정확한 결합 에너지 계산\n",
        "\n",
        "4. ML- or XTB-역장이 도입된 시스템의 결합 에너지 계산\n",
        "\n",
        "----------------------------------------\n",
        "\n",
        "각 탭을 눌러 펼치고 **하나씩** 셀을 실행해주세요!"
      ],
      "metadata": {
        "id": "tYJdP4FzFPCx"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 0. 구글 드라이브 연결 & 계산 준비"
      ],
      "metadata": {
        "id": "JnMbY1wFxzQF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "assert gpu_info.find('command not found') <0, \"Please change to GPU runtime!\"\n",
        "\n",
        "#@markdown 런타임이 재시작됩니다. 놀라지 마시고 기다려주세요.\n",
        "\n",
        "!pip install -q condacolab\n",
        "import condacolab\n",
        "condacolab.install()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "9aJoPllQ1OkF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown 구글 드라이브 내에 생성될 작업 폴더 이름을 입력해주세요. <br/>\n",
        "#@markdown 구글 드라이브 접근 권한을 승인해주세요. Colab 세션이 중간에 끊기더라도 수 GB에 달하는 trajectory 파일을 안전하게 백업하기 위함입니다.\n",
        "\n",
        "import subprocess\n",
        "\n",
        "print('Installing all required python packages...', end='')\n",
        "subprocess.run(['mamba', 'install', '-c', 'pytorch', '-c', 'nvidia', '-c', 'conda-forge',\n",
        "                'pytorch', 'pytorch-cuda=12.1', 'openmm=8', 'openmm-torch', 'openmm-ml',\n",
        "                'openmm-xtb', 'nnpops', 'espaloma', 'openmmforcefields', 'scikit-learn',\n",
        "                'scipy', 'ipywidgets=7', 'nglview', 'parmed', 'rdkit', 'pdbfixer', 'openbabel',\n",
        "                'plotly', 'mdtraj', 'lxml', '--yes'])\n",
        "subprocess.run(['pip', 'install', 'mace-torch', '--yes'])\n",
        "subprocess.run(['apt-get', 'install', 'python3-openbabel'])\n",
        "subprocess.run(['pip', 'install', '--no-deps', 'plip'])\n",
        "print('done.')\n",
        "\n",
        "from pdbfixer import PDBFixer\n",
        "from openmm import *\n",
        "from openmm.app import *\n",
        "from openmm.unit import *\n",
        "from openmmxtb import XtbForce\n",
        "from openmmml import MLPotential\n",
        "from openff.toolkit import Molecule\n",
        "from openmmforcefields.generators import GAFFTemplateGenerator\n",
        "from openmmforcefields.generators import SMIRNOFFTemplateGenerator\n",
        "from openmmforcefields.generators import EspalomaTemplateGenerator\n",
        "from openmm.app.internal import customgbforces\n",
        "from rdkit import Chem\n",
        "from rdkit.Chem import AllChem\n",
        "import io\n",
        "import os\n",
        "import locale\n",
        "import sys\n",
        "import glob\n",
        "import copy\n",
        "import warnings\n",
        "import threading\n",
        "import numpy as np\n",
        "import parmed as pmd\n",
        "import nglview as nv\n",
        "from time import sleep\n",
        "from tqdm.notebook import tqdm\n",
        "from IPython.display import display, clear_output\n",
        "import mdtraj as md\n",
        "from sklearn.cluster import AgglomerativeClustering\n",
        "from sklearn.decomposition import PCA\n",
        "from sklearn.neighbors import KernelDensity\n",
        "from sklearn.model_selection import GridSearchCV\n",
        "from collections import Counter\n",
        "import openbabel\n",
        "from plip.structure.preparation import PDBComplex\n",
        "from plip.exchange.report import BindingSiteReport\n",
        "from scipy.spatial.transform import Rotation\n",
        "from sklearn.covariance import EmpiricalCovariance\n",
        "from collections import Counter\n",
        "import ipywidgets as widgets\n",
        "import plotly.graph_objects as go\n",
        "from plotly.subplots import make_subplots\n",
        "import plotly.io as pio\n",
        "pio.renderers.default = 'colab'\n",
        "from google.colab import drive, output, files\n",
        "output.enable_custom_widget_manager()\n",
        "\n",
        "\n",
        "def navigator(working_directory)->str:\n",
        "    work = \"/content/drive/MyDrive/\"+working_directory\n",
        "\n",
        "    if not os.path.isdir(work):\n",
        "        print(f'{working_directory}란 이름의 폴더가 없습니다.')\n",
        "        os.mkdir(work)\n",
        "        print(f'구글 드라이브 내에 {working_directory}란 이름의 새 폴더를 생성하였습니다.')\n",
        "        print('1-1로 가세요.')\n",
        "\n",
        "    elif os.path.exists(work+'/production_final.xtc'):\n",
        "        print(f'{working_directory} 폴더 내에 이미 끝난 MD trajectory가 있습니다.')\n",
        "        print('3-3을 실행해 더 MD 시뮬레이션을 이어갈 수 있습니다.')\n",
        "        print('4-1로 가서 결과 분석을 수행할 수 있습니다.')\n",
        "        print('5-1로 가서 결합 에너지 분석을 수행할 수 있습니다.')\n",
        "\n",
        "    elif os.path.exists(work+'/production_0.xtc'):\n",
        "        print('production run 수행 중이었군요.')\n",
        "        print('3-3으로 가서 이어가세요.')\n",
        "\n",
        "    elif os.path.exists(work+'/npt_0.xtc'):\n",
        "        print('NPT 평형 단계를 수행 중이셨군요.')\n",
        "        print('3-2로 가서 계속하세요.')\n",
        "\n",
        "    elif os.path.exists(work+'/nvt_0.xtc'):\n",
        "        print('NVT 평형 단계를 수행 중이셨군요.')\n",
        "        print('3-1로 가서 계속하세요.')\n",
        "\n",
        "    elif os.path.exists(work+'/start.xml'):\n",
        "        print('시뮬레이션 객체가 생성되었지만, 계산이 시작되진 않았습니다.')\n",
        "        print('3-1로 가서 NVT 평형을 시작하세요..')\n",
        "\n",
        "    elif os.path.exists(work+'/starting_structure.pdb'):\n",
        "        print('시작 구조를 생성하였지만, 시뮬레이션 객체가 생성되지 않았습니다.')\n",
        "        print('2-1로 가세요..')\n",
        "\n",
        "    else:\n",
        "        print(f'{working_directory}란 폴더가 구글 드라이브 내에 존재하지만, 비어있습니다.')\n",
        "        print('1-1로 가세요.')\n",
        "\n",
        "    work += '/'\n",
        "    return work\n",
        "\n",
        "\n",
        "def read_settings(work:str)->dict:\n",
        "    \"\"\"\n",
        "    Read the parameters from the settings.txt file.\n",
        "    \"\"\"\n",
        "    if not os.path.exists(work+'settings.txt'):\n",
        "        return None\n",
        "\n",
        "    with open(work+'settings.txt', 'r') as f:\n",
        "        settings = {}\n",
        "        for line in f.readlines():\n",
        "            if line.strip().startswith('Forcefield'):\n",
        "                settings['forcefield'] = line.split(\":\")[1].strip()\n",
        "            elif line.strip().startswith('Watermodel'):\n",
        "                settings['watermodel'] = line.split(\":\")[1].strip()\n",
        "            elif line.strip().startswith('Constraints'):\n",
        "                settings['constraints'] = line.split(\":\")[1].strip()\n",
        "            elif line.strip().startswith('Rigidwater'):\n",
        "                settings['rigidwater'] = bool(line.split(\":\")[1].strip())\n",
        "            elif line.strip().startswith('Hydrogen mass'):\n",
        "                settings['hydrogen mass'] = float(line.split(\":\")[1].strip())*amu\n",
        "            elif line.strip().startswith('Precision'):\n",
        "                settings['precision'] = line.split(\":\")[1].strip()\n",
        "            elif line.strip().startswith('Temperature'):\n",
        "                settings['temperature'] = (float(line.split(\":\")[1].strip())+273.15)*kelvin\n",
        "            elif line.strip().startswith('Pressure'):\n",
        "                settings['pressure'] = float(line.split(\":\")[1].strip())*bar\n",
        "            elif line.strip().startswith('Time step'):\n",
        "                settings['stepsize'] = float(line.split(\":\")[1].strip())*femtosecond\n",
        "            elif line.strip().startswith('Ionic strength'):\n",
        "                settings['ionic strength'] = float(line.split(\":\")[1].strip())*molar\n",
        "            elif line.strip().startswith('Periodic Boundary Condition'):\n",
        "                settings['PBC'] = True if line.split(\":\")[1].strip()==\"True\" else False\n",
        "            elif line.strip().startswith('Padding'):\n",
        "                settings['padding'] = float(line.split(\":\")[1].strip())*nanometer\n",
        "            elif line.strip().startswith('Cation'):\n",
        "                settings['cation'] = line.split(\":\")[1].strip()\n",
        "            elif line.strip().startswith('Anion'):\n",
        "                settings['anion'] = line.split(\":\")[1].strip()\n",
        "            elif line.strip().startswith('Save interval'):\n",
        "                settings['save interval'] = int(line.split(\":\")[1].strip())*picosecond\n",
        "            elif line.strip().startswith('Restraint'):\n",
        "                settings['restraint'] = int(line.split(\":\")[1].strip())\n",
        "\n",
        "    return settings\n",
        "\n",
        "\n",
        "def read_non_standard_residues_list(work)->list:\n",
        "    \"\"\"\n",
        "    Read the nonstandard residue list from the non_standard_residues.txt file.\n",
        "    \"\"\"\n",
        "    if not os.path.exists(work+'non_standard_residues.txt'):\n",
        "        return []\n",
        "\n",
        "    with open(work+'non_standard_residues.txt', 'r') as f:\n",
        "        nsres = []\n",
        "        for line in f.readlines():\n",
        "            if not line.startswith('#'):\n",
        "                data = line.strip().split( )\n",
        "                if len(data) > 1:\n",
        "                    if data[1] == 'ligand':\n",
        "                        lig = True\n",
        "                    elif data[1] == 'residue':\n",
        "                        lig = False\n",
        "                    nsres.append({'name':data[0], 'lig':lig, 'ff':data[2]})\n",
        "                else:\n",
        "                    nsres.append({'name':data[0]})\n",
        "    return nsres\n",
        "\n",
        "def determine_forcefield(forcefield, watermodel)->tuple:\n",
        "    \"\"\"\n",
        "    Assign the macromolecular forcefield & water model from the choice.\n",
        "    \"\"\"\n",
        "    if forcefield == \"Amber14\":\n",
        "        ff = ForceField('amber14-all.xml')\n",
        "        if watermodel == \"(explicit) TIP3P\":\n",
        "            ff.loadFile('amber14/tip3p.xml')\n",
        "            water = 'tip3p'\n",
        "        elif watermodel == \"(explicit) SPC/E\":\n",
        "            ff.loadFile('amber14/spce.xml')\n",
        "            water = 'spce'\n",
        "        elif watermodel == \"(explicit) TIP4P-Ew\":\n",
        "            ff.loadFile('amber14/tip4pew.xml')\n",
        "            water = 'tip4pew'\n",
        "        elif watermodel == \"(explicit) TIP5P\":\n",
        "            print(\"ERROR : AMBER 역장은 TIP5P 용매 모델을 지원하지 않습니다!\")\n",
        "            print(\"다른 선택지를 골라주세요.\")\n",
        "            return None, \"\"\n",
        "        elif watermodel == \"(explicit) OPC\":\n",
        "            ff.loadFile('amber14/opc.xml')\n",
        "            water = 'tip4pew'\n",
        "        elif watermodel == \"(explicit) OPC3\":\n",
        "            ff.loadFile('amber14/opc3.xml')\n",
        "            water = 'tip3p'\n",
        "        elif watermodel == \"(implicit) HCT (igb=1)\":\n",
        "            ff.loadFile('implicit/hct.xml')\n",
        "            water = \"\"\n",
        "        elif watermodel == \"(implicit) OBC1 (igb=2)\":\n",
        "            ff.loadFile('implicit/obc1.xml')\n",
        "            water = \"\"\n",
        "        elif watermodel == \"(implicit) OBC2 (igb=5)\":\n",
        "            ff.loadFile('implicit/obc2.xml')\n",
        "            water = \"\"\n",
        "        elif watermodel == \"(implicit) GBn (igb=7)\":\n",
        "            ff.loadFile('implicit/gbn.xml')\n",
        "            water = \"\"\n",
        "        elif watermodel == \"(implicit) GBn2 (igb=8)\":\n",
        "            ff.loadFile('implicit/gbn2.xml')\n",
        "            water = \"\"\n",
        "        else:\n",
        "            print('Error in reading AMBER forcefield parameters.')\n",
        "            return None, \"\"\n",
        "\n",
        "    elif forcefield == \"CHARMM36\":\n",
        "        if watermodel == \"(explicit) TIP3P\":\n",
        "            ff = ForceField('charmm36.xml', 'charmm36/water.xml')\n",
        "            water = 'tip3p'\n",
        "        elif watermodel == \"(explicit) SPC/E\":\n",
        "            ff = ForceField('charmm36.xml', 'charmm36/spce.xml')\n",
        "            water = 'spce'\n",
        "        elif watermodel == \"(explicit) TIP4P-Ew\":\n",
        "            ff = ForceField('charmm36.xml', 'charmm36/tip4pew.xml')\n",
        "            water = 'tip4pew'\n",
        "        elif watermodel == \"(explicit) TIP5P\":\n",
        "            ff = ForceField('charmm36.xml', 'charmm36/tip5p.xml')\n",
        "            water = 'tip5p'\n",
        "        elif watermodel == \"(explicit) OPC\":\n",
        "            print(\"ERROR : CHARMM 역장은 OPC 용매 모델을 지원하지 않습니다.\")\n",
        "            print(\"다른 선택지를 골라주세요.\")\n",
        "            return None, \"\"\n",
        "        elif watermodel == \"(explicit) OPC3\":\n",
        "            print(\"ERROR : CHARMM 역장은 OPC 용매 모델을 지원하지 않습니다.\")\n",
        "            print(\"다른 선택지를 골라주세요.\")\n",
        "            return None, \"\"\n",
        "        elif watermodel == \"(implicit) HCT (igb=1)\":\n",
        "            ff = ForceField('charmm36.xml', 'implicit/hct.xml')\n",
        "            water = ''\n",
        "        elif watermodel == \"(implicit) OBC1 (igb=2)\":\n",
        "            ff = ForceField('charmm36.xml', 'implicit/obc1.xml')\n",
        "            water = ''\n",
        "        elif watermodel == \"(implicit) OBC2 (igb=5)\":\n",
        "            ff = ForceField('charmm36.xml', 'implicit/obc2.xml')\n",
        "            water = ''\n",
        "        elif watermodel == \"(implicit) GBn (igb=7)\":\n",
        "            ff = ForceField('charmm36.xml', 'implicit/gbn.xml')\n",
        "            water = ''\n",
        "        elif watermodel == \"(implicit) GBn2 (igb=8)\":\n",
        "            ff = ForceField('charmm36.xml', 'implicit/gbn2.xml')\n",
        "            water = ''\n",
        "        else:\n",
        "            print(\"Error in reading the CHARMM forcefield parameters\")\n",
        "            return None, \"\"\n",
        "\n",
        "    else:\n",
        "        return None, \"\"\n",
        "\n",
        "    return ff, water\n",
        "\n",
        "\n",
        "class SimulationStream(io.StringIO):\n",
        "    \"\"\"\n",
        "    StringIO's child class to capture stdout simulation output.\n",
        "    \"\"\"\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.string = \"\"\n",
        "\n",
        "    def write(self, string):\n",
        "        self.string += string\n",
        "\n",
        "    def read(self, size=-1)->str:\n",
        "        if size < 0:\n",
        "            result = self.string\n",
        "            self.string = \"\"\n",
        "        else:\n",
        "            result = self.string[:size]\n",
        "            self.string = self.string[size:]\n",
        "        return result\n",
        "\n",
        "\n",
        "def apply_restraints(work, simulation, restraint):\n",
        "    \"\"\"\n",
        "    Apply a harmonic positional restraint potential to the heavy atoms.\n",
        "    \"\"\"\n",
        "    if restraint > 0:\n",
        "        print('위치 좌표별 restraint 생성 중...', end='')\n",
        "        posres = CustomExternalForce(\"k*periodicdistance(x, y, z, x0, y0, z0)^2\")\n",
        "        posres.addGlobalParameter(\"k\", restraint)\n",
        "        posres.addPerParticleParameter(\"x0\")\n",
        "        posres.addPerParticleParameter(\"y0\")\n",
        "        posres.addPerParticleParameter(\"z0\")\n",
        "\n",
        "        pdb = PDBFile(work+'minimized.pdb')\n",
        "        heavy = []\n",
        "        for i, a in enumerate(pdb.topology.atoms()):\n",
        "            if a.element.symbol != 'H':\n",
        "                heavy.append(i)\n",
        "\n",
        "        positions = simulation.context.getState(getPositions=True).getPositions(asNumpy=True)\n",
        "        for i in range(np.shape(positions)[0]):\n",
        "            if i in heavy:\n",
        "                posres.addParticle(i, positions[i])\n",
        "        simulation.system.addForce(posres)\n",
        "        print('완료.')\n",
        "    return simulation\n",
        "\n",
        "\n",
        "\n",
        "def restore_simulation_from_state(work, statename, step, stage):\n",
        "    \"\"\"\n",
        "    Restore a simulation object from a state file.\n",
        "    \"\"\"\n",
        "    s = read_settings(work)\n",
        "    pdb = PDBFile(work+'minimized.pdb')\n",
        "    system = XmlSerializer.deserialize(open(work+'system.xml').read())\n",
        "    integrator = LangevinMiddleIntegrator(s['temperature'], 1/picosecond, s['stepsize'])\n",
        "    platform = Platform.getPlatformByName('CUDA')\n",
        "    properties = {'Precision': s['precision']}\n",
        "    simulation = Simulation(pdb.topology, system, integrator, platform, properties)\n",
        "    simulation.context.setState(XmlSerializer.deserialize(open(statename).read()))\n",
        "    change_flag = False # When the context is changed, reinitialize it\n",
        "    if stage in ['npt', 'production'] and s['PBC']:\n",
        "        barostat = MonteCarloBarostat(s['pressure'], s['temperature'])\n",
        "        simulation.system.addForce(barostat)\n",
        "        change_flag = True\n",
        "    if stage != 'production' and stage != 'start':\n",
        "        simulation = apply_restraints(work, simulation, s['restraint'])\n",
        "        change_flag = True\n",
        "    if change_flag:\n",
        "        simulation.context.reinitialize(preserveState=True)\n",
        "\n",
        "    simulation.context.setStepCount(step)\n",
        "    time = step*s['stepsize']/picosecond\n",
        "    simulation.context.setTime(time)\n",
        "\n",
        "    return simulation\n",
        "\n",
        "\n",
        "def recall_simulation(work):\n",
        "    \"\"\"\n",
        "    Recall a simulation object from the working directory.\n",
        "    \"\"\"\n",
        "    if len(glob.glob(work+'production_step_*.xml')) > 0:\n",
        "        stage = \"production\"\n",
        "        for f in sorted(glob.glob(work+'production_step_*.xml'), reverse=True):\n",
        "            step = int(f.split('.')[-2].split('_')[-1])\n",
        "            simulation = restore_simulation_from_state(work, f, step, 'production')\n",
        "            break\n",
        "\n",
        "    elif len(glob.glob(work+'npt_step_*.xml')) > 0:\n",
        "        stage = \"npt\"\n",
        "        for f in sorted(glob.glob(work+'npt_step_*.xml'), reverse=True):\n",
        "            step = int(f.split('.')[-2].split('_')[-1])\n",
        "            simulation = restore_simulation_from_state(work, f, step, 'npt')\n",
        "            break\n",
        "\n",
        "    elif len(glob.glob(work+'nvt_step_*.xml')) > 0:\n",
        "        stage = \"nvt\"\n",
        "        for f in sorted(glob.glob(work+'nvt_step_*.xml'), reverse=True):\n",
        "            step = int(f.split('.')[-2].split('_')[-1])\n",
        "            simulation = restore_simulation_from_state(work, f, step, 'nvt')\n",
        "            break\n",
        "\n",
        "    elif len(glob.glob(work+'start.xml')) > 0:\n",
        "        stage = \"start\"\n",
        "        print('시뮬레이션 시작...')\n",
        "        simulation = restore_simulation_from_state(work, work+'start.xml', 0, 'start')\n",
        "\n",
        "    else:\n",
        "        print('시뮬레이션 객체를 불러오는데 실패했습니다.')\n",
        "        return None, \"\"\n",
        "\n",
        "    print('done.')\n",
        "    return simulation, stage\n",
        "\n",
        "\n",
        "def run_simulation(work, stage, simulation, stream, length, save_period, backup_period):\n",
        "    \"\"\"\n",
        "    Run the simulation object for length picoseconds.\n",
        "    \"\"\"\n",
        "\n",
        "    s = read_settings(work)\n",
        "\n",
        "    runtime = 0.0\n",
        "    time = simulation.context.getTime().value_in_unit(picosecond)\n",
        "\n",
        "    while runtime < length:\n",
        "        step = simulation.context.getStepCount()\n",
        "        log = StateDataReporter(stream, 50, step=True, potentialEnergy=True,\n",
        "                                temperature=True, speed=True, time=True,\n",
        "                                density=s['PBC'])\n",
        "\n",
        "        i = 0\n",
        "        for f in glob.glob(work+stage+'_*.xtc'):\n",
        "            i += 1\n",
        "\n",
        "        save_freq = int(save_period/s['stepsize'])\n",
        "        xtc = XTCReporter(work+stage+f'_{i}.xtc', save_freq, False, s['PBC'])\n",
        "\n",
        "        simulation.reporters.clear()\n",
        "        simulation.reporters.append(log)\n",
        "        simulation.reporters.append(xtc)\n",
        "\n",
        "        for i in range(backup_period*60):\n",
        "            simulation.runForClockTime(1*second)\n",
        "            newStep = simulation.context.getStepCount()\n",
        "            now = simulation.context.getTime().value_in_unit(picosecond)\n",
        "            runtime = now - time\n",
        "            if runtime >= length:\n",
        "                break\n",
        "            lines = stream.read()\n",
        "            with open(work+stage+'_log.txt', 'a') as f:\n",
        "                f.write(lines)\n",
        "            line = lines.split('\\n')\n",
        "            if len(line) < 2:\n",
        "                continue\n",
        "            data = line[-2].split(',')\n",
        "            print('', end='\\r')\n",
        "            if len(data) < 6:\n",
        "                speed = float(data[4])\n",
        "                density = ''\n",
        "            else:\n",
        "                speed = float(data[5])\n",
        "                density = f'밀도: {float(data[4]):.4f} g/mL, '\n",
        "            if stage != 'production':\n",
        "                T = f'시간: {float(data[1]):.1f} ps, '\n",
        "            else:\n",
        "                T = f'시간: {float(data[1])/1000:.2f} ns, '\n",
        "\n",
        "            if speed > 0:\n",
        "                d, hs = divmod((length - runtime)/speed/1000*24, 24)\n",
        "                h, mins = divmod(hs*60, 60)\n",
        "                min, secs = divmod(mins*60, 60)\n",
        "                sec = int(secs)\n",
        "                remaining = f'{str(int(d))+\":\" if d > 0 else \"\"}{int(h):0>2}:{int(min):0>2}:{sec:0>2} 남음, '\n",
        "            else:\n",
        "                remaining = '--:--:-- 남음, '\n",
        "\n",
        "            print(f'{remaining}step수: {data[0]}, {T}온도: {(float(data[3])-273.15):.2f} C, {density}계산속도: {speed:.1f} ns/day', end='')\n",
        "\n",
        "        state = simulation.context.getState(getPositions=True, getVelocities=True)\n",
        "        with open(work+stage+'_step_'+str(newStep)+'.xml', 'w') as f:\n",
        "            f.write(XmlSerializer.serialize(state))\n",
        "        os.remove(work+stage+'_step_'+str(step)+'.xml')\n",
        "\n",
        "    print('\\n계산 완료!')\n",
        "    return\n",
        "\n",
        "\n",
        "drive.mount('/content/drive')\n",
        "working_directory = \"test\" #@param {type:\"string\"}\n",
        "work = navigator(working_directory)"
      ],
      "metadata": {
        "id": "Uh2F1zMkUnFz",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 1. PDB 시작 파일 준비"
      ],
      "metadata": {
        "id": "xoA5kbSMqkv5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 1-1. PDB 파일 수선\n",
        "\n",
        "assert 'work' in globals(), \"0번 셀을 실행해주세요.\"\n",
        "\n",
        "#@markdown 4자리 PDB ID를 입력해주세요.\n",
        "#@markdown PDB 파일을 업로드 하고 싶으면 아래 칸을 비워주세요.\n",
        "PDB_ID = \"1B0Y\" #@param {type:\"string\"}\n",
        "#@markdown [RCSB Chemical Component Dictionary](https://www.wwpdb.org/data/ccd)에 정의된 소분자/비정규 residue의 경우,\n",
        "#@markdown residue 이름을 입력해주세요. 여러개는 공백 없이 쉼표로 구분해주세요.\n",
        "#@markdown <br/>**해당 residue 이름이 PDB 파일 내의 residue 이름과 같아야 합니다!**\n",
        "non_standard_residues = \"SF4\" #@param {type:\"string\"}\n",
        "#@markdown RCSB Dictionary에도 없는 residue의 경우, 다음 셀을 실행해 SDF 파일을 업로드 해주세요.<br/>\n",
        "#@markdown PDB 파일 내 물 분자를 없애고 싶으면 다음 체크 박스를 선택해주세요.\n",
        "remove_crystallographic_water = True #@param {type:\"boolean\"}\n",
        "\n",
        "amino_acids = ['ALA', 'CYS', 'ASP', 'GLU', 'PHE', 'GLY', 'HIS', 'ILE', 'LYS', 'LEU',\n",
        "               'MET', 'ASN', 'PRO', 'GLN', 'ARG', 'SER', 'THR', 'VAL', 'TRP', 'TYR']\n",
        "nucleic_acids = ['A', 'T', 'G', 'C', 'U', 'I', 'DA', 'DT', 'DG', 'DC', 'DI']\n",
        "\n",
        "if non_standard_residues == \"\":\n",
        "    nsres = []\n",
        "else:\n",
        "    nsres = non_standard_residues.split(',')\n",
        "\n",
        "with open(work+'non_standard_residues.txt', 'w') as f:\n",
        "    f.write('# User-defined Non-standard Residues\\n')\n",
        "    for res in nsres:\n",
        "        f.write(res.upper()+'\\n')\n",
        "\n",
        "if PDB_ID == \"\":\n",
        "    print(\"PDB 파일을 업로드해주세요.\")\n",
        "    pdbfile = files.upload()\n",
        "    try:\n",
        "        fixer = PDBFixer(filename=next(iter(pdbfile)))\n",
        "    except:\n",
        "        print('PDB 파일 입력 오류.')\n",
        "else:\n",
        "    try:\n",
        "        fixer = PDBFixer(pdbid=PDB_ID.upper())\n",
        "    except:\n",
        "        print('RCSB에서 PDB ID 검색에 실패했습니다.')\n",
        "\n",
        "print('빠진 원자들을 채우는중...', end='')\n",
        "fixer.findMissingResidues()\n",
        "terminal_missing_res = []\n",
        "chains = list(fixer.topology.chains())\n",
        "for key in fixer.missingResidues.keys():\n",
        "    if key[1] == 0 or key[1] == len(list(chains[key[0]].residues())):\n",
        "        terminal_missing_res.append(key)\n",
        "for key in terminal_missing_res:\n",
        "    del fixer.missingResidues[key]\n",
        "fixer.findMissingAtoms()\n",
        "fixer.addMissingAtoms()\n",
        "\n",
        "model = Modeller(fixer.topology, fixer.positions)\n",
        "res_to_del = []\n",
        "for res in model.topology.residues():\n",
        "    if remove_crystallographic_water:\n",
        "        if res.name not in amino_acids+nucleic_acids+nsres:\n",
        "            res_to_del.append(res)\n",
        "    else:\n",
        "        if res.name not in (amino_acids+nucleic_acids+nsres+['HOH']):\n",
        "            res_to_del.append(res)\n",
        "\n",
        "model.delete(res_to_del)\n",
        "\n",
        "print('완료.')\n",
        "\n",
        "with open(work+\"fixed.pdb\", 'w') as f:\n",
        "    PDBFile.writeFile(model.topology, model.positions, f)\n",
        "\n",
        "custom_info = {}\n",
        "\n",
        "view1 = nv.NGLWidget()\n",
        "view1._set_size('750px','500px')\n",
        "view1.add_component(nv.FileStructure(work+\"fixed.pdb\"))\n",
        "view1.add_licorice()\n",
        "display(view1)"
      ],
      "metadata": {
        "id": "5QASovvjuj2l",
        "collapsed": true,
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title (선택) Custom ligand/residue를 위한 SDF 파일 업로드하기\n",
        "\n",
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "assert len(nsres) > 0, \"비정규 residue 이름이 위 셀에서 입력되지 않았습니다!\"\n",
        "\n",
        "print(\"SDF 파일(들)을 업로드해주세요.\")\n",
        "ligfiles = files.upload()\n",
        "\n",
        "lignames = list(ligfiles.keys())\n",
        "lig_select = widgets.Dropdown(options=lignames, value=lignames[0], description='Custom residue')\n",
        "view0 = nv.NGLWidget()\n",
        "view0._set_size('400px','300px')\n",
        "\n",
        "for lig in lignames:\n",
        "    custom_info[lig] = {}\n",
        "custom_info[lig_select.value]['id'] = view0.add_component(nv.FileStructure(lig_select.value))\n",
        "\n",
        "res_select = widgets.Dropdown(options=nsres, value=nsres[0], description='Choose residue:')\n",
        "custom_info[lig_select.value]['name'] = res_select.value\n",
        "\n",
        "def update_mol2(change): # callback function to interactively update viewer\n",
        "    view0.remove_component(custom_info[change.old]['id'])\n",
        "    custom_info[change.new]['id'] = view0.add_component(nv.FileStructure(change.new))\n",
        "    view0.center()\n",
        "    custom_info[change.new]['name'] = res_select.value\n",
        "\n",
        "def update_name(change):\n",
        "    custom_info[lig_select.value]['name'] = change.new\n",
        "\n",
        "lig_select.observe(update_mol2, names='value')\n",
        "res_select.observe(update_name, names='value')\n",
        "\n",
        "print('업로드한 파일별로 residue 이름을 지정해주세요.')\n",
        "\n",
        "display(lig_select, view0, res_select)"
      ],
      "metadata": {
        "id": "LWF6fBt3OqU7",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 1-2. PDB 파일에 수소 원자 더하기\n",
        "\n",
        "assert 'work' in globals(), \"0번 셀부터 실행해주세요..\"\n",
        "assert os.path.exists(work+'fixed.pdb'), \"1-1 부터 실행해주세요.\"\n",
        "\n",
        "#@markdown 수소 원자를 더하는데 사용할 도구를 선택해주세요: <br/>\n",
        "#@markdown * PDB2PQR - PropKa를 이용해 pKa를 예측, 수소를 붙임.\n",
        "#@markdown * REDUCE - 중성 pH에 대해서만 작동. 리간드에도 붙여줌\n",
        "#@markdown * OpenBabel - 이론상 모든 분자에 수소를 붙여줄 수 있어야 하지만 버그가 많음.\n",
        "#@markdown * PDBFixer - 정규 아미노산에 대해서만 수소 붙여줌.\n",
        "#@markdown * Custom - 수소가 붙여진 PDB 파일을 업로드 하세요.\n",
        "#@markdown [DelPhipKa](http://compbio.clemson.edu/pka_webserver/),\n",
        "#@markdown [H++](http://newbiophysics.cs.vt.edu/H++/),\n",
        "#@markdown [PypKa](https://pypka.org/run-pypka/) 등의 외부 도구 이용 가능.\n",
        "\n",
        "protonation_method = \"REDUCE\" #@param ['PDB2PQR', 'REDUCE', 'OpenBabel', 'PDBFixer', 'Custom']\n",
        "\n",
        "#@markdown 수소를 붙여줄 기준 pH를 정하세요.\n",
        "pH = 7.4 #@param {type:\"slider\", min:0.0, max:14.0, step:0.1}\n",
        "\n",
        "for key, value in custom_info.items():\n",
        "    subprocess.run(f'cp {key} {work}{value[\"name\"]}.sdf', shell=True)\n",
        "\n",
        "\n",
        "def protonate_w_pdbfixer(work, removeH, pH)->None:\n",
        "    \"\"\"\n",
        "    Add missing hydrogen atoms to the PDB file using PDBFixer.\n",
        "    \"\"\"\n",
        "    print('PDBFixer로 수소를 더하는 중...', end='')\n",
        "    pdb = PDBFile(work+'fixed.pdb')\n",
        "    model = Modeller(pdb.topology, pdb.positions)\n",
        "\n",
        "    if removeH:\n",
        "        Hs = []\n",
        "        for a in model.topology.atoms():\n",
        "            if a.element.symbol == 'H':\n",
        "                Hs.append(a)\n",
        "        model.delete(Hs)\n",
        "\n",
        "    try:\n",
        "        model.addHydrogens(pH=pH)\n",
        "        with open(work+'starting_structure.pdb', 'w') as f:\n",
        "            PDBFile.writeFile(model.topology, model.positions, f)\n",
        "    except Exception as e:\n",
        "        print(e)\n",
        "    print('완료.')\n",
        "    return\n",
        "\n",
        "\n",
        "def protonate_manually(work)->None:\n",
        "    \"\"\"\n",
        "    Enable uploading of custom PDB file.\n",
        "    \"\"\"\n",
        "    print('빠진 부분이 수선된 PDB 파일을 다운로드 받으세요.')\n",
        "    files.download(work+'fixed.pdb')\n",
        "    print('이제 수소가 붙여진 PDB 파일을 업로드 하세요.')\n",
        "    fixed_wH = files.upload()\n",
        "    filename = next(iter(fixed_wH))\n",
        "    subprocess.run(f'cp {filename} {work}starting_structure.pdb', shell=True)\n",
        "    return\n",
        "\n",
        "\n",
        "def protonate_w_openbabel(work, nsres, pH)->None:\n",
        "    \"\"\"\n",
        "    Add missing hydrogen atoms to the PDB file using OpenBabel.\n",
        "    Since OpenBabel names all the hydrogen as 'H', we need to rename them.\n",
        "    \"\"\"\n",
        "    print('OpenBabel로 수소를 붙이는 중...', end='')\n",
        "\n",
        "    # save atom names before OpenBabel butchers them\n",
        "    pdb = PDBFile(work+'fixed.pdb')\n",
        "    nsatoms = []\n",
        "    for res in pdb.topology.residues():\n",
        "        if res.name in nsres:\n",
        "            atoms = []\n",
        "            for a in res.atoms():\n",
        "                atoms.append(a.name)\n",
        "            nsatoms.append({'name':res.name, 'chain':res.chain.id, 'index':res.id, 'atoms':atoms})\n",
        "\n",
        "    locale.getpreferredencoding = lambda: \"UTF-8\"\n",
        "    subprocess.run(f'obabel {work}fixed.pdb -O tmp.pdb -xk -p {pH}', shell=True)\n",
        "\n",
        "    with open('tmp.pdb', 'r') as f: # just number Hs so that OpenMM can discriminate them\n",
        "        lines = f.readlines()\n",
        "        chain = ''\n",
        "        resi = -100\n",
        "        headlines = ''\n",
        "        bodylines = []\n",
        "        taillines = ''\n",
        "        headflag = True\n",
        "        for line in lines:\n",
        "            if line.startswith('ATOM') or line.startswith('HETATM'):\n",
        "                headflag = False\n",
        "                if line[17:20] == 'UNL': # ligands/ncAAs\n",
        "                    if line[21] != chain: # reset index\n",
        "                        chain = line[21]\n",
        "                        resi = -100\n",
        "                    if int(line[22:26]) > resi:\n",
        "                        resi = int(line[22:26])\n",
        "                        count = 0 # 0-based index of nsres\n",
        "                        Hcount = 1 # 1-based index of Hs in the residue\n",
        "                    if line[77] == 'H':\n",
        "                        name = 'H'+str(Hcount)\n",
        "                        Hcount += 1\n",
        "                    else:\n",
        "                        name = nsatoms[resi-1]['atoms'][count]\n",
        "                        count += 1\n",
        "                    resname = nsatoms[resi-1]['name']\n",
        "                    chainid = nsatoms[resi-1]['chain']\n",
        "                    resindex = nsatoms[resi-1]['index']\n",
        "                    newinfo = f'{name:^4} {resname:>3} {chainid}{resindex:>4}'\n",
        "                elif line[77] == 'H': # element symbol\n",
        "                    if line[21] != chain: # chain ID\n",
        "                        chain = line[21]\n",
        "                        resi = -100\n",
        "                    if int(line[22:26]) > resi:\n",
        "                        resi = int(line[22:26])\n",
        "                        Hcount = 1\n",
        "                    name = 'H'+str(Hcount) # count number of Hs in a residue\n",
        "                    Hcount += 1\n",
        "                    newinfo = f'{name:^4}'+line[16:26]\n",
        "                else:\n",
        "                    newinfo = line[12:26]\n",
        "                newline = line[:12]\n",
        "                newline += newinfo\n",
        "                newline += line[26:]\n",
        "                bodylines.append(newline)\n",
        "            elif headflag:\n",
        "                headlines += line\n",
        "            else:\n",
        "                taillines += line\n",
        "\n",
        "    with open('tmp2.pdb', 'w') as f:\n",
        "        f.write(headlines)\n",
        "        for line in sorted(bodylines, key=lambda x: (x[21], int(x[22:26]))):\n",
        "            f.write(line)  # sort by chainID then residue index\n",
        "        f.write(taillines)\n",
        "\n",
        "    warnings.filterwarnings('ignore')\n",
        "    pdb2 = PDBFile('tmp2.pdb')\n",
        "\n",
        "    bonds_wH = []\n",
        "    for bond in pdb2.topology.bonds():\n",
        "        if bond[0].element.symbol == 'H':\n",
        "            bonds_wH.append([bond[1], bond[0]])\n",
        "        elif bond[1].element.symbol == 'H':\n",
        "            bonds_wH.append([bond[0], bond[1]])\n",
        "        else:\n",
        "            continue\n",
        "\n",
        "    counter = 4\n",
        "    for bond in bonds_wH:\n",
        "        multi = [b[0] for b in bonds_wH].count(bond[0]) # number of equivalent Hs\n",
        "        bond[1].residue = bond[0].residue\n",
        "        if counter > multi or counter == 0:\n",
        "            counter = multi\n",
        "        if bond[0].name == 'N':\n",
        "            if multi == 3: # N-terminus\n",
        "                bond[1].name = 'H'+str(counter)\n",
        "                counter -= 1\n",
        "            else:          # amide proton\n",
        "                bond[1].name = 'H'\n",
        "                counter -= 1\n",
        "        else:\n",
        "            if multi == 1:\n",
        "                bond[1].name = 'H'+bond[0].name[1:]\n",
        "                counter -= 1\n",
        "            elif multi == 2: # AMBER atom naming scheme (HB2, HB3 connected to CB)\n",
        "                if bond[0].element.symbol == 'C':\n",
        "                    bond[1].name = 'H'+bond[0].name[1:]+str(counter+1)\n",
        "                    counter -= 1\n",
        "                else:\n",
        "                    bond[1].name = 'H'+bond[0].name[1:]+str(counter)\n",
        "                    counter -= 1\n",
        "            elif multi == 3:\n",
        "                bond[1].name = 'H'+bond[0].name[1:]+str(counter)\n",
        "                counter -= 1\n",
        "\n",
        "    with open(work+'starting_structure.pdb', 'w') as f:\n",
        "        PDBFile.writeFile(pdb2.topology, pdb2.positions, f)\n",
        "    print('done.')\n",
        "    return\n",
        "\n",
        "\n",
        "\n",
        "def protonate_w_pdb2pqr(work, removeH, pH)->None:\n",
        "    \"\"\"\n",
        "    Add missing hydrogen atoms to the PDB file using PDB2PQR.\n",
        "    \"\"\"\n",
        "    print('PDB2PQR로 수소를 붙이는 중...', end='')\n",
        "\n",
        "    if removeH:\n",
        "        pdb = PDBFile(work+'fixed.pdb')\n",
        "        model = Modeller(pdb.topology, pdb.positions)\n",
        "        Hs = []\n",
        "        for a in model.topology.atoms():\n",
        "            if a.element.symbol == 'H':\n",
        "                Hs.append(a)\n",
        "        model.delete(Hs)\n",
        "\n",
        "        with open('fixed-H.pdb', 'w') as f:\n",
        "            PDBFile.writeFile(model.topology, model.positions, f)\n",
        "\n",
        "        params = ' fixed-H.pdb fixed.pqr'\n",
        "\n",
        "    else:\n",
        "        params = f' {work}fixed.pdb fixed.pqr'\n",
        "\n",
        "\n",
        "    params += f' --ff AMBER --ffout AMBER --pdb-output {work}starting_structure.pdb'\n",
        "    params += f' --titration-state-method propka --with-ph {pH} --pH {pH} --drop-water'\n",
        "    subprocess.run(\"pdb2pqr\"+params, shell=True)\n",
        "    print('done.')\n",
        "    return\n",
        "\n",
        "\n",
        "\n",
        "def get_parmed_atom_from_name(structure, res_top, name):\n",
        "    for i, res in enumerate(structure.topology.residues()):\n",
        "        if res.name == res_top.name and res.index == res_top.index:\n",
        "            residue = structure.residues[i]\n",
        "            break\n",
        "    for atom in residue.atoms:\n",
        "        if atom.name == name:\n",
        "            return atom\n",
        "\n",
        "def get_parmed_res_from_top(structure, res_top):\n",
        "    for i, res in enumerate(structure.topology.residues()):\n",
        "        if res.name == res_top.name and res.index == res_top.index:\n",
        "            return structure.residues[i]\n",
        "\n",
        "def find_coordinate(structure, atom):\n",
        "    for i, a in enumerate(structure.atoms):\n",
        "        if a is atom:\n",
        "            return np.array([structure[i].xx, structure[i].xy, structure[i].xz])\n",
        "\n",
        "def add_missed_proton(structure, residue, coords)->None:\n",
        "    \"\"\"\n",
        "    Rarely, REDUCE program omits amide proton of some residues.\n",
        "    This function adds them back.\n",
        "    \"\"\"\n",
        "    H = pmd.Atom(name='H', atomic_number=1, type='H', mass=1.008, charge=0.0)\n",
        "\n",
        "    pmd_res = get_parmed_res_from_top(structure, residue)\n",
        "\n",
        "    structure.add_atom_to_residue(H, pmd_res)\n",
        "\n",
        "    for i, a in enumerate(structure.atoms):\n",
        "        if a is H:\n",
        "            structure[i].xx = coords[0]\n",
        "            structure[i].xy = coords[1]\n",
        "            structure[i].xz = coords[2]\n",
        "\n",
        "    structure.assign_bonds()\n",
        "\n",
        "    print(f'Added a missing amide proton at {residue}!')\n",
        "\n",
        "    return\n",
        "\n",
        "\n",
        "\n",
        "def protonate_w_reduce(work, removeH)->None:\n",
        "    \"\"\"\n",
        "    Add missing hydrogen atoms to the PDB file using REDUCE.\n",
        "    \"\"\"\n",
        "    print('Warning: REDUCE에서 기준 pH는 7.4로 고정됩니다.')\n",
        "    print('REDUCE로 수소를 더하는 중...', end='')\n",
        "\n",
        "    if removeH:\n",
        "        Hflag = '-DROP_HYDROGENS_ON_ATOM_RECORDS '\n",
        "    else:\n",
        "        Hflag = ''\n",
        "    subprocess.run(f'reduce -BUILD {Hflag}{work}fixed.pdb > reduced.pdb', shell=True)\n",
        "    print('done.')\n",
        "\n",
        "    struc = pmd.load_file(\"reduced.pdb\")\n",
        "    struc.assign_bonds()\n",
        "\n",
        "    for chain in struc.topology.chains():\n",
        "        for i, res in enumerate(chain.residues()):\n",
        "            if i == 0 or res.name == 'PRO':\n",
        "                C = get_parmed_atom_from_name(struc, res, 'C')\n",
        "                O = get_parmed_atom_from_name(struc, res, 'O')\n",
        "                continue\n",
        "            N = get_parmed_atom_from_name(struc, res, 'N')\n",
        "\n",
        "            # In most cases, N-H is antiparallel to C=O\n",
        "            a_names = [atom.name for atom in res.atoms()]\n",
        "            if 'N' in a_names and 'CA' in a_names and 'C' in a_names and 'O' in a_names and 'H' not in a_names:\n",
        "                v_c = find_coordinate(struc, C) - find_coordinate(struc, O)\n",
        "                pos_H = find_coordinate(struc, N) + 1.01*v_c/np.linalg.norm(v_c)\n",
        "                add_missed_proton(struc, res, pos_H)\n",
        "\n",
        "            C = get_parmed_atom_from_name(struc, res, 'C')\n",
        "            O = get_parmed_atom_from_name(struc, res, 'O')\n",
        "\n",
        "    with open(work+'starting_structure.pdb', 'w') as f:\n",
        "        PDBFile.writeFile(struc.topology, struc.positions, f)\n",
        "    return\n",
        "\n",
        "\n",
        "#@markdown Check below box to ignore the hydrogen atoms currently present in your PDB.\n",
        "ignore_current_Hs = True #@param {type:\"boolean\"}\n",
        "\n",
        "if protonation_method == \"PDBFixer\":\n",
        "    protonate_w_pdbfixer(work, ignore_current_Hs, pH)\n",
        "\n",
        "elif protonation_method == \"OpenBabel\":\n",
        "    print('OpenBabel 설치 중...', end='')\n",
        "    subprocess.run(['mamba', 'install', '-c', 'conda-forge', 'openbabel'])\n",
        "    print('완료.')\n",
        "    protonate_w_openbabel(work, nsres, pH)\n",
        "\n",
        "elif protonation_method == \"PDB2PQR\":\n",
        "    print('PDB2PQR 설치 중...', end='')\n",
        "    subprocess.run(['pip install pdb2pqr'])\n",
        "    print('완료.')\n",
        "    protonate_w_pdb2pqr(work, ignore_current_Hs, pH)\n",
        "\n",
        "elif protonation_method == 'REDUCE':\n",
        "    print('REDUCE 설치 중...', end='')\n",
        "    subprocess.run(['mamba', 'install', '-c', 'bioconda', 'reduce'])\n",
        "    print('완료.')\n",
        "    protonate_w_reduce(work, ignore_current_Hs)\n",
        "\n",
        "else:\n",
        "    protonate_manually(work)\n",
        "\n",
        "\n",
        "\n",
        "#@markdown Select the center of focus to zoom in/out.<br/>\n",
        "#@markdown You may also long-click an atom to center zoom and translate molecule with a right drag.\n",
        "\n",
        "to_lookat = []\n",
        "\n",
        "pdb = PDBFile(work+'starting_structure.pdb')\n",
        "\n",
        "for res in pdb.topology.residues():\n",
        "    if res.name in nsres:\n",
        "        to_lookat.append(res)\n",
        "\n",
        "# Interactively shows the protonated structure\n",
        "select_res = widgets.Dropdown(options=['whole structure']+to_lookat, value='whole structure', description='Focus at', continuous_update=True)\n",
        "view2 = nv.NGLWidget()\n",
        "view2._set_size('750px','500px')\n",
        "try:\n",
        "    view2.add_component(nv.FileStructure(work+\"starting_structure.pdb\"))\n",
        "except:\n",
        "    print(f\"수소를 더하는 중 {protonation_method}에서 에러가 발생했습니다.\")\n",
        "\n",
        "view2.add_licorice()\n",
        "view2.center()\n",
        "\n",
        "def center_of_view(change): # callback function for the update\n",
        "    view2.center(f'{change.new.id} and :{change.new.chain.id}' if change.new != 'whole structure' else None)\n",
        "\n",
        "select_res.observe(center_of_view, names='value')\n",
        "\n",
        "display(select_res, view2)\n",
        "\n",
        "print(\"수소가 붙은 구조는 작업폴더 내 'starting_structure.pdb'에도 저장되었습니다.\")\n",
        "print(\"수소가 제대로 붙었는지, 특히 비정규 residue와 그 주변에 대해서 꼭 확인해주세요!\")"
      ],
      "metadata": {
        "id": "SYA6pxmcgyti",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 2. 시뮬레이션 객체 생성"
      ],
      "metadata": {
        "id": "g0Ou4jzVqqi5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 2-1. 역장 선택\n",
        "assert 'work' in globals(), \"0번 셀부터 실행해주세요.\"\n",
        "\n",
        "def identify_new_res(fftype, forcefield, model, non_standard_residue_list)->list:\n",
        "    \"\"\"\n",
        "    Returns a list of nonstandard residues that\n",
        "    are not present in the given macromolecular forcefield.\n",
        "    \"\"\"\n",
        "    if len(non_standard_residue_list) == 0:\n",
        "        return []\n",
        "\n",
        "    # Prepare copied modeller objects with a single nonstandard residue\n",
        "    models = {}\n",
        "    for res in non_standard_residue_list:\n",
        "        models[res] = copy.deepcopy(model)\n",
        "    for key, value in models.items():\n",
        "        other_res = []\n",
        "        for res in value.topology.residues():\n",
        "            if res.name in non_standard_residue_list and res.name != key:\n",
        "                other_res.append(res)\n",
        "        if len(other_res) > 0:\n",
        "            value.delete(other_res)\n",
        "\n",
        "    new_res = []\n",
        "\n",
        "    for key, value in models.items():\n",
        "        try:\n",
        "            if fftype == 'Amber14':\n",
        "                value.addExtraParticles(forcefield)\n",
        "            elif fftype == 'CHARMM36': #due to patches, CHARMM forcefield gives errors with addExtraParticles()\n",
        "                forcefield.createSystem(value.topology)\n",
        "            print(f'{fftype} 역장 내에서 {key}에 해당하는 template를 찾았습니다!')\n",
        "\n",
        "        except ValueError: # Usually raised when no matching residue templates are found.\n",
        "            #print(str(e)[:41])\n",
        "            # Trim other residues and save first appearing structure of the nonstandard residue\n",
        "            res_to_del = []\n",
        "            counter = 0\n",
        "            for res in value.topology.residues():\n",
        "                if res.name != key or counter > 0:\n",
        "                    res_to_del.append(res)\n",
        "                else:\n",
        "                    counter += 1\n",
        "            value.delete(res_to_del)\n",
        "            with open(f'{key}.pdb', 'w') as f:\n",
        "                PDBFile.writeFile(value.topology, value.positions, f)\n",
        "            new_res.append(key)\n",
        "    return new_res\n",
        "\n",
        "\n",
        "def isLigand(model, name)->bool:\n",
        "    \"\"\"\n",
        "    Check whether the given residue is connected to the other residues.\n",
        "    If not, it is considered as a ligand.\n",
        "    \"\"\"\n",
        "    for res in model.topology.residues():\n",
        "        if res.name == name:\n",
        "            if len(list(res.external_bonds())) > 0:\n",
        "                return False\n",
        "            else:\n",
        "                return True\n",
        "\n",
        "def scan_elements(model, name)->bool:\n",
        "    \"\"\"\n",
        "    Check whether the given residue contains the elements that\n",
        "    are not supported by GAFF/SMIRNOFF/espaloma/MACE.\n",
        "    \"\"\"\n",
        "    gaff_elements = ['H', 'C', 'N', 'O', 'S', 'P', 'F', 'Cl', 'Br', 'I']\n",
        "\n",
        "    for res in model.topology.residues():\n",
        "        if res.name == name:\n",
        "            residue = res\n",
        "            break\n",
        "\n",
        "    elements = []\n",
        "    for atom in residue.atoms():\n",
        "        if atom.element.symbol not in elements:\n",
        "            elements.append(atom.element.symbol)\n",
        "\n",
        "    elements = list(set(elements))\n",
        "\n",
        "    for ele in elements:\n",
        "        if ele not in gaff_elements:\n",
        "            return False\n",
        "\n",
        "    return True\n",
        "\n",
        "\n",
        "#@markdown 생체고분자에 적용할 역장을 골라주세요.\n",
        "forcefield = 'Amber14' #@param ['Amber14', 'CHARMM36']\n",
        "\n",
        "#@markdown 용매 모델을 골라주세요.\n",
        "watermodel = \"(implicit) GBn2 (igb=8)\" #@param [\"(explicit) TIP3P\", \"(explicit) SPC/E\", \"(explicit) TIP4P-Ew\", \"(explicit) TIP5P\", \"(explicit) OPC\", \"(explicit) OPC3\", \"(implicit) HCT (igb=1)\", \"(implicit) OBC1 (igb=2)\", \"(implicit) OBC2 (igb=5)\", \"(implicit) GBn (igb=7)\", \"(implicit) GBn2 (igb=8)\"]\n",
        "\n",
        "#@markdown 리간드나 비정규 residue에 대해서는 소분자용 역장을 선택해야 합니다.\n",
        "#@markdown * GFN-FF : 반경험적 양자역학 역장 (화학 반응 계산도 가능, 느림)\n",
        "#@markdown * MACE-off : 기계학습 역장\n",
        "#@markdown * GAFF : Ambertools의 Generalized Amber Forcefield\n",
        "#@markdown * Sage : OpenForceField의 Smirnoff\n",
        "#@markdown * espaloma : 신경망을 통해 학습된 리간드용 역장\n",
        "#@markdown * Custom : 새로운 residue에 대한 정보를 담은 XML파일을 업로드할 수 있습니다.\n",
        "\n",
        "ff, water = determine_forcefield(forcefield, watermodel)\n",
        "\n",
        "assert ff is not None, \"역장이 만들어지지 않았습니다.\"\n",
        "\n",
        "pdb = PDBFile(work+'starting_structure.pdb')\n",
        "model = Modeller(pdb.topology, pdb.positions)\n",
        "\n",
        "nsres = [a['name'] for a in read_non_standard_residues_list(work)]\n",
        "uaas = identify_new_res(forcefield, ff, model, nsres)\n",
        "uaa_info = {}\n",
        "uaa_atom_info = []\n",
        "\n",
        "\n",
        "if len(uaas) > 0:\n",
        "    print(f'{forcefield} 역장 데이터베이스에 없는 residue를 발견하였습니다.')\n",
        "    print('소분자용 역장을 선택해주세요.')\n",
        "    if len(uaas) > 1:\n",
        "        print('다른 residue들에 대해서도 역장을 선택해주세요.')\n",
        "\n",
        "    for uaa in uaas:\n",
        "        lig_flag = isLigand(model, uaa)\n",
        "        viable_ffs = ['GFN-FF'] # GFN-FF is applicable to all elements on the Periodic Table.\n",
        "        if scan_elements(model, uaa):\n",
        "            viable_ffs = ['MACE-off23'] + viable_ffs # Machine Learning forcefields\n",
        "        if scan_elements(model, uaa) and lig_flag:\n",
        "            viable_ffs = ['GAFF-2.11', 'Sage-2.2.1'] + viable_ffs # small molecule forcefields\n",
        "        viable_ffs.append('Custom')\n",
        "        uaa_info[uaa] = {'name':uaa, 'file':f'{uaa}.pdb', 'lig':lig_flag, 'ff':'GFN-FF', 'ffs':viable_ffs}\n",
        "\n",
        "    pdb_select = widgets.Dropdown(options=uaas, value=uaas[0], description='새 residue', continuous_update=True)\n",
        "    view3 = nv.NGLWidget()\n",
        "    view3._set_size('400px','300px')\n",
        "    # save NGLView component id to remove it later\n",
        "    uaa_info[pdb_select.value]['id'] = view3.add_component(nv.FileStructure(uaa_info[pdb_select.value]['file']))\n",
        "    ff_select = widgets.Dropdown(options=uaa_info[uaas[0]]['ffs'], value='GFN-FF', description='역장', continuous_update=True)\n",
        "\n",
        "    def update_pdb(change): # callback function to interactively update viewer\n",
        "        view3.remove_component(uaa_info[change.old]['id'])\n",
        "        uaa_info[change.new]['id'] = view3.add_component(nv.FileStructure(uaa_info[change.new]['file']))\n",
        "        view3.center(change.new)\n",
        "        ff_select.options = uaa_info[change.new]['ffs']\n",
        "        ff_select.value = uaa_info[change.new]['ff']\n",
        "\n",
        "    def update_ff(change): # callback function to interactively update residue infos\n",
        "        uaa_info[pdb_select.value]['ff'] = change.new\n",
        "        ffs =[]\n",
        "        for info in uaa_info.values():\n",
        "            ffs.append(info['ff'])\n",
        "\n",
        "    pdb_select.observe(update_pdb, names='value')\n",
        "    ff_select.observe(update_ff, names='value')\n",
        "\n",
        "    display(pdb_select, view3, ff_select)\n",
        "\n",
        "else:\n",
        "    print('모든 residue가 역장 데이터베이스에 존재합니다.')"
      ],
      "metadata": {
        "id": "yN-yWp8R4aoe",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 2-2. 비정규 residue 주변 완충영역 설정\n",
        "assert 'work' in globals(), \"0번 셀부터 실행해주세요.\"\n",
        "assert 'uaas' in globals(), \"2-1부터 실행해주세요.\"\n",
        "\n",
        "#@markdown ML- 또는 XTB- 역장만으로 계산되는 비정규 residue의 원자들은 Lennard-Jones 매개변수가 0으로 설정됩니다.<br/>\n",
        "#@markdown 따라서 이들과 직접적으로 맞부딪히는 원자들까지 완충영역으로 설정해줘야 오차를 최소화할 수 있습니다. <br/>\n",
        "#@markdown 기준 반경을 조절해서 완충영역의 크기를 설정하세요.<br/>\n",
        "#@markdown **이 셀을 연속으로 두 번 이상 실행하지 마세요!**\n",
        "\n",
        "def make_xml_for_custom_residues(fftype, ff, model, residue_infos)->bool:\n",
        "    \"\"\"\n",
        "    Generate a template XML file for custom residues.\n",
        "    \"\"\"\n",
        "    ncAA = []\n",
        "    for value in residue_infos.values():\n",
        "        if value['ff'] == 'Custom':\n",
        "            ncAA.append(value['name'])\n",
        "    if len(ncAA) == 0:\n",
        "        return False\n",
        "\n",
        "    templates, residues = ff.generateTemplatesForUnmatchedResidues(model.topology)\n",
        "\n",
        "    with open('template.xml', 'w') as f:\n",
        "        f.write(\"<ForceField>\\n\")\n",
        "        f.write(\" <AtomTypes>\\n\")\n",
        "        f.write(\" </AtomTypes>\\n\")\n",
        "        f.write(\" <Residues>\\n\")\n",
        "        for i, temp in enumerate(templates):\n",
        "            if temp.name in ncAA:\n",
        "                f.write(f'  <Residue name=\"{temp.name}\">\\n')\n",
        "                for atom in residues[i].atoms():\n",
        "                    f.write(f'   <Atom name=\"{atom.name}\" type=\"Enter_atom_type\"/>\\n')\n",
        "                for bond in residues[i].internal_bonds():\n",
        "                    f.write(f'   <Bond atomName1=\"{bond[0].name}\" atomName2=\"{bond[1].name}\"/>\\n')\n",
        "                for bond in residues[i].external_bonds():\n",
        "                    if bond[0] in list(residues[i].atoms()):\n",
        "                        atom = bond[0]\n",
        "                    elif bond[1] in list(residues[i].atoms()):\n",
        "                        atom = bond[1]\n",
        "                    f.write(f'   <ExternalBond atomName=\"{atom.name}\"/>\\n')\n",
        "                f.write('  </Residue>\\n')\n",
        "        f.write(\" </Residues>\\n\")\n",
        "        f.write(\" <HarmonicBondForce>\\n\")\n",
        "        f.write(\" </HarmonicBondForce>\\n\")\n",
        "        f.write(\" <HarmonicAngleForce>\\n\")\n",
        "        f.write(\" </HarmonicAngleForce>\\n\")\n",
        "        if fftype == 'CHARMM36':\n",
        "            f.write(' <AmoebaUreyBradleyForce>\\n')\n",
        "            f.write(\" </AmoebaUreyBradleyForce>\\n\")\n",
        "        if fftype == 'Amber14':\n",
        "            f.write(' <PeriodicTorsionForce ordering=\"amber\">\\n')\n",
        "        elif fftype == 'CHARMM36':\n",
        "            f.write(' <PeriodicTorsionForce ordering=\"charmm\">\\n')\n",
        "        f.write(\" </PeriodicTorsionForce>\\n\")\n",
        "        if fftype == 'CHARMM36':\n",
        "            f.write(' <CMAPTorsionForce>\\n')\n",
        "            f.write(\" </CMAPTorsionForce>\\n\")\n",
        "        if fftype == 'Amber14':\n",
        "            f.write(' <NonbondedForce coulomb14scale=\"0.8333333333333334\" lj14scale=\"0.5\">\\n')\n",
        "        elif fftype == 'CHARMM36':\n",
        "            f.write(' <NonbondedForce coulomb14scale=\"1.0\" lj14scale=\"1.0\" useDispersionCorrection=\"False\">\\n')\n",
        "        f.write('  <UseAttributeFromResidue name=\"charge\"/>\\n')\n",
        "        f.write(' </NonbondedForce>\\n')\n",
        "        if fftype == 'CHARMM36':\n",
        "            f.write(' <LennardJonesForce lj14scale=\"1.0\" useDispersionCorrection=\"False\">\\n')\n",
        "            f.write(' </LennardJonesForce>\\n')\n",
        "        f.write(\"</ForceField>\\n\")\n",
        "\n",
        "    files.download('template.xml')\n",
        "    sleep(5)\n",
        "    return True\n",
        "\n",
        "\n",
        "def find_terminal_form(pdb)->str:\n",
        "\n",
        "    mol = Chem.MolFromPDBFile(pdb, removeHs=False)\n",
        "\n",
        "    backbone = Chem.MolFromSmarts('[NH1][CH1]C(=O)')\n",
        "    backatoms = list(mol.GetSubstructMatch(backbone))\n",
        "    if len(backatoms) == 0:\n",
        "        return \"\"\n",
        "\n",
        "    nterm = Chem.MolFromSmarts('[NH3+][CH1]C(=O)')\n",
        "    nterm_atoms = list(mol.GetSubstructMatch(nterm))\n",
        "    if len(nterm_atoms) > 0:\n",
        "        return 'N'\n",
        "\n",
        "    cterm = Chem.MolFromSmarts('[NH1][CH1]C(=O)[O-]')\n",
        "    cterm_atoms = list(mol.GetSubstructMatch(cterm))\n",
        "    if len(cterm_atoms) > 0:\n",
        "        return 'C'\n",
        "\n",
        "    return '-'\n",
        "\n",
        "\n",
        "def assign_partial_charges(name, residue_infos)->dict:\n",
        "    charges = {}\n",
        "    pdb = residue_infos[name]['file']\n",
        "    subprocess.run(f'obabel {pdb} -O {name}.mol2 --partialcharge mmff94', shell=True)\n",
        "\n",
        "    with open(f'{name}.mol2', 'r') as f:\n",
        "        lines = f.readlines()\n",
        "        atom_flag = False\n",
        "        for line in lines:\n",
        "            if line.strip().startswith('@<TRIPOS>ATOM'):\n",
        "                atom_flag = True\n",
        "                continue\n",
        "            elif line.strip().startswith('@<TRIPOS>BOND'):\n",
        "                atom_flag = False\n",
        "                break\n",
        "            if atom_flag:\n",
        "                data = line.strip().split( )\n",
        "                charges[data[1]] = round(float(data[-1]), 4)\n",
        "\n",
        "    term = find_terminal_form(pdb)\n",
        "    sum = 0.0\n",
        "    for q in charges.values():\n",
        "        sum += q\n",
        "    offset = float(int(sum))-sum\n",
        "\n",
        "    if term == \"\":\n",
        "        for key, value in charges.items():\n",
        "            value += offset/len(charges)\n",
        "\n",
        "    elif term == \"N\":\n",
        "        before = 0.0\n",
        "        for key, value in charges.items():\n",
        "            if key == 'N':\n",
        "                before += value\n",
        "                value = -0.3\n",
        "            elif key == 'H1' or key == 'H2' or key == 'H3':\n",
        "                before += value\n",
        "                value = 0.33\n",
        "            elif key == 'CA':\n",
        "                before += value\n",
        "                value = 0.21\n",
        "            elif key == 'HA':\n",
        "                before += value\n",
        "                value = 0.1\n",
        "            elif key == 'C':\n",
        "                before += value\n",
        "                value = 0.51\n",
        "            elif key == 'O':\n",
        "                before += value\n",
        "                value = -0.51\n",
        "        offset += (before - 1) # N-terminus has +1 charge\n",
        "\n",
        "        for key, value in charges.items():\n",
        "            if key not in ['N', 'H1', 'H2', 'H3', 'CA', 'HA', 'C', 'O']:\n",
        "                value += offset/(len(charges)-8)\n",
        "\n",
        "    elif term == \"C\":\n",
        "        before = 0.0\n",
        "        for key, value in charges.items():\n",
        "            if key == 'C':\n",
        "                before += value\n",
        "                value = 0.34\n",
        "            elif key == 'O':\n",
        "                before += value\n",
        "                value = -0.67\n",
        "            elif key == 'OXT':\n",
        "                before += value\n",
        "                value = -0.67\n",
        "        offset += (before + 2) # C-terminus has -1 charge & formal charge for amide N is -1\n",
        "\n",
        "        for key, value in charges.items():\n",
        "            if key not in ['C', 'O', 'OXT']:\n",
        "                value += offset/(len(charges)-3)\n",
        "\n",
        "    elif term == '-':\n",
        "        before = 0.0\n",
        "        for key, value in charges.items():\n",
        "            if key == 'N':\n",
        "                before += value\n",
        "                value = -0.47\n",
        "            elif key == 'H':\n",
        "                before += value\n",
        "                value = 0.31\n",
        "            elif key == 'CA':\n",
        "                before += value\n",
        "                value = 0.07\n",
        "            elif key == 'HA':\n",
        "                before += value\n",
        "                value = 0.09\n",
        "            elif key == 'C':\n",
        "                before += value\n",
        "                value = 0.51\n",
        "            elif key == 'O':\n",
        "                before += value\n",
        "                value = -0.51\n",
        "        offset += (before + 1) # formal charge for amide N is -1, so correct it.\n",
        "\n",
        "        for key, value in charges.items():\n",
        "            if key not in ['N', 'H', 'CA', 'HA', 'C', 'O']:\n",
        "                value += offset/(len(charges)-6)\n",
        "    return charges\n",
        "\n",
        "\"\"\"\n",
        "def assign_LJ_parameters(fftype, ff, model, atomtype, element)->tuple:\n",
        "\n",
        "    resname = atomtype.split('_')[0]\n",
        "    atomname = atomtype.split('_')[1]\n",
        "\n",
        "    amber14LJ = {'C':(0.25104, 0.402), 'O':(0.79496, 0.333), 'N':(0.54392, 0.386)}\n",
        "    # parameters from\n",
        "    # M. Freindorf, Y. Shao, T. R. Furiani, and J. Kong.\n",
        "    # \"Lennard-Jones parameters for the combined QM/MM method using the B3LYP/6-31G*/AMBER potential\"\n",
        "    # J. Comput. Chem. 2005, 26(12): 1270-1278.\n",
        "    # https://doi.org/10.1002/jcc.20264\n",
        "\n",
        "    gaff2AB = {'S':(1.9825, 0.2824), 'P':(2.0732, 0.2295), 'F':(1.7029, 0.0832), 'Cl':(1.9452, 0.2638),\n",
        "                 'Br':(2.0275, 0.3932), 'I':(2.1558, 0.4955)}\n",
        "    amberions = {'Al':'AL', 'Ag':'Ag', 'Ba':'BA', 'Br':'BR', 'Be':'Be', 'Ca':'CA', 'Cd':'CD',\n",
        "                 'Cl':'CL', 'Co':'CO', 'Cs':'CS', 'Cu':'CU', 'Dy':'Dy', 'Er':'Er',\n",
        "                 'F':'F', 'Gd':'GD3', 'Hg':'HG', 'Hf':'Hf',\n",
        "                 'In':'IN', 'I':'IOD', 'K':'K', 'La':'LA', 'Li':'LI', 'Lu':'LU', 'Mg':'MG',\n",
        "                 'Mn':'MN', 'Na':'NA', 'Ni':'NI', 'Nd':'Nd', 'Pb':'PB', 'Pd':'PD', 'Pr':'PR',\n",
        "                 'Pt':'PT', 'Pu':'Pu', 'Rb':'RB', 'Ra':'Ra', 'Sm':'Sm', 'Sr':'SR', 'Sn':'Sn',\n",
        "                 'Tb':'TB', 'Th':'Th', 'Tl':'Tl', 'Tm':'Tm', 'U':'U4+', 'V':'V2+', 'Y':'Y',\n",
        "                 'Yb':'YB2', 'Zn':'ZN', 'Zr':'Zr'}\n",
        "    charmm36LJ = {'N':(0.8368, 0.329632525712), 'F':(0.56484, 0.290432982114), 'Si':(2.5104, 0.391995435982),\n",
        "                  'Al':(0.4460144, 0.039377723342), 'P':(2.44764, 0.3830864488), 'Cl':(1.435112, 0.34032331033),\n",
        "                  'Se':(2.606632, 0.387540942391), 'Br':(2.25936, 0.356359487256), 'I':(2.17568, 0.399122625727),\n",
        "                  'Be':(0.4250944, 0.007127189745), 'Sc':(0.589944, 0.195445360786), 'Ti':(1.7016328, 0.344011631023),\n",
        "                  'V':(0.6732056, 0.196656983042), 'Cr':(0.6974728, 0.188727984451), 'Mn':(0.7050040, 0.214474957405),\n",
        "                  'Fe':(0.667348, 0.192451941093), 'Co':(0.6439176, 0.165279530189), 'Ni':(0.5999856, 0.139033653953),\n",
        "                  'Cu':(0.7748768, 0.153590939007), 'Ga':(0.5045904, 0.040428983829), 'Sr':(0.9572992, 0.330683786199),\n",
        "                  'Y':(0.6928704, 0.251839249644), 'Rh':(0.537644, 0.082194315736), 'Pd':(0.7288528, 0.165190440318),\n",
        "                  'Ag':(1.200808, 0.216363662688), 'In':(0.6163032, 0.154357111905), 'Sn':(0.7861736, 0.295422014935),\n",
        "                  'La':(0.7932864, 0.315859231529), 'Ce':(0.7702744, 0.304954631219), 'Pr':(0.7627432, 0.295618012653),\n",
        "                  'Nd':(0.7497728, 0.288633366703), 'Pm':(0.74266, 0.29527947114), 'Sm':(1.0242432,0.330790694046),\n",
        "                  'Eu':(0.9907712,0.327922000173), 'Gd':(0.7209032, 0.26816051416), 'Tb':(0.7083512,0.262815121851),\n",
        "                  'Dy':(0.6999832,0.257362821696), 'Ho':(0.6928704, 0.247794569464), 'Er':(0.6845024,0.243625163463),\n",
        "                  'Tm':(0.677808, 0.240293202257), 'Yb':(0.8593936, 0.286691207498), 'Lu':(0.6644192, 0.24075646959),\n",
        "                  'Pt':(0.6824104, 0.142027073646), 'Au':(1.4907592, 0.112502690127), 'Hg':(0.8593936, 0.21057282102),\n",
        "                  'Tl':(0.6815736, 0.154891651136), 'Pb':(1.0016496, 0.313239989298), 'Bi':(0.77404, 0.243785525232),\n",
        "                  'Ra':(1.2631496,0.372377846208), 'U':(0.7932864, 0.30311937986), 'Pu':(0.7702744,0.297506717936)}\n",
        "    charmmions = {'Li':'LIT', 'Na':'SOD', 'Mg':'MG', 'K':'POT', 'Ca':'CAL', 'Rb':'RUB',\n",
        "                  'Cs':'CES', 'Ba':'BAR', 'Zn':'ZN', 'Cd':'CAD', 'Cl':'CLA'}\n",
        "\n",
        "    if fftype == 'Amber14':\n",
        "        if element == 'H':\n",
        "            for res in model.topology.residues():\n",
        "                if res.name == resname:\n",
        "                    for bond in res.bonds():\n",
        "                        if bond[0].name == atomname:\n",
        "                            parent = bond[1].element.symbol\n",
        "                            break\n",
        "                        elif bond[1].name == atomname:\n",
        "                            parent = bond[0].element.symbol\n",
        "                            break\n",
        "                        else:\n",
        "                            continue\n",
        "                    if parent == 'C': # nonpolar H\n",
        "                        lj = (0.12552, 0.222)\n",
        "                    else:             # polar H\n",
        "                        lj = (0.12552, 0.106)\n",
        "                    break\n",
        "            return lj\n",
        "        elif element in gaff2AB: # parm99.dat in Amber forcefield\n",
        "            A = gaff2AB[element][0]\n",
        "            B = gaff2AB[element][1]\n",
        "            sigma = ((A/B)**(1/6))/10  # sigma in Angstom\n",
        "            epsilon = B/4/(sigma**6)*4.184 # epsilon in kcal/mol\n",
        "            return (epsilon, sigma)\n",
        "        elif element in amber14LJ:\n",
        "            return amber14LJ[element]\n",
        "        elif element in amberions: # use parameters in ions definition\n",
        "            top = Topology()\n",
        "            chain = top.addChain(id='X')\n",
        "            residue = top.addResidue(amberions[element], chain)\n",
        "            top.addAtom(element, Element.getBySymbol(element), residue)\n",
        "            system = ff.createSystem(top, residueTemplates=dict(residue = amberions[element]))\n",
        "            for f in system.getForces():\n",
        "                if f.getName() == 'NonbondedForce':\n",
        "                    q, s, e = f.getParticleParameters(0)\n",
        "                    break\n",
        "            return (e,s)\n",
        "    elif fftype == 'CHARMM36':\n",
        "        if element == 'H':\n",
        "            for res in model.topology.residues():\n",
        "                if res.name == resname:\n",
        "                    for bond in res.bonds():\n",
        "                        if bond[0].name == atomname:\n",
        "                            parent = bond[1].element.symbol\n",
        "                            break\n",
        "                        elif bond[1].name == atomname:\n",
        "                            parent = bond[0].element.symbol\n",
        "                            break\n",
        "                        else:\n",
        "                            continue\n",
        "                    if parent == 'C': # nonpolar H\n",
        "                        lj = (0.192464, 0.040001352445)\n",
        "                    else:             # polar H\n",
        "                        lj = (0.092048, 0.235197261589)\n",
        "                    return lj\n",
        "        elif element == 'C':\n",
        "            for res in model.topology.residues():\n",
        "                if res.name == resname:\n",
        "                    count = 0\n",
        "                    for bond in res.bonds():\n",
        "                        if bond[0].name == atomname or bond[1].name == atomname:\n",
        "                            count += 1\n",
        "                    if count < 4: # aromatic / carbonyl C\n",
        "                        lj = (0.29288, 0.355005321205)\n",
        "                    else:         # aliphatic C\n",
        "                        lj = (0.08368, 0.405358916754)\n",
        "                    return lj\n",
        "        elif element == 'S':\n",
        "            for res in model.topology.residues():\n",
        "                if res.name == resname:\n",
        "                    count = 0\n",
        "                    for bond in res.bonds():\n",
        "                        if bond[0].name == atomname or bond[1].name == atomname:\n",
        "                            count += 1\n",
        "                    if count > 1: # thioether S\n",
        "                        lj = (1.8828, 0.356359487256)\n",
        "                    else:         # thiolate S\n",
        "                        lj = (1.96648, 0.391995435982)\n",
        "                    return lj\n",
        "        elif element == 'O':\n",
        "            for res in model.topology.residues():\n",
        "                if res.name == resname:\n",
        "                    count = 0\n",
        "                    for bond in res.bonds():\n",
        "                        if bond[0].name == atomname or bond[1].name == atomname:\n",
        "                            count += 1\n",
        "                    if count > 1: # ether O\n",
        "                        lj = (0.4184, 0.293996576986)\n",
        "                    else:         # anionic O\n",
        "                        lj = (0.50208, 0.302905564168)\n",
        "                    return lj\n",
        "        elif element in charmm36LJ:\n",
        "            return charmm36LJ[element]\n",
        "\n",
        "        elif element in charmmions: # use parameters in ions definition\n",
        "            top = Topology()\n",
        "            chain = top.addChain(id='X')\n",
        "            residue = top.addResidue(charmmions[element], chain)\n",
        "            top.addAtom(element, Element.getBySymbol(element), residue)\n",
        "            system = ff.createSystem(top)\n",
        "            for f in system.getForces():\n",
        "                if f.getName() in ['NonbondedForce', 'LennardJonesForce']:\n",
        "                    q, s, e = f.getParticleParameters(0)\n",
        "                    break\n",
        "            return (e,s)\n",
        "\n",
        "    # For elements not defined in AMBER/CHARMM, use OpenKIM Universal LJ parameters\n",
        "    # DOI : https://doi.org/10.25950/962b4967\n",
        "    if not os.path.exists('/content/LennardJones612_UniversalShifted.params'):\n",
        "        subprocess.run('wget https://openkim.org/files/MO_959249795837_003/LennardJones612_UniversalShifted.params', shell=True)\n",
        "    with open('LennardJones612_UniversalShifted.params', 'r') as f:\n",
        "        lines = f.readlines()\n",
        "        for line in lines:\n",
        "            if line and not line.startswith('#'):\n",
        "                data = line.strip().split( )\n",
        "                if len(data) == 5 and data[0] == element:\n",
        "                    epsilon = float(data[3])\n",
        "                    sigma = float(data[4])\n",
        "                    break\n",
        "    return (epsilon*96.485, sigma/10) # eV to kJ/mol, Anstrom to nm\n",
        "\"\"\"\n",
        "\n",
        "\n",
        "def assign_parameters(fftype, ff, model, residue_infos)->list:\n",
        "    \"\"\"\n",
        "    Assign those novel atom / residue types to the macromolecular forcefield.\n",
        "    \"\"\"\n",
        "    ambertype = {'N':'protein-N', 'CA':'protein-CX', 'C':'protein-C', 'O':'protein-O', 'H':'protein-H', 'HA':'protein-H1'}\n",
        "    charmmtype = {'N':'NH1', 'CA':'CT1', 'C':'C', 'O':'O', 'H':'H', 'HA':'HB1'}\n",
        "\n",
        "    templates, residues = ff.generateTemplatesForUnmatchedResidues(model.topology)\n",
        "    ncAA = []\n",
        "    for value in residue_infos.values():\n",
        "        if value['ff'] in ['GFN-FF', 'MACE-off23']:\n",
        "            ncAA.append(value['name'])\n",
        "    if len(ncAA) == 0:\n",
        "        return []\n",
        "\n",
        "    new_type = []\n",
        "    unk_index = 0\n",
        "    uaa_res = []\n",
        "    lj_parms = {}\n",
        "\n",
        "    for i, temp in enumerate(templates):\n",
        "        if temp.name not in ncAA:\n",
        "            for bond in residues[i].external_bonds(): # Search for normal atoms connected to the nonstandard atoms\n",
        "                if bond[0].residue is residues[i]:    # e.g. metal coordinating nitrogen in Histidine\n",
        "                    ext_atom = bond[1]\n",
        "                    int_atom = bond[0]\n",
        "                elif bond[1].residue is residues[i]:\n",
        "                    ext_atom = bond[0]\n",
        "                    int_atom = bond[1]\n",
        "                else:\n",
        "                    continue\n",
        "                if int_atom.name not in ambertype.keys(): # Remove bond from the topology\n",
        "                    if ext_atom.residue.name in [value['name'] for value in residue_infos.values()]:\n",
        "                        print(f'WARNING: {int_atom.residue}의 {int_atom.name} 원자와 {ext_atom.residue}의 {ext_atom.name}원자 사이 결합이 존재합니다.')\n",
        "                        print(f'{int_atom.residue}도 완충 영역에 포함시키세요.')\n",
        "                        model.delete(bond)\n",
        "            continue # some normal residues (especially terminal ones) show up in this list, so skip them\n",
        "\n",
        "        lig_flag = residue_infos[temp.name]['lig']\n",
        "\n",
        "        if residue_infos[temp.name]['ff'] != 'GFN-FF':\n",
        "            charges = assign_partial_charges(temp.name, residue_infos)\n",
        "\n",
        "        for atom in temp.atoms:\n",
        "            if residue_infos[temp.name]['ff'] == 'GFN-FF':\n",
        "                atom.parameters = dict(charge=0.0)\n",
        "            else:\n",
        "                atom.parameters = dict(charge=charges[atom.name])\n",
        "\n",
        "            if not lig_flag and fftype == 'Amber14' and atom.name in ambertype.keys():\n",
        "                name = atom.name\n",
        "                atom.type = ambertype[atom.name] # For backbone atoms, follow AMBER scheme\n",
        "\n",
        "            elif not lig_flag and fftype == 'CHARMM36' and atom.name in charmmtype.keys():\n",
        "                name = atom.name\n",
        "                atom.type = charmmtype[atom.name] # For backbone atoms, follow CHARMM scheme\n",
        "\n",
        "            else:\n",
        "                name = temp.name+'_'+atom.name # For other atoms, register new atom type\n",
        "                atom.type = name\n",
        "                mass = atom.element.mass.value_in_unit(amu)\n",
        "                element = atom.element.symbol\n",
        "                ff.registerAtomType({'name':name, 'class':name, 'mass':mass, 'element':element})\n",
        "                new_type.append(name)\n",
        "                #lj_parms[name] = assign_LJ_parameters(fftype, ff, model, name, element)\n",
        "\n",
        "        try:\n",
        "            ff.registerResidueTemplate(temp)\n",
        "        except:\n",
        "            # search for duplicate residue name, and if present, rename it\n",
        "            print(f'{temp.name}란 이름의 residue가 이미 {fftype} 역장 안에 있네요.')\n",
        "            newName = 'UK'+str(unk_index)\n",
        "            print(f'Residue 이름을 {temp.name}에서 {newName}로 변경합니다.')\n",
        "            residues[i].name = newName\n",
        "            ncAA[ncAA.index(temp.name)] = newName\n",
        "            temp.name = newName\n",
        "            ff.registerResidueTemplate(temp)\n",
        "            unk_index += 1\n",
        "\n",
        "        for res in model.topology.residues():\n",
        "            if res.name == temp.name:\n",
        "                uaa_res.append(res)\n",
        "\n",
        "\n",
        "    with open('userff.xml', 'w') as f: # Add custom xml file to set LJ parameters to zeros\n",
        "        f.write(\"<ForceField>\\n\")\n",
        "        if fftype == \"Amber14\":\n",
        "            f.write(' <NonbondedForce coulomb14scale=\"0.8333333333333334\" lj14scale=\"0.5\">\\n')\n",
        "            f.write('  <UseAttributeFromResidue name=\"charge\"/>\\n')\n",
        "            for at in new_type:\n",
        "                f.write(f'  <Atom class=\"{at}\" sigma=\"1.0\" epsilon=\"0.0\"/>\\n')\n",
        "        elif fftype == \"CHARMM36\":\n",
        "            f.write(' <NonbondedForce coulomb14scale=\"1.0\" lj14scale=\"1.0\" useDispersionCorrection=\"False\">\\n')\n",
        "            f.write('  <UseAttributeFromResidue name=\"charge\"/>\\n')\n",
        "            for at in new_type:\n",
        "                f.write(f'  <Atom class=\"{at}\" epsilon=\"0.0\" sigma=\"1.0\"/>\\n')\n",
        "        f.write(' </NonbondedForce>\\n')\n",
        "        if fftype == 'CHARMM36':\n",
        "            f.write(' <LennardJonesForce lj14scale=\"1.0\" useDispersionCorrection=\"False\">\\n')\n",
        "            for at in new_type:\n",
        "                f.write(f'  <Atom epsilon=\"0.0\" sigma=\"1.0\" type=\"{at}\"/>\\n')\n",
        "                #f.write(f'  <Atom epsilon=\"0.0\" epsilon14=\"0.0\" sigma=\"1.0\" sigma14=\"1.0\" type=\"{at}\"/>\\n')\n",
        "            f.write(' </LennardJonesForce>\\n')\n",
        "        f.write(\"</ForceField>\\n\")\n",
        "\n",
        "    ff.loadFile('userff.xml')\n",
        "    return uaa_res\n",
        "\n",
        "\n",
        "def update_qm_region(model, atomlist)->int:\n",
        "    \"\"\"\n",
        "    Similar to ONIOM, the MM/XTB or MM/ML boundary must be at isotropic bonds.\n",
        "    If not placed at carbon-carbon bond, expand the XTB/ML region.\n",
        "    \"\"\"\n",
        "    atoms_to_add = []\n",
        "    counter = 0\n",
        "    for atom1 in atomlist:\n",
        "        for bond in model.topology.bonds():\n",
        "            if bond[0] == atom1:\n",
        "                atom2 = bond[1]\n",
        "            elif bond[1] == atom1:\n",
        "                atom2 = bond[0]\n",
        "            else:\n",
        "                continue\n",
        "            if atom2 in atomlist:\n",
        "                continue\n",
        "            elif atom1.element.symbol != 'C' or atom2.element.symbol != 'C':\n",
        "                atoms_to_add.append(atom2)\n",
        "                counter += 1\n",
        "            elif atom1.residue.name in ['HIS', 'HIP', 'HID', 'HIE', 'PHE', 'TYR', 'TRP']:\n",
        "                if atom1.name not in  ['C', 'O', 'N', 'H', 'CA', 'HA', 'CB', 'HB1', 'HB2', 'HB3']:\n",
        "                    if atom2.name != 'CB':\n",
        "                        atoms_to_add.append(atom2)\n",
        "                        counter += 1\n",
        "    atomlist += atoms_to_add\n",
        "    return counter\n",
        "\n",
        "\n",
        "def find_neighbors(model, res_atoms, cutoff)->list:\n",
        "    \"\"\"\n",
        "    Find atoms within a cutoff distance from the given residue.\n",
        "    \"\"\"\n",
        "    contact_atoms = []\n",
        "    for i in res_atoms:\n",
        "        v1 = model.positions[i]\n",
        "        for j, v2 in enumerate(model.positions):\n",
        "            if j in res_atoms:\n",
        "                continue\n",
        "            if np.linalg.norm(np.array([v1.x-v2.x, v1.y-v2.y, v1.z-v2.z])) < cutoff:\n",
        "                contact_atoms.append(j)\n",
        "\n",
        "    qm_atoms = []\n",
        "    for atom in model.topology.atoms():\n",
        "        if atom.index in set(res_atoms+contact_atoms):\n",
        "            qm_atoms.append(atom)\n",
        "\n",
        "    counter = -1\n",
        "    while counter != 0:\n",
        "        counter = update_qm_region(model, qm_atoms)\n",
        "\n",
        "    return [a.index for a in qm_atoms]\n",
        "\n",
        "\n",
        "if len(uaas) == 0:\n",
        "    print('No need to run this cell!')\n",
        "else:\n",
        "    with open(work+'non_standard_residues.txt', 'w') as f:\n",
        "        f.write('# User-defined Non-standard Residues\\n')\n",
        "        f.write('#name\\t#type\\t#forcefield\\n')\n",
        "        for uaa in uaa_info.values():\n",
        "            f.write(f\"{uaa['name']}\\t{'ligand' if uaa['lig'] else 'residue'}\\t{uaa['ff']}\\n\")\n",
        "\n",
        "    if make_xml_for_custom_residues(forcefield, ff, model, uaa_info):\n",
        "        print('역장 정의를 위한 XML 파일을 업로드 해주세요.')\n",
        "        custom_xml = files.upload()\n",
        "        ff.loadFile(next(iter(custom_xml)))\n",
        "\n",
        "    uaa_res = assign_parameters(forcefield, ff, model, uaa_info)\n",
        "\n",
        "    if len(uaa_res) != 0:\n",
        "        uaa_res_names = [f'{res.name}_{res.index}{res.chain}' for res in uaa_res]\n",
        "        # Interactively determine the group of atoms to be treated differently\n",
        "        res_select = widgets.Dropdown(options=uaa_res_names, value=uaa_res_names[0], description='residue', continuous_update=True)\n",
        "        uaa_atom_info = []\n",
        "        for res in uaa_res:\n",
        "            atoms = []\n",
        "            for a in res.atoms():\n",
        "                atoms.append(a.index)\n",
        "            uaa_atom_info.append({'res':res, 'atoms':atoms, 'cutoff':0.0, 'neighbor':find_neighbors(model, atoms, 0.5), 'charge':0, 'spin':1})\n",
        "\n",
        "        cutoff_select = widgets.FloatSlider(value=0.0, min=0.0, max=1.5, step=0.1, description='cutoff(nm)', continuous_update=True)\n",
        "        charge_input = widgets.IntText(value=0, description='Total charge:', disabled=False, continuous_update=True)\n",
        "        spin_input = widgets.IntText(value=1, description='Spin :', disabled=False, continuous_update=True)\n",
        "\n",
        "        view4 = nv.NGLWidget()\n",
        "        view4._set_size('750px','500px')\n",
        "        view4.add_component(nv.FileStructure(work+\"starting_structure.pdb\"), defaultRepresentation=False)\n",
        "        view4.add_cartoon(\"protein or nucleic\")\n",
        "        view4.add_line()\n",
        "        uaa_atom_info[0]['cutoff'] = cutoff_select.value\n",
        "        neighbor = find_neighbors(model, uaa_atom_info[0]['atoms'], cutoff_select.value)\n",
        "        uaa_atom_info[0]['neighbor'] = neighbor\n",
        "        uaa_atom_info[0]['charge'] = charge_input.value\n",
        "        uaa_atom_info[0]['spin'] = spin_input.value\n",
        "        neighborhood = '@'+','.join([str(i) for i in neighbor])\n",
        "        view4.add_ball_and_stick(neighborhood)\n",
        "        view4.center(neighborhood)\n",
        "\n",
        "        def update_res(change):\n",
        "            res_index = uaa_res_names.index(change.new)\n",
        "            view4.clear_representations()\n",
        "            view4.add_cartoon(\"protein or nucleic\")\n",
        "            view4.add_line()\n",
        "            uaa_atom_info[res_index]['charge'] = charge_input.value\n",
        "            uaa_atom_info[res_index]['spin'] = spin_input.value\n",
        "            neighbor = find_neighbors(model, uaa_atom_info[res_index]['atoms'], cutoff_select.value)\n",
        "            neighborhood = '@'+','.join([str(i) for i in neighbor])\n",
        "            view4.add_ball_and_stick(neighborhood)\n",
        "            view4.center(neighborhood)\n",
        "\n",
        "        def update_cutoff(change):\n",
        "            res_index = uaa_res_names.index(res_select.value)\n",
        "            view4.clear_representations()\n",
        "            view4.add_cartoon(\"protein or nucleic\")\n",
        "            view4.add_line(\"not water\")\n",
        "            uaa_atom_info[res_index]['cutoff'] = change.new\n",
        "            neighbor = find_neighbors(model, uaa_atom_info[res_index]['atoms'], change.new)\n",
        "            neighborhood = '@'+','.join([str(i) for i in neighbor])\n",
        "            view4.add_ball_and_stick(neighborhood)\n",
        "            view4.center(neighborhood)\n",
        "            uaa_atom_info[res_index]['neighbor'] = neighbor\n",
        "            uaa_atom_info[res_index]['charge'] = charge_input.value\n",
        "            uaa_atom_info[res_index]['spin'] = spin_input.value\n",
        "\n",
        "        def update_charge(change):\n",
        "            res_index = uaa_res_names.index(res_select.value)\n",
        "            uaa_atom_info[res_index]['charge'] = change.new\n",
        "\n",
        "        def update_spin(change):\n",
        "            res_index = uaa_res_names.index(res_select.value)\n",
        "            uaa_atom_info[res_index]['spin'] = change.new\n",
        "\n",
        "        res_select.observe(update_res, names='value')\n",
        "        cutoff_select.observe(update_cutoff, names='value')\n",
        "        charge_input.observe(update_charge, names='value')\n",
        "        spin_input.observe(update_spin, names='value')\n",
        "\n",
        "        display(res_select, cutoff_select, view4, charge_input, spin_input)\n",
        "    else:\n",
        "        for uaa in uaa_info.values():\n",
        "            assert uaa['lig'], '이 셀을 두번 이상 실행하셨나요? 그렇다면 2-1부터 다시 실행하세요.'\n",
        "        print('비리간드 residue가 발견되지 않았습니다.')"
      ],
      "metadata": {
        "collapsed": true,
        "id": "GaqRQQMN629M",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 2-3. 시뮬레이션 준비"
      ],
      "metadata": {
        "id": "A5kqyHcPryj3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "assert 'work' in globals(), \"0번 셀부터 실행해주세요.\"\n",
        "assert 'uaas' in globals(), \"2-1부터 실행해주세요.\"\n",
        "\n",
        "#@markdown 용액의 이온세기 (mol/L)를 지정해주세요.\n",
        "ion_conc = 0.15 #@param {type:\"slider\", min:0.0, max:2.0, step:0.05}\n",
        "#@markdown 명시적 용매모델의 경우, 아래 변수들도 지정해주세요.\n",
        "\n",
        "#@markdown 물층을 몇 nm 두께로 덮을건지, 이온 종류를 뭐로 할지 설정하세요.\n",
        "padding = 1.2 #@param {type:\"slider\", min:1.0, max:5.0, step:0.1}\n",
        "cation = \"Na+\" #@param [\"Li+\", \"Na+\", \"K+\", \"Rb+\", \"Cs+\"]\n",
        "anion = \"Cl-\" #@param [\"F-\", \"Cl-\", \"Br-\", \"I-\"]\n",
        "#@markdown 계산 속도를 높이려면 다음 constraint들을 설정하세요.\n",
        "rigidWater = True #@param {type:\"boolean\"}\n",
        "constraints = \"HBonds\" #@param [\"None\", \"HBonds\", \"AllBonds\", \"HAngles\"]\n",
        "hydrogenMass = 2 #@param {type:\"slider\", min:1.0, max:4.0, step:0.1}\n",
        "#@markdown 계산 정밀도를 설정하세요: (빠름) single -> mixed -> double (정밀)\n",
        "precision = \"single\" #@param [\"single\", \"mixed\", \"double\"]\n",
        "#@markdown step의 크기를 fs 단위로 설정하세요.\n",
        "step_size = 2.5 #@param {type:\"slider\", min:0.1, max:4, step:0.1}\n",
        "#@markdown **섭씨 온도**를 설정하세요.\n",
        "temperature = 25 #@param {type:\"slider\", min:0.0, max:100.0, step:1.0}\n",
        "\n",
        "assert forcefield != \"CHARMM36\" or padding >= 1.2, \"물층의 두께는 1.2 nm보다 두꺼워야 합니다!\"\n",
        "assert all(p['ff'] != 'MACE-off23' for p in uaa_info.values()) or water != '', 'ML 역장은 묵시적 용매 환경을 지원하지 않습니다.'\n",
        "assert hydrogenMass > 1.0 or rigidWater or constraints != 'None' or step_size <= 0.5, f\"경고 : 아무 제약 조건 없는 상태에서 {step_size} fs의 step 크기는 너무 큽니다.\"\n",
        "assert hydrogenMass > 1.0 or constaints != 'HBonds' or step_size <= 2.0, f\"경고: {step_size} fs의 step 크기는 너무 큽니다. 시뮬레이션이 터질 수 있습니다.\"\n",
        "\n",
        "\n",
        "def get_openff_ligand(work, ligand_info):\n",
        "    \"\"\"\n",
        "    Register ligand template to the macromolecular forcefield.\n",
        "    \"\"\"\n",
        "    name = ligand_info[\"name\"]\n",
        "\n",
        "    if os.path.exists(work+name+'.sdf'):\n",
        "        sdf_file = work + name+'.sdf'\n",
        "    elif os.path.exists(name+'_ideal.sdf'):\n",
        "        sdf_file = name+'_ideal.sdf'\n",
        "    else:\n",
        "        subprocess.run(f'wget https://files.rcsb.org/ligands/download/{name}_ideal.sdf', shell=True)\n",
        "        sdf_file = f'{name}_ideal.sdf'\n",
        "\n",
        "    SDF = Chem.rdmolfiles.SDMolSupplier(sdf_file)[0]\n",
        "    PDB = AllChem.MolFromPDBFile(f'{name}.pdb')\n",
        "    ligand_noH = AllChem.AssignBondOrdersFromTemplate(SDF, PDB)\n",
        "    ligand = AllChem.MolFromPDBFile(f'{name}.pdb', removeHs=False)\n",
        "\n",
        "    for bond in ligand_noH.GetBonds():\n",
        "        pos1 = ligand_noH.GetConformer().GetAtomPosition(bond.GetBeginAtomIdx())\n",
        "        pos2 = ligand_noH.GetConformer().GetAtomPosition(bond.GetEndAtomIdx())\n",
        "        for i in range(ligand.GetNumAtoms()):\n",
        "            pos = ligand.GetConformer().GetAtomPosition(i)\n",
        "            if all(pos[j] == pos1[j] for j in range(3)):\n",
        "                aid1 = i\n",
        "            elif all(pos[j] == pos2[j] for j in range(3)):\n",
        "                aid2 = i\n",
        "            else:\n",
        "                continue\n",
        "        ligand.GetBondBetweenAtoms(aid1, aid2).SetBondType(bond.GetBondType())\n",
        "\n",
        "    Chem.SanitizeMol(ligand)\n",
        "\n",
        "    Chem.rdmolops.AssignStereochemistryFrom3D(ligand)\n",
        "    openff_ligand = Molecule.from_rdkit(ligand, hydrogens_are_explicit=True)\n",
        "    openff_ligand.name = name\n",
        "    openff_ligand.assign_partial_charges('mmff94')\n",
        "\n",
        "    with open(f'{work}{name}.json', 'w') as f:\n",
        "        f.write(openff_ligand.to_json())\n",
        "\n",
        "    return openff_ligand\n",
        "\n",
        "gaff = []\n",
        "smirnoff = []\n",
        "espaloma = []\n",
        "\n",
        "for key, value in uaa_info.items():\n",
        "    if value['ff'] == 'GAFF-2.11':\n",
        "        gaff.append(get_openff_ligand(work, value))\n",
        "    elif value['ff'] == 'Sage-2.2.1':\n",
        "        smirnoff.append(get_openff_ligand(work, value))\n",
        "    elif value['ff'] == 'espaloma-0.3.2':\n",
        "        espaloma.append(get_openff_ligand(work, value))\n",
        "\n",
        "if len(gaff) > 0:\n",
        "    gaff_g = GAFFTemplateGenerator(molecules=gaff)\n",
        "    ff.registerTemplateGenerator(gaff_g.generator)\n",
        "if len(smirnoff) > 0:\n",
        "    smirnoff_g = SMIRNOFFTemplateGenerator(molecules=smirnoff)\n",
        "    ff.registerTemplateGenerator(smirnoff_g.generator)\n",
        "if len(espaloma) > 0:\n",
        "    espaloma_g = EspalomaTemplateGenerator(molecules=espaloma)\n",
        "    ff.registerTemplateGenerator(espaloma_g.generator)\n",
        "\n",
        "if water != '':\n",
        "    print(\"용매를 더하는중...\", end='')\n",
        "    model.addSolvent(ff, model=water, padding=padding*nanometers, positiveIon=cation,\n",
        "                     negativeIon=anion, ionicStrength=ion_conc*molar, neutralize=True)\n",
        "    print(\"완료.\")\n",
        "\n",
        "    print(\"계 생성 중...\", end='')\n",
        "\n",
        "    if forcefield == \"Amber14\":\n",
        "        system = ff.createSystem(model.topology, nonbondedMethod=PME,\n",
        "                                 nonbondedCutoff=1.0*nanometer, constraints=constraints,\n",
        "                                 rigidWater=rigidWater, removeCMMotion=True,\n",
        "                                 hydrogenMass=hydrogenMass*amu)\n",
        "    elif forcefield == \"CHARMM36\":\n",
        "        # CHARMM forcefield recommends the use of switch distance at 1.0 nm and cutoff distance at 1.2 nm.\n",
        "        system = ff.createSystem(model.topology, nonbondedMethod=PME,\n",
        "                                 nonbondedCutoff=1.2*nanometer, constraints=constraints,\n",
        "                                 rigidWater=rigidWater, removeCMMotion=True,\n",
        "                                 switchDistance=1.0*nanometer,\n",
        "                                 hydrogenMass=float(hydrogenMass)*amu)\n",
        "\n",
        "else:\n",
        "    print(\"계 생성 중...\", end='')\n",
        "\n",
        "    kappa = 367.434915*sqrt(float(ion_conc)/78.5/(int(temperature)+273.15))\n",
        "\n",
        "    if forcefield == \"Amber14\":\n",
        "        system = ff.createSystem(model.topology, nonbondedMethod=CutoffNonPeriodic,implicitSolventKappa=kappa,\n",
        "                                 nonbondedCutoff=1.0*nanometer, constraints=constraints, hydrogenMass=float(hydrogenMass)*amu)\n",
        "\n",
        "    elif forcefield == \"CHARMM36\":\n",
        "        system = ff.createSystem(model.topology, nonbondedMethod=CutoffNonPeriodic,implicitSolventKappa=kappa,\n",
        "                                 nonbondedCutoff=1.2*nanometer, switchDistance=1.0*nanometer, constraints=constraints,\n",
        "                                 hydrogenMass=float(hydrogenMass)*amu)\n",
        "\n",
        "print('완료.')\n",
        "\n",
        "\n",
        "def register_ml_residues(system, topology, pbc, residue_infos, atom_infos):\n",
        "    \"\"\"\n",
        "    Create MM-ML or MM-XTB mixed system\n",
        "    \"\"\"\n",
        "\n",
        "    xtb_atoms = []\n",
        "    total_charge = 0\n",
        "    total_e = 0\n",
        "    ml_atoms = []\n",
        "\n",
        "    for key1, value1 in residue_infos.items():\n",
        "        for res in atom_infos:\n",
        "            if res['res'].name == key1:\n",
        "                if value1['ff'] == 'GFN-FF':\n",
        "                    xtb_atoms.extend(res['neighbor'])\n",
        "                    total_charge += res['charge']\n",
        "                    total_e += float(res['spin']-1)/2\n",
        "                elif value1['ff'] == 'MACE-off23':\n",
        "                    ml_atoms.extend(res['neighbor'])\n",
        "\n",
        "    xtb_atoms = set(xtb_atoms)\n",
        "    ml_atoms = set(ml_atoms)\n",
        "\n",
        "    element_list = []\n",
        "\n",
        "    for a in list(xtb_atoms):\n",
        "        for atom in topology.atoms():\n",
        "            if atom.index == a:\n",
        "                element_list.append(atom.element.atomic_number)\n",
        "\n",
        "    if len(xtb_atoms&ml_atoms)>0:\n",
        "        print('하나의 원자가 두가지 종류 이상의 역장으로 계산될 수 없습니다.')\n",
        "        print('완충영역이 겹치지 않게 기준반경을 줄이거나 역장 종류를 통일하세요.')\n",
        "        return None\n",
        "\n",
        "    if len(ml_atoms) > 0:\n",
        "        potential = MLPotential('mace-off23-small')\n",
        "        print(\"MM-ML 혼합 계 생성 중...\", end='')\n",
        "        warnings.filterwarnings('ignore')\n",
        "        system2 = potential.createMixedSystem(topology, system, ml_atoms, removeConstraints=False, precision='single')\n",
        "        print(\"완료.\")\n",
        "    else:\n",
        "        system2 = system\n",
        "\n",
        "    if len(xtb_atoms) > 0:\n",
        "        print(\"MM-XTB 혼합 계 생성 중...\", end='')\n",
        "        system2.addForce(XtbForce(XtbForce.GFNFF, total_charge, int(total_e/2+1),\n",
        "                                  pbc, xtb_atoms, element_list))\n",
        "        print(\"완료.\")\n",
        "\n",
        "    return system2\n",
        "\n",
        "\n",
        "def save_settings(work, forcefield, watermodel, ion_conc, pbc, cation, anion, padding, rigidWater, constraints, hydrogenMass, precision, step_size, temperature):\n",
        "    \"\"\"\n",
        "    Save simulation settings to settings.txt file.\n",
        "    \"\"\"\n",
        "    settings = f\"\"\"\n",
        "    Forcefield : {forcefield}\n",
        "    Watermodel : {watermodel}\n",
        "    Ionic strength (in mol/L) : {ion_conc}\n",
        "    Periodic Boundary Condition : {pbc}\n",
        "    Cation : {cation}\n",
        "    Anion : {anion}\n",
        "    Padding (in nm) : {padding}\n",
        "    Rigidwater : {rigidWater}\n",
        "    Constraints : {constraints}\n",
        "    Hydrogen mass (in amu) : {hydrogenMass}\n",
        "    Precision : {precision}\n",
        "    Time step (in fs) : {step_size}\n",
        "    Temperature (in Celcius) : {temperature}\n",
        "    \"\"\"\n",
        "    with open(work+\"settings.txt\", 'a') as f:\n",
        "        f.write(settings)\n",
        "    return\n",
        "\n",
        "\n",
        "if water == '':\n",
        "    pbc = False\n",
        "else:\n",
        "    pbc = True\n",
        "\n",
        "final_system = register_ml_residues(system, model.topology, pbc, uaa_info, uaa_atom_info)\n",
        "\n",
        "assert final_system is not None, \"계가 생성되지 않았습니다!\"\n",
        "\n",
        "\n",
        "with open(work+'system.xml', 'w') as f:\n",
        "    f.write(XmlSerializer.serialize(final_system))\n",
        "\n",
        "print(\"시뮬레이션 생성 중...\", end='')\n",
        "integrator = LangevinMiddleIntegrator((int(temperature)+273.15)*kelvin, 1/picosecond, float(step_size)*femtosecond)\n",
        "\n",
        "with open(work+'integrator.xml', 'w') as f:\n",
        "    f.write(XmlSerializer.serialize(integrator))\n",
        "\n",
        "platform = Platform.getPlatformByName('CUDA')\n",
        "properties = {'Precision': precision}\n",
        "simulation = Simulation(model.topology, final_system, integrator, platform, properties)\n",
        "print(\"done.\")\n",
        "\n",
        "print(\"에너지 최적화 중...\", end='')\n",
        "simulation.context.setPositions(model.positions)\n",
        "simulation.minimizeEnergy(tolerance=0.001*kilojoules/(nanometer*mole))\n",
        "positions = simulation.context.getState(getPositions=True).getPositions()\n",
        "\n",
        "with open(work+\"minimized.pdb\", 'w') as f:\n",
        "    PDBFile.writeFile(simulation.topology, positions, f)\n",
        "print(\"done.\")\n",
        "\n",
        "simulation.context.setVelocitiesToTemperature((int(temperature)+273.15)*kelvin)\n",
        "simulation.context.setStepCount(0)\n",
        "simulation.context.setTime(0)\n",
        "state = simulation.context.getState(getPositions=True, getVelocities=True)\n",
        "with open(work+\"start.xml\", 'w') as f:\n",
        "    f.write(XmlSerializer.serialize(state))\n",
        "stage = 'start'\n",
        "\n",
        "save_settings(work, forcefield, watermodel, ion_conc, pbc, cation, anion, padding, rigidWater, constraints, hydrogenMass, precision, step_size, temperature)\n",
        "print(\"시뮬레이션 준비가 끝났습니다!\")\n",
        "\n",
        "view3 = nv.NGLWidget()\n",
        "view3._set_size('750px','500px')\n",
        "view3.add_component(nv.FileStructure(work+\"minimized.pdb\"))\n",
        "view3.add_licorice()\n",
        "if water != '':\n",
        "    view3.add_point(\"water\")\n",
        "    view3.add_spacefill(\"ion\")\n",
        "view3.center()\n",
        "view3"
      ],
      "metadata": {
        "id": "SCYS8pxLys66",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 3. 시뮬레이션 실행"
      ],
      "metadata": {
        "id": "7Xf2GSRAqxDg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 3-1. NVT 평형\n",
        "assert 'work' in globals(), \"0번 셀부터 실행해주세요.\"\n",
        "\n",
        "#@markdown N ps만큼 (추가로) 평형을 시행합니다.\n",
        "eq_time = 50 #@param {type:\"slider\", min:0.0, max:1000.0, step:10.0}\n",
        "#@markdown 실제 N 분 마다 구글 드라이브에 시뮬레이션을 백업합니다.\n",
        "#@markdown 무료 세션을 사용 중이시라면 이 값을 작게 설정하세요.\n",
        "backup_interval = 10 #@param {type:\"slider\", min:5, max:60, step:5}\n",
        "#@markdown 수소가 아닌 원자들에 대해 위치 restraint를 걸 경우 조화진동자의 용수철상수 (kJ/mol*nm^2)를 설정하세요.\n",
        "positional_restraint = 0 #@param {type:\"slider\", min:0, max:3000, step:100}\n",
        "\n",
        "simulation, stage = recall_simulation(work)\n",
        "\n",
        "assert stage != 'npt', 'NPT 평형이 이미 시행되었습니다.'\n",
        "assert stage != 'production', 'Production run이 이미 시행되었습니다.'\n",
        "assert eq_time > 0, '평형 시간을 0으로 설정했습니다.'\n",
        "\n",
        "if stage == 'start':\n",
        "    simulation = apply_restraints(work, simulation, int(positional_restraint))\n",
        "    print('NVT 평형을 시작합니다!')\n",
        "    print(f'{eq_time} ps 만큼 NVT 평형을 시행합니다.')\n",
        "    stage = 'nvt'\n",
        "    state = simulation.context.getState(getPositions=True, getVelocities=True)\n",
        "    with open(work+'nvt_step_0.xml', 'w') as f:\n",
        "        f.write(XmlSerializer.serialize(state))\n",
        "    assert os.path.exists(work+'nvt_step_0.xml'), 'State data not recorded.'\n",
        "\n",
        "else:\n",
        "    print(f'{simulation.context.getTime()} ps 만큼 NVT 평형이 이뤄졌습니다.')\n",
        "    print(f'{eq_time} ps 만큼 NVT 평형을 더 시행합니다.')\n",
        "\n",
        "subprocess.run(f'echo \"\\tRestraint (NVT, kJ/mol*nm^2) : {positional_restraint}\" >> {work}settings.txt', shell=True)\n",
        "stream = SimulationStream()\n",
        "f = open(work+\"nvt_log.txt\", 'w')\n",
        "f.close()\n",
        "run_simulation(work, stage, simulation, stream, float(eq_time), 5*picoseconds, int(backup_interval))"
      ],
      "metadata": {
        "id": "y_rBhzYx6K3m",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "assert os.path.exists(work+'nvt_log.txt'), \"3-1을 먼저 실행하세요.\"\n",
        "\n",
        "#@markdown 온도 수렴 여부 확인\n",
        "time = []\n",
        "temp = []\n",
        "s = read_settings(work)\n",
        "\n",
        "with open(work+\"nvt_log.txt\", 'r') as f:\n",
        "    for line in f.readlines():\n",
        "        if line.startswith(\"#\"):\n",
        "            continue\n",
        "        line = line.strip().split(',')\n",
        "        step = float(line[0])\n",
        "        time.append(step*s['stepsize']/picoseconds)\n",
        "        temp.append(round(float(line[3]),2))\n",
        "\n",
        "target_temp = round(s['temperature']/kelvin, 2)\n",
        "\n",
        "fig1 = go.Figure()\n",
        "fig1.add_trace(go.Scatter(x=time, y=temp, mode='lines', name='Temperature'))\n",
        "fig1.add_hline(y=target_temp, annotation_text=\"target T : \"+str(target_temp)+\" K\")\n",
        "fig1.update_layout(title='NVT equilibration', xaxis_title='Time (ps)', yaxis_title='Temperature (K)', width=750, height=500)\n",
        "fig1.show()\n",
        "\n",
        "print(\"계의 온도가 지정된 값으로 수렴하는지 확인하세요.\")\n",
        "print(\"그렇지 않다면 평형단계를 더 오래 수행하세요.\")\n",
        "print(\"수렴했다면 다음 단계로 넘어가도 좋습니다.\")"
      ],
      "metadata": {
        "id": "Ut9NEwrxnPA9",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 3-2. NPT 평형\n",
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "\n",
        "#@markdown N ps만큼 평형을 (더) 실행합니다.\n",
        "eq_time = 50 #@param {type:\"slider\", min:0.0, max:1000.0, step:10.0}\n",
        "#@markdown bar 단위의 압력을 설정하세요.\n",
        "pressure = 1 #@param {type:\"slider\", min:0.0, max:100.0, step:1.0}\n",
        "#@markdown N 분 (실제시간)마다 구글 드라이브에 데이터를 백업합니다.\n",
        "#@markdown 무료 세션을 실행중인 경우 이 값을 작게 해주세요.\n",
        "backup_interval = 10 #@param {type:\"slider\", min:5, max:60, step:5}\n",
        "#@markdown 수소가 아닌 원자에 대한 위치 restraint 용수철상수 (kJ/mol*nm^2)\n",
        "positional_restraint = 0 #@param {type:\"slider\", min:0, max:3000, step:100}\n",
        "\n",
        "\n",
        "s = read_settings(work)\n",
        "\n",
        "simulation, stage = recall_simulation(work)\n",
        "\n",
        "if not s['PBC']:\n",
        "    print('경계조건이 없습니다. 압력의 정의가 불가능합니다.')\n",
        "\n",
        "elif stage == 'start':\n",
        "    print('NVT 평형부터 시행하세요.')\n",
        "\n",
        "elif stage == 'production':\n",
        "    print('production run이 시행되었습니다. 평형 단계는 더이상 수행할 수 없습니다')\n",
        "\n",
        "elif eq_time == 0:\n",
        "    print('평형 시간을 0으로 설정했습니다.')\n",
        "\n",
        "else:\n",
        "    if stage == 'nvt':\n",
        "        barostat = MonteCarloBarostat(int(pressure)*bar, s['temperature'])\n",
        "        simulation.system.addForce(barostat)\n",
        "        simulation.context.reinitialize(preserveState=True)\n",
        "        simulation.context.setTime(0)\n",
        "        simulation.context.setStepCount(0)\n",
        "        state = simulation.context.getState(getPositions=True, getVelocities=True)\n",
        "        with open(work+\"npt_step_0.xml\", 'w') as f:\n",
        "            f.write(XmlSerializer.serialize(state))\n",
        "        simulation.saveCheckpoint(work+\"npt.chk\")\n",
        "        print('NPT 평형을 시작합니다!')\n",
        "        print(f'{eq_time} ps 만큼 NPT 평형을 시행합니다.')\n",
        "        stage = 'npt'\n",
        "    elif stage == 'npt':\n",
        "        print(f'{simulation.context.getTime()} ps만큼 NPT 평형이 실행되었습니다.')\n",
        "        print(f'{eq_time} ps 더 NPT 평형을 수행합니다.')\n",
        "\n",
        "    subprocess.run(f'echo \"\\tPressure (in bar) : {pressure}\" >> {work}settings.txt', shell=True)\n",
        "    subprocess.run(f'echo \"\\tRestraint (NPT, kJ/mol*nm^2) : {positional_restraint}\" >> {work}settings.txt', shell=True)\n",
        "\n",
        "    stream = SimulationStream()\n",
        "    f = open(work+\"npt_log.txt\", 'w')\n",
        "    f.close()\n",
        "    run_simulation(work, stage, simulation, stream, float(eq_time), 5*picoseconds, int(backup_interval))"
      ],
      "metadata": {
        "cellView": "form",
        "id": "WZzkMi5eMgqh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "assert 'work' in globals(), \"0번 셀부터 실행해주세요.\"\n",
        "assert os.path.exists(work+'nvt_log.txt'), \"3-2부터 실행해주세요.\"\n",
        "\n",
        "#@markdown 온도 & 밀도 수렴 여부 확인\n",
        "time = []\n",
        "temp = []\n",
        "dens = []\n",
        "\n",
        "with open(work+\"npt_log.txt\", 'r') as f:\n",
        "    for line in f.readlines():\n",
        "        if line.startswith(\"#\"):\n",
        "            continue\n",
        "        line = line.strip().split(',')\n",
        "        step = float(line[0])\n",
        "        time.append(step*s['stepsize']/picoseconds)\n",
        "        temp.append(round(float(line[3]),2))\n",
        "        dens.append(round(float(line[5]),5))\n",
        "\n",
        "target_temp = round(s['temperature']/kelvin,2)\n",
        "\n",
        "fig2 = make_subplots(specs=[[{\"secondary_y\":True}]])\n",
        "fig2.add_trace(go.Scatter(x=time, y=temp, mode='lines', name='Temperature'), secondary_y=False)\n",
        "fig2.add_hline(y=target_temp, annotation_text=\"target T : \"+str(target_temp)+\" K\")\n",
        "fig2.add_trace(go.Scatter(x=time, y=dens, mode='lines', name='Density'), secondary_y=True)\n",
        "fig2.update_layout(title='NPT equilibration', width=750, height=500)\n",
        "fig2.update_xaxes(title_text=\"Time (ps)\")\n",
        "fig2.update_yaxes(title_text=\"Temperature (K)\", secondary_y=False)\n",
        "fig2.update_yaxes(title_text=\"Density (g/mL)\", secondary_y=True)\n",
        "fig2.show()\n",
        "\n",
        "print(\"계의 밀도와 온도가 일정한 값으로 수렴하는지 확인해주세요.\")\n",
        "print(\"만일 아니라면 3-2를 더 실행하세요.\")\n",
        "print(\"수렴했다면 production run으로 진행해도 좋습니다.\")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "2SEEPsZyqQph"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 3-3. Production Run\n",
        "assert 'work' in globals(), \"0번 셀부터 실행해주세요.\"\n",
        "\n",
        "#@markdown N ns 동안 production run을 진행합니다.\n",
        "run_time = 5 #@param {type:\"slider\", min:1, max:500.0, step:1}\n",
        "#@markdown 매 N ps 마다 trajectory를 저장합니다.\n",
        "save_interval = 5 #@param {type:\"slider\", min:1, max:50, step:1}\n",
        "#@markdown N 분 (실제시간)마다 구글 드라이브에 데이터를 백업합니다.\n",
        "#@markdown 무료 세션을 사용중이시라면 이 값을 작게 설정하세요.\n",
        "backup_interval = 10 #@param {type:\"slider\", min:5, max:60, step:5}\n",
        "\n",
        "s = read_settings(work)\n",
        "\n",
        "simulation, stage = recall_simulation(work)\n",
        "\n",
        "if stage != 'production':\n",
        "    simulation.context.setTime(0)\n",
        "    simulation.context.setStepCount(0)\n",
        "    state = simulation.context.getState(getPositions=True, getVelocities=True)\n",
        "    with open(work+\"production_step_0.xml\", 'w') as f:\n",
        "        f.write(XmlSerializer.serialize(state))\n",
        "    simulation = restore_simulation_from_state(work, work+'production_step_0.xml', 0, 'production')\n",
        "    if s['PBC']:\n",
        "        barostat = MonteCarloBarostat(s['pressure'], s['temperature'])\n",
        "        simulation.system.addForce(barostat)\n",
        "        simulation.context.reinitialize(preserveState=True)\n",
        "\n",
        "    simulation.saveCheckpoint(work+\"production.chk\")\n",
        "    print('MD production run을 시작합니다!')\n",
        "    print(f'{run_time} ns 만큼 계산을 수행합니다.')\n",
        "    stage = 'production'\n",
        "else:\n",
        "    print(f'MD 시뮬레이션이 {round(simulation.context.getTime().value_in_unit(nanosecond), 1)} ns 만큼 이루어졌습니다.')\n",
        "    print(f'{run_time} ns 더 계산을 수행합니다.')\n",
        "\n",
        "subprocess.run(f'echo \"\\tSave interval (in ps) : {save_interval}\" >> {work}settings.txt', shell=True)\n",
        "stream = SimulationStream()\n",
        "f = open(work+\"production_log.txt\", 'w')\n",
        "f.close()\n",
        "run_simulation(work, stage, simulation, stream, int(run_time)*1000, int(save_interval)*picoseconds, int(backup_interval))"
      ],
      "metadata": {
        "cellView": "form",
        "id": "jT4Ii0UGMrjI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "assert 'work' in globals(), \"0번 셀을 먼저 실행해주세요.\"\n",
        "\n",
        "#@markdown 시뮬레이션 중간 점검 : snapshot 확인\n",
        "view_water_box = False #@param {type:\"boolean\"}\n",
        "s = read_settings(work)\n",
        "\n",
        "view4 = nv.NGLWidget()\n",
        "view4._set_size('750px','500px')\n",
        "view4.add_component(nv.FileStructure(work+\"minimized.pdb\"), defaultRepresentation=False)\n",
        "view4.add_cartoon(\"protein or nucleic\")\n",
        "view4.add_simplified_base(\"nucleic\")\n",
        "view4.add_licorice(\"protein or nucleic or ligand\")\n",
        "if s['PBC'] and view_water_box:\n",
        "    view4.add_point(\"water\")\n",
        "    view4.add_spacefill(\"ion\")\n",
        "state = simulation.context.getState(getPositions=True, enforcePeriodicBox=s['PBC'])\n",
        "view4.set_coordinates({0:state.getPositions(asNumpy=True)*10})\n",
        "view4.update_representation(component=0)\n",
        "view4.center()\n",
        "display(view4)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "5CtPtIP2v2K1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 4. 결과 분석"
      ],
      "metadata": {
        "id": "ktFrmmAnqyp1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 4-1. Trajectory 편집 및 시각화\n",
        "assert 'work' in globals(), \"0번 셀부터 실행해주세요.\"\n",
        "\n",
        "#@markdown 이 셀은 4-1~3에서 생성된 trajectory 파일들을 이어붙이고 편집해줍니다. <br/>\n",
        "#@markdown 더 부드러운 trajectory 영상을 원하시면 아래 체크 박스를 선택하세요. <br/>\n",
        "#@markdown _final.xtc로 끝나는 파일을 구글 드라이브에서 다운로드 하세요.\n",
        "\n",
        "smooth_trajectory = False #@param {type:\"boolean\"}\n",
        "\n",
        "def wrap_trajectories(work, stage):\n",
        "    \"\"\"\n",
        "    1. Merge the trajectory XTC files into one\n",
        "    2. Wrap around PBC box if necessary\n",
        "    3. Align the trajectory to the first frame\n",
        "    4. Save the processed trajectory to the XTC and PDB files.\n",
        "    \"\"\"\n",
        "    s = read_settings(work)\n",
        "    print(f'{stage} trajectory 파일들을 읽는 중...', end='')\n",
        "    trajs = []\n",
        "    for t in sorted(glob.glob(work+stage+'_*.xtc')):\n",
        "        if t != work+stage+'_final.xtc':\n",
        "            trajs.append(md.load(t, top=work+'minimized.pdb'))\n",
        "    trj = md.join(trajs)\n",
        "\n",
        "    if s['PBC']:\n",
        "        trj.image_molecules(inplace=True, anchor_molecules=trj.topology.guess_anchor_molecules())\n",
        "\n",
        "    trj.save_hdf5(work+stage+'.h5')\n",
        "\n",
        "    alphas = trj.topology.select_atom_indices('alpha')\n",
        "    if len(alphas) > 0:\n",
        "        traj2 = trj.superpose(trj, 0, atom_indices=trj.topology.select_atom_indices('alpha'))\n",
        "    else:\n",
        "        nucleic = 'resname A or resname T or resname G or resname C or resname U or resname DA or resname DT or resname DG or resname DC'\n",
        "        traj2 = trj.superpose(trj, 0, atom_indices=trj.topology.select(f'({nucleic}) and (symbol == \"P\" or symbol == \"O\" or symbol == \"N\")'))\n",
        "\n",
        "    if smooth_trajectory:\n",
        "        filter_width = max(12*picoseconds/s['save interval'], 1)\n",
        "        if filter_width > 3:\n",
        "            traj2.smooth(int(filter_width), order=3, inplace=True)\n",
        "\n",
        "    traj2.save_xtc(work+stage+'_final.xtc')\n",
        "\n",
        "    print('완료.')\n",
        "\n",
        "    return\n",
        "\n",
        "\n",
        "to_view = {}\n",
        "if os.path.exists(work+'nvt_0.xtc'):\n",
        "    wrap_trajectories(work, 'nvt')\n",
        "if os.path.exists(work+'npt_0.xtc'):\n",
        "    wrap_trajectories(work, 'npt')\n",
        "if os.path.exists(work+'production_0.xtc'):\n",
        "    wrap_trajectories(work, 'production')\n",
        "\n",
        "if os.path.exists(work+'nvt_final.xtc'):\n",
        "    to_view['nvt'] = md.load(work+'nvt_final.xtc', top=work+'minimized.pdb')\n",
        "if os.path.exists(work+'npt_final.xtc'):\n",
        "    to_view['npt'] = md.load(work+'npt_final.xtc', top=work+'minimized.pdb')\n",
        "if os.path.exists(work+'production_final.xtc'):\n",
        "    to_view['production'] = md.load(work+'production_final.xtc', top=work+'minimized.pdb')\n",
        "\n",
        "assert len(to_view) > 0, \"Trajectory XTC 파일이 없습니다!\"\n",
        "\n",
        "choice = widgets.Dropdown(options=to_view.keys(), value='production',\n",
        "                          description='trajectory to view',\n",
        "                          style={'description_width': 'initial'},\n",
        "                          layout=widgets.Layout(width='750px'),\n",
        "                          continuous_update=True)\n",
        "\n",
        "def update_choice(change):\n",
        "    with view_traj:\n",
        "        clear_output(wait=True)\n",
        "        view = nv.show_mdtraj(to_view[change.new], defaultRepresentation=False)\n",
        "        view._set_size('750px','500px')\n",
        "        view.add_cartoon()\n",
        "        view.add_simplified_base(\"nucleic\")\n",
        "        view.add_ball_and_stick(\"ligand\")\n",
        "        view.center()\n",
        "        display(view)\n",
        "\n",
        "choice.observe(update_choice, names='value')\n",
        "\n",
        "view_traj = widgets.Output()\n",
        "\n",
        "with view_traj:\n",
        "    view = nv.show_mdtraj(to_view['production'], defaultRepresentation=False)\n",
        "    view._set_size('750px','500px')\n",
        "    view.add_cartoon()\n",
        "    view.add_simplified_base(\"nucleic\")\n",
        "    view.add_ball_and_stick(\"ligand\")\n",
        "    view.center()\n",
        "    display(view)\n",
        "\n",
        "print('NGLViewer로 볼 trajectory를 선택하세요.')\n",
        "print('로딩이 좀 오래 걸리니 기다려주세요.')\n",
        "display(choice)\n",
        "display(view_traj)"
      ],
      "metadata": {
        "id": "UPwZPr-83OnW",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 4-2. 프레임별 clustering\n",
        "assert 'work' in globals(), \"0번 셀부터 실행해주세요.\"\n",
        "assert os.path.exists(work+'production.h5'), \"4-1을 먼저 실행해주세요.\"\n",
        "\n",
        "if 'traj' not in globals():\n",
        "    traj = md.load(work+'production.h5')\n",
        "\n",
        "#@markdown RMSD 기준 값 (nm)을 정하면 이보다 작은 RMSD를 갖는 구조들이 한 덩어리로 묶입니다.\n",
        "rmsd_threshold = 0.15 #@param {type:\"slider\", min:0.05, max:1.0, step:0.05}\n",
        "\n",
        "#@markdown 각 뭉치들의 대표 구조 (중심 구조)들이 \"clustered.pdb\"의 다른 프레임에 저장됩니다.\n",
        "rmsd_matrix = np.zeros((traj.n_frames, traj.n_frames))\n",
        "\n",
        "for i in tqdm(range(traj.n_frames)):\n",
        "    rmsd_matrix[i] = md.rmsd(traj, traj, frame=i, atom_indices=traj.topology.select_atom_indices('heavy'))\n",
        "\n",
        "clustering = AgglomerativeClustering(n_clusters=None, distance_threshold=rmsd_threshold, metric='precomputed', linkage='average')\n",
        "cluster_labels = clustering.fit_predict(rmsd_matrix)\n",
        "\n",
        "n_clusters = len(set(cluster_labels))\n",
        "cluster_sizes = Counter(cluster_labels)\n",
        "sorted_clusters = sorted(cluster_sizes.items(), key=lambda x:x[1], reverse=True)\n",
        "\n",
        "print(f\"Total {n_clusters} clusters found!\")\n",
        "\n",
        "centroids = []\n",
        "for cluster, size in sorted_clusters:\n",
        "    cluster_frames = np.where(cluster_labels == cluster)[0]\n",
        "    centroid_idx = cluster_frames[np.argmin([np.sum(rmsd_matrix[frame, cluster_frames]) for frame in cluster_frames])]\n",
        "    print(f\"Cluser {cluster} (size = {size})'s centroid structure : frame {centroid_idx}\")\n",
        "    centroids.append((cluster, size, centroid_idx))\n",
        "\n",
        "cent_traj = md.join([traj[idx] for _, _, idx in centroids])\n",
        "cent_traj.save_pdb(work+\"clustered.pdb\")\n",
        "\n",
        "view6 = nv.show_mdtraj(cent_traj, defaultRepresentation=False)\n",
        "view6._set_size('750px','500px')\n",
        "view6.add_cartoon()\n",
        "view6.add_licorice()\n",
        "view6.add_simplified_base(\"nucleic\")\n",
        "view6.add_ball_and_stick(\"ligand\")\n",
        "view6.center()\n",
        "view6.player.delay = 3000\n",
        "display(view6)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "p27P31aeNR5t"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 4-3. Root Mean Squared Deviation\n",
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "assert os.path.exists(work+'production.h5'), \"4-1을 먼저 실행하세요.\"\n",
        "\n",
        "if 'traj' not in globals():\n",
        "    traj = md.load(work+'production.h5')\n",
        "\n",
        "#@markdown Trajectory의 첫 프레임과 비교했을 때 backbone 원자들의 좌표로 계산한 RMSD값 변화를 나타낸 그래프입니다.\n",
        "\n",
        "nsres = read_non_standard_residues_list(work)\n",
        "rmsd = md.rmsd(traj, traj, 0, atom_indices=traj.topology.select('backbone'))\n",
        "\n",
        "lig_select = widgets.Dropdown(options=['whole structure']+[f\"{a['name']} and rest\" for a in nsres],\n",
        "                              value='whole structure', description='calculate RMSDs for',\n",
        "                              style={'description_width': 'initial'},\n",
        "                              layout=widgets.Layout(width='750px'),\n",
        "                              continuous_update=True)\n",
        "\n",
        "def update_RMSD_chart(change):\n",
        "    with view_chart:\n",
        "        clear_output(wait=True)\n",
        "        if change.new == 'whole structure':\n",
        "            fig = go.Figure()\n",
        "            fig.add_trace(go.Scatter(x=sorted(traj.time/1000), y=rmsd, mode='lines', name='RMSD', connectgaps=False))\n",
        "            fig.update_layout(title='whole structure backbone RMSD', xaxis_title='Time (ns)',\n",
        "                               yaxis_title='RMSD (nm)', width=750, height=500)\n",
        "        else:\n",
        "            lig_name = change.new.split( )[0]\n",
        "            lig_rmsd = md.rmsd(traj, traj, 0, atom_indices=traj.topology.select(f'resname {lig_name}'))\n",
        "            rec_rmsd = md.rmsd(traj, traj, 0, atom_indices=traj.topology.select(f'(resname != {lig_name}) and backbone'))\n",
        "            fig = make_subplots(rows=2, cols=1, subplot_titles=(\"Receptor backbone RMSD\", f\"Ligand {lig_name} all-atom RMSD\"), vertical_spacing=0.2)\n",
        "            fig.add_trace(go.Scatter(x=sorted(traj.time/1000), y=rec_rmsd, mode='lines', name='Receptor RMSD', connectgaps=False), row=1, col=1)\n",
        "            fig.add_trace(go.Scatter(x=sorted(traj.time/1000), y=lig_rmsd, mode='lines', name=f'Ligand {lig_name} RMSD', connectgaps=False), row=2, col=1)\n",
        "            fig.update_xaxes(title='Time (ns)', showticklabels=True, row=1, col=1)\n",
        "            fig.update_xaxes(title='Time (ns)', row=2, col=1)\n",
        "            fig.update_yaxes(title='RMSD (nm)', row=1, col=1)\n",
        "            fig.update_yaxes(title='RMSD (nm)', row=2, col=1)\n",
        "            fig.update_layout(width=750, height=800, showlegend=False)\n",
        "        display(fig)\n",
        "\n",
        "\n",
        "lig_select.observe(update_RMSD_chart, names='value')\n",
        "\n",
        "view_chart = widgets.Output()\n",
        "\n",
        "with view_chart:\n",
        "    fig = go.Figure()\n",
        "    fig.add_trace(go.Scatter(x=sorted(traj.time/1000), y=rmsd, mode='lines', name='RMSD', connectgaps=False))\n",
        "    fig.update_layout(title='whole structure backbone RMSD', xaxis_title='Time (ns)',\n",
        "                       yaxis_title='RMSD (nm)', width=750, height=500)\n",
        "    display(fig)\n",
        "\n",
        "display(lig_select, view_chart)"
      ],
      "metadata": {
        "id": "8HF0OqTqAbhf",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 4-4. Radius of Gyration\n",
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "assert os.path.exists(work+'production.h5'), \"4-1을 먼저 실행하세요.\"\n",
        "\n",
        "if 'traj' not in globals():\n",
        "    traj = md.load(work+'production.h5')\n",
        "\n",
        "#@markdown 단백질/핵산 분자의 회전 반경 변화를 나타냅니다.\n",
        "\n",
        "rg = md.compute_rg(traj)\n",
        "\n",
        "print('Autoscale을 눌러 확대하세요.')\n",
        "\n",
        "fig4 = go.Figure()\n",
        "fig4.add_trace(go.Scatter(x=traj.time/1000, y=rg, mode='lines', name='RoG'))\n",
        "fig4.update_yaxes(range=[0, rg.max()*1.1])\n",
        "fig4.update_layout(title='Radius of Gyration', xaxis_title='Time (ns)', yaxis_title='Radius of Gyration (nm)', width=750, height=500)\n",
        "fig4.show()"
      ],
      "metadata": {
        "id": "usLGlIz-AcSS",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 4-5. Free Energy Landscape\n",
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "assert os.path.exists(work+'production.h5'), \"4-1을 먼저 실행하세요.\"\n",
        "\n",
        "#@markdown 특정 RMSD/RoG의 조합을 갖는 구조의 존재확률로부터 자유에너지 분포를 계산합니다.\n",
        "\n",
        "if 'traj' not in globals():\n",
        "    traj = md.load(work+'production.h5')\n",
        "\n",
        "rmsd = md.rmsd(traj, traj, 0)\n",
        "rg = md.compute_rg(traj)\n",
        "s = read_settings(work)\n",
        "X = np.vstack([rmsd, rg]).T\n",
        "\n",
        "# Automatic bandwidth sweep with 5-fold cross-validation\n",
        "bandwidths = np.logspace(-1, 1, 20)\n",
        "grid = GridSearchCV(KernelDensity(kernel='gaussian'), {'bandwidth': bandwidths}, cv=5)\n",
        "grid.fit(X)\n",
        "\n",
        "kde = grid.best_estimator_\n",
        "xi, yi = np.mgrid[rmsd.min():rmsd.max():100j, rg.min():rg.max():100j]\n",
        "xy_sample = np.vstack([xi.ravel(), yi.ravel()]).T\n",
        "z = np.exp(kde.score_samples(xy_sample))\n",
        "zi = z.reshape(xi.shape)\n",
        "\n",
        "Fi = -np.log(zi)\n",
        "Fi -= np.max(Fi) # zero at maximum\n",
        "Fi *= (8.31446261815324*s['temperature'].value_in_unit(kelvin)/1000) # convert from kT to kJ/mol\n",
        "\n",
        "fig5 = go.Figure(data=[go.Surface(z=Fi, x=xi, y=yi, colorbar=dict(title='kJ/mol'),\n",
        "                                  hovertemplate='RMSD: %{x:.3f} nm<br>Rg: %{y:.3f} nm<br>FE: %{z:.1f} kJ/mol<extra></extra>')])\n",
        "\n",
        "fig5.update_layout(title='Relative Free Energy Landscape',\n",
        "                     scene = dict(xaxis_title='RMSD (nm)', yaxis_title='Rg (nm)', zaxis_title='Free Energy (kJ/mol)'),\n",
        "                     width=750, height=750, autosize=True)\n",
        "\n",
        "fig5.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "KjAZprHVEdr5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 4-6. Residue-wise Root Mean Square Fluctuation\n",
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "assert os.path.exists(work+'production.h5'), \"4-1을 먼저 실행하세요.\"\n",
        "\n",
        "if 'traj' not in globals():\n",
        "    traj = md.load(work+'production.h5')\n",
        "\n",
        "#@markdown Residue별 모든 원자들의 좌표로부터 계산한 RMSF 값을 보여줍니다.\n",
        "\n",
        "#@markdown RMSF가 큰 residue는 더 이리저리 많이 움직인다는 것을 뜻합니다.\n",
        "\n",
        "prot_indices = traj.topology.select(\"not water and (mass 11 to 33) and (symbol != 'Na')\")\n",
        "traj2 = traj.superpose(traj, 0, atom_indices=prot_indices)\n",
        "traj2.atom_slice(atom_indices=prot_indices, inplace=True)\n",
        "\n",
        "atom_rmsf = np.sqrt(3*np.mean(np.square(traj2.xyz - np.mean(traj2.xyz, axis=0)), axis=0))\n",
        "\n",
        "res = []\n",
        "res_rmsf = []\n",
        "for c in traj.topology.chains:\n",
        "    chain = c.chain_id\n",
        "    for i, r in enumerate(c.residues):\n",
        "        res.append(f'{\"\" if chain is None else f\"{chain}:\"}{r.name.capitalize()}{i+1}')\n",
        "        atom_idx = [atom.index for atom in r.atoms]\n",
        "        res_rmsf.append(np.mean(atom_rmsf[atom_idx]))\n",
        "\n",
        "fig6 = go.Figure()\n",
        "fig6.add_trace(go.Scatter(x=np.arange(len(res)), y=res_rmsf, mode='lines', name='RMSF', text=res))\n",
        "fig6.update_traces(hovertemplate='Residue: %{text}<br>RMSF: %{y:.3f} nm<extra></extra>')\n",
        "fig6.update_yaxes(range=[0, max(res_rmsf)*1.1])\n",
        "fig6.update_layout(title='Residue-wise heavy atom RMSF', xaxis_title='Residue index', yaxis_title='RMSF (nm)',\n",
        "                   hovermode='x unified', width=750, height=500)\n",
        "fig6.show()"
      ],
      "metadata": {
        "id": "YgV67PHfAexa",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 4-7. Ramachandran Map\n",
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "assert os.path.exists(work+'production.h5'), \"4-1을 먼저 실행하세요.\"\n",
        "\n",
        "if 'traj' not in globals():\n",
        "    traj = md.load(work+'production.h5')\n",
        "\n",
        "#@markdown 단백질의 backbone 원자들이 이루는 이면각들인 phi와 psi 값을 보여주는 Ramachandran diagram을 그려줍니다.\n",
        "\n",
        "#@markdown 큰 단백질의 경우 로딩이 오래 걸릴 수 있습니다.\n",
        "\n",
        "stride = int(traj.n_frames/50)\n",
        "\n",
        "if stride < 2:\n",
        "    traj2 = traj\n",
        "    stride = 1\n",
        "else:\n",
        "    traj2 = traj[::stride]\n",
        "\n",
        "s = read_settings(work)\n",
        "\n",
        "phi = md.compute_phi(traj2, periodic=s['PBC'])[1]\n",
        "psi = md.compute_psi(traj2, periodic=s['PBC'])[1]\n",
        "\n",
        "res = []\n",
        "for c in traj.topology.chains:\n",
        "    chain = c.chain_id\n",
        "    for i, r in enumerate(c.residues):\n",
        "        res.append(f'{\"\" if chain is None else f\"{chain}:\"}{r.name.capitalize()}{i+1}')\n",
        "\n",
        "fig8 = go.Figure()\n",
        "colors = np.linspace(0, traj.n_frames*s['save interval'].value_in_unit(nanosecond), traj2.n_frames)\n",
        "\n",
        "fig8.add_trace(go.Scatter(x=phi.flatten(), y=psi.flatten(), mode='markers', marker=dict(color=np.repeat(colors, phi.shape[1]), colorscale='Viridis', colorbar=dict(title='Time (ns)'), showscale=True),\n",
        "                          text=[f'Time: {(i//phi.shape[1])*stride*s[\"save interval\"]}<br>Residue: {res[i%phi.shape[1]]}<br>Phi: {phi[i//phi.shape[1], i%phi.shape[1]]:.2f}<br>Psi: {psi[i//phi.shape[1], i%phi.shape[1]]:.2f}'\n",
        "                                for i in range(phi.size)], hoverinfo='text'))\n",
        "\n",
        "fig8.update_layout(title='Time-dependent Ramachandran Map', xaxis_title='Phi (radians)', yaxis_title='Psi (radians)', xaxis_range=[-np.pi, np.pi], yaxis_range=[-np.pi, np.pi], width=750, height=500)\n",
        "fig8.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "_tPxFFanA3mf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 4-8. Solvent Accessible Surface Area\n",
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "assert os.path.exists(work+'production.h5'), \"4-1을 먼저 실행하세요.\"\n",
        "\n",
        "if 'traj' not in globals():\n",
        "    traj = md.load(work+'production.h5')\n",
        "\n",
        "#@markdown 주어진 residue가 용매와 접촉하는 면적을 보여줍니다.\n",
        "\n",
        "#@markdown 보고 싶은 residue의 번호를 입력하세요.\n",
        "\n",
        "residue_index = \"140\" #@param {type:\"string\"}\n",
        "res_idx = int(residue_index)-1\n",
        "\n",
        "assert res_idx < traj.topology.n_residues, \"Invalid residue index!\"\n",
        "\n",
        "sasa = md.shrake_rupley(traj, mode='residue')\n",
        "\n",
        "for c in traj.topology.chains:\n",
        "    chain = c.chain_id\n",
        "    for i, r in enumerate(c.residues):\n",
        "        if r.index == res_idx:\n",
        "            name = f'{\"\" if chain is None else f\"{chain}:\"}{r.name.capitalize()}{i+1}'\n",
        "            break\n",
        "\n",
        "fig9 = go.Figure()\n",
        "fig9.add_trace(go.Scatter(x=traj.time/1000, y=sasa[:,res_idx], mode='lines', name='SASA'))\n",
        "fig9.update_layout(title=f'Solvent Accessible Surface Area of the residue {name}',\n",
        "                   xaxis_title='Time (ns)', yaxis_title='SASA (nm^2)', width=750, height=500)\n",
        "fig9.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "dCZiVq5DBAn9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 4-9. Residue간 수소 결합 분석\n",
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "assert os.path.exists(work+'production.h5'), \"4-1을 먼저 실행하세요.\"\n",
        "\n",
        "#@markdown 주어진 residue 번호들을 갖는 두 residue 사이의 수소 결합 갯수의 변화를 보여줍니다.\n",
        "\n",
        "#@markdown 하나만 입력하면 residue_1과 물 분자들 사이의 수소결합 갯수의 변화를 보여줍니다.\n",
        "\n",
        "traj = md.load(work+'production.h5')\n",
        "\n",
        "residue_1 = \"118\" #@param {type:\"string\"}\n",
        "residue_2 = \"\" #@param {type:\"string\"}\n",
        "\n",
        "residue_1 = int(residue_1) -1\n",
        "if residue_2 != \"\":\n",
        "    water_flag = False\n",
        "    residue_2 = int(residue_2) -1\n",
        "else:\n",
        "    water_flag = True\n",
        "\n",
        "assert residue_1 != residue_2, \"동일한 residue 번호가 입력되었습니다!\"\n",
        "\n",
        "s = read_settings(work)\n",
        "\n",
        "topol = traj.topology.to_openmm()\n",
        "\n",
        "donors = []\n",
        "acceptors = []\n",
        "\n",
        "for a in topol.atoms():\n",
        "    if a.element.symbol in ['N', 'O', 'F']:\n",
        "        if a.residue.index == residue_1:\n",
        "            donors.append(a.index)\n",
        "        elif a.residue.index == residue_2:\n",
        "            acceptors.append(a.index)\n",
        "        elif water_flag and a.residue.name == 'HOH':\n",
        "            acceptors.append(a.index)\n",
        "\n",
        "donor_Hs = {}\n",
        "acceptor_Hs = {}\n",
        "\n",
        "for b in topol.bonds():\n",
        "    if b[0].index in donors and b[1].element.symbol == 'H':\n",
        "        donor_Hs[b[1].index] = b[0].index\n",
        "    elif b[1].index in donors and b[0].element.symbol == 'H':\n",
        "        donor_Hs[b[0].index] = b[1].index\n",
        "    elif b[0].index in acceptors and b[1].element.symbol == 'H':\n",
        "        acceptor_Hs[b[1].index] = b[0].index\n",
        "    elif b[1].index in acceptors and b[0].element.symbol == 'H':\n",
        "        acceptor_Hs[b[0].index] = b[1].index\n",
        "\n",
        "H_bonds = []\n",
        "\n",
        "for key, value in donor_Hs.items():\n",
        "    for a in acceptors:\n",
        "        H_bonds.append([value, key, a])\n",
        "for key, value in acceptor_Hs.items():\n",
        "    for d in donors:\n",
        "        H_bonds.append([value, key, d])\n",
        "\n",
        "assert len(H_bonds) > 0, \"수소결합이 가능한 작용기가 없습니다.\"\n",
        "\n",
        "angles = md.compute_angles(traj, np.array(H_bonds), s['PBC'], True)\n",
        "distances = md.compute_distances(traj, np.array([[h[1], h[2]] for h in H_bonds]), s['PBC'], True)\n",
        "\n",
        "hbonds = np.sum(np.where(angles > 2*np.pi/3, 1, 0)*np.where(distances < 0.25, 1, 0), axis=1)\n",
        "\n",
        "for c in traj.topology.chains:\n",
        "    chain = c.chain_id\n",
        "    for i, r in enumerate(c.residues):\n",
        "        if r.index == residue_1:\n",
        "            name1 = f'{\"\" if chain is None else f\"{chain}:\"}{r.name.capitalize()}{i+1}'\n",
        "        elif not water_flag and r.index == residue_2:\n",
        "            name2 = f'{\"\" if chain is None else f\"{chain}:\"}{r.name.capitalize()}{i+1}'\n",
        "        else:\n",
        "            continue\n",
        "\n",
        "fig10 = go.Figure()\n",
        "fig10.add_trace(go.Scatter(x=traj.time/1000, y=hbonds, mode='lines', name='H-bond'))\n",
        "fig10.update_layout(xaxis_title='Time (ns)', yaxis_title='Number of H-bonds', width=750, height=500)\n",
        "if water_flag:\n",
        "    fig10.update_layout(title=f'Number of Hydrogen bonds between {name1} and water molecules')\n",
        "else:\n",
        "    fig10.update_layout(title=f'Number of Hydrogen Bonds between {name1} and {name2}')\n",
        "fig10.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "Ab0Le5IaS8BG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 4-10. DSSP를 통한 2차구조 분석\n",
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "assert os.path.exists(work+'production.h5'), \"4-1을 먼저 실행하세요.\"\n",
        "\n",
        "if 'traj' not in globals():\n",
        "    traj = md.load(work+'production.h5')\n",
        "\n",
        "#@markdown Define Secondary Structure of Proteins 알고리즘을 이용해 각 residue가 갖는 2차구조를 보여줍니다.\n",
        "\n",
        "dssp = md.compute_dssp(traj)\n",
        "dssp_map = {'H':0, 'E':1, 'C':2}\n",
        "dssp_num = np.array([[dssp_map.get(ss,2) for ss in frame] for frame in dssp])\n",
        "\n",
        "s = read_settings(work)\n",
        "\n",
        "fig11 = go.Figure(data=go.Heatmap(z=dssp_num.T, x=traj.time/1000, y=np.arange(traj.n_residues),\n",
        "                                 colorscale=[[0,'red'],[0.33, 'red'], [0.33, 'yellow'], [0.66, 'yellow'], [0.66, 'blue'], [1, 'blue']],\n",
        "                                 colorbar=dict(tickvals=[0,1,2],ticktext=['Helix', 'Sheet', 'Coil'])))\n",
        "fig11.update_layout(title='Time-dependent Secondary Structure (DSSP)', xaxis_title='Time (ns)', yaxis_title='Residue index', width=750, height=1000)\n",
        "fig11.data[0].hovertext = [[f'Time : {i*s[\"save interval\"]}<br>Residue: {j}<br>SS: {dssp[i][j]}' for i in range(traj.n_frames)] for j in range(traj.n_residues)]\n",
        "fig11.data[0].hoverinfo = 'text'\n",
        "fig11.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "2c8tedqCBDbN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 4-11. Dynamic Cross-Correlation Matrix\n",
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "assert os.path.exists(work+'production.h5'), \"4-1을 먼저 실행하세요.\"\n",
        "\n",
        "if 'traj' not in globals():\n",
        "    traj = md.load(work+'production.h5')\n",
        "\n",
        "#@markdown 각 residue쌍의 알파탄소의 운동의 상관관계를 그린 표입니다.\n",
        "\n",
        "#@markdown 파랑색 (+1)에 가까울 수록 양의 상관관계가 큰 것이고, 두 residue가 같은 방향으로 움직이는 경향이 있다는 뜻입니다.\n",
        "\n",
        "#@markdown 빨강색 (-1)에 가까울 수록 음의 상관관계가 큰 것이고, 두 residue가 반대 방향으로 움직이는 경향이 있다는 뜻입니다.\n",
        "\n",
        "#@markdown 흰색 (0)에 가까울 수록 상관관계가 없는 것이고, 두 residue가 독립적으로 움직인다는 뜻입니다.\n",
        "\n",
        "ca_idx = traj.topology.select_atom_indices('alpha')\n",
        "mean_pos = np.mean(traj.xyz[:, ca_idx, :], axis=0)\n",
        "fluct = traj.xyz[:, ca_idx, :] - mean_pos\n",
        "\n",
        "dccm = np.zeros((traj.n_residues, traj.n_residues))\n",
        "\n",
        "for i in range(len(ca_idx)):\n",
        "    for j in range(len(ca_idx)):\n",
        "        numerator = np.mean(np.sum(fluct[:,i,:]*fluct[:,j,:], axis=1))\n",
        "        denominator = np.sqrt(np.mean(np.sum(fluct[:,i,:]**2, axis=1))*np.mean(np.sum(fluct[:,j,:]**2, axis=1)))\n",
        "        dccm[i,j] = numerator/denominator\n",
        "\n",
        "fig12 = go.Figure(data=go.Heatmap(z=dccm, colorscale='RdBu', zmid=0, colorbar=dict(title='Correlation')))\n",
        "fig12.update_layout(title='Dynamic Cross-Correlation Matrix', xaxis_title='Residue index', yaxis_title='Residue index', width=750, height=750)\n",
        "fig12.data[0].hovertext = [[f'Residue i: {i+1}<br>Residue j: {j+1}<br>Correlation: {dccm[i,j]:.2f}' for j in range(dccm.shape[1])] for i in range(dccm.shape[0])]\n",
        "fig12.data[0].hoverinfo = 'text'\n",
        "fig12.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "V0qemYvRBF5e"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 4-12. 주성분분석 (Principal Component Analysis)"
      ],
      "metadata": {
        "id": "EBe114g84Lq9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 2차원 주성분분석\n",
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "assert os.path.exists(work+'production.h5'), \"4-1을 먼저 실행하세요.\"\n",
        "\n",
        "if 'traj' not in globals():\n",
        "    traj = md.load(work+'production.h5')\n",
        "\n",
        "#@markdown 원자들의 집합적인 운동을 직교벡터 성분들로 추출해서 가장 고윳값이 큰 벡터들을 주성분이라고 부릅니다. 고윳값은 그 주성분이 갖는 설명력이라고 할 수 있습니다.\n",
        "\n",
        "prot_indices = traj.topology.select(\"not water and (mass 11 to 33) and (symbol != 'Na')\")\n",
        "traj2 = traj.superpose(traj, 0, atom_indices=prot_indices)\n",
        "traj2.atom_slice(atom_indices=prot_indices, inplace=True)\n",
        "pca = PCA(n_components=2)\n",
        "reduced_cart = pca.fit_transform(traj2.xyz.reshape(traj2.n_frames, traj2.n_atoms*3))\n",
        "var_explained = pca.explained_variance_ratio_*100\n",
        "\n",
        "s = read_settings(work)\n",
        "\n",
        "fig13 = go.Figure()\n",
        "fig13.add_trace(go.Scatter(x=reduced_cart[:,0], y=reduced_cart[:,1], mode='markers',\n",
        "                          marker=dict(size=5, color=np.arange(len(reduced_cart))*s[\"save interval\"].value_in_unit(nanosecond), colorscale='Viridis', colorbar=dict(title='Time (ns)'), showscale=True),\n",
        "                          text=[f'Time: {i*s[\"save interval\"]}' for i in range(len(reduced_cart))], hoverinfo='text', name='PCA'))\n",
        "fig13.update_layout(title='Two Dimensional Principal Component Analysis', xaxis_title=f'PC1 ({var_explained[0]:.2f}%)', yaxis_title=f'PC2 ({var_explained[1]:.2f}%)', width=750, height=500)\n",
        "fig13.show()"
      ],
      "metadata": {
        "id": "PgPSAoZ2Apnr",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 3차원 주성분분석\n",
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "assert os.path.exists(work+'production.h5'), \"4-1을 먼저 실행하세요.\"\n",
        "\n",
        "if 'traj' not in globals():\n",
        "    traj = md.load(work+'production.h5')\n",
        "\n",
        "prot_indices = traj.topology.select(\"not water and (mass 11 to 33) and (symbol != 'Na')\")\n",
        "traj2 = traj.superpose(traj, 0, atom_indices=prot_indices)\n",
        "traj2.atom_slice(atom_indices=prot_indices, inplace=True)\n",
        "pca = PCA(n_components=3)\n",
        "reduced_cart = pca.fit_transform(traj2.xyz.reshape(traj2.n_frames, traj2.n_atoms*3))\n",
        "var_explained = pca.explained_variance_ratio_*100\n",
        "\n",
        "s = read_settings(work)\n",
        "s['save interval'] = 10*picosecond\n",
        "\n",
        "fig14 = go.Figure()\n",
        "fig14.add_trace(go.Scatter3d(x=reduced_cart[:,0], y=reduced_cart[:,1], z=reduced_cart[:,2], mode='markers',\n",
        "                          marker=dict(size=5, color=s[\"save interval\"].value_in_unit(nanosecond)*np.arange(len(reduced_cart)), colorscale='Viridis', opacity=0.8, colorbar=dict(title='Time (ns)'), showscale=True),\n",
        "                          text=[f'Time: {i*s[\"save interval\"]}' for i in range(len(reduced_cart))], hoverinfo='text', name='PCA'))\n",
        "fig14.update_layout(title='Three Dimensional Principal Component Analysis', scene=dict(xaxis_title=f'PC1 ({var_explained[0]:.2f}%)', yaxis_title=f'PC2 ({var_explained[1]:.2f}%)', zaxis_title=f'PC3 ({var_explained[2]:.2f}%)'), width=750, height=750)\n",
        "fig14.show()"
      ],
      "metadata": {
        "id": "Kov9UUH6AsEX",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 주성분별 시각화\n",
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "assert os.path.exists(work+'production.h5'), \"4-1을 먼저 실행하세요.\"\n",
        "\n",
        "if 'traj' not in globals():\n",
        "    traj = md.load(work+'production.h5')\n",
        "\n",
        "#@markdown 각 주성분이 어떤 운동을 나타내는지 시각화해줍니다.\n",
        "\n",
        "#@markdown 보고 싶은 주성분의 번호를 골라보세요.\n",
        "component_to_view = \"1\" #@param [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]\n",
        "#@markdown 운동의 크기를 지정해주세요. 클 수록 눈에 잘 띄지만 너무 커지면 구조가 안뜰 수 있습니다.\n",
        "scale = 3 #@param {type:\"slider\", min:1, max:10, step:1}\n",
        "\n",
        "nsres = read_non_standard_residues_list(work)\n",
        "nsres_to_view = \"\"\n",
        "for i, r in enumerate(nsres):\n",
        "    nsres_to_view += f\"(resname {r['name']})\"\n",
        "    if i < len(nsres)-1:\n",
        "        nsres_to_view += \" or \"\n",
        "\n",
        "ligand_indices = traj.topology.select(f'{nsres_to_view}')\n",
        "prot_indices = traj.topology.select(\"not water and (mass 11 to 33) and (symbol != 'Na')\")\n",
        "view_indices = np.array(sorted(set(list(ligand_indices)+list(prot_indices))))\n",
        "traj2 = traj.superpose(traj, 0, atom_indices=prot_indices)\n",
        "traj2.atom_slice(atom_indices=view_indices, inplace=True)\n",
        "coords = traj2.xyz.reshape(traj2.n_frames, -1)\n",
        "\n",
        "pca = PCA()\n",
        "pca_results = pca.fit_transform(coords)\n",
        "top_pcs = pca.components_[:10]\n",
        "\n",
        "pc_trajs = top_pcs.reshape(10, traj2.n_atoms, 3)\n",
        "sine_wave = np.sin(np.linspace(0, 2*np.pi, num=50))\n",
        "\n",
        "pc_traj = md.Trajectory(np.zeros((50, traj2.n_atoms, 3)), traj2.topology)\n",
        "\n",
        "for i in range(50):\n",
        "    pc_traj.xyz[i] = traj2.xyz[0] + scale*sine_wave[i]*pc_trajs[int(component_to_view)-1]\n",
        "\n",
        "view6 = nv.show_mdtraj(pc_traj, defaultRepresentation=False)\n",
        "view6._set_size('750px','500px')\n",
        "view6.add_cartoon()\n",
        "view6.add_simplified_base(\"nucleic\")\n",
        "view6.add_ball_and_stick(\"ligand\")\n",
        "view6.center()\n",
        "view6.player.delay = 500\n",
        "display(view6)"
      ],
      "metadata": {
        "id": "f_jIQg9KAwI-",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 5. 분자간 결합 분석"
      ],
      "metadata": {
        "id": "1EGMmauoHQmT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 5-1. PLIP을 통한 상호작용 시각화\n",
        "assert 'work' in globals(), \"0번 셀부터 실행해주세요.\"\n",
        "assert os.path.exists(work+'starting_structure.pdb'), \"starting_structure.pdb 파일이 없습니다!\"\n",
        "assert os.path.exists(work+'production.h5'), \"결과 trajectory가 없습니다. 4-1을 실행하세요.\"\n",
        "assert os.path.exists(work+'clustered.pdb'), \"4-2를 실행하세요.\"\n",
        "\n",
        "traj = md.load(work+'production.h5')\n",
        "\n",
        "#@markdown Protein-ligand 상호작용을 보여줍니다. <br/>\n",
        "#@markdown * H-bonds : 빨강\n",
        "#@markdown * Hydrophobic : 회색\n",
        "#@markdown * Salt bridge : 노랑\n",
        "#@markdown * Water bridge : 청록\n",
        "#@markdown * Pi-pi stacking : 파랑\n",
        "#@markdown * Pi-cation : 보라\n",
        "#@markdown * Halogen bonds : 초록\n",
        "#@markdown * Metal coordination : 하늘\n",
        "\n",
        "clusterPDB = PDBComplex()\n",
        "clusterPDB.load_pdb(work+'clustered.pdb')\n",
        "clusterPDB.analyze()\n",
        "\n",
        "ligands = {}\n",
        "for ligand in clusterPDB.ligands:\n",
        "    clusterPDB.characterize_complex(ligand)\n",
        "\n",
        "for key, value in sorted(clusterPDB.interaction_sets.items()):\n",
        "    ligands[key] = BindingSiteReport(value)\n",
        "\n",
        "interaction_types = ['hbond', 'hydrophobic', 'saltbridge', 'waterbridge', 'pistacking', 'pication', 'halogen', 'metal']\n",
        "color = [(1,0,0), (0.5,0.5,0.5), (1,1,0), (0,1,1), (0,0,1), (1,0,1), (0,1,0), (0.9,0.5,0.9)]\n",
        "\n",
        "view_interaction = widgets.Output()\n",
        "\n",
        "with view_interaction:\n",
        "    view = nv.NGLWidget()\n",
        "    view.add_component(nv.FileStructure(work+'clustered.pdb'), defaultRepresentation=False)\n",
        "    view._set_size('750px','500px')\n",
        "    view.add_cartoon()\n",
        "    view.add_simplified_base(\"nucleic\")\n",
        "    view.add_licorice(\"ligand\")\n",
        "    view.center()\n",
        "    display(view)\n",
        "\n",
        "def center_ligand(change): # callback function for the update\n",
        "    with view_interaction:\n",
        "        clear_output(wait=True)\n",
        "        view = nv.NGLWidget()\n",
        "        view.add_component(nv.FileStructure(work+'clustered.pdb'), defaultRepresentation=False)\n",
        "        view._set_size('750px','500px')\n",
        "        view.add_cartoon()\n",
        "        view.add_simplified_base(\"nucleic\")\n",
        "        view.add_licorice(\"ligand\")\n",
        "        if change.new == 'none':\n",
        "            view.center()\n",
        "        else:\n",
        "            intres = []\n",
        "            for i, inttype in enumerate(interaction_types):\n",
        "                intdata = getattr(ligands[change.new], inttype+'_info')\n",
        "                for interaction in intdata:\n",
        "                    intres.append(f'({interaction[0]} and {interaction[1]} and :{interaction[2]})')\n",
        "                    view.shape.add_cylinder(interaction[-2], interaction[-1], color[i], [0.05], inttype)\n",
        "            view.add_licorice(\" or \".join(intres))\n",
        "            resdata = change.new.split(':')\n",
        "            view.center(f'{resdata[2]} and :{resdata[1]} and {resdata[0]}')\n",
        "        display(view)\n",
        "\n",
        "print('리간드를 새로 선택한 후에는 NGLView 위젯을 기다려주세요. ')\n",
        "\n",
        "if len(ligands) > 0:\n",
        "    select_lig = widgets.Dropdown(options=['none']+list(ligands.keys()), value='none', description='Ligand : ', continuous_update=True)\n",
        "    select_lig.observe(center_ligand, names='value')\n",
        "    display(select_lig, view_interaction)\n",
        "else:\n",
        "    display(view_interaction)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "QFxU3Eqx3XlG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 5-2. MM/GBSA 준비\n",
        "assert 'work' in globals(), \"0번 셀부터 실행해주세요.\"\n",
        "#@markdown Ligand와 receptor의 chain 이름을 입력해주세요. 여러개는 공백없이 쉼표로 구분해주세요. <br/>\n",
        "#@markdown Ligand만 입력하면 나머지 chain들은 자동으로 receptor로 간주됩니다.\n",
        "ligand_chain = \"C\" #@param {type:\"string\"}\n",
        "receptor_chain = \"\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Production trajectory의 시작부분을 계산에 포함하지 않으려면 아래 항목을 0보다 크게 하세요.\n",
        "skip_initial_fraction = 0 #@param {type:\"slider\", min:0, max:0.9, step:0.1}\n",
        "#@markdown N번째 프레임마다 계산에 활용할지 설정하세요.\n",
        "stride = 1  #@param {type:\"slider\", min:1, max:20, step:1}\n",
        "\n",
        "pdb = PDBFile(work+'starting_structure.pdb')\n",
        "\n",
        "lig_a = []\n",
        "rec_a = []\n",
        "com_a = []\n",
        "elements = []\n",
        "\n",
        "lig_present = False\n",
        "rec_present = False\n",
        "lig_r = []\n",
        "rec_r = []\n",
        "\n",
        "for a in pdb.topology.atoms():\n",
        "    if a.residue.chain.id in ligand_chain.split(','):\n",
        "        lig_present = True\n",
        "        lig_r.append(a.residue.index)\n",
        "        lig_a.append(a.index)\n",
        "    elif receptor_chain == \"\" or a.residue.chain.id in receptor_chain.split(','):\n",
        "        rec_present = True\n",
        "        rec_r.append(a.residue.index)\n",
        "        rec_a.append(a.index)\n",
        "    else:\n",
        "        continue\n",
        "    com_a.append(a.index)\n",
        "    elements.append(a.element.symbol)\n",
        "\n",
        "for e in set(elements):\n",
        "    assert e in ['H', 'C', 'N', 'O', 'F', 'P', 'S', 'Si', 'Cl'], f'{e}에 대한 bondi parameter가 없습니다!'\n",
        "assert lig_present, f'{ligand_chain}에 해당하는 분자가 없습니다!'\n",
        "assert rec_present, f'{receptor_chain}에 해당하는 분자가 없습니다!'\n",
        "\n",
        "\n",
        "s = read_settings(work)\n",
        "ff, wm = determine_forcefield(s['forcefield'], s['watermodel'])\n",
        "\n",
        "nsres = read_non_standard_residues_list(work)\n",
        "smffs = []\n",
        "for res in nsres:\n",
        "    if res['lig']:\n",
        "        smffs.append(res['ff'])\n",
        "\n",
        "esp_flag = False\n",
        "if 'espaloma-0.3.2' in set(smffs):\n",
        "    esp_flag = True\n",
        "\n",
        "if esp_flag:\n",
        "    print('Full package를 설치하세요!')"
      ],
      "metadata": {
        "id": "K0ftSkx5Kb9-",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 5-3. MM/GBSA 실행\n",
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "assert 'ff' in globals(), \"5-2를 먼저 실행하세요.\"\n",
        "\n",
        "#@markdown Generalized Born 모델을 선택하세요.\n",
        "GB_model = \"GBn2 (igb=8)\" #@param [\"HCT (igb=1)\", \"OBC1 (igb=2)\", \"OBC2 (igb=5)\", \"GBn (igb=7)\", \"GBn2 (igb=8)\"]\n",
        "#@markdown 용매와 용질의 유전상수를 정해주세요.<br/>\n",
        "#@markdown 용매의 권장값은 78이고 용질의 권장값은 2~5 입니다. <br/>\n",
        "#@markdown 용질의 유전상수를 높여서 상호작용 에너지를 줄일 수 있습니다.\n",
        "solvent_dielectric = 78 #@param {type:\"slider\", min:1, max:80, step:1}\n",
        "solute_dielectric = 5 #@param {type:\"slider\", min:1, max:80, step:1}\n",
        "\n",
        "assert solvent_dielectric > solute_dielectric, \"용매의 유전상수는 용질의 유전상수보다 커야합니다!\"\n",
        "\n",
        "gaff_flag = False\n",
        "smirnoff_flag = False\n",
        "esp_flag = False\n",
        "\n",
        "gaff = []\n",
        "smirnoff = []\n",
        "esp = []\n",
        "for res in nsres:\n",
        "    if res['ff'] == 'GAFF-2.11':\n",
        "        gaff_flag = True\n",
        "        with open(f'{work}{res[\"name\"]}.json', 'r') as f:\n",
        "            gaff.append(Molecule.from_json(f.read()))\n",
        "    elif res['ff'] == 'Sage-2.2.1':\n",
        "        smirnoff_flag = True\n",
        "        with open(f'{work}{res[\"name\"]}.json', 'r') as f:\n",
        "            smirnoff.append(Molecule.from_json(f.read()))\n",
        "    elif res['ff'] == 'espaloma-0.3.2':\n",
        "        esp_flag = True\n",
        "        with open(f'{work}{res[\"name\"]}.json', 'r') as f:\n",
        "            esp.append(Molecule.from_json(f.read()))\n",
        "\n",
        "if gaff_flag:\n",
        "    g = GAFFTemplateGenerator(molecules=gaff)\n",
        "    ff.registerTemplateGenerator(g.generator)\n",
        "if smirnoff_flag:\n",
        "    g = SMIRNOFFTemplateGenerator(molecules=smirnoff)\n",
        "    ff.registerTemplateGenerator(g.generator)\n",
        "if esp_flag:\n",
        "    g = EspalomaTemplateGenerator(molecules=esp)\n",
        "    ff.registerTemplateGenerator(g.generator)\n",
        "\n",
        "def create_GBSA_system(trajectory, forcefield, fftype, eps_solv, eps_solut, kappa):\n",
        "\n",
        "    top = trajectory.topology.to_openmm(trajectory)\n",
        "\n",
        "    for chain in top.chains():\n",
        "        a1 = None\n",
        "        a2 = None\n",
        "        for i, res in enumerate(chain.residues()):\n",
        "            if i == 0 and res.name == 'PRO':\n",
        "                for atom in res.atoms():\n",
        "                    if atom.name == 'H2':\n",
        "                        a2 = atom\n",
        "                    elif atom.name == 'N':\n",
        "                        a1 = atom\n",
        "                    else:\n",
        "                        continue\n",
        "            else:\n",
        "                break\n",
        "        if a1 is not None and a2 is not None:\n",
        "            top.addBond(a1, a2)\n",
        "\n",
        "    system = forcefield.createSystem(top, nonbondedMethod=NoCutoff)\n",
        "\n",
        "    charges = []\n",
        "\n",
        "    for force in system.getForces():\n",
        "        if isinstance(force, NonbondedForce):\n",
        "            force.setForceGroup(1)\n",
        "            force.addGlobalParameter(f'coul_scale', 1)\n",
        "            force.addGlobalParameter(f'lj_scale', 1)\n",
        "            for i in range(force.getNumParticles()):\n",
        "                q, s, e = force.getParticleParameters(i)\n",
        "                charges.append(q)\n",
        "                force.setParticleParameters(i, 0, s, 0)\n",
        "                force.addParticleParameterOffset(f'coul_scale', i, q, 0, 0)\n",
        "                force.addParticleParameterOffset(f'lj_scale', i, 0, 0, e)\n",
        "            for i in range(force.getNumExceptions()):\n",
        "                p1, p2, q_p, s, e = force.getExceptionParameters(i)\n",
        "                force.setExceptionParameters(i, p1, p2, 0, s, 0)\n",
        "\n",
        "        elif isinstance(force, CustomNonbondedForce):\n",
        "            force.setForceGroup(2)\n",
        "            for i in range(force.getNumExceptions()):\n",
        "                p1, p2, q_p, s, e = force.getExceptionParameters(i)\n",
        "                force.setExceptionParameters(i, p1, p2, 0, s, 0)\n",
        "        else:\n",
        "            force.setForceGroup(0)\n",
        "\n",
        "    args = {'soluteDielectric': eps_solut, 'solventDielectric': eps_solv,\n",
        "            'kappa': kappa, 'SA':'ACE'}\n",
        "\n",
        "    if fftype == \"HCT (igb=1)\":\n",
        "        gb = customgbforces.GBSAHCTForce(**args)\n",
        "        params = customgbforces.GBSAHCTForce.getStandardParameters(top)\n",
        "    elif fftype == \"OBC1 (igb=2)\":\n",
        "        gb = customgbforces.GBSAOBC1Force(**args)\n",
        "        params = customgbforces.GBSAOBC1Force.getStandardParameters(top)\n",
        "    elif fftype == \"OBC2 (igb=5)\":\n",
        "        gb = customgbforces.GBSAOBC2Force(**args)\n",
        "        params = customgbforces.GBSAOBC2Force.getStandardParameters(top)\n",
        "    elif fftype == \"GBn (igb=7)\":\n",
        "        gb = customgbforces.GBSAGBnForce(**args)\n",
        "        parmas = customgbforces.GBSAGBnForce.getStandardParameters(top)\n",
        "    elif fftype == \"GBn2 (igb=8)\":\n",
        "        gb = customgbforces.GBSAGBn2Force(**args)\n",
        "        params = customgbforces.GBSAGBn2Force.getStandardParameters(top)\n",
        "    else:\n",
        "        raise ValueError\n",
        "\n",
        "    gb.setForceGroup(3)\n",
        "    for i, p in enumerate(params):\n",
        "        if fftype in ['GBn (igb=7)', 'GBn2 (igb=8)']:\n",
        "            gb.addParticle([charges[i], p[0], p[1], p[2], p[3], p[4]])\n",
        "        else:\n",
        "            gb.addParticle([charges[i], p[0], p[1]])\n",
        "\n",
        "    gb.finalize()\n",
        "    gb.setNonbondedMethod(CustomNonbondedForce.NoCutoff)\n",
        "    system.addForce(gb)\n",
        "\n",
        "    return system\n",
        "\n",
        "\n",
        "def calculate_energies(work, name, trajectory, forcefield, fftype, eps_solv, eps_solut):\n",
        "\n",
        "    print(f'{name} 분석 중 :')\n",
        "\n",
        "    s = read_settings(work)\n",
        "    temperature = s['temperature']/kelvin\n",
        "    ion_conc = s['ionic strength']/molar\n",
        "    kappa = 367.434915*sqrt(float(ion_conc)/78.5/(temperature))\n",
        "\n",
        "    system = create_GBSA_system(trajectory, forcefield, fftype, eps_solv, eps_solut, kappa)\n",
        "    integrator = VerletIntegrator(1*femtosecond)\n",
        "    context = Context(system, integrator)\n",
        "\n",
        "    coulombic = []\n",
        "    LJ = []\n",
        "    GBSA = []\n",
        "\n",
        "    for i in tqdm(range(trajectory.n_frames)):\n",
        "        context.setPositions(trajectory.openmm_positions(i))\n",
        "        context.setParameter('coul_scale', eps_solut**(-1/2))\n",
        "        context.setParameter('lj_scale', 0)\n",
        "        coulombic.append(context.getState(getEnergy=True, groups={1}).getPotentialEnergy().value_in_unit(kilojoules_per_mole))\n",
        "        context.setParameter('coul_scale', 0)\n",
        "        context.setParameter('lj_scale', 1/eps_solut)\n",
        "        LJ.append(context.getState(getEnergy=True, groups={1,2}).getPotentialEnergy().value_in_unit(kilojoules_per_mole))\n",
        "        GBSA.append(context.getState(getEnergy=True, groups={3}).getPotentialEnergy().value_in_unit(kilojoules_per_mole))\n",
        "\n",
        "    return np.array(coulombic), np.array(LJ), np.array(GBSA)\n",
        "\n",
        "\n",
        "def derive_G_K(dH, dH_std, dS, dS_std, temperature):\n",
        "    dG = dH - temperature*dS/1000\n",
        "    dG_std = np.sqrt(dH_std**2+(dS_std*temperature/1000)**2)\n",
        "    Kd = np.exp(dG*1000/8.31446261815324/temperature)\n",
        "    Kd_std = Kd*1000/8.31446261815324/temperature*dG_std\n",
        "    return dG, dG_std, Kd, Kd_std\n",
        "\n",
        "\n",
        "rough_traj = md.load(work+'production.h5')\n",
        "start_point = int(skip_initial_fraction*rough_traj.n_frames)\n",
        "\n",
        "lig_traj = rough_traj[start_point::stride].atom_slice(atom_indices=lig_a)\n",
        "rec_traj = rough_traj[start_point::stride].atom_slice(atom_indices=rec_a)\n",
        "com_traj = rough_traj[start_point::stride].atom_slice(atom_indices=com_a)\n",
        "\n",
        "lig_coul, lig_lj, lig_gb = calculate_energies(work, 'ligand', lig_traj, ff, GB_model, solvent_dielectric, solute_dielectric)\n",
        "rec_coul, rec_lj, rec_gb = calculate_energies(work, 'receptor', rec_traj, ff, GB_model, solvent_dielectric, solute_dielectric)\n",
        "com_coul, com_lj, com_gb = calculate_energies(work, 'complex', com_traj, ff, GB_model, solvent_dielectric, solute_dielectric)\n",
        "\n",
        "couls = com_coul - lig_coul - rec_coul\n",
        "ljs = com_lj - lig_lj - rec_lj\n",
        "gbsas = com_gb - lig_gb - rec_gb\n",
        "dHs = couls + ljs + gbsas\n",
        "dH = np.mean(dHs)\n",
        "\n",
        "fig15 = go.Figure()\n",
        "fig15.add_trace(go.Scatter(x=lig_traj.time/1000, y=dHs, mode='lines', name='delta H'))\n",
        "fig15.add_trace(go.Scatter(x=lig_traj.time/1000, y=couls, mode='lines', name='Coulombic E'))\n",
        "fig15.add_trace(go.Scatter(x=lig_traj.time/1000, y=ljs, mode='lines', name='Lennard-Jones E'))\n",
        "fig15.add_trace(go.Scatter(x=lig_traj.time/1000, y=gbsas, mode='lines', name='Solvation E'))\n",
        "fig15.update_layout(title='Binding Energy', xaxis_title='Time (ns)', yaxis_title='Energy (kJ/mol)', width=750, height=500)\n",
        "\n",
        "dE = np.mean(couls+ljs)\n",
        "dE_std = np.std(couls+ljs)\n",
        "dH_std = np.sqrt(np.var(couls+ljs)+np.var(gbsas))\n",
        "\n",
        "fig16 = go.Figure()\n",
        "x = ['ΔE', 'polar', 'nonpolar', 'GBSA', 'ΔH']\n",
        "y = [dE, np.mean(couls), np.mean(ljs), np.mean(gbsas), dH]\n",
        "colors = ['#1f77b4', '#17becf', '#17becf', '#9467bd', '#d62728']\n",
        "err = dict(type='data', array=[dE_std, np.std(couls), np.std(ljs), np.std(gbsas), dH_std])\n",
        "fig16.add_trace(go.Bar(x=x, y=y, error_y=err, marker_color=colors))\n",
        "fig16.update_layout(title=\"Energy Contributions\", xaxis_title='Type', yaxis_title='Energy (kJ/mol)', width=750, height=500)\n",
        "\n",
        "display(fig15, fig16)\n",
        "\n",
        "print(f'Interaction Energy    : {dE:>+7.2f} +/- {dE_std:>5.2f} kJ/mol')\n",
        "print(f'Polar contribution    : {np.mean(couls):>+7.2f} +/- {np.std(couls):>5.2f} kJ/mol')\n",
        "print(f'Nonpolar contribution : {np.mean(ljs):>+7.2f} +/- {np.std(ljs):>5.2f} kJ/mol')\n",
        "print(f'GBSA Solvation Energy : {np.mean(gbsas):>+7.2f} +/- {np.std(gbsas):>5.2f} kJ/mol')\n",
        "print(f'Binding Enthalpy      : {dH:>+7.2f} +/- {dH_std:>5.2f} kJ/mol')"
      ],
      "metadata": {
        "id": "SB6-5YsX5KYs",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 5-4. Interaction Entropy & C2 Entropy 계산\n",
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "assert 'couls' in globals(), \"5-2를 먼저 실행하세요.\"\n",
        "\n",
        "#@markdown 10^N 번 만큼 부트스트랩 샘플링을 수행합니다.\n",
        "bootstrap_power = 4 #@param {type:\"slider\", min:3, max:8, step:1}\n",
        "\n",
        "def calculate_interaction_entropy(interaction_energies, temperature):\n",
        "    \"\"\"\n",
        "    Calculate the interaction entropy using the method from Duan et al.\n",
        "    https://doi.org/10.1021/jacs.6b02682\n",
        "    \"\"\"\n",
        "    R = 8.31446261815324  # Gas constant in J/mol/K\n",
        "    beta = 1 / (R * temperature)\n",
        "    avg_energy = np.mean(interaction_energies)\n",
        "    delta_e = interaction_energies - avg_energy\n",
        "    exp_beta_delta_e = np.exp(beta * delta_e *1000) # kJ/mol to J/mol\n",
        "\n",
        "    interaction_entropy = - R * np.log(np.mean(exp_beta_delta_e))\n",
        "    return interaction_entropy\n",
        "\n",
        "def calculate_c2_entropy(interaction_energies, temperature):\n",
        "    \"\"\"\n",
        "    Calculate the C2 entropy using the method from Menzer et al.\n",
        "    https://doi.org/10.1021/acs.jctc.8b00418\n",
        "    \"\"\"\n",
        "    R = 8.31446261815324  # Gas constant in J/mol/K\n",
        "    beta = 1 / (R * temperature)\n",
        "    variance = np.var(interaction_energies*1000)\n",
        "\n",
        "    c2_entropy = -0.5 * beta * variance / temperature # kJ/mol to J/mol\n",
        "    return c2_entropy\n",
        "\n",
        "\n",
        "def bootstrap_entropy_error(interaction_energies, temperature, n_bootstrap):\n",
        "\n",
        "    n_samples = len(interaction_energies)\n",
        "    interaction_entropies = []\n",
        "    c2_entropies = []\n",
        "\n",
        "    print(f'Running with n_bootstrap : {n_bootstrap}...')\n",
        "\n",
        "    for _ in tqdm(range(n_bootstrap)):\n",
        "        bootstrap_sample = np.random.choice(interaction_energies, size=n_samples, replace=True)\n",
        "\n",
        "        # Calculate entropies for this sample\n",
        "        interaction_entropy = calculate_interaction_entropy(bootstrap_sample, temperature)\n",
        "        c2_entropy = calculate_c2_entropy(bootstrap_sample, temperature)\n",
        "\n",
        "        interaction_entropies.append(interaction_entropy)\n",
        "        c2_entropies.append(c2_entropy)\n",
        "\n",
        "    # Calculate mean and standard error\n",
        "    IE = np.mean(interaction_entropies)\n",
        "    IE_std = np.std(interaction_entropies, ddof=1)\n",
        "\n",
        "    c2 = np.mean(c2_entropies)\n",
        "    c2_std = np.std(c2_entropies, ddof=1)\n",
        "\n",
        "    return IE, IE_std, c2, c2_std\n",
        "\n",
        "temperature = s['temperature']/kelvin\n",
        "\n",
        "# set from https://doi.org/10.1021/acs.jctc.1c00374\n",
        "if dE_std > 15:\n",
        "    print('Warning : Interaction Entropy 값을 신뢰하기 어렵습니다!')\n",
        "if dE_std > 25:\n",
        "    print('Warning : C2 Entropy 값을 신뢰하기 어렵습니다!')\n",
        "\n",
        "IE, IE_std, c2, c2_std = bootstrap_entropy_error(couls+ljs, temperature, 10**bootstrap_power)\n",
        "\n",
        "print(f'Interaction Entropy : {IE:>+7.2f} +/- {IE_std:>3.2f} J/mol/K')\n",
        "print(f'C2 Entropy          : {c2:>+7.2f} +/- {c2_std:>3.2f} J/mol/K')\n",
        "\n",
        "dG_IE, dG_IE_std, Kd_IE, Kd_IE_std = derive_G_K(dH, dH_std, IE, IE_std, temperature)\n",
        "dG_c2, dG_c2_std, Kd_c2, Kd_c2_std = derive_G_K(dH, dH_std, c2, c2_std, temperature)\n",
        "\n",
        "print(f'Binding Free Energy (ΔG) from IE : {dG_IE:>+7.2f} +/- {dG_IE_std:>3.2f} kJ/mol')\n",
        "print(f'Binding Free Energy (ΔG) from C2 : {dG_c2:>+7.2f} +/- {dG_c2_std:>3.2f} kJ/mol')\n",
        "print(f'Dissociation Constant (Kd) from IE: {Kd_IE:>5.2e} +/- {Kd_IE_std:>3.2e} M')\n",
        "print(f'Dissociation Constant (Kd) from C2: {Kd_c2:>5.2e} +/- {Kd_c2_std:>3.2e} M')\n",
        "\n",
        "\n",
        "fig17 = go.Figure()\n",
        "fig17.add_trace(go.Bar(name='IE', x=['ΔH', '-TΔS', 'ΔG'], y=[dH, -temperature*IE/1000, dG_IE],\n",
        "                        error_y=dict(type='data', array=[dH_std, -temperature*IE_std/1000, dG_IE_std])))\n",
        "fig17.add_trace(go.Bar(name='c2', x=['ΔH', '-TΔS', 'ΔG'], y=[dH, -temperature*c2/1000, dG_c2],\n",
        "                        error_y=dict(type='data', array=[dH_std, -temperature*c2_std/1000, dG_c2_std])))\n",
        "fig17.update_layout(title='Binding Free Energy', xaxis_title='Components', yaxis_title='Energy (kJ/mol)', barmode='group', width=750, height=500)\n",
        "display(fig17)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "xiOXqrXq66DM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 5-5. (Optional) 준조화근사를 이용한 엔트로피 계산\n",
        "#@markdown **WARNING: 이 방법은 무지무지 오래 걸려요!!**\n",
        "\n",
        "assert 'work' in globals(), \"0번 셀부터 실행하세요.\"\n",
        "assert 'temperature' in globals(), \"5-2를 먼저 실행하세요.\"\n",
        "\n",
        "#@markdown 리간드의 표준 농도 (mol/L)를 정하세요.\n",
        "reference_concentration = 1 #@param {type:\"string\"}\n",
        "bootstrap_power = 3 #@param {type:\"slider\", min:3, max:10, step:1}\n",
        "\n",
        "def estimate_rot_entropy_loss(lig_traj, temperature):\n",
        "    print(\"리간드의 회전 엔트로피 감소분을 계산 중...\", end='')\n",
        "\n",
        "    R = 8.31446261815324  # Gas constant in J/mol/K\n",
        "    N_A = 6.022e23  # Avogadro's number\n",
        "    h = 6.62607004e-34  # Planck's constant in J*s\n",
        "\n",
        "    inertia_tensor = md.compute_inertia_tensor(lig_traj)\n",
        "    moment, axes = np.linalg.eigh(inertia_tensor)\n",
        "    quaternions = Rotation.from_matrix(axes).as_quat()\n",
        "\n",
        "    kde_rot = KernelDensity(kernel='gaussian', bandwidth='scott').fit(quaternions)\n",
        "    S_bound = -R * kde_rot.score_samples(quaternions)\n",
        "    S_bound += 0.5*R*np.log(np.prod(moment.mean(axis=0)))\n",
        "\n",
        "    principal_axes = []\n",
        "    for i in range(lig_traj.n_frames):\n",
        "        I_A, I_B, I_C = inertia_tensor[i,:,:].diagonal()\n",
        "        principal_axes.append(I_A*I_B*I_C)\n",
        "\n",
        "    S_free = []\n",
        "    for i in range(lig_traj.n_frames):\n",
        "        S_free.append(R*(1.5*np.log(temperature)+np.log(principal_axes[i])+1.5*np.log(8*np.pi**2 * R/N_A/h**2)+1.5))\n",
        "\n",
        "    dS = S_bound - np.array(S_free)\n",
        "\n",
        "    print(\"done.\")\n",
        "\n",
        "    return np.mean(dS), np.std(dS)\n",
        "\n",
        "\n",
        "def estimate_trans_entropy_loss(lig_traj, rec_traj, temperature, reference_concentration):\n",
        "    R = 8.31446261815324  # Gas constant in J/mol/K\n",
        "    N_A = 6.022e23  # Avogadro's number\n",
        "\n",
        "    print(f'기준 농도는 {reference_concentration} mol/L 입니다.')\n",
        "    print(\"리간드의 병진 엔트로피 감소분 계산 중...\", end='')\n",
        "    # Calculate free volume based on reference concentration\n",
        "    ref_v  = 1000 * 1e21/ (reference_concentration * N_A)\n",
        "\n",
        "    free_volume = np.zeros(lig_traj.n_frames)\n",
        "\n",
        "    for i in range(lig_traj.n_frames):\n",
        "        free_volume[i] = ref_v - 4/3*np.pi*(md.compute_rg(lig_traj)[i]**3)\n",
        "\n",
        "    lig_com = md.compute_center_of_mass(lig_traj)\n",
        "    rec_com = md.compute_center_of_mass(rec_traj)\n",
        "    bound_volume = 4/3 * np.pi\n",
        "    for i in range(3):\n",
        "        bound_volume *= np.var((lig_com-rec_com)[:,i])\n",
        "\n",
        "    dS_trans = R * np.log(np.full(lig_traj.n_frames, bound_volume) / free_volume)\n",
        "\n",
        "    print(\"done.\")\n",
        "\n",
        "    return np.mean(dS_trans), np.std(dS_trans)\n",
        "\n",
        "\n",
        "def autocorrelation_time(trajectory):\n",
        "    n_frames, n_atoms, _ = trajectory.shape\n",
        "    mean_positions = np.mean(trajectory, axis=0)\n",
        "\n",
        "    displacements = trajectory - mean_positions\n",
        "    sd = np.sum(displacements**2, axis=(1, 2))\n",
        "\n",
        "    msd = np.mean(sd)\n",
        "\n",
        "    y = sd - msd\n",
        "    acorr = np.correlate(y, y, mode='full')[n_frames-1:]\n",
        "    acorr /= acorr[0]\n",
        "    return np.sum(acorr)\n",
        "\n",
        "def extract_independent_frames(coordinates):\n",
        "    \"\"\"\n",
        "    Extract independent frames based on dynamic autocorrelation times.\n",
        "    \"\"\"\n",
        "    n_frames, n_atoms, n_dims = coordinates.shape\n",
        "    independent_indices = []\n",
        "    last_index = 0\n",
        "\n",
        "    while last_index < n_frames:\n",
        "        independent_indices.append(last_index)\n",
        "        remaining_data = coordinates[last_index:,:,:]\n",
        "        if remaining_data.shape[0] > 1:  # Ensure we have enough data to calculate autocorrelation\n",
        "            try:\n",
        "                lag = int(np.ceil(autocorrelation_time(remaining_data)))\n",
        "            except:\n",
        "                break\n",
        "\n",
        "        last_index += lag\n",
        "\n",
        "    return coordinates[independent_indices]\n",
        "\n",
        "def bootstrap_sample(data, n_samples):\n",
        "    indices = np.random.choice(data.shape[0], size=n_samples, replace=True)\n",
        "    return data[indices,:,:]\n",
        "\n",
        "def quasi_harmonic(name, coordinates, n_bootstrap):\n",
        "    R = 8.31446261815324  # Gas constant in J/mol/K\n",
        "\n",
        "    print(f\"준조화근사를 통한 {name}의 진동 엔트로피 계산 중...\")\n",
        "    independent_coords = extract_independent_frames(coordinates)\n",
        "    print(f\"총 {independent_coords.shape[0]} 개의 독립 프레임을 선택했습니다.\")\n",
        "    bootstrapped_coords = bootstrap_sample(independent_coords, n_bootstrap)\n",
        "    reshaped_coords = bootstrapped_coords.reshape(bootstrapped_coords.shape[0], -1)\n",
        "    print(f'공분산행렬 계산중...', end='')\n",
        "    cov_estimator = EmpiricalCovariance()\n",
        "    cov_estimator.fit(reshaped_coords)\n",
        "    cov_matrix = cov_estimator.covariance_\n",
        "\n",
        "    S_corr = 0.5 * np.log(np.linalg.det(2 * np.pi * np.e * cov_matrix))\n",
        "\n",
        "    print('완료.')\n",
        "\n",
        "    return R * S_corr\n",
        "\n",
        "\n",
        "n_bootstrap = 10**bootstrap_power\n",
        "print(f'{n_bootstrap}회 부트스트랩 샘플링을 진행합니다.')\n",
        "\n",
        "dS_rot, dS_rot_std = estimate_rot_entropy_loss(lig_traj, temperature)\n",
        "dS_trans, dS_trans_std = estimate_trans_entropy_loss(lig_traj, rec_traj, temperature, float(reference_concentration))\n",
        "\n",
        "qh_lig_S = quasi_harmonic('ligand', lig_traj.xyz, n_bootstrap)\n",
        "qh_rec_S = quasi_harmonic('receptor', rec_traj.xyz, n_bootstrap)\n",
        "qh_com_S = quasi_harmonic('complex', com_traj.xyz, n_bootstrap)\n",
        "\n",
        "qh_dS_vib = com_S - lig_S - rec_S\n",
        "\n",
        "qh = dS_vib + dS_rot + dS_trans\n",
        "qh_std = np.sqrt(dS_rot_std**2 + dS_trans_std**2)\n",
        "\n",
        "print(f'Estimated Binding Entropy : {qh:.2f} +/- {qh_std:.2f} J/mol/K')\n"
      ],
      "metadata": {
        "id": "MxEOi2oC6FjX",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 사용된 패키지들"
      ],
      "metadata": {
        "id": "B8BgKwrGEKsT"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "* REDUCE\n",
        "> J. M. Word, S. C. Lovell, J. S. Richardson, and D. C. Richardson. (1999) \"Asparagine and glutamine: using hydrogen atom contacts in the choice of side-chain amide orientation\" J. Mol. Biol. 285(4): 1735-1747. DOI:[10.1006/jmbi.1998.2401](https://doi.org/10.1006/jmbi.1998.2401)\n",
        "\n",
        "*  RDKit\n",
        "> RDKit: Open-source cheminformatics. https://www.rdkit.org\n",
        "\n",
        "* OpenBabel\n",
        "> N. M. O'Boyle, M. Banck, C. A. James, C. Morley, T. Vandermeersch, and G. R. Hutchison (2011) \"Open Babel: An open chemical toolbox\" J. Cheminf. 3:33. DOI:[10.1186/1758-2946-3-33](https://doi.org/10.1186/1758-2946-3-33)\n",
        "\n",
        "* PDB2PQR\n",
        "> T. J. Dolinsky, J. E. Nielsen, J. A. McCammon, and N. A. Baker. (2004) \"PDB2PQR: an automated pipeline for the setup, execution, and analysis of Poisson-Boltzmann electrostatics calculations.\" Nucleic Acids Res. 32: W665-667. DOI:[10.1093/nar/gkh381](https://doi.org/10.1093/nar/gkh381)\n",
        "\n",
        "* PropKa\n",
        "> H. Li, A. D. Robertson, and J. H. Jensen. (2005) \"Very Fast Empirical Prediction and Rationalization of Protein pKa Values.\" Proteins, 61: 704-721. DOI:[10.1002/prot.20660](https://doi.org/10.1002/prot.20660)\n",
        "\n",
        "*  OpenMM\n",
        "> P. Eastman, J. Swails, J. D. Chodera, R. T. McGibbon, Y. Zhao, K. A. Beauchamp, L.-P. Wang, A. C. Simmonett, M. P. Harrigan, C. D. Stern, R. P. Wiewiora, B. R. Brooks, and V. S. Pande. (2017) \"OpenMM 7: Rapid development of high performance algorithms for molecular dynamics.” PLOS Comp. Biol. 13(7): e1005659. DOI:[10.1371/journal.pcbi.1005659](https://doi.org/10.1371/journal.pcbi.1005659)\n",
        "\n",
        "*  NGLViewer\n",
        "> H. Nguyen, D. A. Case, and A. S. Rose. (2018) \"NGLview - Interactive molecular graphics for Jupyter notebooks\" Bioinformatics 34(7): 1241-1242. DOI:[10.1093/bioinformatics/btx789](https://doi.org/10.1093/bioinformatics/btx789)\n",
        "\n",
        "* MDTraj\n",
        "> R. T. McGibbon, K. A. Beauchamp, M. P. Harrigan, C. Klein, J. M. Swails, C. X. Hernández, C. R. Schwantes, L-P. Wang, T. J. Lane, and V. S. Pande (2011) \"MDTraj: A Modern Open Library for the Analysis of Molecular Dynamics Trajectories\" Biophys. J. 109(8): 1528-1532. DOI:[10.1016/j.bpj.2015.08.015](https://doi.org/10.1016/j.bpj.2015.08.015)\n",
        "\n",
        "* NumPy\n",
        "> C. R. Harris, K. J. Millman, S. J. van der Walt et al. (2020) \"Array programming with NumPy\" Nature 585, 357-362. DOI:[10.1038/s41586-020-2649-2](https://doi.org/10.1038/s41586-020-2649-2)\n",
        "\n",
        "* Scikit-learn\n",
        "> F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, David Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay (2011) \"Scikit-learn: Machine Learning in Python\" JMLR 12(85): 2825-2830. (https://jmlr.csail.mit.edu/papers/v12/pedregosa11a.html)\n",
        "\n",
        "* Plotly\n",
        "> Plotly Technologies Inc. (https://plot.ly)\n",
        "\n",
        "* ipywidgets\n",
        "> Jupyter widgets community (https://github.com/jupyter-widgets/ipywidgets)\n",
        "\n",
        "* PLIP\n",
        "> S. Salentin, S. Schreiber, V. J. Haupt, M. F. Adasme, and M. Schroeder (2015) \"PLIP: fully automated protein-ligand interaction profiler\" Nucleic Acids Research 43(W1): W443-W447. DOI:[10.1093/nar/gkv315](https://doi.org/10.1093/nar/gkv315)"
      ],
      "metadata": {
        "id": "TkzwxKqOBLZb"
      }
    }
  ]
}