{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pulling Simulations\n",
    "\n",
    "<font color='red'>Pulling Simulations is available in OpenMichroM versions >= 1.0.7</font> \n",
    "\n",
    "March 19, 2023.  Tutorial written by Ben Ruben (benruben@g.harvard.edu)\n",
    "\n",
    "In this tutoral, we will demonstrate how to simulate a constant-force and constant-distance pulling of mitotic chromosomes using the openMiChroM library.  These methods were used in the paper \"Structural Reogranization and Relaxation Dynamics of Axially Stressed Chromosomes,\" published in the Biophysical Journal (DOI:https://doi.org/10.1016/j.bpj.2023.03.029).  This tutorial constains two parts:\n",
    "\n",
    "##### Part 1: Constant-Force Pulling\n",
    "1. Load a polymer from a condensed starting conformation and define the Homopolymer, Ideal Chromosome, and Type-Type potentials\n",
    "2. Define and add the \"Pin\" and \"Slide\" pulling potentials and equilibrate.\n",
    "3. Define a constant pulling forces which act on the centers of mass of the \"pull groups\" defined as the first and last 50 beads of the polymer.\n",
    "4. Apply constant pulling force and measure end-to-end distance under force.  Save pulled structures for following CD simulations\n",
    "5. Release constnat force and measure extension as the chromosome retracts.\n",
    "\n",
    "##### Part 2: Constant-Distance Pulling\n",
    "In this section, we deminstrate a constant-distane pulling experiment which measured the force-extension curve of the chromosome.  Frames generated in thne constant-force pulling simulation are used as starting structures for simulations.  After the simulations, we plot a rough force-extension curve."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from OpenMiChroM.ChromDynamics import MiChroM\n",
    "\n",
    "import time\n",
    "import numpy as np\n",
    "\n",
    "from sys import stdout, argv\n",
    "import numpy as np\n",
    "from six import string_types\n",
    "import os\n",
    "import time\n",
    "import random\n",
    "import h5py\n",
    "from scipy.spatial import distance\n",
    "import scipy as sp\n",
    "import itertools\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.rcParams.update({'font.size': 16, 'figure.figsize': [6.0, 5.0]})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Part 0: Setting Up Input Files\n",
    "\n",
    "The folder [input](https://github.com/junioreif/OpenMiChroM/tree/main/Tutorials/Chromosome_Pulling_Tutorial/input) contains two files: \n",
    "-DT40_chr7.eigen contains the A/B type labels of each locus in the chromosome.\n",
    "-lambdas_ic_15m contains the ideal chromosome interaction parameters learned for the mitotic DT40 chromosome.\n",
    "\n",
    "We will create two additional input files used to set up the simulations.  The first is a text file which contains the type labels, constructed from the .eigen file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "awk '{if ($1 < 0) print \"B1\"; else print \"A1\"}' input/DT40_chr7.eigen | cat -n > input/DT40_chr7.txt\n",
    "\n",
    "#Show the first 10 lines of the text file:\n",
    "head input/DT40_chr7.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we create a text file containing a table of the type-type interaction parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "echo \"A1,B1\n",
    "-1.644455115938274231e-02, -2.043214285797510181e-03\n",
    "-2.043214285797510181e-03 ,-1.512813001184262077e-02\" > input/lambdas_types_15m\n",
    "\n",
    "cat input/lambdas_types_15m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Auxilary Function Used to Read Coordinates from PDB Files.\n",
    "def load_Coords_PDB(filename):\n",
    "        print('Coords Only')\n",
    "        aFile = open(filename,'r')\n",
    "        pos = aFile.read().splitlines()\n",
    "        x = []\n",
    "        y = []\n",
    "        z = []\n",
    "\n",
    "        for t in range(len(pos)):\n",
    "            pos[t] = pos[t].split()\n",
    "            if pos[t][0] == 'ATOM':\n",
    "                x.append(float(pos[t][5]))\n",
    "                y.append(float(pos[t][6]))\n",
    "                z.append(float(pos[t][7]))\n",
    "\n",
    "        return np.vstack([x,y,z]).T"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Part 1: Constant-Force Simulation (Stress-Relaxation Experiment)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Simulation Parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "force = 5\n",
    "\n",
    "#platform = 'cuda'#Use cuda for nvidia GPUs.  \n",
    "#Note: GPU must support 64-bit atomic calculations to use CustomCentroidBondForce class in openmm\n",
    "platform = 'cpu' #use cpu if gpu supporting 64-bit atomic graphics is not available.  \n",
    "\n",
    "blockSize = 10 #Pull Coordinate is recorded every block, whicih is every 10 steps\n",
    "bs = blockSize\n",
    "numBlocks_nat = 2*10**3 #Number of simulation blocks to simulate before applying pulling force\n",
    "numBlocks_pull = 10**4 #Number of simulation blocks to run under constant pulling force\n",
    "numBlocks_release = 2*10**3 #Number of simulation blocks to run after releasing constant force.\n",
    "\n",
    "blocksPerFrame = 500 #Positions of all beads are recorded every frame\n",
    "\n",
    "CDInits = []#List of filenames to be used later for CD simulations"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Function that retrieves pull coordinate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#This function retrieves the pull coordinate, defined as the distance between the centers of mass of the pull groups, along the x-axis.\n",
    "def getPCoord(group1, group2, positions):\n",
    "    first_centroid = np.mean(positions[group1], axis=0) \n",
    "    second_centroid = np.mean(positions[group2], axis=0)\n",
    "\n",
    "    ## calculate r0 distance between groups\n",
    "\n",
    "    return np.sqrt( (first_centroid[0] - second_centroid[0])**2 )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Functions to create pin, slide, and pull forces"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use center of mass pulling actng on the pull groups, defined as the first and last 50 beads in the polymer:\n",
    "\n",
    "\\begin{equation} \\label{COMs}\n",
    "\\begin{split}\n",
    "\\mathbf{R}_{left} = &(x_l, y_l, z_l) = \\frac{1}{50}\\sum_{i=1}^{50} \\mathbf{r}_i \\\\\n",
    "\\mathbf{R}_{right} = &(x_r, y_r, z_r) = \\frac{1}{50}\\sum_{i=N-49}^{N} \\mathbf{r}_i\n",
    "\\end{split}\n",
    "\\end{equation}\n",
    "\n",
    "We define a pull coordinate $\\xi = x_r - x_l$, as the chromosome's end-to-end x-distance.  In chromosome pulling experiments, the ends of the chromosome are held by micropipettes and one is moved along a linear track.  To mimic this process, we introduce orientation constraints:\n",
    "\\begin{equation} \\label{PinSlide_Main}\n",
    "\\begin{split}\n",
    "&\\mathbf{U}_{pin} = \\frac{1}{2} k_r (x_l^2 + y_l^2 + z_l^2) \\\\\n",
    "&\\mathbf{U}_{slide} = \\frac{1}{2} k_r (y_r^2 + z_r^2)\n",
    "\\end{split}\n",
    "\\end{equation}\n",
    "$\\mathbf{U}_{pin}$ restrains the center of mass of the left pull group to the origin.  $\\mathbf{U}_{slide}$ restrains the right pull group to the x-axis, but allows it to slide along the x-axis.  A large value $k_r = 10^5 \\epsilon/\\sigma^2$ is used so that $\\mathbf{R}_{left} \\approx 0$ and $y_r, z_r \\approx 0$.  We use two methods to subject the simulated chromosomes to axial strain.  In constant-force (CF) pulling, a linear potential is used to apply a constant elongating force to the pull coordinate $\\xi$.  In constant-distance (CD) pulling, a harmonic potential with a very strong spring constant is used to constrain $\\xi$ to a very small window around a chosen reference distance $\\xi_0$.\n",
    "\\begin{equation} \\label{CFCDPotentials}\n",
    "\\begin{split}\n",
    "&\\mathbf{U}_{CF} = - F \\xi\\\\\n",
    "&\\mathbf{U}_{CD} = \\frac{1}{2} k_p (\\xi -\\xi_0)^2\n",
    "\\end{split}\n",
    "\\end{equation}\n",
    "A large value of $k_p = 10^5 \\epsilon/\\sigma^2$ is chosen so that $\\xi \\approx \\xi_0$ during CD sampling.\n",
    "\n",
    "These potentials are implemented in the functions in the cell below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#This function is used in Constant-Distance pulling experiments.  It defines a harmonic restraining potential on the pull coordinate.\n",
    "def harmonic_pull_force(group1, group2, r0, kp):\n",
    "    import simtk.openmm as openmm\n",
    "    \n",
    "    #r0 = r0 * units.meter * 1e-9\n",
    "    pullequation = \"0.5 * kpull * ((x2-x1)- rp)^2\" #Enforces x2>x1\n",
    "\n",
    "    pullforce = openmm.CustomCentroidBondForce(2, pullequation)\n",
    "\n",
    "    pullforce.addGlobalParameter('kpull', kp)\n",
    "    pullforce.addGlobalParameter('rp', r0)\n",
    "    pullforce.addGroup(group1)\n",
    "    pullforce.addGroup(group2)\n",
    "    pullforce.addBond([0,1])\n",
    "    #pullforce.setForceGroup(8)\n",
    "    return(pullforce)\n",
    "\n",
    "#This function is used in Constant-Force pulling.  It applies a constant force between the centers of mass of the pull groups.  Written as a potential energy, this is a linear function of the pull coordinate.\n",
    "def constant_pull_force(group1, group2, f):\n",
    "    import simtk.openmm as openmm\n",
    "    \n",
    "    #Pulls x2 to the right and x1 to the left\n",
    "    pullequation = \"-1*f*(x2-x1)\"\n",
    "\n",
    "    pullforce = openmm.CustomCentroidBondForce(2, pullequation)\n",
    "\n",
    "    pullforce.addGlobalParameter('f', f)\n",
    "    pullforce.addGroup(group1)\n",
    "    pullforce.addGroup(group2)\n",
    "    pullforce.addBond([0,1])\n",
    "    return(pullforce)\n",
    "\n",
    "#This potential energy keeps the left pull group near the origin of the coordinate system.\n",
    "def pin_force(group1, kpin=100):\n",
    "    import simtk.openmm as openmm\n",
    "    \n",
    "    #r0 = r0 * units.meter * 1e-9\n",
    "    pullequation = \"0.5 * kpin * (x1^2+y1^2+z1^2)\"\n",
    "\n",
    "    pullforce = openmm.CustomCentroidBondForce(1, pullequation)\n",
    "\n",
    "    print(str(kpin))\n",
    "    pullforce.addGlobalParameter('kpin', kpin)\n",
    "    pullforce.addGroup(group1)\n",
    "    pullforce.addBond([0])\n",
    "    return(pullforce)\n",
    "\n",
    "#This potential energy keeps the right pull group near the positive x-axis.\n",
    "def slide_force(group1, x_min=0, kslide = 100):\n",
    "    import simtk.openmm as openmm\n",
    "    \n",
    "    pullequation = \"0.5 * kslide * ((step(x_min-x1)*((x_min-x1)^2))+y1^2+z1^2)\"\n",
    "\n",
    "    pullforce = openmm.CustomCentroidBondForce(1, pullequation)\n",
    "\n",
    "    pullforce.addGlobalParameter('kslide', kslide)\n",
    "    pullforce.addGlobalParameter('x_min', x_min)\n",
    "    pullforce.addGroup(group1)\n",
    "    pullforce.addBond([0])\n",
    "    return(pullforce)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Sets up simulation and loads initial structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "##Start openMiChroM library\n",
    "sim = MiChroM(name='sim', temperature=120)\n",
    "sim.setup(platform=platform, integrator=\"Langevin\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Folder to save outputs\n",
    "sim.saveFolder('output_files')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Creates an initial state\n",
    "mypol = sim.createSpringSpiral(ChromSeq='input/DT40_chr7.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Load initial conformation into the system.  Uses pre-oriented and equilibrated input structure\n",
    "sim.loadStructure(mypol, center=False)\n",
    "dataToLoad = load_Coords_PDB('condensed.pdb')\n",
    "sim.loadStructure(dataToLoad, center = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Adds Homopolymer forces\n",
    "sim.addFENEBonds(kfb=30.0)\n",
    "sim.addAngles(ka=1.0)\n",
    "sim.addRepulsiveSoftCore(Ecut=4.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Adds type-type potential for mitotic chromosome\n",
    "#type_lambs = sim_aux.getlambfromfile('input/types-DT40-15m')\n",
    "sim.addCustomTypes(mu=1.51, rc = 2.12, TypesTable='input/lambdas_types_15m')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Adds Ideal Chromosome potential for mitotic chromosome\n",
    "sim.addCustomIC(mu=1.51, rc = 2.12, dinit=3, dend=735, IClist = 'input/lambdas_ic_15m')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Define pull groups."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = [ index for index in range(len(sim.getPositions())) ]\n",
    "g1 = index[:50] #Left pull group is the first 50 beads\n",
    "g2 = index[-50:] #Right pull group is the last 50 beads\n",
    "\n",
    "positions = sim.getPositions() #get the position of each bead\n",
    "\n",
    "first_centroid = np.mean(positions[g1], axis=0) \n",
    "print(\"Location of the left pull group:\")\n",
    "print(first_centroid)\n",
    "\n",
    "second_centroid = np.mean(positions[g2], axis=0)\n",
    "print(\"\\nLocation of the right pull group:\")\n",
    "print(second_centroid)\n",
    "\n",
    "## calculate r0 distance between groups\n",
    "\n",
    "r0 = np.sqrt( (first_centroid[0] - second_centroid[0])**2)\n",
    "\n",
    "print(\"\\nInitial distance between groups is r = {}\".format(r0))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### WARNING! -- Orient Initial Structures\n",
    "The initial structure used here has been pre-oriented so that the left pull group is very near to the origin and the right pull group is very near to the positive x-axis.  This is so that the pin and slide potentials do not generate very large forces on the initial structure.  When performing simulations with an arbitrary starting structure, make sure to rotate the coordinates so that it aligns with the restraining potentials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pforce_sim = constant_pull_force(g1,g2, 0)#We first set the pull force to 0\n",
    "sim.forceDict[\"pulling\"] =  pforce_sim\n",
    "\n",
    "pin_force_sim = pin_force(g1, kpin=100) #here we using x_pos = 0 and kp = 100\n",
    "sim.forceDict[\"pin\"] = pin_force_sim\n",
    "\n",
    "# Pin right end to positive x-axis\n",
    "slide_force_sim = slide_force(g2, kslide=100) #here we using x_pos = 0 and kp = 100\n",
    "sim.forceDict[\"slide\"] = slide_force_sim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#define name to save .cndb\n",
    "sim.initStorage('CF_traj', mode='w')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Matrix to save time\n",
    "time_record = []\n",
    "\n",
    "#Matrix to save pull coordinate\n",
    "pcoord = []\n",
    "\n",
    "#Matrix to save frame labels\n",
    "flabels = []\n",
    "\n",
    "#Append Initial pcoord and time:\n",
    "time_record.append(0)\n",
    "pcoord.append(getPCoord(g1, g2, sim.getPositions()))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Native Simulation\n",
    "We first simulate without adding a pulling force."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#run simulation\n",
    "time1 = time.time()\n",
    "\n",
    "#Save Initial Configuration\n",
    "#sim.save()\n",
    "flabels.append(sim.step)\n",
    "\n",
    "nb = numBlocks_nat\n",
    "\n",
    "#Start Simulation\n",
    "for t in range(1, nb+1):\n",
    "    \n",
    "    sim.runSimBlock(bs, increment=True) #relaxamento\n",
    "    \n",
    "    #Recurd Pulling Data Every Block!\n",
    "    time_record.append(sim.timestep*t*bs)\n",
    "    pcoord.append(getPCoord(g1, g2, sim.getPositions()))\n",
    "    \n",
    "    if t% blocksPerFrame == 0:\n",
    "        ##save trajectory\n",
    "        sim.saveStructure()\n",
    "        flabels.append(sim.step)\n",
    "time2 = time.time()\n",
    "print('This run took {:.2f} seconds'.format(time2-time1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Save zero force structure to ndb file\n",
    "sim.saveStructure(mode = 'pdb')\n",
    "CDInits.append('output_files/'+sim.name +\"_0_block%d.\" % sim.step + 'pdb')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Constant Force"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Update the force parameter to the specified pulling force.\n",
    "sim.context.setParameter('f', force)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#run simulation\n",
    "time1 = time.time()\n",
    "\n",
    "nb = numBlocks_pull"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Start Simulation\n",
    "for t in range(1, nb+1):\n",
    "    \n",
    "    sim.runSimBlock(bs, increment=True) #relaxamento\n",
    "    \n",
    "    #Recurd Pulling Data Every Block!\n",
    "    time_record.append(sim.timestep*t*bs)\n",
    "    pcoord.append(getPCoord(g1, g2, sim.getPositions()))\n",
    "    \n",
    "    if t% blocksPerFrame == 0:\n",
    "        ##save trajectory\n",
    "        sim.saveStructure()\n",
    "        #Save pdb structures throughout pulling \n",
    "        sim.saveStructure(mode = 'pdb')\n",
    "        CDInits.append('output_files/'+sim.name +\"_0_block%d.\" % sim.step + 'pdb')\n",
    "        flabels.append(sim.step)\n",
    "time2 = time.time()\n",
    "print('This run took {:.2f} seconds'.format(time2-time1))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Release Simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Update the force parameter to the specified pulling force.\n",
    "sim.context.setParameter('f', 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#run simulation\n",
    "time1 = time.time()\n",
    "\n",
    "nb = numBlocks_release\n",
    "\n",
    "#Start Simulation\n",
    "for t in range(1, nb+1):\n",
    "    \n",
    "    sim.runSimBlock(bs, increment=True) #relaxamento\n",
    "    \n",
    "    #Recurd Pulling Data Every Block!\n",
    "    time_record.append(sim.timestep*t*bs)\n",
    "    pcoord.append(getPCoord(g1, g2, sim.getPositions()))\n",
    "    \n",
    "    if t% blocksPerFrame == 0:\n",
    "        sim.saveStructure()\n",
    "        flabels.append(sim.step)\n",
    "time2 = time.time()\n",
    "print('This run took {:.2f} seconds'.format(time2-time1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#close storage file\n",
    "sim.storage[0].close()\n",
    "\n",
    "#save last conformation in pdb\n",
    "sim.saveStructure(mode = 'pdb')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Save Time and Pull Coordinate to output file.\n",
    "coordinfo = np.vstack([time_record, pcoord]).T\n",
    "np.savetxt('output_files/CF_Pull_Coord.txt', coordinfo)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Plot Pull Coordinate\n",
    "We may now plot the end-to-end distance as a function of time during the constant-force experiment.  The vertical lines indicate when force is applied and released."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_times = sim.timestep*np.array(range(len(pcoord)))\n",
    "all_pcoord = pcoord\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(all_times, all_pcoord)\n",
    "ax.set_xlabel(r\"Time ($\\tau$)\")\n",
    "ax.set_ylabel(r\"End to End Distance ($\\sigma$)\")\n",
    "ax.set_title('Force = '+ str(force) + r\" $\\epsilon/\\sigma$\")\n",
    "ax.axvline(x=all_times[numBlocks_nat], ls='--', color = 'r')\n",
    "ax.axvline(x = all_times[numBlocks_nat+numBlocks_pull], color = 'b')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Part 2: Constant-Distance Pulling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "numDists = 5\n",
    "kpull = 10**5 #A very strong spring constant is used to restrain the end-to-end distance to a small window.\n",
    "numBlocks_CD = 10000 #number of blocks to run each CD simulation\n",
    "\n",
    "xiVals = [] #Array to store pull coordinates used for pulling\n",
    "forceVals = []#Array to store measured forces"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each iteration of the below loop runs a constant-distance pulling simulation at the distances corresponding to the starting structures saved from the stress relaxation simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CDInits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for distNum in range(0, numDists+1):\n",
    "    \n",
    "    print('Starting Simulation ' + str(distNum))\n",
    "    #Sets up a new simulation\n",
    "    sim = MiChroM(name='sim', temperature=120)\n",
    "    sim.setup(platform=platform, integrator=\"Langevin\")\n",
    "    \n",
    "    #Folder to save outputs\n",
    "    sim.saveFolder('output_files_CD_'+str(distNum))\n",
    "    \n",
    "    #Creates an initial state\n",
    "    mypol = sim.createSpringSpiral(ChromSeq='input/DT40_chr7.txt')\n",
    "    \n",
    "    \n",
    "    #Load initial conformation into the system.  Uses outputs of CF Simulation as starting points.\n",
    "    sim.loadStructure(mypol, center=True)\n",
    "    coords = load_Coords_PDB(CDInits[distNum])\n",
    "    sim.loadStructure(coords)#Sets coordinates to initial structure.\n",
    "\n",
    "    #Adds Homopolymer forces\n",
    "    sim.addFENEBonds(kfb=30.0)\n",
    "    sim.addAngles(ka=1.0)\n",
    "    sim.addRepulsiveSoftCore(Ecut=4.0)\n",
    "\n",
    "    #Adds type-type potential for mitotic chromosome\n",
    "    sim.addCustomTypes(mu=1.51, rc = 2.12, TypesTable='input/lambdas_types_15m')\n",
    "\n",
    "    #Adds Ideal Chromosome potential for mitotic chromosome\n",
    "    sim.addCustomIC(mu=1.51, rc = 2.12, dinit=3, dend=735, IClist = 'input/lambdas_ic_15m')\n",
    "    \n",
    "    #calculate the initial pull coordinate:\n",
    "    positions = sim.getPositions()\n",
    "    \n",
    "    first_centroid = np.mean(positions[g1], axis=0) \n",
    "    print(\"Location of the left pull group:\")\n",
    "    print(first_centroid)\n",
    "\n",
    "    second_centroid = np.mean(positions[g2], axis=0)\n",
    "    print(\"\\nLocation of the right pull group:\")\n",
    "    print(second_centroid)\n",
    "\n",
    "    ## calculate r0 distance between groups\n",
    "    r0 = np.sqrt( (first_centroid[0] - second_centroid[0])**2)  \n",
    "    print('Initial Distance Is: ' + str(r0))\n",
    "    \n",
    "    #Add pulling forcers\n",
    "    pforce_sim = harmonic_pull_force(g1,g2, r0, kpull)\n",
    "    sim.forceDict[\"pulling\"] = pforce_sim\n",
    "    \n",
    "    pin_force_sim = pin_force(g1, kpin=100) #here we using x_pos = 0 and kp = 100\n",
    "    sim.forceDict[\"pin\"] = pin_force_sim\n",
    "\n",
    "    # Pin right end to positive x-axis\n",
    "    slide_force_sim = slide_force(g2, kslide=100) #here we using x_pos = 0 and kp = 100\n",
    "    sim.forceDict[\"slide\"] = slide_force_sim\n",
    "    \n",
    "    pcoord = []\n",
    "    \n",
    "    #run simulation\n",
    "    time1 = time.time()\n",
    "\n",
    "    #Save Initial Configuration\n",
    "    #sim.save()\n",
    "    flabels.append(sim.step)\n",
    "\n",
    "    nb = numBlocks_CD\n",
    "\n",
    "    #Start Simulation\n",
    "    for t in range(1, nb+1):\n",
    "\n",
    "        sim.runSimBlock(bs, increment=True) #relaxamento\n",
    "\n",
    "        #Recurd Pulling Data Every Block!\n",
    "        pcoord.append(getPCoord(g1, g2, sim.getPositions()))\n",
    "\n",
    "        #if want print forces\n",
    "        if t % 1000 == 0:\n",
    "            sim.printForces()\n",
    "            print(\"Radius of Gyration: {:.2f}\\tBlock: {}/{}\".format(sim.RG(), t,nb))\n",
    "    time2 = time.time()\n",
    "    print('This run took {:.2f} seconds'.format(time2-time1))\n",
    "    \n",
    "    meanDist = np.mean(pcoord)\n",
    "    meanForce = kpull*(r0-meanDist)\n",
    "    \n",
    "    #Add to forces\n",
    "    forceVals.append(meanForce)\n",
    "    \n",
    "    #Add to XiVals:\n",
    "    xiVals.append(meanDist)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now plot a force-extension curve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.scatter(xiVals, forceVals)\n",
    "ax.set_xlabel(r'Extension $(\\sigma)$')\n",
    "ax.set_ylabel(r'Force $(\\epsilon/\\sigma)$')"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}