nbs/_example_prokaryotic.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp simulations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DiffPALM – Example usage on prokaryotic datasets\n",
    "\n",
    "> API details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from diffpalm.core import DiffPALM\n",
    "from diffpalm.msa_parsing import read_msa\n",
    "from diffpalm.datasets import generate_dataset, dataset_tokenizer\n",
    "\n",
    "from pathlib import Path\n",
    "from datetime import datetime\n",
    "import pickle\n",
    "\n",
    "\n",
    "def save_parameters(parameters_all, filepath):\n",
    "    \"\"\"Saves the parameters dictionary\"\"\"\n",
    "    for name, parameters in parameters_all.items():\n",
    "        with open(filepath / f\"{name}.pkl\", \"wb\") as f:\n",
    "            pickle.dump(parameters, f)\n",
    "        with open(filepath / f\"{name}.csv\", \"w\") as f:\n",
    "            for key in parameters.keys():\n",
    "                f.write(\"%s, %s\\n\" % (key, parameters[key]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RESULTS_DIR = Path(input(prompt=\"Insert path to directory where results will be stored (default: 'Results'): \") or \"Results\")\n",
    "RESULTS_DIR.mkdir(exist_ok=True)\n",
    "\n",
    "DEVICE = input(prompt=\"Choose PyTorch device (default: 'cuda'): \") or \"cuda\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load one of the two prokaryotic datasets used in our paper: HK-RR and MALG-MALK.\n",
    "\n",
    "`get_species_name` extracts species names from the FASTA header.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PROKARYOTIC DATASETS\n",
    "\n",
    "msa_data = [read_msa(\"data/HK-RR/HK_in_Concat_nnn.fasta\", -1),\n",
    "            read_msa(\"data/HK-RR/RR_in_Concat_nnn.fasta\", -1)]\n",
    "get_species_name = (lambda strn: strn.split(\"|\")[1])\n",
    "\n",
    "# msa_data = [read_msa(\"data/MALG-MALK/MALG_cov75_hmmsearch_extr5000_withLast_b.fasta\", -1),\n",
    "#             read_msa(\"data/MALG-MALK/MALK_cov75_hmmsearch_extr5000_withLast_b.fasta\", -1)]\n",
    "# get_species_name = (lambda strn: strn.split(\"_\")[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters_dataset = {\n",
    "    \"N\": 50,  # Average number of sequences in the input\n",
    "    \"pos\": 0,  # Size of the context pairs to use as positive example \n",
    "    \"max_size\": 100,  # Max size of species MSAs (if same as N there is no limit on size)\n",
    "    \"NUMPY_SEED\": 10,\n",
    "    \"NUMPY_SEED_OTHER\": 11,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset, species_sizes = generate_dataset(\n",
    "    parameters_dataset, msa_data, get_species_name=get_species_name\n",
    ")\n",
    "tokenized_dataset = dataset_tokenizer(dataset, device=DEVICE)\n",
    "\n",
    "left_msa, right_msa = tokenized_dataset[\"msa\"][\"left\"], tokenized_dataset[\"msa\"][\"right\"]\n",
    "positive_examples = tokenized_dataset[\"positive_examples\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train single block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EPOCHS = 100\n",
    "TORCH_SEED = 100\n",
    "\n",
    "parameters_init = {\n",
    "    \"device\": DEVICE,\n",
    "    \"p_mask\": 0.7,\n",
    "    \"random_seed\": TORCH_SEED\n",
    "}\n",
    "\n",
    "parameters_train = {\n",
    "    \"std_init\": 0.,\n",
    "    \"scheduler_name\": \"ReduceLROnPlateau\",\n",
    "    \"scheduler_kwargs\": {\"mode\": \"min\", \"factor\": 0.8, \"patience\": 20},\n",
    "    \"optimizer_name\": \"Adadelta\",\n",
    "    \"optimizer_kwargs\": {\"lr\": 9, \"weight_decay\": 1e-1},\n",
    "    \"tau\": 1.,\n",
    "    \"n_sink_iter\": 10,\n",
    "    \"batch_size\": 1,\n",
    "    \"epochs\": EPOCHS,\n",
    "    \"noise\": True,\n",
    "    \"noise_factor\": 0.1,  # If noise_std is False, this is just the std of the noise\n",
    "    \"noise_scheduler\": True,\n",
    "    \"noise_std\": True,\n",
    "    \"use_rand_perm\": True,\n",
    "}\n",
    "\n",
    "parameters_target_loss = {\n",
    "    \"batch_size\": 200\n",
    "}\n",
    "\n",
    "parameters_all = {\n",
    "    \"init\": parameters_init,\n",
    "    \"target_loss\": parameters_target_loss,\n",
    "    \"train\": parameters_train,\n",
    "    \"dataset\": parameters_dataset\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dpalm = DiffPALM(species_sizes, **parameters_init)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "date = datetime.now().strftime(\"%Y_%m_%d-%H:%M:%S\")\n",
    "output_dir = RESULTS_DIR / date\n",
    "output_dir.mkdir()\n",
    "\n",
    "save_parameters(parameters_all, output_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When `save_all_figs=True`, a figure is saved and shown after each gradient step, illustrating the current state of the optimization. This slows the overall optimization down and may create memory leakage issues. Set `save_all_figs=False` to only have the figure saved and shown after the last gradient step."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The plotting function is able to show the number of correctly predicted pairs because the ground truth pairs are known. The model assumes that the input pairs are already correctly matched (i.e. the correct matching matrix is a diagonal matrix) because in the HK-RR and MALG-MALK datasets the sequences are are already ordered with the correct matches in the same position of the MSA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tar_loss = dpalm.target_loss(\n",
    "    left_msa,\n",
    "    right_msa,\n",
    "    positive_examples=positive_examples,\n",
    "    **parameters_target_loss\n",
    ")\n",
    "\n",
    "(losses,\n",
    " list_scheduler,\n",
    " shuffled_indexes,\n",
    " mat_perm,\n",
    " mat_gs,\n",
    " list_log_alpha) = dpalm.train(\n",
    "    left_msa,\n",
    "    right_msa,\n",
    "    positive_examples=positive_examples,\n",
    "    tar_loss=np.mean(tar_loss),\n",
    "    output_dir=output_dir,\n",
    "    save_all_figs=True,\n",
    "    **parameters_train,\n",
    ")\n",
    "\n",
    "results = {\n",
    "    \"trainng_results\": (losses, list_scheduler, shuffled_indexes, [mat_perm, mat_gs], list_log_alpha),\n",
    "    \"target_loss\": tar_loss,\n",
    "    \"species_sizes\": species_sizes\n",
    "}"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}