{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multiple Chromosome Simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial should take between 20 to 30 minutes of reading and performing simulations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Chromatin Dynamics Simulations on Chromosome 10 and Chromosome 11 of GM12878 Cell Line"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note**: This tutorial will be running in OpenMiChroM version 1.1.0 or greater. Please ensure you have the correct version installed before proceeding."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first step is to import the **OpenMiChroM** module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from OpenMiChroM.ChromDynamics import MiChroM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`OpenMiChroM` allows the simulation of multiples chromosomes. In this tutorial, the multiple chain simulation will be performed using the chromosomes 10 and 11 of the human GM12878 cell line.The system is generated based on the collapsed structure of each chromosome. This collapse step is similar to the one presented in the [single chromosome simulation tutorial](https://open-michrom.readthedocs.io/en/latest/Tutorials/Tutorial_Single_Chromosome.html). For each chromosome, the collapse simulation should be performed individually."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Chromosome 10 Collapse Simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "sim_chr10 = MiChroM(name=\"chr10\", temperature=1.0, timeStep=0.01)\n",
    "sim_chr10.setup(platform=\"opencl\")\n",
    "sim_chr10.saveFolder('output_nucleus')\n",
    "chr10 = sim_chr10.initStructure(mode='random',ChromSeq='Tutorials/Chromosomes_simulations/inputs/chr10_beads.txt', isRing=False)\n",
    "sim_chr10.loadStructure(chr10, center=True)\n",
    "\n",
    "sim_chr10.addHomoPolymerForces()\n",
    "sim_chr10.addFlatBottomHarmonic()\n",
    "\n",
    "sim_chr10.addTypetoType(mu=3.22, rc=1.78)\n",
    "sim_chr10.addIdealChromosome(mu=3.22, rc=1.78, dinit=3, dend=500)\n",
    "\n",
    "sim_chr10.createSimulation()\n",
    "\n",
    "print(\"Performing chr10 simulation...\")\n",
    "sim_chr10.run(nsteps=10**6, report=True, interval=10**5)\n",
    "\n",
    "sim_chr10.saveStructure(filename=\"chr10\", mode=\"ndb\")\n",
    "del sim_chr10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Chromosome 11 Collapse Simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "sim_chr11 = MiChroM(name=\"chr11\", temperature=1.0, timeStep=0.01)\n",
    "sim_chr11.setup(platform=\"opencl\")\n",
    "sim_chr11.saveFolder('output_nucleus')\n",
    "chr11 = sim_chr11.initStructure(mode='random',ChromSeq='Tutorials/Chromosomes_simulations/inputs/chr11_beads.txt', isRing=False)\n",
    "sim_chr11.loadStructure(chr11, center=True)\n",
    "\n",
    "sim_chr11.addHomoPolymerForces()\n",
    "sim_chr11.addFlatBottomHarmonic()\n",
    "\n",
    "sim_chr11.addTypetoType(mu=3.22, rc=1.78)\n",
    "sim_chr11.addIdealChromosome(mu=3.22, rc=1.78, dinit=3, dend=500)\n",
    "\n",
    "sim_chr11.createSimulation()\n",
    "\n",
    "print(\"Performing chr11 simulation...\")\n",
    "sim_chr11.run(nsteps=10**6, report=True, interval=10**5)\n",
    "\n",
    "sim_chr11.saveStructure(filename=\"chr11\", mode=\"ndb\")\n",
    "del sim_chr11"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Multiple Chromosome Simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the collapse simulations are done, the collapsed structures of each chromosome should be included in the same simulation system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_nucleus = MiChroM(name=\"nucleus\", temperature=1.0, timeStep=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_nucleus.setup(platform=\"opencl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_nucleus.saveFolder('output_nucleus')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function `initStructure()` receives a list of files and saves the positions of the collapsed chromosomes in the variable `initial_conf`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_conf = sim_nucleus.initStructure(\n",
    "    CoordFiles=['output_nucleus/chr10.ndb',\n",
    "                'output_nucleus/chr11.ndb']\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`initial_conf` contains the coordinates of all beads for each chromosome. When used to load coordinate's files (`.ndb`,`.pdb`, or `.gro`), the function `initStructure` prints the variable `chains`, which shows the information of each chromosome.\n",
    "\n",
    "For each chain, we have (`start`,`end`,`is_ring`). `start` is the index of the first bead of the chain, `end` is the index for the last bead and `is_ring` indicates whether the first and last bead are connect or not. `is_ring=0` represents a open chromosome chain and `is_ring=1` represents a circular polymer (used for simulating bacteria genome, for example)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_nucleus.chains"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before loading the structures in the simulation context, it is necessary to spatially distribute the chromosome chains. This step is essential to guarantee no overlap between the chains in the initial configuration. It also helps randomizing the initial condition when simulating different replicas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_conf = sim_nucleus.setFibPosition(initial_conf, factor=1.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This function distributes the center of mass of each chain in a spherical shell, according to the Fibonacci Sphere Algorithm. The argument `factor` sets the radius of the spherical shell, in comparison to the radius of the nucleus (see function [documentation](https://open-michrom.readthedocs.io/en/latest/OpenMiChroM.html?highlight=setFibPosition#OpenMiChroM.ChromDynamics.MiChroM.setFibPosition)). For simulations with just a few chromosomes, `factor=1.5` should be fine."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can load the chromosomes with adjusted positions in the simulation context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_nucleus.loadStructure(initial_conf, center=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The initial configuration of the chromosome structures can be saved in the [.ndb file format](https://ndb.rice.edu/ndb-format). The file is stored in the path set in `saveFolder`. We advise saving the structure and double-checking if there is no overlap between the chains.\n",
    "\n",
    "When having multiple chains in the simulation context, the `saveStructure` function will save each chromosome structure in a different file. Each chain is associated with an index starting from 0. This index follows the order of addition of the chains in the `initStructure` function. In this tutorial, chromosome 10 receives index 0, and chromosome 11 index 1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next step is to add the force field in the simulation object `sim_nucleus`.\n",
    "\n",
    "In this tutorial, the forces can be divided into two sets:\n",
    "\n",
    "**MiChroM Homopolymer (Bonded) Potentials**  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_nucleus.addFENEBonds()\n",
    "sim_nucleus.addAngles()\n",
    "sim_nucleus.addSelfAvoidance()\n",
    "sim_nucleus.addFlatBottomHarmonic(nRad=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**MiChroM Non-Bonded Potentials**\n",
    "\n",
    "The `addTypetoType` interaction is independent of the chromosome chain and only depends on the chromatin subcompartment annotation for the interacting loci. This potential is added for both chains simultaneously."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_nucleus.addTypetoType(mu=3.22, rc=1.78)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On the other hand, for the Ideal Chromossome potential, we add the forces for each chromosome separately. The function `addMultiChainIC` receives the chromosome chain index and adds the IC potential.\n",
    "\n",
    "The chromosome chain information can be obtained from the aforementioned variable `chains`. Note that the Ideal Chromosome potential was applied from the genomic distance $d$ = 3 to $d$ = 500. These cutoff values can be adjusted based on the chromosome length and user needs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_nucleus.addMultiChainIC(chainIndex=0, mu=3.22, rc=1.78, dinit=3, dend=500)\n",
    "sim_nucleus.addMultiChainIC(chainIndex=1, mu=3.22, rc=1.78, dinit=3, dend=500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_nucleus.createSimulation()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**The simulation setup is complete!**\n",
    "\n",
    "Before running the production simulation, it is necessary to run a collapse and equilibration simulation as the chains are initially apart. The `addFlatBottomHarmonic` function used above includes in the simulation a harmonic potential to drive the chains together."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the parameters of the collapse simulation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "collapse_steps = 10**6\n",
    "production_steps = 10**6 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_nucleus.run(nsteps=collapse_steps, report=True, interval=10**5, totalSteps=(collapse_steps+production_steps))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After the collapse, we should remove the harmonic potential and add the nucleus confinement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_nucleus.removeFlatBottomHarmonic()\n",
    "sim_nucleus.addAdditionalForce(sim_nucleus.addSphericalConfinementLJ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function `addAdditionalForce` adds a force to the system after the system has been already initialized. The initialization happens when the function `runSimBlock` is executed for the first time after the addition of the forces. In this case, we are using the function to add the Lennard-Jones spherical confinement (`addSphericalConfinementLJ`) in the system after removing the harmonic potential."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we create the reporters to save the simulation infos. There are 3 types of reporters:\n",
    "\n",
    "**statistics**: Attaches a reporter to collect simulation statistics such as step number, radius of gyration (RG), total energy, potential energy, kinetic energy, and temperature.\n",
    "\n",
    "**trajectory**:  Attaches a reporter to save trajectory data (xyz per bead per chain) during the simulation. The file format to save the trajectory data. Options are 'cndb', 'swb','ndb', 'pdb', 'gro', 'xyz'. (Default: 'cndb')\n",
    "\n",
    "**energy components**: Saves energy components per force group to a separate file named 'energyComponents.txt' in the simulation folder. Requires that statistics is True\n",
    "\n",
    "set the number of steps interval we will save this information, here I chose 1000 steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_nucleus.createReporters(statistics=True, traj=True, trajFormat=\"cndb\", energyComponents=True, interval=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "sim_nucleus.run(nsteps=production_steps, report=True, interval=10**5, totalSteps=(collapse_steps+production_steps))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To visualize the chromosome's 3D structures in the standard visualization softwares for macromolecules, there are available scripts for converting the `ndb`/`cndb` file formats to `.pdb` and `.gro`. For details, please check the [Nucleome Data Bank](https://ndb.rice.edu/ndb-format).\n",
    "\n",
    "The `ndb` plugin for visualizing the chromatin dynamics trajectories in VMD/Chimera/Pymol is under development."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_nucleus.saveStructure(mode=\"ndb\")\n",
    "sim_nucleus.saveStructure(mode=\"gro\")\n",
    "sim_nucleus.saveStructure(mode=\"pdb\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "openmm8.2",
   "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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}