{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial about fluopy - distributed simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we outline parallel computing with fluopy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install [ray](https://docs.ray.io) for distributing multiple simulations on different cores or machines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import ray\n",
    "\n",
    "import fluopy\n",
    "import fluopy.fluorophores as fl\n",
    "import fluopy.simulation as si\n",
    "import fluopy.transitions as tr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fluopy.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define a fluorophore and transition set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fluorophore = fl.Fluorophore(name=\"cy5_dna\", position=[0, 0])\n",
    "fluorophore_system = fl.FluorophoreSystem(fluorophores=[fluorophore])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transitions = fluorophore_system.load_transitions(\n",
    "    summarize=True,\n",
    "    irradiance=2,\n",
    "    wavelength=640,\n",
    "    bleaching=True,\n",
    "    energy_transfer=False,\n",
    "    dstorm=False,\n",
    ")\n",
    "transition_set = tr.TransitionSet(transitions, fluorophore_system)\n",
    "transition_set.finalize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transition_set.transition_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run a simulation in parallel using ray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.init()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "\n",
    "@ray.remote\n",
    "def worker(transition_set, seed):\n",
    "    simulation = si.Simulation(transition_set)\n",
    "    simulation.run(start_at=None, size=1e6, end_time=None, seed=seed)\n",
    "    return simulation\n",
    "\n",
    "\n",
    "n_processes = 10\n",
    "ss = np.random.SeedSequence()\n",
    "child_seeds = ss.spawn(n_processes)\n",
    "\n",
    "futures = [\n",
    "    worker.remote(transition_set=transition_set, seed=seed) for seed in child_seeds\n",
    "]\n",
    "simulations = ray.get(futures)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.shutdown()  # might be needed to reinitialize ray if executed multiple times"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Extract and present data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_times = [simulation_.time_series[-1] for simulation_ in simulations]\n",
    "len(final_times)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(final_times)\n",
    "plt.gca().set(\n",
    "    title=f\"The mean bleaching time of {len(simulations)} simulations is {np.mean(final_times):0.2f} s\",\n",
    "    xlabel=\"bleaching time / s\",\n",
    "    ylabel=\"counts\",\n",
    ");"
   ]
  }
 ],
 "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.12.8"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}