{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8f70aab6-4bd7-4037-802b-58f060d35ec9",
   "metadata": {},
   "source": [
    "## Tutorials: Chromosome Chromosome Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "485057a3-4d7c-4981-a754-5273dbc7ffa0",
   "metadata": {},
   "source": [
    "### Import the necessary packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7fbf57dc-ce2f-4720-bc4e-75de146108d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import subprocess\n",
    "import sys\n",
    "import copy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc6ff132-cb12-4b4e-b2f7-e2ca9b421cc8",
   "metadata": {},
   "source": [
    "### Important parameters\n",
    "We defined the important parameters in the next blocks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ed5cfd82-e410-435b-a60c-c13734a4133a",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_number = 1 #The index of the current iteration\n",
    "N_replicas = 1 #The number of the total replicas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ead03225-6bd3-410e-995c-2f78416e268e",
   "metadata": {},
   "outputs": [],
   "source": [
    "nAtom  = 60642 #The number of the atoms in our system\n",
    "nIdeal = 2490-1 #The number of the ideal parameters\n",
    "ncompt = 3 #The number of the compartments\n",
    "nChrom = 22 #The number of the chromosomes including inter-chromosomal interactions, excluding the sex chromosome \n",
    "nInter = int(nChrom*(nChrom-1)/2)\n",
    "ncv    = int(nIdeal + ncompt*(ncompt+1)/2 + nInter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "95244d5f-061c-4bbf-872b-d95a6dbd98b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "chop_ideals             = 1000\n",
    "\n",
    "## Adam optimizer parameters\n",
    "m_dw                    = np.loadtxt('adam_chr_chr_param/%02d/mdw.txt'%(run_number-1))\n",
    "v_dw                    = np.loadtxt('adam_chr_chr_param/%02d/vdw.txt'%(run_number-1))\n",
    "m_db                    = np.loadtxt('adam_chr_chr_param/%02d/mdb.txt'%(run_number-1))\n",
    "v_db                    = np.loadtxt('adam_chr_chr_param/%02d/vdb.txt'%(run_number-1))\n",
    "beta1                   = 0.9\n",
    "beta2                   = 0.999\n",
    "epsilon                 = 1e-8\n",
    "t                       = int(np.loadtxt('adam_chr_chr_param/%02d/t.txt'%(run_number-1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "67aa6151-04a8-42d6-8601-d69bb1918b94",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_cv                = 1 #The starting index of the parameters that we hope to optimize\n",
    "end_cv                  = 2726 #The ending index of the parameters that we hope to optimize\n",
    "\n",
    "old_iter                = run_number-1\n",
    "\n",
    "cvInd   = np.zeros((ncv, ), dtype=float)\n",
    "irun    = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c56836b4-0aaa-48b9-b701-d55d79858d5c",
   "metadata": {},
   "source": [
    "We normalize the contact probabilities firstly, and people can have their own path for the contact probs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ffbc8837-4396-4579-bd70-b2450488bac9",
   "metadata": {},
   "outputs": [],
   "source": [
    "for replica in range(1,N_replicas+1,1): #We show the situation when the N_replicas = 1\n",
    "    #If simulation didn't complete\n",
    "    if os.path.exists(\"contact_prob/contact_prob.txt\"):\n",
    "        cvInd   += np.loadtxt(\"contact_prob/contact_prob.txt\")\n",
    "        irun    += np.loadtxt(\"contact_prob/nframes.txt\")\n",
    "cvInd /= irun"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c051ff9-1bc0-4ee3-beb7-4c9bbb01308c",
   "metadata": {},
   "source": [
    "We load the experimental results and scale them to the same magnitude of the simulated contact probabilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dfbe437b-c19c-43bc-b8a3-8dc6ee31c486",
   "metadata": {},
   "outputs": [],
   "source": [
    "expt        = np.loadtxt(\"expt_constraints_HFF_100KB.txt\")\n",
    "expt        = cvInd[0]/expt[0]*expt #We scale the experimental contacts to the same magnitude of simulated results\n",
    "\n",
    "grad        = -cvInd + expt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a2e1fdc-ebb9-44a1-8cdc-27f9df78face",
   "metadata": {},
   "source": [
    "### Adam optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2b4ff4fb-83ef-46c5-90cf-4166876bb5ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "## START TO DO THE ADAM OPTIMIZATION\n",
    "## momentum beta 1\n",
    "# *** weights *** #\n",
    "m_dw        = beta1*m_dw + (1-beta1)*grad\n",
    "# *** biases *** #\n",
    "m_db        = beta1*m_db + (1-beta1)*grad\n",
    "## rms beta 2\n",
    "# *** weights *** #\n",
    "v_dw        = beta2*v_dw + (1-beta2)*(grad**2)\n",
    "# *** biases *** #\n",
    "v_db        = beta2*v_db + (1-beta2)*grad"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dff16b0a-4381-46e8-b841-73bceb1f6f7a",
   "metadata": {},
   "source": [
    "We save the updated parameters in Adam for the next iteration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "eb821b7d-e4bb-4a2e-af83-7ab3d37a5e9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "subprocess.call([\"mkdir -p adam_chr_chr_param/%02d\"%(run_number)],shell=True,stdout=subprocess.PIPE)\n",
    "np.savetxt('adam_chr_chr_param/%02d/mdw.txt'%(run_number), m_dw.reshape((-1,1)), fmt='%15.12e')\n",
    "np.savetxt('adam_chr_chr_param/%02d/vdw.txt'%(run_number), v_dw.reshape((-1,1)), fmt='%15.12e')\n",
    "np.savetxt('adam_chr_chr_param/%02d/mdb.txt'%(run_number), m_db.reshape((-1,1)), fmt='%15.12e')\n",
    "np.savetxt('adam_chr_chr_param/%02d/vdb.txt'%(run_number), v_db.reshape((-1,1)), fmt='%15.12e')\n",
    "np.savetxt('adam_chr_chr_param/%02d/t.txt'%(run_number), np.array([t+1]).reshape((-1,1)), fmt='%d')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e78aee5d-8097-4c31-bf75-80059c9b904b",
   "metadata": {},
   "outputs": [],
   "source": [
    "## bias correction\n",
    "m_dw_corr   = m_dw/(1-beta1**t)\n",
    "m_db_corr   = m_db/(1-beta1**t)\n",
    "v_dw_corr   = v_dw/(1-beta2**t)\n",
    "v_db_corr   = v_db/(1-beta2**t)\n",
    "\n",
    "dalpha     = m_dw_corr/(np.sqrt(v_dw_corr)+epsilon)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9abc0e39-5f44-4e2e-bfd5-f66ba2ba27e4",
   "metadata": {},
   "source": [
    "### Start to update the parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "157b07ef-5433-43c9-af48-4012be987d95",
   "metadata": {},
   "outputs": [],
   "source": [
    "#We set all the default learning rates as 0.001\n",
    "eta1           = 0.001\n",
    "eta2           = 0.001\n",
    "eta3           = 0.001\n",
    "\n",
    "#Load the old potential\n",
    "ideal_old         = np.loadtxt('potential/%02d/ideal_param_file.txt'%(run_number-1), usecols=[1])\n",
    "compt_old         = np.loadtxt('potential/%02d/compt_param_file.txt'%(run_number-1), usecols=[2])\n",
    "inter_old         = np.loadtxt('potential/%02d/interchr_param_file.txt'%(run_number-1), usecols=[2])\n",
    "\n",
    "ideal_new = copy.deepcopy(ideal_old)\n",
    "compt_new = copy.deepcopy(compt_old)\n",
    "inter_new = copy.deepcopy(inter_old)\n",
    "\n",
    "ideal_new[2:chop_ideals+1]          -= dalpha[1:chop_ideals]*eta1\n",
    "compt_new[np.array([0,1,2,4,5,7])]  -= dalpha[nIdeal:nIdeal+ncompt*(ncompt+1)//2]*eta2 #The 6 positions represent the 6 type-type interactions\n",
    "inter_new                           -= dalpha[nIdeal+ncompt*(ncompt+1)//2:]*eta3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e08f96d5-e58c-4545-9bba-c6f235a90b2b",
   "metadata": {},
   "source": [
    "Save all the updated parameters as the next iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "54b21a76-cdca-4f7d-96e1-ab94357166d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Save the new potential\n",
    "subprocess.call([\"mkdir -p potential/%02d\"%(run_number)],shell=True,stdout=subprocess.PIPE)\n",
    "\n",
    "ideal_idx = np.loadtxt('potential/%02d/ideal_param_file.txt'%(run_number-1), usecols=[0], dtype='int')\n",
    "compt_idx = np.loadtxt('potential/%02d/compt_param_file.txt'%(run_number-1), usecols=[0,1], dtype='int')\n",
    "inter_idx = np.loadtxt('potential/%02d/interchr_param_file.txt'%(run_number-1), usecols=[0,1], dtype='int')\n",
    "\n",
    "np.savetxt('potential/%02d/ideal_param_file.txt'%run_number,\n",
    "           np.hstack((ideal_idx.reshape((-1,1)), ideal_new.reshape((-1,1)))), fmt='%d %.6f')\n",
    "np.savetxt('potential/%02d/compt_param_file.txt'%run_number,\n",
    "           np.hstack((compt_idx, compt_new.reshape((-1,1)))), fmt='%d %d %.6f')\n",
    "np.savetxt('potential/%02d/interchr_param_file.txt'%run_number,\n",
    "           np.hstack((inter_idx, inter_new.reshape((-1,1)))), fmt='%d %d %.6f')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}