{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Batch image flattening with 2D plane using regression"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Import packages and find all the *.nid files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\\\\saphir\\nanosurf\\Marketing\\Projects\\Jupyter Notebooks\\2D Plane Removal with 4 *.nid files\n"
     ]
    }
   ],
   "source": [
    "from NSFopen.read import read as afmreader\n",
    "\n",
    "import os\n",
    "from os import listdir\n",
    "from os.path import isfile, join\n",
    "\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "import collections\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "from mpl_toolkits.axes_grid1.colorbar import colorbar\n",
    "%matplotlib inline\n",
    "\n",
    "font = {'size': 12}\n",
    "matplotlib.rc('font', **font)\n",
    "\n",
    "dirpath = os.getcwd()\n",
    "onlyfiles = [f for f in listdir(dirpath) if isfile(join(dirpath, f)) and f.split(\".\")[-1] == \"nid\"]\n",
    "filenum = len(onlyfiles)\n",
    "print(dirpath, \"with\", filenum,\"*.nid files\")\n",
    "print(onlyfiles)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plane fitting and median search functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit2DPlane(matrix, xaxis, yaxis):\n",
    "    m = len(xaxis)\n",
    "    X1, Y1 = np.meshgrid(xaxis, yaxis)\n",
    "    X = np.hstack(( np.reshape(X1, (m*m, 1)) , np.reshape(Y1, (m*m, 1)) ))\n",
    "    X = np.hstack(( np.ones((m*m, 1)) , X ))\n",
    "    YY = np.reshape(matrix, (m*m, 1))\n",
    "\n",
    "    theta = np.dot(np.dot( np.linalg.pinv(np.dot(X.transpose(), X)), X.transpose()), YY)\n",
    "    plane = np.reshape(np.dot(X, theta), (m, m))\n",
    "    \n",
    "    return plane\n",
    "\n",
    "def findLimits(matrix, strength):\n",
    "    array1D = matrix.flatten()\n",
    "    mean = np.mean(array1D)\n",
    "    stdev = np.std(array1D)\n",
    "    armin = mean - strength/2*stdev\n",
    "    armax = mean + strength/2*stdev\n",
    "    return armin, armax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define what to plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "signal_to_plot = 'Topography' # 'Amplitude', 'Z-Axis', 'Z-Axis Sensor' there could be more, i.e. 2nd lock-in signal, etc.\n",
    "signal_to_plot1 = 'Z-Axis' \n",
    "scan_direction = 'Forward' # 'Forward', 'Backward', '2nd Forward', '2nd Backward'\n",
    "scan_to_plot = 'Image' # 'Image', 'Spec' or 'Sweep'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Build X, Y axes and read data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "grid1.nid\n"
     ]
    }
   ],
   "source": [
    "# We store the images and respective X and Y axes in ordered dictionaries\n",
    "data_dict = collections.OrderedDict()\n",
    "xaxis_dict = collections.OrderedDict()\n",
    "yaxis_dict = collections.OrderedDict()\n",
    "data_levelled = collections.OrderedDict()\n",
    "\n",
    "# Read the data into the dictionaries, depending on the version of the file,\n",
    "# there could be various identifiers, i.e. 'Topography' or 'Z-Axis'.\n",
    "for filename in onlyfiles:\n",
    "    try:\n",
    "        rawdata0 = afmreader(filename, verbose=False).data\n",
    "        rawdata = rawdata0[scan_to_plot][scan_direction][signal_to_plot]*1e9\n",
    "    except:\n",
    "        print(filename) # Print filenames with outdated (old) identifiers\n",
    "        rawdata0 = afmreader(filename, verbose=False).data\n",
    "        rawdata = rawdata0[scan_to_plot][scan_direction][signal_to_plot1]*1e9\n",
    "    \n",
    "    param = afmreader(filename).param\n",
    "    xrange = param['X']['range'][0]\n",
    "    yrange = param['Y']['range'][0]\n",
    "    \n",
    "    key = filename.split(\".nid\")[0]\n",
    "    xaxis = np.linspace(0, xrange*1E6, len(rawdata[0]))\n",
    "    xaxis_dict[key] = xaxis\n",
    "    \n",
    "    yaxis = np.linspace(0, yrange*1E6, len(rawdata))\n",
    "    yaxis_dict[key] = yaxis\n",
    "    \n",
    "    data = np.array(rawdata, dtype = 'float')\n",
    "    data_dict[key] = data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Apply plane fitting to all the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key in data_dict.keys():\n",
    "    data = data_dict[key]\n",
    "    xaxis = xaxis_dict[key]\n",
    "    yaxis = yaxis_dict[key]\n",
    "    \n",
    "    plane = fit2DPlane(data, xaxis, yaxis)  \n",
    "    data_levelled[key] = data-plane"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exclude rogue points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key in data_levelled.keys():\n",
    "    data = data_levelled[key]\n",
    "    armin, armax = findLimits(data, 3) # 3 -means the range of the z-axis on the colorbar is 3 standard deviations.\n",
    "    data_levelled[key] = np.clip(data, armin, armax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Prepare the figure and plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 530428 characters]\n",
      "text/plain": [
       "<Figure size 1440x720 with 16 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "im = collections.OrderedDict() # Create a dictionary for the plots\n",
    "cb = collections.OrderedDict()\n",
    "cmap = \"copper\" # colormap\n",
    "xlbl = 'X (um)'\n",
    "ylbl = 'Y (um)'\n",
    "axnum = 0\n",
    "\n",
    "plot_columns = 4\n",
    "plot_rows = math.ceil(len(data_dict)*2/plot_columns)\n",
    "\n",
    "fig, axarr = plt.subplots(plot_rows, plot_columns, figsize=(5*plot_columns,5*plot_rows))\n",
    "fig.tight_layout(pad=5.0)\n",
    "\n",
    "for ax in axarr.flat:\n",
    "    if axnum <= filenum*2-1:\n",
    "        dict_key = onlyfiles[math.floor(axnum/2)].split(\".nid\")[0]\n",
    "\n",
    "        if axnum%2 == 0:\n",
    "            im[str(axnum)] = ax.pcolormesh(xaxis_dict[dict_key], yaxis_dict[dict_key], data_dict[dict_key], cmap = cmap)\n",
    "            ax.set_title(dict_key)\n",
    "        if axnum%2 == 1:\n",
    "            im[str(axnum)] = ax.pcolormesh(xaxis_dict[dict_key], yaxis_dict[dict_key], data_levelled[dict_key], cmap = cmap)\n",
    "            ax.set_title(dict_key + ' levelled')\n",
    "\n",
    "        ax.set(xlabel=xlbl, ylabel=ylbl)\n",
    "        ax_divider = make_axes_locatable(ax)\n",
    "        cax = ax_divider.append_axes(\"right\", size=\"5%\", pad=\"2%\")\n",
    "        cb[str(axnum)] = colorbar(im[str(axnum)], cax=cax)\n",
    "        #cb[str(axnum)].ax.set_ylabel(\"Z (nm)\")\n",
    "        axnum += 1"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.1 64-bit",
   "language": "python",
   "name": "python38164bit13404dc2f00c48a39c0eb9d561458b9c"
  },
  "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
}
