filament_mp_v2.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/roncv/MPss/blob/main/filament_mp_v2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "##MPss - Filament & Membrane Protein\n",
        "\n",
        "Protocol for associatiation of protofilaments to a lipid bilayer with an embedded (trans)membrane protein.\n",
        "\n",
        "The protein structures are converted to  [Martini 3](https://www.nature.com/articles/s41592-021-01098-3) coarse-grained (CG) representations using [Martinize](https://github.com/marrink-lab/vermouth-martinize). The membrane-spanning regions and protein orientation is predicted using [Memembed](https://doi.org/10.1186/1471-2105-14-276). A lipid bilayer is constructed using [Insane](https://doi.org/10.1021/acs.jctc.5b00209). The system is briefly equilibrated using [GROMACS](https://doi.org/10.1016/j.softx.2015.06.001).\n",
        "\n",
        "Change the settings below then click **Runtime → Run all**.  You will be prompted to upload a Soluble Protein .pdb file to simulate.\n",
        "\n",
        "**NB:** To generate filaments the .pdb file supplied should contain crystal lattice vectors. "
      ],
      "metadata": {
        "id": "wJx9txs6NPux"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "gLmjlu0Tp-4k"
      },
      "outputs": [],
      "source": [
        "#@title Define Settings\n",
        "#@markdown #### Set Membrane Type:\n",
        "MembraneType = \"POPE:POPG:CARDIOLIPIN\" #@param [\"POPC\",\"POPE:POPG\",\"POPE:POPG:CARDIOLIPIN\"]\n",
        "PE_PG_ratio = 3.5 #@param {type:\"slider\", min:1, max:10, step:0.1}\n",
        "CARD_conc = 0.05 #@param {type:\"slider\", min:0.0, max:0.3, step:0.01}\n",
        "if CARD_conc == 0.0:\n",
        "    card = False\n",
        "#@markdown ---\n",
        "#@markdown #### Define Orientation of the Membrane Protein:\n",
        "Nterminus = \"in\" #@param [\"in\",\"out\"]\n",
        "#@markdown ---\n",
        "#@markdown #### Set number of Repeats\n",
        "Repeats = 5 #@param {type:\"slider\", min:1, max:25, step:1}\n",
        "#@markdown ---\n",
        "#@markdown #### Rotate Protein:\n",
        "Rotate = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown #### Set Box Dimensions:\n",
        "Box_Width = 30 #@param {type:\"slider\", min:8, max:30, step:1}\n",
        "Box_Height = 22 #@param {type:\"slider\", min:12, max:30, step:1}\n",
        "#@markdown #### Add padding to the ends of filaments:\n",
        "Box_Length_ext = 10 #@param {type:\"slider\", min:0, max:20, step:1}\n",
        "#@markdown ---\n",
        "#@markdown #### Create Filament:\n",
        "Filament = True #@param {type:\"boolean\"}\n",
        "Units = 4 #@param {type:\"slider\", min:1, max:20, step:1}\n",
        "#Continuous = False #@param {type:\"boolean\"}\n",
        "#@markdown ---\n",
        "#@markdown #### Define Distance between Protein and Membrane (nm):\n",
        "Embed = 2 #@param {type:\"slider\", min:-3, max:4, step:0.5}\n",
        "#@markdown ---\n",
        "#@markdown #### Output:\n",
        "Experiment_Name = \"fil_mp\" #@param {type:\"string\"}\n",
        "Download = True #@param {type:\"boolean\"}\n",
        "Zip_Name = \"fil_mp\" #@param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "MsMspm6Ln2SK",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Upload Protein Structure\n",
        "# Upload protein .pdb structure\n",
        "\n",
        "import glob\n",
        "import shutil\n",
        "import subprocess\n",
        "from random import randrange\n",
        "import os\n",
        "import sys\n",
        "from decimal import *\n",
        "import random\n",
        "\n",
        "!python3 -m pip install py3dmol\n",
        "!python3 -m pip install colorama\n",
        "\n",
        "import py3Dmol\n",
        "from colorama import Fore\n",
        "from google.colab import files\n",
        "sys.path.append('/usr/local/lib/python3.7/site-packages/')\n",
        "\n",
        "def upload_proteins(dirname='filament_mp'):\n",
        "    \"\"\"\n",
        "    ...\n",
        "    \"\"\"\n",
        "    os.chdir('/content/')\n",
        "    print(Fore.BLUE + \"\\nUpload Membrane Protein Structure:\\n\")\n",
        "    upload1 = files.upload() # request upload of .pdb - membrane\n",
        "\n",
        "    print(Fore.BLUE + \"\\nUpload Soluble Protein Structure:\\n\")\n",
        "    upload2 = files.upload() # request upload of .pdb - soluble\n",
        "\n",
        "    # Arrange working directories\n",
        "    filename1 = next(iter(upload1))\n",
        "    name1 = os.path.splitext(filename1)[0]\n",
        "    filename2 = next(iter(upload2))\n",
        "    name2 = os.path.splitext(filename2)[0]\n",
        "\n",
        "    name = f\"{name1}-{name2}-{dirname}\" # combine names\n",
        "    working_dir = f'/content/{name}/'\n",
        "\n",
        "    if os.path.exists(working_dir):\n",
        "        shutil.rmtree(working_dir)\n",
        "    os.makedirs(working_dir)\n",
        "\n",
        "    # Protein 1 - Membrane\n",
        "    print(Fore.GREEN + \"\\nMembrane Protein:\\n\")\n",
        "    os.rename(filename1, working_dir + filename1)\n",
        "    os.chdir(working_dir)\n",
        "    mol1 = open(working_dir + filename1, 'r').read()\n",
        "    mview = py3Dmol.view(width=500,height=280)\n",
        "    mview.addModel(mol1,'pdb')\n",
        "    mview.setStyle({'cartoon':{'color':'spectrum'}})\n",
        "    mview.setBackgroundColor('0xffffff')\n",
        "    mview.zoomTo()\n",
        "    mview.show()\n",
        "\n",
        "    # Protein 2 - Soluble\n",
        "    print(Fore.GREEN + \"\\nSoluble Protein (Filament):\\n\")\n",
        "    os.chdir('/content/')\n",
        "    os.rename(filename2, working_dir + filename2)\n",
        "    os.chdir(working_dir)\n",
        "    mol2 = open(working_dir + filename2, 'r').read()\n",
        "    mview = py3Dmol.view(width=500,height=280)\n",
        "    mview.addModel(mol2,'pdb')\n",
        "    mview.setStyle({'cartoon':{'color':'spectrum'}})\n",
        "    mview.setBackgroundColor('0xffffff')\n",
        "    mview.zoomTo()\n",
        "    mview.show()\n",
        "\n",
        "    return working_dir, name, filename1, filename2\n",
        "\n",
        "working_dir, name, filename1, filename2 = upload_proteins(dirname=Experiment_Name) # usage\n",
        "\n",
        "# Directory organsiation...\n",
        "os.chdir(working_dir)\n",
        "for file in glob.glob(r'/content/martini*.itp'): # martini .itp files\n",
        "    print(file)\n",
        "    shutil.copy(file, working_dir)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "OQYnJxxCnPzc",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Install Dependencies\n",
        "# Install Dependencies\n",
        "%%capture\n",
        "os.chdir('/content/')\n",
        "if not os.path.isdir(\"/content/memembed/\"):\n",
        "  !apt-get update -y\n",
        "  !apt-get install dssp\n",
        "  !git clone https://github.com/timnugent/memembed\n",
        "  %cd memembed\n",
        "  !make\n",
        "  %cd ../\n",
        "  !python3 -m pip install pdb2pqr\n",
        "  !python3 -m pip install vermouth\n",
        "  !python3 -m pip install GromacsWrapper\n",
        "  !python3 -m pip install MDAnalysis\n",
        "  !git clone https://github.com/pstansfeld/cg2at\n",
        "  !wget https://github.com/pstansfeld/MemProtMD/raw/main/martini_v300.zip\n",
        "  !unzip -o martini_v300.zip\n",
        "  !wget 'http://cgmartini.nl/images/applications/water/water.gro'\n",
        "  !wget https://github.com/pstansfeld/MemProtMD/raw/main/insane3.py\n",
        "  !wget https://github.com/pstansfeld/MemProtMD/raw/main/gromacs.zip\n",
        "  !unzip -o gromacs.zip\n",
        "  !ln -s /content/content/gromacs/bin/gmx /usr/bin/gmx\n",
        "\n",
        "import gromacs\n",
        "import MDAnalysis\n",
        "from MDAnalysis import *\n",
        "from MDAnalysis.lib.distances import self_capped_distance\n",
        "from MDAnalysis.lib.distances import capped_distance\n",
        "import numpy as np\n",
        "from numpy import *\n",
        "import pandas as pd\n",
        "pd.options.display.float_format = '{:.5f}'.format"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0ib-qDQA1T1Q",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Get lipid concentrations\n",
        "# Function to get clean ratio of lipids for a given concentration of CARD\n",
        "def get_lipid_ratio(pe_pg_ratio=3.5, pg_conc=None, total=10, card_conc=0.1, tolerance=0.1):\n",
        "    \"\"\"\n",
        "    Function to get clean a ratio of lipids for a given concentration of cardiolipin.\n",
        "\n",
        "    Parameters\n",
        "    \n",
        "    pe_pg_ratio : Ratio of PE:PG in PE:PG or PE:PG:CARD membrane.\n",
        "    pg_conc : Float between 0 and 1. Conc. of pg within a PE:PG only membrane. Only takes 2.d.p.\n",
        "    total : Target ratio sum total. May exceed depending on error tolerance.\n",
        "    card_conc : Concentration of CARD in a PE:PG:PE membrane.\n",
        "    tolerance : Error tolerance for ratio.\n",
        "    \"\"\"\n",
        "    \n",
        "    # No CARD\n",
        "    if card_conc == 0:\n",
        "        if pg_conc is None:\n",
        "            pg_conc = 1 - (pe_pg_ratio / (pe_pg_ratio+1))\n",
        "            #print(f\"PG conc: {round(pg_conc, 2)}\")\n",
        "        elif pg_conc == 0:\n",
        "            #print(f\"PG conc: {round(pg_conc, 2)}\")\n",
        "            ratio = [1, 0, 0]\n",
        "            return ratio\n",
        "\n",
        "        #if tolerance == 0: # maybe ignore this when conc is given\n",
        "        pg_conc = Decimal(f\"{pg_conc}\")\n",
        "        pg_conc = round(pg_conc, 2)\n",
        "        #print(f\"PG conc: {pg_conc}\")\n",
        "        pg_int = int(pg_conc * 100)\n",
        "        pe_int = 100 - pg_int\n",
        "        pe_pg_ratio = round(pe_int/pg_int, 1)\n",
        "        #print(f\"PE:PG ratio - {pe_pg_ratio} : 1\")\n",
        "        ratio = [pe_int, pg_int, 0]\n",
        "\n",
        "        # Use HCF, LCM, GCD ¿to utilise totals? and minimise this value ***\n",
        "        # https://softwareengineering.stackexchange.com/questions/416627/how-to-find-the-nearest-integer-number-that-when-multiplied-by-a-set-of-decimals\n",
        "    \n",
        "    # With CARD\n",
        "    else:\n",
        "        card_conc = Decimal(f\"{card_conc}\")\n",
        "\n",
        "        # CARD concentration\n",
        "        card_flt = card_conc * total # float result\n",
        "        card_rnd = int(round(card_conc * total, 0)) # rounded result\n",
        "        card_err = np.sqrt(np.square(card_flt-card_rnd)) # difference\n",
        "\n",
        "        # Fit to tolerance\n",
        "        r = 0 # counter\n",
        "        if float(card_err) <= tolerance:\n",
        "            pass\n",
        "        else:\n",
        "            while float(card_err) > tolerance:\n",
        "                total+=1; r+=1 # counters\n",
        "                card_flt = card_conc * total\n",
        "                card_rnd = int(round(card_conc * total, 0))\n",
        "                card_err = np.sqrt(np.square(card_flt-card_rnd))\n",
        "        # Set CARD count\n",
        "        card = card_rnd\n",
        "\n",
        "        # PE + PG conc\n",
        "        pe = int(2*(total//(2*(pe_pg_ratio+1)))*pe_pg_ratio)\n",
        "        pg = total - (pe+card_rnd)\n",
        "        ratio = [pe, pg, card]\n",
        "\n",
        "    return ratio\n",
        "\n",
        "#ratio = get_lipid_ratio(pe_pg_ratio=PE_PG_ratio, total=10, card_conc=CARD_conc, tolerance=0.01) # usage\n",
        "#ratio = get_lipid_ratio(pe_pg_ratio=3.5, pg_conc=0.3, total=10, card_conc=0) # usage\n",
        "#ratio"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YS8dQRR50OUi",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Functions\n",
        "# title Get area per lipid\n",
        "def get_area_per_lipid(ratio):\n",
        "    \"\"\"\n",
        "    Function to get the area per lipid for a given lipid ratio.\n",
        "    \"\"\"\n",
        "    # Constants - area per lipid\n",
        "    pe_apl = 0.63\n",
        "    pg_apl = 0.67\n",
        "    cl_apl = 1.23\n",
        "\n",
        "    # Weighted densities\n",
        "    pe_wd = ratio[0]/sum(ratio)*pe_apl\n",
        "    pg_wd = ratio[1]/sum(ratio)*pg_apl\n",
        "    cl_wd = ratio[2]/sum(ratio)*cl_apl\n",
        "\n",
        "    # Sum up for avg. density\n",
        "    avg_apl = sum([pe_wd,pg_wd,cl_wd])\n",
        "\n",
        "    return avg_apl\n",
        "\n",
        "#get_area_per_lipid(ratio) # usage\n",
        "\n",
        "\n",
        "# ANSI Escape Code styling\n",
        "def esc(codes, string:str):\n",
        "    \"\"\"\n",
        "    Function to generate ANSI escape codes for customizing cell output.\n",
        "    \"\"\"\n",
        "    if isinstance(codes, list): # process list format\n",
        "        # Generate text with input string, then format as ANSI escape code\n",
        "        return f\"\\033[{';'.join(map(str, codes))}m{string}\\033[0m\"\n",
        "    if isinstance(codes, int): # process as an integer\n",
        "        return f'\\033[{codes}m{string}\\033[0m'\n",
        "    else:\n",
        "        raise TypeError(\"ANSI codes must be provided as a list or integer\")\n",
        "\n",
        "\n",
        "# Get Cell Dimensions\n",
        "def get_dimensions(structure):\n",
        "    \"\"\"\n",
        "    Function to get dimensions of a system's cell from a .pdb file.\n",
        "    \"\"\"\n",
        "    # Create universe with structure file...\n",
        "    u = MDAnalysis.Universe(structure)\n",
        "    x = round(float(u.dimensions[0]/10)) # in angstroms\n",
        "    y = round(float(u.dimensions[1]/10))\n",
        "    z = round(float(u.dimensions[2]/10))\n",
        "    cube = [x,y,z]\n",
        "\n",
        "    return x, y, z, cube\n",
        "    \n",
        "#x, y, z, cube = get_dimensions('CG-system-extended.gro') # usage\n",
        "\n",
        "\n",
        "# Format topology file\n",
        "def format_topology(topology, Units=1):\n",
        "    \"\"\"\n",
        "    Format a topology file to include all .itp files.\n",
        "    \"\"\"\n",
        "    replacements = {'Protein        1\\n':f'ProteinA      {Units}' +'\\n','NA+':'NA',\n",
        "                    'CL-':'CL',\n",
        "                    '#include \"martini_v3.itp\"':'#include \"martini_v3.0.0.itp\"\\n#include \"martini_v3.0.0_ions_v1.itp\"\\n#include \"martini_v3.0.0_solvents_v1.itp\"\\n#include \"martini_v3.0.0_phospholipids_v1.itp\"\\n'}\n",
        "    lines = []\n",
        "    with open(topology) as infile:\n",
        "        for line in infile:\n",
        "            for src, target in replacements.items():\n",
        "                line = line.replace(src, target)\n",
        "            lines.append(line)\n",
        "    with open(topology, 'w') as outfile:\n",
        "        for line in lines:\n",
        "            outfile.write(line)\n",
        "\n",
        "#format_topology('topol_test2.top') # usage\n",
        "\n",
        "\n",
        "# Position protein\n",
        "def position_protein_insane(structure, selection='name BB or name SC*',\n",
        "                            monolayer_width=1.95):\n",
        "    \"\"\"\n",
        "    Finds the COM and minimum z-coordinate of protein structure given. Difference \n",
        "    in values is added to the average monolayer width. Made to work with the insane3.py \n",
        "    -center and -dm flags, where the functions output used as the -dm value results in \n",
        "    the protein just touching the membrane. To embed the protein more or less, increase \n",
        "    or decrease the value accordingly.\n",
        "\n",
        "    Distances in nm.\n",
        "    \"\"\"\n",
        "    # Create universe\n",
        "    u = MDAnalysis.Universe(structure)\n",
        "\n",
        "    # Protein dimensions in z-axis\n",
        "    z_coords_p = u.select_atoms(f'{selection}').positions[:, 2] # select protein\n",
        "    p_min_z = float(z_coords_p.min())/10.0\n",
        "    com_p = (u.select_atoms(f'{selection}').center_of_mass()/10.0)[2]\n",
        "    adjustment = (com_p - p_min_z) + monolayer_width # in nm\n",
        "\n",
        "    return adjustment\n",
        "\n",
        "#position_protein_insane('protein-em.pdb') # usage\n",
        "\n",
        "\n",
        "# Topology file entries\n",
        "def insane_topadd(top_add, og_top):\n",
        "    \"\"\"\n",
        "    Function to add new entries to a topology file.\n",
        "    \"\"\"\n",
        "    with open(og_top, \"a\") as output_file:\n",
        "        with open(top_add) as input_data:\n",
        "            # Find additional entries for topology\n",
        "            for line in input_data:\n",
        "                if line.strip() == 'Protein        1':\n",
        "                    break\n",
        "            for line in input_data:  # read + write them to original file\n",
        "                output_file.write(line.rstrip('\\n').strip()+\"\\n\")\n",
        "\n",
        "#insane_topadd('topol_ext.top', 'topol1.top') # usage\n",
        "\n",
        "\n",
        "# Interrupt cell execution at specified opint\n",
        "class StopExecution(Exception):\n",
        "    def _render_traceback_(self):\n",
        "        pass\n",
        "# Function call for cell interruption\n",
        "def exit_c(): raise StopExecution\n",
        "\n",
        "#exit_c() # usage\n",
        "\n",
        "\n",
        "# Write EM mdp file\n",
        "def mdp_em(filename='em'):\n",
        "    \"\"\"\n",
        "    Write an .mdp file for EM\n",
        "    \"\"\"\n",
        "    with open(f'{filename}.mdp','w') as em:\n",
        "                em.write('integrator = steep\\nnsteps = 5000\\nemtol = 100\\nemstep = 0.001')\n",
        "\n",
        "#mdp_em() # usage\n",
        "\n",
        "\n",
        "# Format .pdb\n",
        "def format_pdb_residues(pdb):\n",
        "    \"\"\"\n",
        "    Format residues of a .pdb.\n",
        "    \"\"\"\n",
        "    with open(pdb, 'r') as file :\n",
        "        filedata = file.read()\n",
        "    filedata = filedata.replace('HSE', 'HIS')\n",
        "    filedata = filedata.replace('HSD', 'HIS')\n",
        "    filedata = filedata.replace('MSE', 'MET')\n",
        "    filedata = filedata.replace(' SE ', ' SD ')\n",
        "    \n",
        "    # Update .pdb\n",
        "    with open(pdb, 'w') as file:\n",
        "        file.write(filedata)\n",
        "\n",
        "#format_pdb_residues(pdb='protein1.pdb') # usage\n",
        "\n",
        "\n",
        "# Write a topology file\n",
        "def write_topology(idx, idx_a, n=1):\n",
        "    \"\"\"\n",
        "    Write Topology file for protein.\n",
        "    \"\"\"\n",
        "    with open('topol.top','w') as top:\n",
        "                top.write('#include \"martini_v3.0.0.itp\"\\n'\n",
        "                            '#include \"martini_v3.0.0_ions_v1.itp\"\\n'\n",
        "                            '#include \"martini_v3.0.0_solvents_v1.itp\"\\n'\n",
        "                            '#include \"martini_v3.0.0_phospholipids_v1.itp\"\\n'\n",
        "                            f'#include \"protein{idx}-cg.itp\"\\n[ system ]\\n[ molecules ]\\nProtein{idx_a} {n}')\n",
        "\n",
        "    os.system(f'cp topol.top topol_protein{idx_a}.top')\n",
        "\n",
        "#write_topology(idx=1, idx_a='A') # usage\n",
        "\n",
        "\n",
        "# Visualise System\n",
        "def visualise_system(file, _type='gro', waters=False, ions=False, dir=working_dir):\n",
        "    \"\"\"\n",
        "    Visualise systems in 3D using Py3Dmol.\n",
        "    \"\"\"\n",
        "    if file is not None: # specified filepath (rel/full)\n",
        "        systems = [file]\n",
        "        _type = file[-3:] # get file extension\n",
        "    else:\n",
        "        # Access first repeat of each experiment - automated\n",
        "        if _type == 'pdb':\n",
        "            systems = glob.glob(os.path.expanduser(f'{dir}*/CG-system1.{_type}'))\n",
        "        else:\n",
        "            systems = glob.glob(os.path.expanduser(f'{dir}*/EQ1/eq.{_type}'))\n",
        "    # Visualise\n",
        "    for i in systems:\n",
        "        print(f'{Fore.BLUE}\\n{i}\\n') # convert to f strings ***\n",
        "        mview = py3Dmol.view(width=600,height=450)\n",
        "        mol1 = open(i, 'r').read()\n",
        "        mview.addModel(mol1,f'{_type}')\n",
        "        mview.setStyle({'cartoon':{'color':'spectrum'}})\n",
        "        mview.setStyle({'atom':'PO1','atom':'PO2','atom':'PO3','atom':'PO4'},\n",
        "                       {'sphere':{'color':'0xE63629', 'radius':2.2,'quality':20}}) # lipids\n",
        "        if waters is True:\n",
        "            mview.setStyle({'atom':'W'},{'sphere':{'color':'0x5014F5','quality':20}}) # waters\n",
        "        if ions is True:\n",
        "                try:\n",
        "                    mview.setStyle({'atom':'CL'},{'sphere':{'color':'0x00E649','radius':1.5,'quality':20}}) # CL\n",
        "                except:\n",
        "                    pass\n",
        "                try:\n",
        "                    mview.setStyle({'atom':'NA'},{'sphere':{'color':'0xE644D5','radius':1.5,'quality':20}}) # NA\n",
        "                except:\n",
        "                    pass\n",
        "        mview.setStyle({'atom':'BB'},\n",
        "                       {'sphere':{'color':'0x5FC8F5', 'radius':3,'quality':20}}) # protein\n",
        "        try:\n",
        "            mview.setStyle({'chain':'B', 'atom':'BB'},\n",
        "                           {'sphere':{'color':'0x17E6DB', 'radius':3,'quality':20}}) # protein\n",
        "        except:\n",
        "            pass\n",
        "        mview.setBackgroundColor('0xffffff')\n",
        "        mview.addUnitCell() # only with .pdb files\n",
        "        mview.zoomTo()\n",
        "        mview.show()\n",
        "        \n",
        "        # Remove temporary files\n",
        "        for file in glob.glob(r'#*'):\n",
        "            os.remove(file)\n",
        "    # Change back to working directory\n",
        "    #os.chdir(dir)\n",
        "    # Remove temporary files\n",
        "    #for file in glob.glob(r'*/#*'):\n",
        "    #    os.remove(file)\n",
        "    \n",
        "#visualise_system('pdb', waters=False, ions=False) # use .pdb to see unit cell"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AL-a83ArtV34",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Predict Membrane Protein Orientations\n",
        "def predict_mp_orientation(Nterminus='in', structure=filename1, o='memembed1', dir=working_dir):\n",
        "    \"\"\"\n",
        "    ...\n",
        "    \"\"\"\n",
        "    os.chdir(dir)\n",
        "    os.system(f'/content/memembed/bin/memembed -n {Nterminus} -o {o}.pdb {structure}')\n",
        "    print(f\"{esc(91, f'periplasmic')}\\n{esc(91, f'oooo   oooo')}\\n¦¦¦¦ {esc([97,100], f' ')} ¦¦¦¦\\n{esc(34, f'oooo   oooo')}\\n{esc(34, f'cytoplasmic')}\\n\")\n",
        "    #mview = py3Dmol.view(width=800,height=400)  \n",
        "    #mol1 = open(f'{dir}/{o}.pdb', 'r').read()\n",
        "    #mview.addModel(mol1)\n",
        "    #mview.setStyle({'cartoon':{'color':'spectrum','quality':20}})\n",
        "    #mview.setStyle({'resn':'DUM'},{'sphere':{'radius':0.4,'quality':20}})\n",
        "    #mview.setBackgroundColor('0xffffff')\n",
        "    #mview.zoomTo()\n",
        "    #mview.show()\n",
        "\n",
        "#predict_mp_orientation(Nterminus=Nterminus, structure=filename1) # usage\n",
        "\n",
        "# \n",
        "def position_membrane_protein(mem_prot, sol_prot, idx=1, Box_Width=Box_Width,\n",
        "                              Box_Height=Box_Width):\n",
        "    \"\"\"\n",
        "    ...\n",
        "    \"\"\"\n",
        "    cube = [Box_Width,Box_Width,Box_Height]\n",
        "\n",
        "    gromacs.make_ndx(f=mem_prot, o=f'index{idx}.ndx', input=('del 0', 'del 1-100','rDUM','q'),backup=False)\n",
        "    gromacs.editconf(f=mem_prot,o=f'centered{idx}.pdb',n=f'index{idx}.ndx',c=True,box=cube,input=(0,0),backup=False)\n",
        "\n",
        "    x, y, z, cube = get_dimensions(f'centered{idx}.pdb')\n",
        "\n",
        "    gromacs.confrms(f2=mem_prot, f1=f'centered{idx}.pdb', one=True, o=f'aligned{idx}.gro', input=(3,3),backup=False)\n",
        "    gromacs.editconf(f=f'aligned{idx}.gro', o=f'protein{idx}.pdb',label='A',resnr=1,n=f'index{idx}.ndx',box=cube,input=(0,0),backup=False)\n",
        "\n",
        "    v = MDAnalysis.Universe(f'aligned{idx}.gro')\n",
        "    dum = v.select_atoms('resname DUM')\n",
        "\n",
        "    #print(\"\\n\", [x, y, z], \"\\n\")\n",
        "    dm = (z/2) - (round(dum.center_of_mass()[2]/10))\n",
        "\n",
        "    return dm, cube\n",
        "\n",
        "#dm, cube = position_membrane_protein(mem_prot='memembed1.pdb', sol_prot=filename2, idx=1) # usage"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pfMr18MiGoJN",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Build Systems\n",
        "def coarse_grain_protein(idx, idx_a, ElasticCutoff=0.7,cube=None,rotation=None,translation=None, Continuous=False, dir=working_dir): # cube=cube - clarity?\n",
        "    # martini .itp files\n",
        "    for file in glob.glob(r'/content/martini_v300/*.itp'):\n",
        "        #print(file)\n",
        "        shutil.copy(file, dir)\n",
        "    \n",
        "    # GENERATE COARSE GRAINED PROTEIN\n",
        "    print(f\"\\nCoarse graining protein {idx}({idx_a})...\")\n",
        "    merge = 'A' \n",
        "    eunit = 'molecule' # 'molecule\n",
        "    os.system(f'martinize2 -f protein{idx}.pdb -dssp mkdssp -ff martini3001 '\n",
        "              f'-x protein{idx}-cg.pdb -o protein{idx}-cg.top -elastic -ef 500 -eu {str(ElasticCutoff)} ' \n",
        "              f'-el 0.5 -ea 0 -ep 0 -merge {merge} -maxwarn 100000 -eunit {eunit}') # see eunit + merge (chains...)\n",
        "    print(\"Done.\")\n",
        "\n",
        "    os.system(f\"sed -e 's/^molecule.*/Protein{idx_a} 1/g' molecule_0.itp | grep -v ';' >  protein{idx}-cg.itp\")\n",
        "\n",
        "    # Continuity restraints - add continuity to filament elastic network \n",
        "    if Continuous == True:\n",
        "        u = MDAnalysis.Universe(f'protein{idx}-cg.pdb')\n",
        "        boxed = u.dimensions\n",
        "\n",
        "        with open(f\"protein{idx}-cg.itp\") as f:\n",
        "            contents1, constraints, contents2 = f.read().partition(\"[ constraints ]\\n\")\n",
        "\n",
        "        # Restraints - add continuity to filament elastic network\n",
        "        U = MDAnalysis.Universe(f'protein{idx}-cg.pdb')\n",
        "        BB = U.select_atoms('name BB')\n",
        "        d1 = pd.DataFrame(np.column_stack(self_capped_distance(BB.positions,max_cutoff=7,min_cutoff=5,box=boxed,method='pkdtree')))\n",
        "        d2 = pd.DataFrame(np.column_stack(self_capped_distance(BB.positions,max_cutoff=7,min_cutoff=5,method='pkdtree')))\n",
        "        zid = pd.DataFrame(BB.ix+1)\n",
        "        d1 = d1.sort_values(by=[0, 1])\n",
        "        d2 = d2.sort_values(by=[0, 1])\n",
        "        df = pd.concat([d1,d2])\n",
        "        df = df.sort_values(by=[0, 1])\n",
        "        df = df.drop_duplicates(keep=False,subset=(0,1))\n",
        "        df[2] = (df[2]/10).map('{:.5f}'.format)\n",
        "        zi = dict(zip(zid.index,zid[0]))\n",
        "        df[3] = df[0].replace(zi).map('{:.0f}'.format)\n",
        "        df[4] = df[1].replace(zi).map('{:.0f}'.format)\n",
        "        df[5] = 1\n",
        "        df[6] = 500\n",
        "        df = df[df[0]<df[1]]\n",
        "        infinity = df.to_csv(columns=(3, 4, 5, 2, 6),sep='\\t', index=False, header=False)    \n",
        "\n",
        "        # New .itp parameters\n",
        "        with open(f\"protein{idx}-cg.itp\", \"w\") as new:\n",
        "            new.write(contents1 + infinity + constraints + contents2)\n",
        "\n",
        "    #print(\"\\nWriting files...\")\n",
        "    file_object = open(f'protein{idx}-cg.itp', 'a')\n",
        "    file_object.write(f'#ifdef POSRES\\n#include \"protein{idx}-posre.itp\"\\n#endif\\n')\n",
        "    file_object.close()\n",
        "\n",
        "    #print(\"\\nRestraints...\") # use of this??\n",
        "    gromacs.genrestr(f=f'protein{idx}-cg.pdb',o=f'protein{idx}-posre.itp',input=(1,1)) # **??\n",
        "    \n",
        "    # EM .mdp \n",
        "    mdp_em()\n",
        "    os.system(f'cp em.mdp em{idx}.mdp')\n",
        "\n",
        "    # Setup and run minimisation\n",
        "    #print(\"\\nTidy file system...\")\n",
        "    try:\n",
        "        os.remove(f'em{idx}*')\n",
        "    except:\n",
        "        pass\n",
        "\n",
        "    if rotation is not None:\n",
        "        #print(\"\\nSingle condition...\")\n",
        "        gromacs.editconf(f=f'protein{idx}-cg.pdb', o=f'protein{idx}-box.pdb', c=True, box=cube, rotate=rotation)\n",
        "    else:\n",
        "        #print(\"\\nNone condition...\")\n",
        "        gromacs.editconf(f=f'protein{idx}-cg.pdb', o=f'protein{idx}-box.pdb', c=True, box=cube)\n",
        "    print(f\"\\nMinimizing protein {idx}({idx_a})...\")\n",
        "    gromacs.grompp(f=f'em{idx}.mdp',o=f'em{idx}.tpr',c=f'protein{idx}-box.pdb',maxwarn='-1',backup=False,v=True)\n",
        "    gromacs.mdrun(deffnm=f'em{idx}', c=f'protein{idx}-em.pdb',backup=False)\n",
        "    if translation is not None:\n",
        "        gromacs.editconf(f=f'protein{idx}-em.pdb', o=f'protein{idx}-em.pdb', translate=translation)\n",
        "    print(\"Done.\\n\")\n",
        "\n",
        "#coarse_grain_protein(idx=1, idx_a='A') # usage\n",
        "\n",
        "def build_proteins(mp, sp, Units=4, rotate_sp=180, Box_Width=Box_Width,\n",
        "                   Box_Height=Box_Height, Box_Length_ext=None, Embed=0,\n",
        "                   Continuous=False, Filament=Filament, dir=working_dir):\n",
        "    \"\"\"\n",
        "    Parameters:\n",
        "\n",
        "    mp: str, path object or file-like object\n",
        "    Membrane protein\n",
        "\n",
        "    sp: str, path object or file-like object\n",
        "    Soluble protein\n",
        "\n",
        "    Units: int\n",
        "    Number of units in the protofilament\n",
        "\n",
        "    rotate_sp: int\n",
        "    Rotate filament (sp), 0 for no rotation\n",
        "    \"\"\"\n",
        "\n",
        "    # Predict membrane orientation of input mp\n",
        "    predict_mp_orientation(Nterminus=Nterminus, structure=mp, o='memembed1', dir=dir) # usage - red is periplasmic\n",
        "\n",
        "    # Get membrane positining\n",
        "    dm, cube = position_membrane_protein(mem_prot='memembed1.pdb', sol_prot=sp, idx=1)\n",
        "\n",
        "    # Format mp residues\n",
        "    format_pdb_residues(pdb='protein1.pdb')\n",
        "\n",
        "    # Write mp topology\n",
        "    write_topology(idx=1, idx_a='A')\n",
        "\n",
        "    # Convert mp to CG rep\n",
        "    coarse_grain_protein(idx=1, idx_a='A', cube=cube, dir=dir)\n",
        "\n",
        "    # Generate protofilaments from soluble protein\n",
        "    if Filament == True:\n",
        "        gromacs.genconf(f=sp,o='ready2.pdb',nbox=[Units,1,1],backup=False,renumber=True)\n",
        "        x, y, z, _ = get_dimensions('ready2.pdb')\n",
        "        if Box_Length_ext is not None:# or Box_Length_ext != 0:\n",
        "            cube = [x+Box_Length_ext, Box_Width, Box_Height] # add x-padding\n",
        "            Continuous = False # prevent continuous filament with x-padding\n",
        "        else:\n",
        "            cube = [x, Box_Width, Box_Height]\n",
        "        gromacs.editconf(f='protein1-em.pdb',o='protein1-em.pdb',c=True,box=cube) # format membrane protein cell size\n",
        "    else:\n",
        "        gromacs.genconf(f=sp,o='ready2.pdb',nbox=[1,1,1],backup=False,renumber=True)  \n",
        "        cube = [Box_Width,Box_Width,Box_Height] \n",
        "\n",
        "    print(f\"\\nDimensions: {cube}\")\n",
        "    \n",
        "    # Make .ndx for filament\n",
        "    gromacs.make_ndx(f='ready2.pdb', o='index2.ndx',input=('del 0','del 1-100','rDUM','q'),backup=False)\n",
        "    # Center in defined box size\n",
        "    gromacs.editconf(f='ready2.pdb',o='protein2.pdb',n='index2.ndx',c=True,box=cube,input=(0,0),backup=False,label='A',resnr=1)\n",
        "    format_pdb_residues(pdb='protein2.pdb') # format residues\n",
        "    write_topology(idx=2, idx_a='B') # write filament .top file\n",
        "\n",
        "    # Convert protofilament to CG rep + rotate if necessary\n",
        "    coarse_grain_protein(idx=2,idx_a='B',rotation=[rotate_sp,0,0],\n",
        "                         cube=cube,dir=dir, Continuous=Continuous)\n",
        "    \n",
        "    # Position soluble protein in box\n",
        "    p_height = position_protein_insane('protein2-em.pdb')\n",
        "    Distance = Embed + p_height\n",
        "    print(f'SP Distance: {Distance}\\n')\n",
        "    gromacs.editconf(f=f'protein2-em.pdb', o=f'protein2-em.pdb', translate=[0,0,-Distance])\n",
        "\n",
        "    # Match cell sizes for both proteins\n",
        "    u = MDAnalysis.Universe('protein2-em.pdb')\n",
        "    cube = list((u.dimensions[:3]/10)) # get cell dimensions of EM'd filament\n",
        "    gromacs.editconf(f='protein1-em.pdb', o='protein1-em.pdb', c=True, box=cube) # apply dimensions to mp\n",
        "\n",
        "    return\n",
        "\n",
        "#build_proteins(mp=filename1, sp=filename2, Units=Units, Continuous=False, Box_Length_ext=Box_Length_ext) # usage"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "DnHS8YMINkKt",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Equillibrate Systems\n",
        "# %%script false --no-raise-error\n",
        "def setup_cg_system(mp, sp, Units=4, rotate_sp=180, rotate_mp=True,\n",
        "                    MembraneType=MembraneType, Repeats=Repeats, ElasticCutoff=0.7,\n",
        "                    NamingConvention=None, Box_Width=Box_Width, Box_Height=Box_Height,\n",
        "                    Box_Length_ext=None, Embed=0, Continuous=Continuous,\n",
        "                    Ratio=None, dir=working_dir):\n",
        "    \"\"\"\n",
        "    Function to set-up coarse-grained membrane-protein systems.\n",
        "    \"\"\"\n",
        "    # rotate mp\n",
        "    # translate mp\n",
        "\n",
        "    print(f'\\n{esc([100, 97], \"INITIALIZE\")}\\n')\n",
        "\n",
        "    # Move to working directory\n",
        "    os.chdir(dir)\n",
        "\n",
        "    # Options\n",
        "    if NamingConvention == \"elastic_cutoff\": # ***\n",
        "        name_nc = f'-eu{ElasticCutoff}'\n",
        "    elif NamingConvention == None:\n",
        "        name_nc = ''\n",
        "    else:\n",
        "        name_nc = NamingConvention\n",
        "\n",
        "    # LIPID COMPOSITION\n",
        "    if isinstance(Ratio, list):\n",
        "        pe, pg, crd = Ratio\n",
        "        if crd == 0:\n",
        "            lipid_name = f'pe{pe}-pg{pg}'\n",
        "            lipid = f'-l POPE:{pe} -l POPG:{pg}'\n",
        "        else:\n",
        "            lipid_name = f'pe{pe}-pg{pg}-c{crd}'\n",
        "            lipid = f'-l POPE:{pe} -l POPG:{pg} -l CARD:{crd}'\n",
        "        apl = get_area_per_lipid(Ratio)\n",
        "        print(f\"Lipid Ratio: {Ratio}\")\n",
        "    else:\n",
        "        raise TypeError(\"Invalid lipid selection!\")\n",
        "\n",
        "    # Organize Directories - Including naming conventions\n",
        "    name_filaments = f\"f{Units}\"\n",
        "    name_embed = f\"eb{Embed}\"\n",
        "    subdir = f\"{dir}{name}-{lipid_name}-{name_filaments}-{name_embed}{name_nc}/\"\n",
        "    print(subdir)\n",
        "    shutil.rmtree(subdir, ignore_errors=True)\n",
        "    os.makedirs(subdir, exist_ok=True)\n",
        "    shutil.copyfile(mp, subdir+mp)\n",
        "    shutil.copyfile(sp, subdir+sp)\n",
        "    os.chdir(subdir)\n",
        "\n",
        "    # BUILD PROTEINS\n",
        "    print(\"\\nBuilding Proteins...\\n\")\n",
        "    build_proteins(mp=mp,sp=sp,Units=Units,Continuous=Continuous,Embed=Embed,\n",
        "                   Box_Length_ext=Box_Length_ext,rotate_sp=rotate_sp,Box_Width=Box_Width,\n",
        "                   Box_Height=Box_Height,Filament=Filament, dir=subdir) # usage\n",
        "\n",
        "    # CONTINUITY\n",
        "    # .mdp settings\n",
        "    if Continuous == True:\n",
        "        pcoupltype = 'anisotropic'\n",
        "        compressibility = '3e-4 3e-4 3e-4 0 0 0'\n",
        "        ref_p = '1.0 1.0 1.0 0 0 0'\n",
        "        # New .itp parameters ** not yet functional\n",
        "        #with open(\"protein-cg.itp\", \"w\") as new:\n",
        "        #    new.write(contents1 + infinity + constraints + contents2)\n",
        "    else:\n",
        "        pcoupltype = 'semiisotropic'\n",
        "        compressibility = '3e-4 3e-4'\n",
        "        ref_p = '1.0 1.0'\n",
        "\n",
        "    # REPEATS\n",
        "    for rep in range(1,Repeats+1):\n",
        "        #rep=str(rep)\n",
        "\n",
        "        print(f\"\\n{esc(1, f'Repeat {rep}')} ({esc(34, f'CG-system{rep}')})\")\n",
        "        os.chdir(subdir)\n",
        "\n",
        "        # Random rotation\n",
        "        if rotate_mp is True:\n",
        "            mp_rotation = random.randint(0,360)\n",
        "            print(f\"\\nRotation {rep}: {mp_rotation}\")\n",
        "            gromacs.editconf(f=f'protein1-em.pdb', o=f'protein1-em{rep}.pdb',\n",
        "                             c=True,rotate=[0,0,mp_rotation])\n",
        "            \n",
        "        # Random translation ***\n",
        "        # ...\n",
        "\n",
        "        # Add lipids\n",
        "        x, y, z, cube = get_dimensions('protein1-em.pdb')\n",
        "        dm = 0 # mp already centered (mmembrane centered also)\n",
        "        print(\"\\nAdding lipids...\")\n",
        "        print(f\"Lipid ratio: {Ratio}\")\n",
        "        print(f\"Area per lipid: {round(apl, 3)}\")\n",
        "        os.system(f'python2.7 /content/insane3.py {lipid} -a {apl} '\n",
        "                  f'-o protein1-system.gro -p topol.top -f protein1-em{rep}.pdb '\n",
        "                  f'-center -ring -x {x} -y {y} -z {z} -dm {dm}')\n",
        "        print(\"Done.\\n\")\n",
        "\n",
        "        #os.system(\"wget 'http://cgmartini.nl/images/applications/water/water.gro'\") # remove ***\n",
        "\n",
        "        # COMBINE PROTEINS + BILAYER\n",
        "        u1 = MDAnalysis.Universe('protein1-system.gro')\n",
        "        u2 = MDAnalysis.Universe(\"protein2-em.pdb\")\n",
        "        multi = MDAnalysis.Merge(u2.atoms, u1.atoms)\n",
        "        multi.atoms.write(\"protein12.pdb\")\n",
        "        \n",
        "        # Format topology\n",
        "        replacements = {'Protein        1\\n':'ProteinB       1\\nProteinA       1\\n',\n",
        "                        '#include \"protein-cg.itp\"':'#include \"protein1-cg.itp\"\\n#include \"protein2-cg.itp\"\\n',\n",
        "                        'NA+':'NA', 'CL-':'CL',\n",
        "                        '#include \"martini_v3.itp\"':'#include \"martini_v3.0.0.itp\"\\n#include \"martini_v3.0.0_ions_v1.itp\"\\n#include \"martini_v3.0.0_solvents_v1.itp\"\\n#include \"martini_v3.0.0_phospholipids_v1.itp\"\\n'}\n",
        "        lines = []\n",
        "\n",
        "        with open('topol.top') as infile:\n",
        "            for line in infile:\n",
        "                for src, target in replacements.items():\n",
        "                    line = line.replace(src, target)\n",
        "                lines.append(line)\n",
        "        with open('topol.top', 'w') as outfile:\n",
        "            for line in lines:\n",
        "                outfile.write(line)\n",
        "        \n",
        "        # SOLVATION\n",
        "        #x, y, z, cube = get_dimensions('protein12.pdb')\n",
        "        print(\"\\nSolvating the system...\")\n",
        "        os.system(f'python2.7 /content/insane3.py -sol W -salt 0.15 '\n",
        "                  f'-o CG-system{rep}.gro -p topol_solv_test.top -f protein12.pdb '\n",
        "                  f'-center -ring -x {x} -y {y} -z {z}')\n",
        "        # Update topology\n",
        "        insane_topadd('topol_solv_test.top', 'topol.top') \n",
        "        os.system(f'cp topol.top topol{rep}.top')\n",
        "        print(\"Done.\")\n",
        "        \n",
        "        # ENERGY MINIMIZATION\n",
        "        print(f\"\\nPreparing CG system {rep} for energy minimization...\")\n",
        "        mdp_em() # generate mdp\n",
        "        gromacs.grompp(f='em.mdp',o=f'em_m{rep}.tpr',c=f'CG-system{rep}.gro',maxwarn='-1',backup=False,v=True) # em naming ***\n",
        "        gromacs.mdrun(deffnm=f'em_m{rep}', c=f'CG-system{rep}.pdb',backup=False)\n",
        "        gromacs.trjconv(f=f'CG-system{rep}.pdb', o=f'CG-system{rep}.pdb', pbc='mol',\n",
        "                        s=f'em_m{rep}.tpr', conect=True, input=[0,0],backup=False)\n",
        "        os.makedirs(f'EQ{rep}', exist_ok=True)\n",
        "        #gromacs.make_ndx(f='CG-system'+ rep +'.pdb', o='index_sys'+ rep +'.ndx', input=('del 0', 'del 1-40', '0|rPOP*','1&!0','!1','del 1','name 1 LIPID','name 2 SOL_ION','q'),backup=False)\n",
        "        gromacs.make_ndx(f=f'CG-system{rep}.pdb', o=f'index_sys{rep}.ndx',\n",
        "                         input=('del 2-100', 'rPOP*|rCAR*','rW*|rION','name 1 PROTEIN',\n",
        "                                'name 2 LIPID','name 3 SOL_ION','q'),\n",
        "                         backup=False)\n",
        "        print(\"System minimzed.\")\n",
        "\n",
        "        # EQUILIBRATION\n",
        "        print(f\"\\nPreparing CG system {rep} for equillibration...\")\n",
        "        with open('cgequil.mdp','w') as md: # create equil .mdp\n",
        "            md.write('define = -DPOSRES\\nintegrator = md\\ntinit = 0.0\\ndt = 0.02\\n'\n",
        "                     'nsteps = 10000\\nnstxout = 0\\nnstvout = 0\\nnstfout = 0\\n'\n",
        "                     'nstlog = 50000\\nnstenergy = 50000\\nnstxout-compressed = 50000\\n'\n",
        "                     'compressed-x-precision = 10000\\nnstlist  = 10\\nns_type  = grid\\n'\n",
        "                     'pbc   = xyz\\ncoulombtype  = Reaction_field\\nrcoulomb_switch = 0.0\\n'\n",
        "                     'rcoulomb  = 1.1\\nepsilon_r  = 15\\nvdw_type  = cutoff\\n'\n",
        "                     'rvdw_switch  = 0.9\\nrvdw   = 1.1\\ncutoff-scheme = verlet\\n'\n",
        "                     'coulomb-modifier = Potential-shift\\nvdw-modifier  = Potential-shift\\n'\n",
        "                     'epsilon_rf  = 0\\nverlet-buffer-tolerance = 0.005\\ntcoupl  = v-rescale\\n'\n",
        "                     'tc-grps  = PROTEIN LIPID SOL_ION\\ntau_t  = 1.0 1.0 1.0\\nref_t  = 310 310 310\\n'\n",
        "                     'Pcoupl  = berendsen\\nPcoupltype  = semiisotropic\\ntau_p  = 12.0\\n' \n",
        "                     'compressibility = 3e-4 3e-4\\nref_p  = 1.0 1.0\\ngen_vel  = yes\\n' \n",
        "                     'gen_temp  = 310\\ngen_seed  = -1\\nconstraints  = none\\n'\n",
        "                     'constraint_algorithm = Lincs\\ncontinuation  = no\\nlincs_order  = 4\\n'\n",
        "                     'lincs_warnangle = 30\\nrefcoord_scaling = com')\n",
        "        gromacs.grompp(f='cgequil.mdp',o=f'EQ{rep}/eq',r=f'CG-system{rep}.pdb',\n",
        "                       c=f'CG-system{rep}.pdb',maxwarn=-1, n=f'index_sys{rep}.ndx',backup=False)\n",
        "        os.chdir(f'{subdir}EQ{rep}')\n",
        "        print(\"Equillibration starting...\")\n",
        "        gromacs.mdrun(deffnm='eq',backup=False, v=True)\n",
        "        os.chdir(subdir)\n",
        "        print(\"Equillibration complete.\")\n",
        "\n",
        "        # VISUALIZE\n",
        "        #if rep == 1:\n",
        "        #    visualise_system(f'EQ{rep}/eq.gro', waters=False, ions=True)\n",
        "\n",
        "        # SETTING UP A PRODUCTION RUN\n",
        "        print(f\"\\nSetting up a production run for CG system {rep}...\")\n",
        "        os.makedirs(f'{subdir}MD{rep}', exist_ok=True)\n",
        "        os.chdir(subdir)\n",
        "        with open(f'{subdir}cgmd.mdp','w') as md: # create production .mdp\n",
        "                    md.write('integrator = md\\ntinit = 0.0\\ndt = 0.02\\nnsteps = 250000000\\n'\n",
        "                            'nstxout = 0\\nnstvout = 0\\nnstfout = 0\\nnstlog = 50000\\nnstenergy = 50000\\n'\n",
        "                            'nstxout-compressed = 50000\\ncompressed-x-precision = 10000\\nnstlist  = 10\\n'\n",
        "                            'ns_type  = grid\\npbc   = xyz\\ncoulombtype  = Reaction_field\\n'\n",
        "                            'rcoulomb_switch = 0.0\\nrcoulomb  = 1.1\\nepsilon_r  = 15\\n'\n",
        "                            'vdw_type  = cutoff\\nrvdw_switch  = 0.9\\nrvdw   = 1.1\\n'\n",
        "                            'cutoff-scheme = verlet\\ncoulomb-modifier = Potential-shift\\n'\n",
        "                            'vdw-modifier  = Potential-shift\\nepsilon_rf  = 0\\n'\n",
        "                            'verlet-buffer-tolerance = 0.005\\ntcoupl  = v-rescale\\n'\n",
        "                            'tc-grps  = PROTEIN LIPID SOL_ION\\ntau_t  = 1.0 1.0 1.0\\n'\n",
        "                            f'ref_t  = 310 310 310\\nPcoupl  = berendsen\\nPcoupltype  = {pcoupltype}\\n'\n",
        "                            f'tau_p  = 12.0\\ncompressibility = {compressibility}\\nref_p  = {ref_p}\\n'\n",
        "                            'gen_vel  = yes\\ngen_temp  = 310\\ngen_seed  = -1\\nconstraints  = none\\n'\n",
        "                            'constraint_algorithm = Lincs\\ncontinuation  = no\\nlincs_order  = 4\\n'\n",
        "                            'lincs_warnangle = 30\\n')\n",
        "\n",
        "        gromacs.grompp(f='cgmd.mdp',o=f'MD{rep}/md',c=f'EQ{rep}/eq.gro',\n",
        "                       maxwarn=-1,n=f'index_sys{rep}.ndx',backup=False)\n",
        "        #shutil.rmtree('MD', ignore_errors=True)\n",
        "        print(\"Production run setup complete.\")\n",
        "        print(f\"\\n{esc(34, f'CG-system{rep}')} {esc(92, f'complete.')}\")\n",
        "\n",
        "# Setp system\n",
        "ratio = get_lipid_ratio(pe_pg_ratio=PE_PG_ratio, total=10, card_conc=CARD_conc, tolerance=0.0001)\n",
        "\n",
        "setup_cg_system(filename1, filename2, Units=Units, rotate_sp=0,\n",
        "                Repeats=Repeats, Box_Width=Box_Width,Box_Height=Box_Height, Box_Length_ext=Box_Length_ext,\n",
        "                Embed=Embed,Continuous=False, Ratio=ratio, \n",
        "                NamingConvention=f'')\n",
        "\n",
        "# LOOP EXAMPLE\n",
        "# concs\n",
        "#concs = [0.2] #[0, 0.05, 0.1, 0.15]#, 0.2]\n",
        "\n",
        "#for i in concs:\n",
        "#    ratio = get_lipid_ratio(pe_pg_ratio=3.5, total=10, card_conc=i, tolerance=0.0001)\n",
        "#    setup_cg_system(filename1, filename2, Units=4, rotate_sp=180,\n",
        "#                    Repeats=5, Box_Width=25,Box_Height=21, Box_Length_ext=20,\n",
        "#                    Embed=2,Continuous=False, Ratio=ratio, \n",
        "#                    NamingConvention=f'-c{i}')\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gLVCu8iqXU4U",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Download Zip File\n",
        "#%%script false --no-raise-error\n",
        "\n",
        "def download(nm=None):\n",
        "    \"\"\"\n",
        "    Download results as a zip file.\n",
        "    \"\"\"\n",
        "    if nm == None:\n",
        "        os.chdir('/content/')\n",
        "        os.system(f'zip -r {name}.zip {name}')\n",
        "        files.download(f'{name}.zip')\n",
        "    else:\n",
        "        os.chdir('/content/')\n",
        "        os.system(f'zip -r {name}-{nm}.zip {name}')\n",
        "        files.download(f'{name}-{nm}.zip')\n",
        "\n",
        "if Download is True:\n",
        "    download(f'{Zip_Name}') # usage"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyONeUw4iTexgq+S7DcsPrEv",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}