{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chroma import api\n",
    "api.register_key(\"2cdade6d058b4fd1b85fa5badb501312\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# unconditional sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chroma import Chroma\n",
    "\n",
    "chroma = Chroma()\n",
    "protein = chroma.sample(chain_lengths=[300])\n",
    "\n",
    "protein.to(\"sample.cif\")\n",
    "\n",
    "display(protein)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# conditional sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from chroma import Chroma, conditioners\n",
    "\n",
    "chroma = Chroma()\n",
    "conditioner = conditioners.SymmetryConditioner(G=\"C_3\", num_chain_neighbors=2)\n",
    "protein = chroma.sample(\n",
    "    chain_lengths=[100],\n",
    "    conditioner=conditioner,\n",
    "    langevin_factor=8,\n",
    "    inverse_temperature=8,\n",
    "    sde_func=\"langevin\",\n",
    "    potts_symmetry_order=conditioner.potts_symmetry_order)\n",
    "\n",
    "protein.to(\"sample-C3.cif\")\n",
    "display(protein)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# design"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Redesign a Protein\n",
    "from chroma import Protein, Chroma\n",
    "chroma = Chroma()\n",
    "\n",
    "protein = Protein('1GFP',device='cuda')\n",
    "protein = chroma.design(protein)\n",
    "\n",
    "\n",
    "protein.to(\"1GFP-redesign.cif\")\n",
    "display(protein)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Redesign"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Redesign a Protein\n",
    "from chroma import Protein, Chroma\n",
    "chroma = Chroma()\n",
    "\n",
    "protein = Protein('sample.cif',device='cuda') # PDB is fine too\n",
    "protein = chroma.design(protein, design_selection=\"resid 20-50 around 5.0\") #  5 angstrom bubble around indices 20-50\n",
    "\n",
    "protein.to(\"my_favorite_protein_redesign.cif\")\n",
    "display(protein)"
   ]
  }
 ],
 "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}