{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/disulfide_design.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YRk_XS2NILil"
      },
      "source": [
        "# disulfide-hallucination\n",
        "For a given length, generate/hallucinate a disulfide-rich sequence that AlphaFold thinks folds into a well structured protein"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-AXy0s_4cKaK"
      },
      "outputs": [],
      "source": [
        "#@title setup\n",
        "%%time\n",
        "import os\n",
        "if not os.path.isdir(\"params\"):\n",
        "  # get code\n",
        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git@v1.1.1\")\n",
        "  # for debugging\n",
        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
        "  # download params\n",
        "  os.system(\"mkdir params\")\n",
        "  os.system(\"apt-get install aria2 -qq\")\n",
        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n",
        "  os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n",
        "\n",
        "import warnings\n",
        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
        "\n",
        "import os\n",
        "from colabdesign import mk_afdesign_model, clear_mem\n",
        "from IPython.display import HTML\n",
        "from google.colab import files\n",
        "import numpy as np\n",
        "\n",
        "def get_pdb(pdb_code=\"\"):\n",
        "  if pdb_code is None or pdb_code == \"\":\n",
        "    upload_dict = files.upload()\n",
        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
        "    return \"tmp.pdb\"\n",
        "  elif os.path.isfile(pdb_code):\n",
        "    return pdb_code\n",
        "  elif len(pdb_code) == 4:\n",
        "    os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
        "    return f\"{pdb_code}.pdb\"\n",
        "  else:\n",
        "    os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
        "    return f\"AF-{pdb_code}-F1-model_v3.pdb\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7JbLyVu4ILil"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "from jax.lax import dynamic_slice\n",
        "import jax.numpy as jnp\n",
        "\n",
        "from colabdesign.af.loss import _get_con_loss\n",
        "def get_con_loss(dgram, dgram_bins, cutoff=None, binary=True,\n",
        "                 num=1, seqsep=0, offset=None):\n",
        "  '''convert distogram into contact loss'''  \n",
        "  x = _get_con_loss(dgram, dgram_bins, cutoff, binary)  \n",
        "  a,b = x.shape\n",
        "  if offset is None:\n",
        "    mask = jnp.abs(jnp.arange(a)[:,None] - jnp.arange(b)[None,:]) >= seqsep\n",
        "  else:\n",
        "    mask = jnp.abs(offset) >= seqsep\n",
        "  x = jnp.sort(jnp.where(mask,x,jnp.nan))\n",
        "  k_mask = (jnp.arange(b) < num) * (jnp.isnan(x) == False)    \n",
        "  return jnp.where(k_mask,x,0.0).sum(-1) / (k_mask.sum(-1) + 1e-8)\n",
        "\n",
        "def generate_disulfide_pattern(L, disulfide_num, min_sep=5):\n",
        "    disulfide_pattern = []\n",
        "    positions = list(range(L))\n",
        "    for n in range(disulfide_num):\n",
        "        for _ in range(100): # try 100 time per postion.\n",
        "            i,j = random.sample(positions, k=2)\n",
        "            if abs(i-j)<=min_sep: continue # set min loop len.\n",
        "            positions.remove(i)\n",
        "            positions.remove(j)\n",
        "            disulfide_pattern.append((i,j))\n",
        "            # check\n",
        "            if _ > 99:\n",
        "                print('Not find good disulfide_pos! exit....')\n",
        "                return 0  # not good pose!\n",
        "            else:\n",
        "                break\n",
        "    sequence_pattern = list('X'*L)\n",
        "    for pair in disulfide_pattern:\n",
        "        for i in pair: sequence_pattern[i] = 'C'\n",
        "\n",
        "    return disulfide_pattern, ''.join(sequence_pattern), L\n",
        "\n",
        "def disulfide_loss(inputs, outputs):\n",
        "  def get_disulfide_loss(dgram, dgram_bins, disulfide_pattern):\n",
        "    '''\n",
        "    Func: simple disulfide loss, make the contacts < 7.0/7.5A.\n",
        "    # see: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC7316719/\n",
        "    params: disulfide_pattern: List[(pos1, pos2), (pos3, pos4)...]\n",
        "    '''\n",
        "    disulfide_loss = 0.0\n",
        "    for pair in disulfide_pattern:\n",
        "      i,j = pair\n",
        "      pair_dgram = dynamic_slice(dgram, (i,j,0), (1,1,len(dgram_bins))) + dynamic_slice(dgram, (j,i,0), (1,1,len(dgram_bins)))\n",
        "      disulfide_loss += get_con_loss(pair_dgram, dgram_bins, cutoff=7.0, binary=False, num=1)\n",
        "    return disulfide_loss.mean()\n",
        "\n",
        "  # add disulfide loss here:\n",
        "  dgram_logits = outputs['distogram']['logits']\n",
        "  dgram_bins = jnp.append(0, outputs[\"distogram\"][\"bin_edges\"])\n",
        "  return {\"disulfide\":get_disulfide_loss(dgram_logits, dgram_bins, inputs['opt']['disulfide_pattern'])}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aU4e-ZkfILim"
      },
      "source": [
        "0. generate_cys_pattern"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vrSatV2sILim"
      },
      "outputs": [],
      "source": [
        "disulfide_pattern, sequence_pattern, L = generate_disulfide_pattern(35, 3)\n",
        "print(disulfide_pattern)\n",
        "print(sequence_pattern)\n",
        "print(L)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jGNGjHeyILim"
      },
      "source": [
        "1. hallucination with disulfide-pattern"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HmsBDNhmILin"
      },
      "outputs": [],
      "source": [
        "# using some disulfide sequence as input:\n",
        "clear_mem()\n",
        "af_model = mk_afdesign_model(protocol=\"hallucination\", loss_callback=disulfide_loss)\n",
        "af_model.opt[\"weights\"][\"disulfide\"] = 1.0\n",
        "af_model.prep_inputs(length=L)\n",
        "\n",
        "print(\"length\",af_model._len)\n",
        "print(\"weights\",af_model.opt[\"weights\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MmBtBMnBILin"
      },
      "outputs": [],
      "source": [
        "# set disulfide_pattern sequence.\n",
        "af_model.restart(seq=sequence_pattern, add_seq=True, rm_aa='C')\n",
        "# set disulfide_pattern:\n",
        "af_model.opt['disulfide_pattern'] = disulfide_pattern\n",
        "# reweight con:\n",
        "af_model.opt[\"weights\"]['con'] = 0.5"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "af_model.design_3stage(50,50,10)"
      ],
      "metadata": {
        "id": "wstI0n4zMydi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ewlj8671ILio"
      },
      "outputs": [],
      "source": [
        "af_model.save_pdb(f\"{af_model.protocol}.pdb\")\n",
        "af_model.plot_pdb(show_sidechains=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iVEqJKaKILio"
      },
      "source": [
        "2. fix-sidechain with pyRosetta"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VGwyvaodILio"
      },
      "outputs": [],
      "source": [
        "from pyrosetta import init, pose_from_pdb, Pose, create_score_function\n",
        "from pyrosetta.rosetta.protocols.relax import FastRelax\n",
        "from pyrosetta.rosetta.core.pack.task import TaskFactory\n",
        "from pyrosetta.rosetta.core.kinematics import MoveMap\n",
        "from pyrosetta.rosetta.core.pack.task.operation import InitializeFromCommandline\n",
        "from pyrosetta.rosetta.protocols.denovo_design import DisulfidizeMover\n",
        "from pyrosetta.rosetta.core.pack.task.operation import RestrictToRepacking\n",
        "from pyrosetta.rosetta.core.select.residue_selector import ChainSelector\n",
        "\n",
        "def fastrelax(pose):\n",
        "    full_score = create_score_function('ref2015')\n",
        "    tf = TaskFactory()\n",
        "    # tf.push_back(InitializeFromCommandline())\n",
        "    tf.push_back(RestrictToRepacking())\n",
        "    mmap = MoveMap()\n",
        "    mmap.set_bb(True)\n",
        "    mmap.set_chi(True)\n",
        "\n",
        "    # FastRelax:\n",
        "    fast_design = FastRelax(full_score, 4)  # repeat:6\n",
        "    fast_design.set_task_factory(tf)\n",
        "    fast_design.set_movemap(mmap)\n",
        "    fast_design.apply(pose)\n",
        "    \n",
        "def build_stapled_pose(pose, cys_pattern):\n",
        "    # rebuild \n",
        "    disulfidizer = DisulfidizeMover()\n",
        "    disulfidizer.set_match_rt_limit(999.0)\n",
        "    disulfidizer.set_max_disulf_score(999.0)\n",
        "    for pair in cys_pattern:\n",
        "        full_score = create_score_function('ref2015')\n",
        "        stapled_pose = Pose().assign(pose)  # pose for disulfides stapled\n",
        "        residue1, residue2 = pair\n",
        "        disulfidizer.make_disulfide(stapled_pose, residue1+1, residue2+1, False, full_score)\n",
        "        # fastdesign mover\n",
        "        pdbnum1 = stapled_pose.pdb_info().pose2pdb(residue1)\n",
        "        pdbnum2 = stapled_pose.pdb_info().pose2pdb(residue2)\n",
        "        print('# OPTIMIZE:  Around These Residues: %s %s' % (pdbnum1, pdbnum2))\n",
        "\n",
        "    # save pdb:\n",
        "    fastrelax(stapled_pose)\n",
        "    \n",
        "    # filter: all stapled?\n",
        "    v = ChainSelector(1).apply(stapled_pose)\n",
        "    l = disulfidizer.find_current_disulfides(stapled_pose, v, v)\n",
        "    print(len([i for i in l]), len(cys_pattern))\n",
        "    if len([i for i in l]) == len(cys_pattern):\n",
        "        return stapled_pose\n",
        "    else:\n",
        "        return 0\n",
        "    "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zhMey-4lILip"
      },
      "outputs": [],
      "source": [
        "# optimze the side-chain\n",
        "init('-mute all')\n",
        "pose = pose_from_pdb(f\"{af_model.protocol}.pdb\")\n",
        "stapled_pose = build_stapled_pose(pose, disulfide_pattern)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "q4qiU9I0QHSz"
      ],
      "name": "disulfide_design.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3.8.10 64-bit",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.10"
    },
    "vscode": {
      "interpreter": {
        "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}