filament_v2_1.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_v2_1.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7a--qtEquUzt"
      },
      "source": [
        "<img src=\"https://github.com/roncv/MemProtMD/raw/main/img/filament_output.png\" height=\"250\" align=\"right\" style=\"height:250px\">\n",
        "\n",
        "##MPss Filament\n",
        "\n",
        "Membrane protein association protocol for allowing a filamentous proteins to -associate with a preformed lipid membrane. \n",
        "\n",
        "The protein structure is converted to [Martini 3](https://www.nature.com/articles/s41592-021-01098-3) coarse-grained (CG) representation using [Martinize](https://github.com/marrink-lab/vermouth-martinize) with a lipid membrane 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. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jsgA0B1Rtgl4",
        "cellView": "form"
      },
      "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 #### Set number of Repeats:\n",
        "Repeats = 5 #@param {type:\"slider\", min:1, max:10, step:1}\n",
        "#@markdown ---\n",
        "#@markdown #### Rotate Protein:\n",
        "Rotate = 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 #### Set Box Dimensions:\n",
        "Box_Width = 20 #@param {type:\"slider\", min:8, max:30, step:1}\n",
        "Box_Height = 21 #@param {type:\"slider\", min:12, max:50, step:1}\n",
        "#@markdown #### Add padding to the ends of filaments:\n",
        "Box_Length_ext = 0 #@param {type:\"slider\", min:0, max:20, step:1}\n",
        "#@markdown ---\n",
        "#@markdown #### Create Filament:\n",
        "Filament = True #@param {type:\"boolean\"}\n",
        "Units = 5 #@param {type:\"slider\", min:1, max:20, step:1}\n",
        "Continuous = True #@param {type:\"boolean\"}\n",
        "#@markdown ---\n",
        "#@markdown #### Output:\n",
        "Download = True #@param {type:\"boolean\"}\n",
        "Zip_Name = \"filament\" #@param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Kt6ShO3gticF",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Upload Protein Structure\n",
        "\n",
        "# Imports + Dependencies\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",
        "\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_protein():\n",
        "    \"\"\"\n",
        "    Upload a .pdb coordinate file and visualise structure.\n",
        "    \"\"\"\n",
        "    os.chdir('/content/')\n",
        "    print(Fore.BLUE + \"\\nUpload Protein Structure:\\n\")\n",
        "    upload = files.upload() # request upload of .pdb\n",
        "\n",
        "    # Arrange working directories\n",
        "    filename = next(iter(upload))\n",
        "    name = os.path.splitext(filename)[0]\n",
        "    working_dir = '/content/' + name +'/'\n",
        "\n",
        "    if os.path.exists(working_dir):\n",
        "        shutil.rmtree(working_dir)\n",
        "    os.makedirs(working_dir)\n",
        "\n",
        "    os.rename(filename, working_dir + filename)\n",
        "    os.chdir(working_dir)\n",
        "    mol1 = open(working_dir + filename, 'r').read()\n",
        "\n",
        "    mview = py3Dmol.view(width=600,height=400)\n",
        "    mview.addModel(mol1,'pdb')\n",
        "    mview.setStyle({'cartoon':{'color':'spectrum'}})\n",
        "    mview.setBackgroundColor('0xffffff')\n",
        "    mview.zoomTo()\n",
        "    mview.show()\n",
        "\n",
        "    return working_dir, name, filename\n",
        "\n",
        "working_dir, name, filename = upload_protein() # usage"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yHpvS0OXteVS",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Install Dependencies\n",
        "%%capture\n",
        "os.chdir('/content/')\n",
        "if not os.path.isdir(\"/content/cg2at/\"):\n",
        "  !apt-get update -y\n",
        "  !apt-get install dssp\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 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",
        "\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": "2DPZxbjotjbh",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Get lipid concentrations\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": "O83c0L4DtmK9",
        "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",
        "    Based on predetermined coefficients for each lipid.\n",
        "\n",
        "    Coefficients:\n",
        "    POPE = 0.63\n",
        "    POPG = 0.67\n",
        "    CARD = 1.23\n",
        "\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 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 styling.\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 given protein structure. 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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1LFAv4qhtyQD",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Set-up Coarse-Grained Membrane Protein Systems\n",
        "%%capture\n",
        "def setup_cg_system(Units=Units, MembraneType=MembraneType, Repeats=Repeats,\n",
        "                    Rotate=Rotate, ElasticCutoff=0.7, NamingConvention=None, Box_Width=Box_Width,\n",
        "                    Box_Height=Box_Height, Box_Length_ext=None, Embed=0, Filament=Filament,\n",
        "                    Continuous=Continuous, pe=7, pg=2, crd=1, Ratio=None):\n",
        "    \"\"\"\n",
        "    Function to set-up coarse-grained membrane-protein systems.\n",
        "    \"\"\"\n",
        "    print(f'\\n{esc([100, 97], \"INITIALIZE\")}\\n')\n",
        "\n",
        "    # Move to working directory\n",
        "    os.chdir(working_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 Ratio == None: # depreciate this\n",
        "        Ratio = [pe, pg, crd]\n",
        "        if crd == 0:\n",
        "            MembraneType = 'POPE:POPG'\n",
        "        # Break in the if statement\n",
        "        if MembraneType == \"POPC\":\n",
        "            lipid_name = 'pc'\n",
        "            lipid = '-l POPC:1'\n",
        "        elif MembraneType == \"POPE:POPG\":\n",
        "            lipid_name = 'pepg'\n",
        "            lipid = f'-l POPE:{pe} -l POPG:{pg}'\n",
        "        elif MembraneType == \"POPE:POPG:CARDIOLIPIN\":\n",
        "            lipid_name = f'pe{pe}-pg{pg}-c{crd}'\n",
        "            lipid = f'-l POPE:{pe} -l POPG:{pg} -l CARD:{crd}'\n",
        "        else:\n",
        "            print('Invalid lipid ratios!')\n",
        "    # Using Ratio instead of individual compositions - PREFERRED\n",
        "    elif 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\"{working_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(filename, subdir+filename)\n",
        "    os.chdir(subdir)\n",
        "\n",
        "    # martini .itp files\n",
        "    for file in glob.glob(r'/content/martini*.itp'):\n",
        "        #print(file)\n",
        "        shutil.copy(file, subdir)\n",
        "\n",
        "    # Adjust cell size based on no. of filaments + Generate filaments\n",
        "    print(\"\\nBuilding cell...\\n\")\n",
        "    if Filament == True:\n",
        "        gromacs.genconf(f=filename,o='ready.pdb',nbox=[Units,1,1],backup=False,renumber=True)\n",
        "        x, y, z, _ = get_dimensions('ready.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",
        "    else:\n",
        "        cube = [Box_Width, Box_Width, Box_Height]\n",
        "        gromacs.genconf(f=filename,o='ready.pdb',nbox=[1,1,1],backup=False,renumber=True)\n",
        "    # Print dimensions\n",
        "    print(f\"\\nCell dimensions: {cube}\")\n",
        "\n",
        "\n",
        "    # Write selected settings to dir # *** UPDATE AT THE END\n",
        "    with open('settings.txt', 'w') as file:\n",
        "            file.write(f'MembraneType = {MembraneType}\\nRepeats = {Repeats}\\n'\n",
        "                       f'Rotate = {Rotate}\\nEmbed = {Embed}\\nBox_Width (y) = {Box_Width}\\n'\n",
        "                       f'Box_Length (x) = {cube[0]}\\nBox_Height (z) = {Box_Height}\\n'\n",
        "                       f'Box_Length_ext (y) = {Box_Length_ext}\\nFilament = {Filament}\\nUnits = {Units}\\n'\n",
        "                       f'ElasticBondUpperCutoff = {ElasticCutoff}\\nRatio = {Ratio}\\nContinuous = {Continuous}')\n",
        "    # \n",
        "    gromacs.make_ndx(f=\"ready.pdb\",o='index.ndx',input=('del 0', 'del 1-100','q'))\n",
        "    gromacs.editconf(f=\"ready.pdb\",o='protein.pdb',n='index.ndx',\n",
        "                     c=True,box=cube,input=(0,0),backup=False,label='A',resnr=1)\n",
        "\n",
        "\n",
        "    # Get dimensions of the box \n",
        "    x, y, z, _ = get_dimensions('protein.pdb')\n",
        "    u = MDAnalysis.Universe('protein.pdb')\n",
        "    boxed = u.dimensions\n",
        "\n",
        "    # Format residues\n",
        "    print(\"\\nFormatting .pdb...\")\n",
        "    with open('protein.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('protein.pdb', 'w') as file:\n",
        "        file.write(filedata)\n",
        "    # Write EM .mdp\n",
        "    with open('em.mdp','w') as em:\n",
        "                em.write('integrator = steep\\nnsteps = 5000\\nemtol = 100\\nemstep = 0.001')\n",
        "    # Write Topology\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-cg.itp\"\\n[ system ]\\n[ molecules ]\\nProteinA 1')\n",
        "    os.system(f'cp topol.top topol_protein.top')\n",
        "\n",
        "    # GENERATE COARSE GRAINED PROTEIN\n",
        "    print(\"\\nCoarse graining the protein...\")\n",
        "    merge = 'A' \n",
        "    eunit = 'molecule' # 'molecule\n",
        "    os.system('martinize2 -f protein.pdb -dssp mkdssp -ff martini3001 '\n",
        "              f'-x protein-cg.pdb -o protein-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(\"sed -e 's/^molecule.*/ProteinA 1/g' molecule_0.itp | grep -v ';' >  protein-cg.itp\")\n",
        "\n",
        "    with open(\"protein-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('protein-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",
        "    # Rotation\n",
        "    if Rotate == True:\n",
        "        gromacs.editconf(f='protein-cg.pdb',o='protein-cg.pdb',\n",
        "                         rotate=[randrange(360),randrange(360),0],box=cube)\n",
        "    else:\n",
        "        gromacs.editconf(f='protein-cg.pdb',o='protein-cg.pdb', box=cube)\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\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",
        "    # ENERGY MINIMIZE PROTEIN\n",
        "    print(\"\\nPreparing CG protein for energy minimization...\")\n",
        "    gromacs.editconf(f='protein-cg.pdb',o='protein-box.pdb',c=True,box=cube)\n",
        "    gromacs.grompp(f='em.mdp',o='em.tpr',c='protein-box.pdb',maxwarn='-1',backup=False,v=True)\n",
        "    gromacs.mdrun(deffnm='em',c='protein-em.pdb',backup=False)\n",
        "\n",
        "    # Get protein height and calculate membrane insertion depth\n",
        "    p_height = position_protein_insane('protein-em.pdb')\n",
        "    Distance = p_height + Embed\n",
        "\n",
        "    # REPEATS\n",
        "    for rep in range(1,Repeats+1):\n",
        "        rep=str(rep)\n",
        "        print(f\"\\n{esc(1, f'Repeat {rep}')} ({esc(34, f'CG-system{rep}')})\")\n",
        "        os.chdir(subdir)\n",
        "\n",
        "        # Add lipids\n",
        "        print(\"\\nAdding lipids...\")\n",
        "        print(f\"Lipid ratio: {Ratio}\")\n",
        "        print(f\"Area per lipid: {round(apl, 3)}\")\n",
        "        if Rotate == True:\n",
        "            os.system(f'python2.7 /content/insane3.py {lipid} -salt 0.15 -sol W -a {apl} '\n",
        "                      f'-o CG-system{rep}.gro -p topol.top -f protein-em.pdb -center '\n",
        "                      f'-rotate random -ring -x {x} -y {y} -z {z} -dm {Distance}')\n",
        "        else:\n",
        "            os.system(f'python2.7 /content/insane3.py {lipid} -salt 0.15 -sol W -a {apl} '\n",
        "                      f'-o CG-system{rep}.gro -p topol.top -f protein-em.pdb '\n",
        "                      f'-center -ring -x {x} -y {y} -z {z} -dm {Distance}')\n",
        "        print(\"Done.\")\n",
        "\n",
        "        # Format topology\n",
        "        format_topology('topol.top')\n",
        "        os.system(f'cp topol.top topol{rep}.top')\n",
        "\n",
        "\n",
        "        # ENERGY MINIMIZATION\n",
        "        print(f\"\\nPreparing CG system {rep} for energy minimization...\")\n",
        "        gromacs.grompp(f='em.mdp', o=f'em{rep}.tpr', c=f'CG-system{rep}.gro',maxwarn='-1',backup=False,v=True)\n",
        "        gromacs.mdrun(deffnm=f'em{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{rep}.tpr',conect=True,input='0',backup=False)\n",
        "        print(\"System minimzed.\")\n",
        "\n",
        "        # EQUILIBRATION\n",
        "        print(f\"\\nPreparing CG system {rep} for equillibration...\")\n",
        "        os.makedirs('EQ'+rep, exist_ok=True)\n",
        "        gromacs.make_ndx(f=f'CG-system{rep}.pdb',o=f'index{rep}.ndx',\n",
        "                         input=('del 0', 'del 1-40', '0|rPOP*','1&!0','!1','del 1','name 1 Lipid','name 2 SOL_ION','q'),\n",
        "                         backup=False)\n",
        "        with open('cgequil.mdp','w') as md: # create equil .mdp\n",
        "                    md.write('integrator = 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\\n')\n",
        "        gromacs.grompp(f='cgequil.mdp',o=f'EQ{rep}/eq',c=f'CG-system{rep}.pdb',\n",
        "                       maxwarn=-1, n=f'index{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",
        "        # SETTING UP A PRODUCTION RUN\n",
        "        print(f\"\\nSetting up a production run for CG system {rep}...\")\n",
        "        os.makedirs(f'MD{rep}', exist_ok=True)\n",
        "        with open('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{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.')}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dw5Z2WVYGi9g",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Set up system\n",
        "\n",
        "# Build and equillibriate system\n",
        "ratio = get_lipid_ratio(pe_pg_ratio=PE_PG_ratio, card_conc=CARD_conc, total=10, tolerance=0.001)\n",
        "\n",
        "setup_cg_system(Units=Units, Repeats=Repeats, Box_Width=Box_Width, Box_Height=Box_Height, Ratio=ratio,\n",
        "                Continuous=Continuous, Embed=Embed, Box_Length_ext=Box_Length_ext, NamingConvention=f'')\n",
        "\n",
        "# EXAMPLE LOOP\n",
        "# With various filament lengths...\n",
        "#fil_lens = [6, 8, 10]\n",
        "\n",
        "#for i, fl in enumerate(fil_lens):\n",
        "    #ratio = get_lipid_ratio(pe_pg_ratio=3.5, pg_conc=0.05, total=10)\n",
        "    #print(f\"Ratio {i}: {ratio}\\n\")\n",
        "    #print(f'-fil{fl}\\n')\n",
        "#    setup_cg_system(Units=fl, Repeats=5, Box_Width=30, Box_Height=20, Ratio=ratio,\n",
        "#                    Continuous=True, Embed=2, Box_Length_ext=None)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1rHUULmIHPpR",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Visualise System\n",
        "def visualise_system(_type='gro', waters=False):\n",
        "    \"\"\"\n",
        "    Visualise systems in 3D using Py3Dmol.\n",
        "    \"\"\"\n",
        "    # Access first repeat of each experiment\n",
        "    if _type == 'pdb':\n",
        "        systems = glob.glob(os.path.expanduser(f'{working_dir}*/CG-system1.{_type}'))\n",
        "    else:\n",
        "        systems = glob.glob(os.path.expanduser(f'{working_dir}*/EQ1/eq.{_type}'))\n",
        "    # Visualise\n",
        "    for i in systems:\n",
        "        print(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':'0xF5FF1D', 'radius':2.2,'quality':20}}) # lipids\n",
        "        if waters is True:\n",
        "            mview.setStyle({'atom':'W'},{'sphere':{'color':'0x5014F5','quality':20}}) # waters\n",
        "        mview.setStyle({'atom':'BB'},\n",
        "                       {'sphere':{'color':'0x5FC8F5', 'radius':3,'quality':20}}) # protein\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(working_dir)\n",
        "    # Remove temporary files\n",
        "    #for file in glob.glob(r'*/#*'):\n",
        "    #    os.remove(file)\n",
        "    \n",
        "#visualise_system('pdb', waters=False) # use .pdb to see unit cell"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xITragcGHiXI",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Download Zip File\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": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.8.8"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}