{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial about fluopy - three state simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we outline a simulation procedure for an ideal three state system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import fluopy\n",
    "import fluopy.analysis as an\n",
    "import fluopy.blinking as bl\n",
    "import fluopy.emissions as em\n",
    "import fluopy.fcs as fcs_p\n",
    "import fluopy.figure as fi\n",
    "import fluopy.fluorophores as fl\n",
    "import fluopy.formulas as fo\n",
    "import fluopy.miscellaneous as mi\n",
    "import fluopy.prediction as pr\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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(seed=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up the photophysical system"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up a single fluorophore representing a three state system (tss)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the transitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transitions = {\n",
    "    \"tss\": [\n",
    "        tr.Transition(tr.TransitionType.EXCITATION, rate=1e9, fluorophore_ids=[0]),\n",
    "        tr.Transition(\n",
    "            tr.TransitionType.FLUORESCENT_EMISSION, rate=1e9, fluorophore_ids=[0]\n",
    "        ),\n",
    "        tr.Transition(\n",
    "            tr.TransitionType.INTERSYSTEM_CROSSING_ST, rate=1e6, fluorophore_ids=[0]\n",
    "        ),\n",
    "        tr.Transition(\n",
    "            tr.TransitionType.S1_S0_TRANSITIONS, rate=1e9, fluorophore_ids=[0]\n",
    "        ),\n",
    "        tr.Transition(\n",
    "            tr.TransitionType.T1_S0_TRANSITIONS, rate=1e6, fluorophore_ids=[0]\n",
    "        ),\n",
    "    ]\n",
    "}\n",
    "pprint(transitions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the fluorophore system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fluorophore = fl.Fluorophore(name=\"tss\", position=[0, 0])\n",
    "fluorophore_system = fl.FluorophoreSystem(fluorophores=[fluorophore])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pprint(vars(fluorophore_system))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Alternatively: define the fluorophore system with FluorophoreData"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fluorophore_data_tss = fluopy.fluo_data.FluorophoreData(\n",
    "    QUANTUM_YIELD=1,\n",
    "    FLUORESCENCE_LIFETIME=2e-9,\n",
    "    S1_QUENCH_RATE=1e9,\n",
    "    ISC_ST_RATE=1e6,\n",
    "    ISC_TS_RATE=1e6,\n",
    "    PHOTOBLEACH_T1_RATE=10,\n",
    ")\n",
    "pprint(fluorophore_data_tss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fluorophore_tss = fl.Fluorophore(\n",
    "    name=\"tss\", position=[0, 0], constants=fluorophore_data_tss\n",
    ")\n",
    "fluorophore_system_tss = fl.FluorophoreSystem(fluorophores=[fluorophore_tss])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pprint(vars(fluorophore_system_tss))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "THIS IS NOT YET WORKING."
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "transitions = fluorophore_system_tss.load_transitions(energy_transfer=False, dstorm=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the transition set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transition_set = tr.TransitionSet(transitions, fluorophore_system)\n",
    "transition_set.finalize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transition_set.plot(graph_type=\"shell\", colors=None, scale=1);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transition_set.transition_df"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "for key, value in transition_set.transitions.items():\n",
    "    pprint(key)\n",
    "    pprint(value)\n",
    "    print(\"---\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make a prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "prediction = pr.Prediction(transition_set)\n",
    "prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction.plot_frequency_transitions(scale=0.5)\n",
    "prediction.plot_frequency_states(scale=0.5)\n",
    "prediction.plot_mean_lifetimes(scale=0.5)\n",
    "prediction.plot_mean_transition_times(scale=0.5)\n",
    "prediction.plot_state_occupations(scale=0.5)\n",
    "prediction.plot_transition_time_distributions(\n",
    "    fluorophore=\"tss\", transition_id=0, scale=0.5\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run a simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simulation = si.Simulation(transition_set)\n",
    "simulation"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "%%time\n",
    "# basic simulation run for a given number of steps (size)\n",
    "simulation.run(start_at=None, size=1e6, end_time=None, seed=rng, use_memmap=None)\n",
    "mi.print_class(simulation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# simulate until it reaches given end_time\n",
    "simulation.run(start_at=None, size=1e6, end_time=1e-2, seed=rng, use_memmap=None)\n",
    "mi.print_class(simulation)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "%%time\n",
    "# if a lot of steps have to be carried out, the arrays may become too large for RAM\n",
    "# in that case, use memory maps store on drive\n",
    "simulation.run(start_at=None, size=1e6, end_time=None, seed=rng,\n",
    "               use_memmap='')\n",
    "mi.print_class(simulation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analyze the simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "analysis = an.Analysis(simulation=simulation)\n",
    "mi.print_class(analysis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "analysis.get_fluorescence_lifetimes(fluorophore=\"tss\")\n",
    "analysis.get_emitting_transition_lifetimes(fluorophore=\"tss\")\n",
    "\n",
    "analysis.plot_frequency_transitions(scale=0.5, prediction=prediction)\n",
    "analysis.plot_frequency_states(scale=0.5, prediction=prediction)\n",
    "analysis.plot_mean_transition_times(scale=0.5, prediction=prediction)\n",
    "analysis.plot_mean_lifetimes(scale=0.5, prediction=prediction)\n",
    "analysis.plot_state_occupations(scale=0.5, prediction=prediction)\n",
    "analysis.plot_lifetime_distributions(\n",
    "    scale=0.5, prediction=prediction, fluorophore=\"tss\", state_identity=1\n",
    ")\n",
    "analysis.plot_transition_time_distributions(\n",
    "    scale=0.5, prediction=prediction, fluorophore=\"tss\", transition_id=0\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulation of experimentally observable (photons per frames) only"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Extract photon emission events from simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "emissions = em.Emissions(frame_time=\"1us\", seed=rng)\n",
    "emissions.extract(simulation=simulation)\n",
    "emissions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Simulate photon emission events"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "%%time\n",
    "emissions = em.Emissions(frame_time=\"5ms\", seed=rng, bandpass=(600, 800))\n",
    "emissions.simulate(transition_set=transition_set, store_time_points=True, frames=200)\n",
    "emissions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Correct for detection efficiency and noise contributions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "emissions.add_photon_collection_objective(p=0.1, seed=rng)  # 1.\n",
    "emissions.add_quantum_efficiency(p=0.9, seed=rng)  # 3.1.\n",
    "emissions.add_poisson_noise(\n",
    "    rate=0.05, seed=rng\n",
    ")  # 3.2. (dark noise), note the frame time\n",
    "emissions.add_emccd_gain(emccd_gain=10, seed=rng)  # 4.\n",
    "emissions.add_gaussian_noise(mean=10, std=1, seed=rng)  # 5. (readout noise)\n",
    "emissions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "emissions = em.Emissions(frame_time=\"5ms\", seed=rng, bandpass=(660, 700))\n",
    "emissions.extract(simulation=approximation)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2.\n",
    "# at this point, the bandpass filter was applied\n",
    "# yet, the effect of photon collection by the objective is missing\n",
    "# the order is not relevant for two consecutive binomial distributions\n",
    "# it is more convenient to apply the bandpass first because it needs the\n",
    "# information about the emitting fluorophore whereas all the other effects are\n",
    "# roughly wavelength independent\n",
    "p_photon_collection = fo.calculate_photon_collection_rate(NA=1.45, n1=1.51)\n",
    "emissions.add_photon_collection_objective(p=p_photon_collection)  # 1.\n",
    "emissions.add_quantum_efficiency(p=0.9)  # 4.1.\n",
    "emissions.add_transmittance(p=0.99)  # 3 (depending on number of components of optical\n",
    "# path, may be applied multiple times)\n",
    "emissions.add_poisson_noise(rate=0.05)  # 4.2. (dark noise), note the frame time\n",
    "emissions.add_emccd_gain(emccd_gain=10)  # 5. (+ multiplicative noise)\n",
    "emissions.add_gaussian_noise(mean=10, std=1, seed=rng)  # 6. (readout noise)\n",
    "# CIC (spurious noise) neglected since low probability to happen in the pixels of\n",
    "# interest\n",
    "emissions.apply_threshold(threshold=100)  # 7 (thresholding)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "mi.print_class(\"emissions\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "emissions.plot_cumulative_events(scale=1)\n",
    "emissions.plot_histogram(scale=1)\n",
    "emissions.plot_time_series(scale=1)\n",
    "\n",
    "# to save the time_series and time_points\n",
    "# emissions.save(path='', name_extension='test')\n",
    "\n",
    "# to load time_series and time_points\n",
    "# emissions.load(path='', name_extension='test')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulation of pulsed excitation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "emissions_tcspc = em.Emissions(frame_time=\"10us\", seed=rng, bandpass=None)\n",
    "lifetimes_DA, lifetimes_D, lifetimes_all, simulation_object = emissions_tcspc.tcspc(\n",
    "    transition_set=transition_set,\n",
    "    number_pulses=1e5,\n",
    "    pulse_duration=1e-11,\n",
    "    time_between_pulses=1e-7,\n",
    "    excitation_rates={\"tss\": 1e11},\n",
    "    size=1e5,\n",
    "    store_time_points=True,\n",
    "    # details = True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "emissions_tcspc.plot_time_series()\n",
    "fi.universal_figure(\n",
    "    type_=\"hist\", data=lifetimes_all, ylabel=\"PD\", density=True, xlabel=\"Lifetime (s)\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fluorescence correlation spectroscopy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Observed fluorescence emission events can be analyzed by a correlation analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fcs = fcs_p.FCS(emissions)\n",
    "list(vars(fcs).keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Autocorrelation of time points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "fcs.autocorrelate_time_points(\n",
    "    exp_min=-20, exp_max=-2, points_per_base=4, base=4, normalize=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mi.print_class(fcs)\n",
    "fcs.plot(normalize_to=None, unit=\"s\", scale=1);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Autocorrelation of time series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fcs.autocorrelate_time_series(log=True, m=4, normalize=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fcs.plot(normalize_to=None, unit=\"s\", scale=1);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Antibunching"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, you can focus on fast time scales in a linear scale and observe antibunching."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sensible to tau_max and bin_width, see coincidence notebook\n",
    "hist, bins = fcs_p.coincidence(\n",
    "    emissions.event_time_points[: int(2e5)], tau_max=1e-8, bin_width=1e-10, seed=rng\n",
    ")\n",
    "fi.universal_figure(\n",
    "    type_=\"line\",\n",
    "    data=[bins, hist],\n",
    "    xlabel=r\"$\\tau$ (s)\",\n",
    "    ylabel=r\"$g^{(2)}(\\tau)$\",\n",
    "    scale=1,\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Blinking"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Emissions from a short simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "emissions = em.Emissions(frame_time=\"200ns\", seed=rng, bandpass=None)\n",
    "emissions.simulate(transition_set=transition_set, store_time_points=False, frames=2000)\n",
    "emissions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "threshold: int = 40"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "emissions.plot_time_series(scale=1)\n",
    "plt.hlines(threshold, xmin=0, xmax=0.0004)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "blinks = bl.Blinking(emissions, threshold=threshold)\n",
    "blinks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mi.print_class(blinks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot a histogram of OFF times\n",
    "blinks.plot(\n",
    "    mode=\"off_histogram\", density=True, display_mean=True, as_time=\"s\", scale=0.5\n",
    ")\n",
    "\n",
    "# plot a histogram of ON times\n",
    "blinks.plot(\n",
    "    mode=\"on_histogram\", density=True, display_mean=True, as_time=\"ms\", scale=0.5\n",
    ")\n",
    "\n",
    "# plot a time series of OFF times\n",
    "blinks.plot(mode=\"off_frame_series\", scale=0.5)\n",
    "\n",
    "# plot a time series of ON times\n",
    "blinks.plot(mode=\"on_frame_series\", scale=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# to get the analytical OFF statistics as the same view, use\n",
    "on_off_times_analytic, on_off_values_analytic = bl.get_analytical_off_statistics(\n",
    "    off_frames=blinks.off_periods_frames,\n",
    "    off_periods=blinks.off_periods,\n",
    "    on_frames=blinks.on_periods_frames,\n",
    "    frame_time=blinks.emissions.parameters[\"frame_time\"],\n",
    ")\n",
    "\n",
    "# plot the analytical OFF statistics (no differentiation between fluorophores)\n",
    "bl.plot_off_statistics(\n",
    "    on_off_times_analytic, on_off_values_analytic, scale=1, title=\"analytical OFF\"\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Emissions from the long simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get more detailed information from a complete simulation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "emissions = em.Emissions(frame_time=\"200ns\", seed=rng, bandpass=None)\n",
    "emissions.extract(simulation=simulation)\n",
    "emissions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "blinks = bl.Blinking(emissions, threshold=threshold)\n",
    "blinks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mi.print_class(blinks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot a histogram of OFF times\n",
    "blinks.plot(\n",
    "    mode=\"off_histogram\", density=True, display_mean=True, as_time=\"s\", scale=0.5\n",
    ")\n",
    "\n",
    "# plot a histogram of ON times\n",
    "blinks.plot(\n",
    "    mode=\"on_histogram\", density=True, display_mean=True, as_time=\"ms\", scale=0.5\n",
    ")\n",
    "\n",
    "# plot a time series of OFF times\n",
    "blinks.plot(mode=\"off_frame_series\", scale=0.5)\n",
    "\n",
    "# plot a time series of ON times\n",
    "blinks.plot(mode=\"on_frame_series\", scale=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get more detailed information from a complete simulation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# to get the analytical OFF statistics as the same view, use\n",
    "on_off_times_analytic, on_off_values_analytic = bl.get_analytical_off_statistics(\n",
    "    off_frames=blinks.off_periods_frames,\n",
    "    off_periods=blinks.off_periods,\n",
    "    on_frames=blinks.on_periods_frames,\n",
    "    frame_time=blinks.emissions.parameters[\"frame_time\"],\n",
    ")\n",
    "\n",
    "# plot the analytical OFF statistics (no differentiation between fluorophores)\n",
    "bl.plot_off_statistics(\n",
    "    on_off_times_analytic, on_off_values_analytic, scale=1, title=\"analytical OFF\"\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.14.4"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}