notebooks/model_notebooks/RFDiffusion.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3298e6f0-7d37-463b-8b1e-3d2aa3782e44",
   "metadata": {},
   "source": [
    "# RFDiffusion\n",
    "\n",
    "RFDiffusion is a state-of-the-art method for computational protein structure design. Many possible use cases are demonstrated in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc2fc3e5-45a9-45eb-b30f-63edd1cd30be",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61a8cfbb-a779-47d2-a3e2-8870dec5409c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from prtm import protein\n",
    "from prtm.models.rfdiffusion import config\n",
    "from prtm.models.rfdiffusion.modeling import RFDiffusionForStructureDesign\n",
    "from prtm.models.rfdiffusion.samplers import UnconditionalSampler, SelfConditioningSampler, ScaffoldedSampler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "556925ca-64ca-49cb-bd37-438f0ed23827",
   "metadata": {},
   "source": [
    "There are 8 different sets of model weights for RFDiffusion. The right weights depend on the use case. When `model_name` is set to `auto` the model weights to load are determined by the input inference parameters. Unlike other pipelines, this means that the model is loaded at runtime instead of at instantiation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37b7fe93-0439-4fbd-ac8f-4b6ec84494d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "designer = RFDiffusionForStructureDesign(model_name=\"auto\", random_seed=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa2a91f3-a0c9-4987-9d1a-394ba0b77e0c",
   "metadata": {},
   "source": [
    "## Unconditional Design\n",
    "\n",
    "First, let's look at the simplest design setup where structures are either unconstrained or only lightly constrained by a guiding potential."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f822b65-8388-490b-89dd-70d1b7f564ca",
   "metadata": {},
   "source": [
    "The `UnconditionalSamplerConfig` is used when there are no reference structures or scaffolds to be used in design. The only argument that must be set for this sampler is `contigmap_params`. `RFDiffusion` uses a `ContigMap` and `contig` string to specify the lengths of generated chains and the number of chains. For example:\n",
    "\n",
    "- [\"50-50\"] will create a single chain of exactly 50 residues.\n",
    "- [\"50-100\"] will create a single chain with anywhere from 50-100 residues.\n",
    "\n",
    "In addition to specifying the length of the protein to be designed, we can provide symmetry and potential constraints to the sampler. Choices for symmetry are:\n",
    "\n",
    "- `tetrahedral`\n",
    "- `octahedral`\n",
    "- `icosahedral`\n",
    "- C symmetries (`C2`, `C3`, ...)\n",
    "- D symmetries (`D2`, `D3`, ...)\n",
    "\n",
    "Examples of using symmetries and potentials are provided in the follow sections of this notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c8d5d74-3504-4317-8b5e-70a36e436ab1",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Random Length\n",
    "\n",
    "Generate a structure with an exact number of residues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37b93317-546e-4a44-a691-65a3e34496bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler_config = config.UnconditionalSamplerConfig(\n",
    "    contigmap_params=config.ContigMap(contigs=[\"50-50\"]),\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb7e4194-0629-465b-8a4c-1e6cce70bf00",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: all residues in the structure are glycines which have no sidechains\n",
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34371a1f-71f4-4b6f-9a21-3642241955b9",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Monomer ROG Potential\n",
    "\n",
    "Add a potential during design to encourage a smaller radius of gyration for the monomer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45dd6c2f-08cd-4de0-9557-2cec38cfdcd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler_config = config.UnconditionalSamplerConfig(\n",
    "    # Variable number of residues from 100 to 200\n",
    "    contigmap_params=config.ContigMap(contigs=[\"100-200\"]),\n",
    "    # Potentials are defined in potentials.py and arguments are delimited by commas and colons\n",
    "    potentials_params=config.PotentialsParams(\n",
    "        guiding_potentials=[\"type:monomer_ROG,weight:1,min_dist:5\"],\n",
    "        guide_scale=2,\n",
    "        guide_decay=\"quadratic\",\n",
    "    ),\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3f4ac42-075a-44ef-8ba3-1a74d505bf79",
   "metadata": {},
   "outputs": [],
   "source": [
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6a23b00-e514-484b-8d8a-8f4512ad1dbd",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Contact Potential\n",
    "\n",
    "Add a potential during design to encourage more contacts in a monomer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba8aca40-3096-483d-a8e9-61ce3d59db1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler_config = config.UnconditionalSamplerConfig(\n",
    "    contigmap_params=config.ContigMap(contigs=[\"100-200\"]),\n",
    "    potentials_params=config.PotentialsParams(guiding_potentials=[\"type:monomer_contacts,weight:0.05\"]),\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0639427d-8a4b-4c6a-81f3-1be1f7f58d35",
   "metadata": {},
   "outputs": [],
   "source": [
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58890cd9-697d-4574-a07e-34cdc43335b7",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Tetrahedral Oligos\n",
    "\n",
    "Generate a protein with tetrahedral symmetry and add a potential to encourage more contacts within and between chains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "238bfe39-63dd-48c2-b83b-3837d8961558",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler_config = config.UnconditionalSamplerConfig(\n",
    "    contigmap_params=config.ContigMap(contigs=[\"240-240\"]),\n",
    "    symmetry_params=config.SymmetryParams(symmetry=\"tetrahedral\"),\n",
    "    potentials_params=config.PotentialsParams(\n",
    "        guiding_potentials=[\"type:olig_contacts,weight_intra:1,weight_inter:0.1\"],\n",
    "        olig_inter_all=True,\n",
    "        olig_intra_all=True,\n",
    "        guide_scale=2,\n",
    "        guide_decay=\"quadratic\",\n",
    "    ),\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce3bff90-5a2c-4561-a6ed-abecb3a44344",
   "metadata": {},
   "outputs": [],
   "source": [
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e724583-549b-4696-a60a-88cc9ea35752",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Cyclic Oligos\n",
    "\n",
    "Generate a protein with 6-fold cyclic symmetry and add a potential to encourage more contacts within and between chains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c61e7856-ec7f-4796-b98a-539a5c1af242",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler_config = config.UnconditionalSamplerConfig(\n",
    "    contigmap_params=config.ContigMap(contigs=[\"90-90\"]),\n",
    "    symmetry_params=config.SymmetryParams(symmetry=\"C6\"),\n",
    "    potentials_params=config.PotentialsParams(\n",
    "        guiding_potentials=[\"type:olig_contacts,weight_intra:1,weight_inter:0.1\"], \n",
    "        olig_intra_all=True, \n",
    "        olig_inter_all=True, \n",
    "        guide_scale=2.0, \n",
    "        guide_decay=\"quadratic\",\n",
    "    ),\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d01f79c5-f1ce-4eaf-aad5-865304b719a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63842c27-2744-4519-b53e-516ec3d9b974",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Dihedral Oligos\n",
    "\n",
    "Generate a protein with dihedral symmetry and add a potential to encourage more contacts within and between chains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e602b3f4-80be-48a3-8a05-143b4d1d0ad8",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler_config = config.UnconditionalSamplerConfig(\n",
    "    contigmap_params=config.ContigMap(contigs=[\"120-120\"]),\n",
    "    symmetry_params=config.SymmetryParams(symmetry=\"D2\"),\n",
    "    potentials_params=config.PotentialsParams(\n",
    "        guiding_potentials=[\"type:olig_contacts,weight_intra:1,weight_inter:0.1\"], \n",
    "        olig_intra_all=True, \n",
    "        olig_inter_all=True, \n",
    "        guide_scale=2.0, \n",
    "        guide_decay=\"quadratic\",\n",
    "    ),\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23875916-feeb-42bf-b2ba-5c54d7a84907",
   "metadata": {},
   "outputs": [],
   "source": [
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2d158d5-402c-4c8b-ab2e-d2a831d976eb",
   "metadata": {},
   "source": [
    "## Self Conditioning Design\n",
    "\n",
    "Now we'll look at a more sophisticated design problem where a reference structure is provided for motifscaffolding and binder design."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fb27f47-36bf-4942-8e04-df764b3ecd6f",
   "metadata": {},
   "source": [
    "The `SelfConditioningSamplerConfig` is used with a reference structure. The reference structure can be any `protein.Protein` object. Like the `UnconditionalSamplerConfig` we can set potentials and symmetries, but now we can specify `contigs` that incorporate residues from a reference structure. For example:\n",
    "\n",
    "- [\"30-50/B10-20/40-70\"] will scaffold on the `B` chain of a provided reference structure with 30-50 residues before and 40-70 residues after the specified segment of chain `B`.\n",
    "- [\"5-15/A10-25/30-40/0 B1-100\"] will scaffold on chain `A` while accounting for residues 1-100 in chain `B`. The `\\0` together with a trailing space denotes a chain break."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca816c97-28cc-4e5e-9054-3123e245384c",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Motifscaffolding\n",
    "\n",
    "Design a structure using a subchain from a reference structure as the scaffold."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56a25a4b-d084-435b-886a-6d21f0d066f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./data/5tpn.pdb', mode=\"r\") as f:\n",
    "    pdb_str = f.read()\n",
    "\n",
    "reference_structure = protein.Protein14.from_pdb_string(pdb_str, parse_hetatom=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fed8240-c281-4c39-bf07-4c902083eeca",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler_config = config.SelfConditioningSamplerConfig(\n",
    "    inference_params=config.InferenceParams(reference_structure=reference_structure),\n",
    "    contigmap_params=config.ContigMap(contigs=[\"10-40/A163-181/10-40\"]),\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e41bb1a8-0cd2-478a-8935-e8e8f7fbff22",
   "metadata": {},
   "outputs": [],
   "source": [
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8c053c2-7a73-42c2-8581-2a946e9875b3",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Motifscaffolding with target\n",
    "\n",
    "Design a structure with respect to a particular target chain by using a subchain from a reference structure as a scaffold."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a41dc79-9a45-42ab-a0b4-554a7d824c5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This could have been done by reinstantiating the pipeline, but the `set_model` method is a convenience\n",
    "# to change the loaded model\n",
    "designer.set_model(\"complex_base\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0fcc437-2f24-476c-a2a7-f4ca548bce21",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./data/1ycr.pdb', mode=\"r\") as f:\n",
    "    pdb_str = f.read()\n",
    "\n",
    "reference_structure = protein.Protein14.from_pdb_string(pdb_str, parse_hetatom=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31173dc0-e113-4a7e-ac7a-6cfbb8409165",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler_config = config.SelfConditioningSamplerConfig(\n",
    "    inference_params=config.InferenceParams(reference_structure=reference_structure),\n",
    "    # Length here denotes the total allowable length range of the generated oligomer\n",
    "    contigmap_params=config.ContigMap(contigs=[\"A25-109/0 0-70/B17-29/0-70\"], length=\"70-120\"),\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed65880f-5b60-478d-94c5-010ca8d5679e",
   "metadata": {},
   "outputs": [],
   "source": [
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc13adb9-d130-466e-aab8-a167e30564eb",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Enzyme\n",
    "\n",
    "Design a structure to bind with an enzyme, use a potential to encourage contacts with the substrate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ebb42d4-32c3-41d8-a981-d0a080ee4f48",
   "metadata": {},
   "outputs": [],
   "source": [
    "designer.set_model(\"active_site\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14d48823-db7a-47bc-925d-cf1bc451004a",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"./data/5an7.pdb\", mode=\"r\") as f:\n",
    "    pdb_str = f.read()\n",
    "    \n",
    "reference_structure = protein.Protein14.from_pdb_string(pdb_str, parse_hetatom=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5cdf3f1-224b-4458-be5e-3532145b0326",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler_config = config.SelfConditioningSamplerConfig(\n",
    "    inference_params=config.InferenceParams(reference_structure=reference_structure),\n",
    "    contigmap_params=config.ContigMap(contigs=[\"10-100/A1083-1083/10-100/A1051-1051/10-100/A1180-1180/10-100\"]),\n",
    "    potentials_params=config.PotentialsParams(\n",
    "        guiding_potentials=[\"type:substrate_contacts,s:1,r_0:8,rep_r_0:5.0,rep_s:2,rep_r_min:1\"], \n",
    "        guide_scale=1,\n",
    "        guide_decay=\"quadratic\",\n",
    "        substrate=\"LLK\",\n",
    "    ),\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a42b6c5-c666-4ed0-a404-530bbe868cf4",
   "metadata": {},
   "outputs": [],
   "source": [
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "396a6605-79e1-4140-9e77-6fa0e0a33451",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Nickel Motif\n",
    "\n",
    "Design a symmetric structure where each chain scaffolds on top of the same reference chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "304f77e4-0537-48c5-a08b-8683040daf5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "designer.set_model(\"base_epoch8\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61133bfe-159d-4317-8c3f-07aaa7b0935f",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./data/nickel_motif.pdb', mode=\"r\") as f:\n",
    "    pdb_str = f.read()\n",
    "    \n",
    "reference_structure = protein.Protein14.from_pdb_string(pdb_str, parse_hetatom=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76c6e3b3-990d-43d4-bc25-3f68c6cbd676",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler_config = config.SelfConditioningSamplerConfig(\n",
    "    inference_params=config.InferenceParams(reference_structure=reference_structure),\n",
    "    contigmap_params=config.ContigMap(contigs=[\"50/A2-4/50/0 50/A7-9/50/0 50/A12-14/50/0 50/A17-19/50/0\"]),\n",
    "    symmetry_params=config.SymmetryParams(symmetry=\"C4\"),\n",
    "    potentials_params=config.PotentialsParams(\n",
    "        guiding_potentials=[\"type:olig_contacts,weight_intra:1,weight_inter:0.06\"],\n",
    "        olig_inter_all=True,\n",
    "        olig_intra_all=True,\n",
    "        guide_scale=2,\n",
    "        guide_decay=\"quadratic\",\n",
    "    ),\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84ce0562-34e1-4759-96eb-ceda0b3ebf9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4823d69-4c7b-4b90-af58-0f4820117f7c",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Insulin PPI\n",
    "\n",
    "Design a structure for a target with hotspot residues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97bc87af-552c-4772-b748-ca248287b7c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "designer.set_model(\"auto\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d9938cd-d6a1-44b6-8ae4-98f50add0417",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./data/insulin_target.pdb', mode=\"r\") as f:\n",
    "    pdb_str = f.read()\n",
    "\n",
    "reference_structure = protein.Protein14.from_pdb_string(pdb_str, parse_hetatom=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74fafd16-ae56-4466-8123-caad0f5b841e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler_config = config.SelfConditioningSamplerConfig(\n",
    "    inference_params=config.InferenceParams(reference_structure=reference_structure),\n",
    "    contigmap_params=config.ContigMap(contigs=[\"A1-150/0 70-100\"]),\n",
    "    ppi_params=config.PPIParams(hotspot_res=[\"A59\", \"A83\", \"A91\"]),\n",
    "    denoiser_params=config.DenoiserParams(noise_scale_ca=0, noise_scale_frame=0),\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dc2cd8d-8a2f-4278-bf63-41f721def283",
   "metadata": {},
   "outputs": [],
   "source": [
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dec611a-5bde-43ed-b0a0-40c73cb53424",
   "metadata": {},
   "source": [
    "### Insulin PPI Beta Model\n",
    "\n",
    "Design a structure for a target with hotspot residues using a model that favors generating structures with beta sheets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c81435ef-9944-4dbc-ae87-a1d098a4f7bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "designer.set_model(\"complex_beta\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b375a514-8a5f-4641-a7e6-367b3e8338e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./data/insulin_target.pdb', mode=\"r\") as f:\n",
    "    pdb_str = f.read()\n",
    "\n",
    "reference_structure = protein.Protein14.from_pdb_string(pdb_str, parse_hetatom=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1932631-f3b7-42f1-9472-9b0182ea6b81",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler_config = config.SelfConditioningSamplerConfig(\n",
    "    inference_params=config.InferenceParams(reference_structure=reference_structure),\n",
    "    contigmap_params=config.ContigMap(contigs=[\"A1-150/0 70-100\"]),\n",
    "    ppi_params=config.PPIParams(hotspot_res=[\"A59\", \"A83\", \"A91\"]),\n",
    "    denoiser_params=config.DenoiserParams(noise_scale_ca=0, noise_scale_frame=0),\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fd1f99a-48cd-453a-8956-69bee4fcf79a",
   "metadata": {},
   "outputs": [],
   "source": [
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f4979bf-5402-4191-9936-b0295b6460f8",
   "metadata": {},
   "source": [
    "## Sequence Inpainting\n",
    "\n",
    "Sequence inpainting uses a reference structure as a scaffold but RFDiffusion is allowed to redesign parts of the scaffold as needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7afaf0b9-a4cb-419c-8a68-fc96d564f7a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "designer.set_model(\"auto\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04bb72c1-aa41-4fd4-b50e-4f7690a9597a",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"./data/5tpn.pdb\", mode=\"r\") as f:\n",
    "    pdb_str = f.read()\n",
    "    \n",
    "reference_structure = protein.Protein14.from_pdb_string(pdb_str, parse_hetatom=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f28ec3c-85af-426f-9189-b91fc6fef5f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler_config = config.SelfConditioningSamplerConfig(\n",
    "    inference_params=config.InferenceParams(reference_structure=reference_structure),\n",
    "    contigmap_params=config.ContigMap(contigs=[\"10-40/A163-181/10-40\"], inpaint_seq=[\"A163-168/A170-171/A179\"]),\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9017260a-4960-47ba-8cb0-848fd2abbb2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d2be0b8-acce-47ad-9596-c2f359073e90",
   "metadata": {},
   "source": [
    "## Partial Diffusion\n",
    "\n",
    "Partial diffusion adds some noise to a reference structure and then denoises the reference structure."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a2460ab-647a-4210-8187-af784337ddc1",
   "metadata": {},
   "source": [
    "### Without Sequence\n",
    "\n",
    "Noise the whole reference structure and then denoise it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a25e543-688d-41b3-9960-6dfebf68e352",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./data/2kl8.pdb', mode=\"r\") as f:\n",
    "    pdb_str = f.read()\n",
    "\n",
    "reference_structure = protein.Protein14.from_pdb_string(pdb_str, parse_hetatom=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a23acc9-f911-49f1-a230-8be87d00c9b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "diffuser_config_override = config.DiffuserConfig(partial_T=10)\n",
    "sampler_config = config.SelfConditioningSamplerConfig(\n",
    "    inference_params=config.InferenceParams(reference_structure=reference_structure),\n",
    "    contigmap_params=config.ContigMap(contigs=[\"79-79\"]),\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config, diffuser_config_override=diffuser_config_override)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92e710d1-a6cd-4cb7-bed3-ac0b5edb384a",
   "metadata": {},
   "outputs": [],
   "source": [
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0804431-f53e-4067-94c7-27f0246cc531",
   "metadata": {},
   "source": [
    "### With Sequence\n",
    "\n",
    "Noise a simplistic structure and peptide binding model and denoise it to make something more plausible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1c972e8-c313-4dc4-907c-c7c32dfaa8ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./data/peptide_complex.pdb', mode=\"r\") as f:\n",
    "    pdb_str = f.read()\n",
    "\n",
    "reference_structure = protein.Protein14.from_pdb_string(pdb_str, parse_hetatom=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d540d5ba-66ad-4b0a-9716-f697970fc2ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "diffuser_config_override = config.DiffuserConfig(partial_T=10)\n",
    "sampler_config = config.SelfConditioningSamplerConfig(\n",
    "    inference_params=config.InferenceParams(reference_structure=reference_structure),\n",
    "    contigmap_params=config.ContigMap(contigs=[\"172-172/0 34-34\"], provide_seq=[\"172-205\"]),\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config, diffuser_config_override=diffuser_config_override)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b121d87-90b0-41e1-8679-17592690890d",
   "metadata": {},
   "outputs": [],
   "source": [
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8117f6d3-9fb5-46cc-8469-85f3f70afa4c",
   "metadata": {},
   "source": [
    "### With Multisequence\n",
    "\n",
    "Noise a simplistic structure and peptide binding model and denoise it to make something more plausible, but this time specify multiple disjoint sequences to hold constant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3da633b5-212e-4230-910f-265f27972ff2",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./data/peptide_complex.pdb', mode=\"r\") as f:\n",
    "    pdb_str = f.read()\n",
    "\n",
    "reference_structure = protein.Protein14.from_pdb_string(pdb_str, parse_hetatom=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f22d063d-088f-4c3b-bd44-97d0a9294fcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "diffuser_config_override = config.DiffuserConfig(partial_T=10)\n",
    "sampler_config = config.SelfConditioningSamplerConfig(\n",
    "    inference_params=config.InferenceParams(reference_structure=reference_structure),\n",
    "    contigmap_params=config.ContigMap(contigs=[\"172-172/0 34-34\"], provide_seq=[\"172-177,200-205\"]),\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config, diffuser_config_override=diffuser_config_override)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2cdc29e-d314-42d0-bafa-c980831bf010",
   "metadata": {},
   "outputs": [],
   "source": [
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72258f0e-fd7d-471c-baa1-fe2c060400ed",
   "metadata": {},
   "source": [
    "## Scaffold Guided\n",
    "\n",
    "Here we can specify a target protein and tell RFdiffusion that we want to do \"scaffoldguided\" diffusion (i.e. we want to specify the fold of the protein)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b12bdb8f-5771-4ab8-931d-f9a413e5d281",
   "metadata": {},
   "source": [
    "The `ScaffoldedSamplerConfig` can be used with or without a reference structure. For scaffoldguided design, a list of scaffold structures is provided and, optionally, a target structure (e.g., binder) can be given. Like other samplers, we can define potentials and symmetries as desired."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acbbc3a0-287f-4feb-b346-d1a72ceab2ff",
   "metadata": {},
   "source": [
    "### TIM Barrel\n",
    "\n",
    "Scaffoldguided design using a `TIM` barrel as the scaffold structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25ae120c-a6f2-4306-9257-e2e4d6085e3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./data/1qys.pdb', mode=\"r\") as f:\n",
    "    pdb_str = f.read()\n",
    "    \n",
    "reference_structure = protein.Protein14.from_pdb_string(pdb_str, parse_hetatom=True)\n",
    "\n",
    "with open('./data/tim10.pdb', mode=\"r\") as f:\n",
    "    scaffold_pdb_str = f.read()\n",
    "\n",
    "scaffold_structure = protein.Protein14.from_pdb_string(scaffold_pdb_str, parse_hetatom=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9984e929-d371-47ed-ac1b-bff9229f1ef6",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sampler_config = config.ScaffoldedSamplerConfig(\n",
    "    inference_params=config.InferenceParams(reference_structure=reference_structure),\n",
    "    denoiser_params=config.DenoiserParams(noise_scale_ca=0.5, noise_scale_frame=0.5),\n",
    "    scaffoldguided_params=config.ScaffoldGuidedParams(\n",
    "        target_structure=None,\n",
    "        target_adj=False,\n",
    "        target_ss=False,\n",
    "        scaffold_structure_list=[scaffold_structure],\n",
    "        sampled_insertion=\"0-5\",\n",
    "        sampled_N=\"0-5\",\n",
    "        sampled_C=\"0-5\",\n",
    "    )\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "268805bb-5a2f-4108-b3fb-3c5ddc732aa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db723652-2f18-4d00-8436-f9f6dbd180db",
   "metadata": {},
   "source": [
    "### PPI Scaffolded\n",
    "\n",
    "Scaffoldguided design with a binder and hotspot residues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9e02157-7dee-43fe-849c-021f3bcdac89",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./data/1qys.pdb', mode=\"r\") as f:\n",
    "    pdb_str = f.read()\n",
    "reference_structure = protein.Protein14.from_pdb_string(pdb_str, parse_hetatom=True)\n",
    "\n",
    "with open('./data/insulin_target.pdb', mode=\"r\") as f:\n",
    "    target_pdb_str = f.read()\n",
    "target_structure = protein.Protein14.from_pdb_string(target_pdb_str, parse_hetatom=True)\n",
    "\n",
    "with open('./data/5L33.pdb', mode=\"r\") as f:\n",
    "    scaffold_pdb_str = f.read()\n",
    "scaffold_structure = protein.Protein14.from_pdb_string(scaffold_pdb_str, parse_hetatom=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcab0cbc-4991-485e-a17e-87cb26c06d94",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sampler_config = config.ScaffoldedSamplerConfig(\n",
    "    inference_params=config.InferenceParams(reference_structure=reference_structure),\n",
    "    denoiser_params=config.DenoiserParams(noise_scale_ca=0, noise_scale_frame=0),\n",
    "    ppi_params=config.PPIParams(hotspot_res=[\"A59\", \"A83\", \"A91\"]),\n",
    "    scaffoldguided_params=config.ScaffoldGuidedParams(\n",
    "        target_structure=target_structure,\n",
    "        target_adj=True,\n",
    "        target_ss=True,\n",
    "        scaffold_structure_list=[scaffold_structure],\n",
    "        sampled_insertion=\"0-5\",\n",
    "        sampled_N=\"0-5\",\n",
    "        sampled_C=\"0-5\",\n",
    "    )\n",
    ")\n",
    "designed_structure, _ = designer(sampler_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3b2b4a0-3031-46a5-a55d-fecd4142a1d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "designed_structure.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fe0f820-6bcc-45cb-9191-4a8469ce5986",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}