{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial about drift analysis and correction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lateral drift correction is useful in most SMLM experiments. To determine the amount of drift a method based on image cross-correlation or an iterative closest point algorithm can be applied.\n",
    "\n",
    "We demonstrate drift analysis and correction on simulated data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.stats as stats\n",
    "\n",
    "import locan as lc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lc.show_versions(system=False, dependencies=False, verbose=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Synthetic data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use synthetic data that follows a Neyman-Scott spatial distribution (blobs). The intensity values are exponentially distributed and the number of localizations per frame follows a Poisson distribution: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(seed=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "intensity_mean = 1000\n",
    "localizations_per_frame_mean = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dat_blob = lc.simulate_Thomas(parent_intensity=1e-4, region=((0, 1000), (0, 1000)), cluster_mu=1000, cluster_std=10, seed=rng)\n",
    "dat_blob.dataframe['intensity'] = stats.expon.rvs(scale=intensity_mean, size=len(dat_blob), loc=500)\n",
    "dat_blob.dataframe['frame'] = lc.simulate_frame_numbers(n_samples=len(dat_blob), lam=localizations_per_frame_mean, seed=rng)\n",
    "\n",
    "dat_blob = lc.LocData.from_dataframe(dataframe=dat_blob.data)\n",
    "\n",
    "print('Data head:')\n",
    "print(dat_blob.data.head(), '\\n')\n",
    "print('Summary:')\n",
    "dat_blob.print_summary()\n",
    "print('Properties:')\n",
    "print(dat_blob.properties)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lc.render_2d(dat_blob, bin_size=10, rescale='equal');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add linear drift"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We add linear drift with a velocity given in length units per frame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dat_blob_with_drift = lc.add_drift(dat_blob, velocity=(0.002, 0.001), seed=rng)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 3, figsize=(15, 5))\n",
    "lc.render_2d(dat_blob_with_drift, ax=axes[0], bin_size=10);\n",
    "lc.render_2d(dat_blob_with_drift, ax=axes[1], bin_size=2, rescale='equal', bin_range=((0, 500),(0, 500)));\n",
    "lc.render_2d_mpl(dat_blob_with_drift, ax=axes[2], other_property='frame', bin_size=2, bin_range=((0, 500),(0, 500)), cmap='viridis');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Estimate RMS errors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Knowing the ground truth, you can define a root mean squared error between the original localization coordinates and those after drift and later after correction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rmse(locdata, other_locdata):\n",
    "    return np.sqrt(np.mean(np.square(np.subtract(locdata.coordinates, other_locdata.coordinates)), axis=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rmse(dat_blob, dat_blob_with_drift).round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Estimate drift"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Drift can be estimated by comparing different chunks of successive localizations using either an \"iterative closest point\" algorithm or a \"cross-correlation\" algorithm. Per default, the icp algorithm is applied."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "drift = lc.Drift(chunk_size=10_000, target='first', method='icp').compute(dat_blob_with_drift)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Transformations to register the different data chunks are represented by a transformation matrix and a transformation offset that together specifiy an affine transformation. The tansformation parameters are kept under the `transformations` attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift.transformations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The parameters can be visualized using the plot function. The matrix in this case is close to the unit matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift.plot(transformation_component='matrix', element=None);\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift.plot(transformation_component='offset', element=None)\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model drift"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A continuous transformation model as function of frame number is estimated by fitting the individual transformation components with the specified fit models. Fit models can be provided as `DriftComponent` or by a string representing standard model functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lmfit.models import ConstantModel, LinearModel, PolynomialModel\n",
    "\n",
    "drift.fit_transformations(slice_data=slice(None), offset_models=(lc.DriftComponent('spline', s=100), 'linear'), verbose=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The fit models are represented as `DriftComponent` and can be accessed through the transformation_models attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift.transformation_models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift.transformation_models['offset'][0].type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift.transformation_models['offset'][0].eval(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each `DriftModel` carries detailed information about the fit under the model_result attribute. In most cases, except splines, this will be a `lmfit.ModelResult` object. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift.transformation_models['offset'][0].model_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift.transformation_models['offset'][1].type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift.transformation_models['offset'][1].model_result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Drift correction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The estimated drift is corrected by applying a transformation on the localization chunks (from_model=False)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "drift.apply_correction(from_model=False);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The same correction can be applied to any other localization dataset."
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "%%time\n",
    "drift.apply_correction(from_model=False, locdata=other_locdata);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(15, 5))\n",
    "lc.render_2d(drift.locdata_corrected, ax=axes[0], bin_size=2, rescale='equal', bin_range=((0, 200),(0, 200)));\n",
    "lc.render_2d_mpl(drift.locdata_corrected, ax=axes[1], other_property='frame', bin_size=2, bin_range=((0, 200),(0, 200)), cmap='viridis');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rmse(dat_blob, drift.locdata_corrected).round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or the estimated drift is corrected by applying a transformation on each individual localization using the drift models (from_model=True)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "drift.apply_correction(from_model=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(15, 5))\n",
    "lc.render_2d(drift.locdata_corrected, ax=axes[0], bin_size=2, rescale='equal', bin_range=((0, 200),(0, 200)));\n",
    "lc.render_2d_mpl(drift.locdata_corrected, ax=axes[1], other_property='frame', bin_size=2, bin_range=((0, 200),(0, 200)), cmap='viridis');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rmse(dat_blob, drift.locdata_corrected).round(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift.locdata_corrected.meta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Drift analysis by a cross-correlation algorithm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The same kind of drift estimation and correction can be applied using the image cross-correlation algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "drift = lc.Drift(chunk_size=10_000, target='first', method='cc').\\\n",
    "        compute(dat_blob_with_drift).\\\n",
    "        fit_transformations(slice_data=slice(None), offset_models=(LinearModel(), LinearModel()), verbose=True).\\\n",
    "        apply_correction(from_model=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(15, 5))\n",
    "lc.render_2d(drift.locdata_corrected, ax=axes[0], bin_size=2, rescale='equal', bin_range=((0, 200),(0, 200)));\n",
    "lc.render_2d_mpl(drift.locdata_corrected, ax=axes[1], other_property='frame', bin_size=2, bin_range=((0, 200),(0, 200)), cmap='viridis');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rmse(dat_blob, drift.locdata_corrected)"
   ]
  }
 ],
 "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.6"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}