{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Processing Data and Making 2D Images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook demonstrates how to open a Nanosurf image file (*.nid) and perform basic post-processing operations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import required modules\n",
    "from NSFopen.read import read\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading Header\n",
      "Reading Data\n",
      "Elapsed Time: 0.04 sec\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# import the data\n",
    "data_file = \"MoS2.nid\"\n",
    "afm = read(data_file)\n",
    "data = afm.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make a function to flatten the image with an optional mask\n",
    "def flatten(data, order=1, mask=[]):\n",
    "    data_out = np.copy(data)  # create copy of data\n",
    "    data_in = np.copy(data)\n",
    "    \n",
    "    if np.any(mask):\n",
    "        data_in[mask] = np.nan\n",
    "    for idx, (out, line) in enumerate(zip(data_out, data_in)):\n",
    "        ix = np.isfinite(line)\n",
    "        \n",
    "        x = np.arange(len(line))\n",
    "        p = np.polyfit(x[ix], line[ix], order)  # fit data to polynomial\n",
    "        y = np.polyval(p, x)\n",
    "        data_out[idx] = out - y  # subtract fit from data\n",
    "    return data_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "Zaxis = data['Image']['Forward']['Z-Axis'] * 1e9  # height data scaled to nanometers\n",
    "\n",
    "Zaxis_ = flatten(Zaxis, order=1)  # flatten data with 1st order polynomial (i.e. line)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# X and Y data is stored in the parameters\n",
    "param = afm.param\n",
    "Xmin = param.X['min'][0] * 1e6  # the minimum of X of channel 0 scaled to um\n",
    "Xmax = param.X['range'][0] * 1e6  # the range of X of channel 0 scaled to um\n",
    "\n",
    "Ymin = param.Y['min'][0] * 1e6  # the minimum of Y for channel 0 scaled to um\n",
    "Ymax = param.Y['range'][0] * 1e6  # the range of Y for channel 0 scaled to um\n",
    "\n",
    "extents = [Xmin, Xmax, Ymin, Ymax]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We want to limit the color range of the plot to not include the upper and lower 2% of data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# caluculate the 2% and 98% percentiles\n",
    "clip = [np.percentile(Zaxis_, percent) for percent in [2, 98]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 338012 characters]\n",
      "text/plain": [
       "<Figure size 1008x720 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# define function to plot data with colorbar\n",
    "def plot2D(data, extents=extents, limits=clip):\n",
    "    plt.figure(figsize=(14, 10))\n",
    "    im = plt.imshow(data,\n",
    "                    extent = extents,\n",
    "                    origin = 'lowerleft',\n",
    "                    clim = limits,\n",
    "                    cmap=cm.afmhot)\n",
    "\n",
    "    plt.xlabel('[$\\mu$m]')\n",
    "    plt.ylabel('[$\\mu$m]')\n",
    "    cb = plt.colorbar(im)\n",
    "    cb.set_label('[nm]')\n",
    "    plt.show()\n",
    "\n",
    "plot2D(Zaxis_, extents=extents, limits=clip)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create mask to ignore data > 2 nm\n",
    "mask = Zaxis_>2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 347228 characters]\n",
      "text/plain": [
       "<Figure size 1008x720 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# flatten data again but using mask\n",
    "Zaxis_2 = flatten(Zaxis_,mask=mask)\n",
    "plot2D(Zaxis_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
