{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import supporting package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "\n",
    "from uncertainties import ufloat\n",
    "from uncertainties import unumpy as unp\n",
    "from uncertainties import umath\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from DataContainer.ReadData import read_hdf5_file\n",
    "from Analyser.ImagingAnalyser import ImageAnalyser\n",
    "from Analyser.FitAnalyser import FitAnalyser\n",
    "from Analyser.FitAnalyser import ThomasFermi2dModel, DensityProfileBEC2dModel, Polylog22dModel\n",
    "from Analyser.FitAnalyser import NewFitModel\n",
    "from ToolFunction.ToolFunction import *\n",
    "\n",
    "from ToolFunction.HomeMadeXarrayFunction import errorbar, dataarray_plot_errorbar\n",
    "xr.plot.dataarray_plot.errorbar = errorbar\n",
    "xr.plot.accessor.DataArrayPlotAccessor.errorbar = dataarray_plot_errorbar\n",
    "\n",
    "imageAnalyser = ImageAnalyser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from matplotlib.colors import ListedColormap, LinearSegmentedColormap\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "plt.rcParams[\"font.family\"] = \"arial\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Ncount_to_atoms():\n",
    "    return 1 / 8.4743e-14 / 0.3725 * 5.86e-6**2 / 0.6606**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "\n",
    "colormap = np.zeros((1024, 3))\n",
    "\n",
    "with open('smooth-cool-warm-table-float-1024.csv', newline='') as csvfile:\n",
    "    spamreader = csv.reader(csvfile, delimiter=' ', quotechar='|')\n",
    "    i = 0\n",
    "    for row in spamreader:\n",
    "        try:\n",
    "            a = row[0].split(',')\n",
    "            colormap[i, 0] = float(a[1])\n",
    "            colormap[i, 1] = float(a[2])\n",
    "            colormap[i, 2] = float(a[3])\n",
    "            i = i + 1\n",
    "        except:\n",
    "            pass\n",
    "\n",
    "colormap = ListedColormap(colormap)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_colors = colormap(np.linspace(0, 1, 7))\n",
    "plot_blue = data_colors[-3]\n",
    "plot_red = data_colors[-1]\n",
    "plot_red_alpha = 1"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start a client for parallel computing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\"> </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px;\">Client</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Client-65877ea1-1120-11ee-bd28-9c7bef43b4fb</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "\n",
       "        <tr>\n",
       "        \n",
       "            <td style=\"text-align: left;\"><strong>Connection method:</strong> Cluster object</td>\n",
       "            <td style=\"text-align: left;\"><strong>Cluster type:</strong> distributed.LocalCluster</td>\n",
       "        \n",
       "        </tr>\n",
       "\n",
       "        \n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:8787/status\" target=\"_blank\">http://127.0.0.1:8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\"></td>\n",
       "            </tr>\n",
       "        \n",
       "\n",
       "        </table>\n",
       "\n",
       "        \n",
       "            <button style=\"margin-bottom: 12px;\" data-commandlinker-command=\"dask:populate-and-launch-layout\" data-commandlinker-args='{\"url\": \"http://127.0.0.1:8787/status\" }'>\n",
       "                Launch dashboard in JupyterLab\n",
       "            </button>\n",
       "        \n",
       "\n",
       "        \n",
       "            <details>\n",
       "            <summary style=\"margin-bottom: 20px;\"><h3 style=\"display: inline;\">Cluster Info</h3></summary>\n",
       "            <div class=\"jp-RenderedHTMLCommon jp-RenderedHTML jp-mod-trusted jp-OutputArea-output\">\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\">\n",
       "    </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px; margin-top: 0px;\">LocalCluster</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">ee7c3b53</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard:</strong> <a href=\"http://127.0.0.1:8787/status\" target=\"_blank\">http://127.0.0.1:8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Workers:</strong> 6\n",
       "                </td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total threads:</strong> 60\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total memory:</strong> 55.88 GiB\n",
       "                </td>\n",
       "            </tr>\n",
       "            \n",
       "            <tr>\n",
       "    <td style=\"text-align: left;\"><strong>Status:</strong> running</td>\n",
       "    <td style=\"text-align: left;\"><strong>Using processes:</strong> True</td>\n",
       "</tr>\n",
       "\n",
       "            \n",
       "        </table>\n",
       "\n",
       "        <details>\n",
       "            <summary style=\"margin-bottom: 20px;\">\n",
       "                <h3 style=\"display: inline;\">Scheduler Info</h3>\n",
       "            </summary>\n",
       "\n",
       "            <div style=\"\">\n",
       "    <div>\n",
       "        <div style=\"width: 24px; height: 24px; background-color: #FFF7E5; border: 3px solid #FF6132; border-radius: 5px; position: absolute;\"> </div>\n",
       "        <div style=\"margin-left: 48px;\">\n",
       "            <h3 style=\"margin-bottom: 0px;\">Scheduler</h3>\n",
       "            <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Scheduler-f17926fd-5020-4803-9692-4e5947c8d4dd</p>\n",
       "            <table style=\"width: 100%; text-align: left;\">\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Comm:</strong> tcp://127.0.0.1:63894\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Workers:</strong> 6\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Dashboard:</strong> <a href=\"http://127.0.0.1:8787/status\" target=\"_blank\">http://127.0.0.1:8787/status</a>\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total threads:</strong> 60\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Started:</strong> Just now\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total memory:</strong> 55.88 GiB\n",
       "                    </td>\n",
       "                </tr>\n",
       "            </table>\n",
       "        </div>\n",
       "    </div>\n",
       "\n",
       "    <details style=\"margin-left: 48px;\">\n",
       "        <summary style=\"margin-bottom: 20px;\">\n",
       "            <h3 style=\"display: inline;\">Workers</h3>\n",
       "        </summary>\n",
       "\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 0</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:63924\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 10\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:63932/status\" target=\"_blank\">http://127.0.0.1:63932/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 9.31 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:63897\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> C:\\Users\\JIANSH~1\\AppData\\Local\\Temp\\dask-worker-space\\worker-oaisnlk0\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 1</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:63928\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 10\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:63941/status\" target=\"_blank\">http://127.0.0.1:63941/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 9.31 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:63898\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> C:\\Users\\JIANSH~1\\AppData\\Local\\Temp\\dask-worker-space\\worker-lgn_00uq\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 2</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:63921\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 10\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:63930/status\" target=\"_blank\">http://127.0.0.1:63930/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 9.31 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:63899\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> C:\\Users\\JIANSH~1\\AppData\\Local\\Temp\\dask-worker-space\\worker-1ozf7xf0\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 3</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:63926\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 10\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:63937/status\" target=\"_blank\">http://127.0.0.1:63937/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 9.31 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:63900\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> C:\\Users\\JIANSH~1\\AppData\\Local\\Temp\\dask-worker-space\\worker-53pv7g96\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 4</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:63927\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 10\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:63938/status\" target=\"_blank\">http://127.0.0.1:63938/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 9.31 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:63901\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> C:\\Users\\JIANSH~1\\AppData\\Local\\Temp\\dask-worker-space\\worker-tr6dxqb7\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 5</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:63925\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 10\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:63934/status\" target=\"_blank\">http://127.0.0.1:63934/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 9.31 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:63902\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> C:\\Users\\JIANSH~1\\AppData\\Local\\Temp\\dask-worker-space\\worker-7qa0mcnj\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "\n",
       "    </details>\n",
       "</div>\n",
       "\n",
       "        </details>\n",
       "    </div>\n",
       "</div>\n",
       "            </details>\n",
       "        \n",
       "\n",
       "    </div>\n",
       "</div>"
      ],
      "text/plain": [
       "<Client: 'tcp://127.0.0.1:63894' processes=6 threads=60, memory=55.88 GiB>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dask.distributed import Client\n",
    "client = Client(n_workers=6, threads_per_worker=10, processes=True, memory_limit='10GB')\n",
    "client"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set global path for experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "groupList = [\n",
    "    \"images/MOT_3D_Camera/in_situ_absorption\",\n",
    "    \"images/ODT_1_Axis_Camera/in_situ_absorption\",\n",
    "    \"images/ODT_2_Axis_Camera/in_situ_absorption\",\n",
    "]\n",
    "\n",
    "dskey = {\n",
    "    \"images/MOT_3D_Camera/in_situ_absorption\": \"camera_0\",\n",
    "    \"images/ODT_1_Axis_Camera/in_situ_absorption\": \"camera_1\",\n",
    "    \"images/ODT_2_Axis_Camera/in_situ_absorption\": \"camera_2\",\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_dir = '//DyLabNAS/Data/'\n",
    "SequenceName = \"Repetition_scan\" + \"/\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MOT optimize"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 09.06.2023"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The detected scaning axes and values are: \n",
      "\n",
      "{'cmot_initial_current': array([0.1 , 0.12, 0.14, 0.16, 0.18, 0.2 , 0.22, 0.24, 0.26, 0.28]), 'initial_freq': array([100.5, 100.9, 101.3, 101.7, 102.1, 102.5, 102.9, 103.3, 103.7,\n",
      "       104.1]), 'runs': array([0., 1., 2.])}\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 15912 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "folderPath = img_dir + SequenceName + '2023/06/09'# get_date()\n",
    "\n",
    "shotNum = \"0010\"\n",
    "filePath = folderPath + \"/\" + shotNum + \"/*.h5\"\n",
    "\n",
    "dataSetDict = {\n",
    "    dskey[groupList[i]]: read_hdf5_file(filePath, groupList[i])\n",
    "    for i in range(len(groupList))\n",
    "}\n",
    "\n",
    "dataSet = dataSetDict[\"camera_0\"]\n",
    "\n",
    "print_scanAxis(dataSet)\n",
    "\n",
    "scanAxis = get_scanAxis(dataSet)\n",
    "\n",
    "dataSet = auto_rechunk(dataSet)\n",
    "\n",
    "dataSet = imageAnalyser.get_absorption_images(dataSet)\n",
    "\n",
    "imageAnalyser.center = (690, 1435)\n",
    "imageAnalyser.span = (1000, 700)\n",
    "imageAnalyser.fraction = (0.1, 0.1)\n",
    "\n",
    "dataSet_cropOD = imageAnalyser.crop_image(dataSet.OD)\n",
    "dataSet_cropOD = imageAnalyser.substract_offset(dataSet_cropOD).load()\n",
    "\n",
    "Ncount = imageAnalyser.get_Ncount(dataSet_cropOD)\n",
    "Ncount_mean = calculate_mean(Ncount)\n",
    "Ncount_std = calculate_std(Ncount)\n",
    "\n",
    "fig = plt.figure()\n",
    "# ax = fig.gca()\n",
    "Ncount_mean.plot.pcolormesh()\n",
    "plt.xlabel('MOT AOM Frequency (MHz)')\n",
    "plt.ylabel('MOT Gradient Coil Current (A)')\n",
    "plt.tight_layout()\n",
    "# plt.grid(visible=1)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "freq = -(Ncount_mean.initial_freq.to_numpy() - 99.969) *2 / 0.136\n",
    "current = Ncount_mean.cmot_initial_current.to_numpy() * 5.17\n",
    "data = Ncount_mean.to_numpy() * Ncount_to_atoms()\n",
    "\n",
    "X, Y = np.meshgrid(freq, current)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 51932 characters]",
      "text/plain": [
       "<Figure size 816x1615.68 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "\n",
    "fig = plt.figure(figsize=(6.8, 6.8/4*3*2*(1 + 16/50)), dpi=120)\n",
    "grid = fig.add_gridspec(4, 1, height_ratios=[2, 50, 10, 50], wspace=0.4, hspace=0.1)\n",
    "\n",
    "filter = ''\n",
    "edgecolor = None\n",
    "\n",
    "colorbar_ticks = np.linspace(0, 2.5e8, 6)\n",
    "\n",
    "plot_axes = plt.subplot(grid[1, 0])\n",
    "\n",
    "im = plot_axes.pcolormesh(X, Y, data, cmap=colormap, edgecolor=edgecolor, vmin=0)\n",
    "\n",
    "colorbar = fig.colorbar(im, ax=plot_axes, ticks=np.array([0, 0.5, 1, 1.5, 2])*1e8)\n",
    "colorbar.formatter.set_powerlimits((0, 0))\n",
    "colorbar.formatter.set_useMathText(True)\n",
    "colorbar.ax.yaxis.set_offset_position('left')\n",
    "colorbar.ax.tick_params(axis='both', labelsize=20)\n",
    "colorbar.ax.yaxis.offsetText.set_fontsize(20)\n",
    "\n",
    "plot_axes.set_xlabel(\"3D-MOT detuning $\\delta_\\mathrm{3D}$ $(\\Gamma_\\mathrm{626})$\", fontsize=20)\n",
    "plot_axes.set_ylabel(\"3D-MOT gradient $b'_\\mathrm{3D}$ (G/cm)\", fontsize=20)\n",
    "colorbar.set_label(\"Loaded atom number $N_\\mathrm{4s}$\", fontsize=20)\n",
    "\n",
    "plot_axes.tick_params(axis='both', which='major', labelsize=20)\n",
    "plot_axes.tick_params(axis='both', which='minor', labelsize=16)\n",
    "plot_axes.xaxis.offsetText.set_fontsize(20)\n",
    "plot_axes.yaxis.offsetText.set_fontsize(20)\n",
    "\n",
    "plt.setp(plot_axes.spines.values(), linewidth=3)\n",
    "plot_axes.xaxis.set_tick_params(width=3)\n",
    "plot_axes.yaxis.set_tick_params(width=3)\n",
    "plot_axes.tick_params(direction='in', length=10)\n",
    "\n",
    "# # plotting.figure[dataExtractor_key].colorbar_min = 0\n",
    "# # plotting.figure[dataExtractor_key].colorbar_max = 2.5e8\n",
    "\n",
    "# plotting.figure[dataExtractor_key].add_pcolormesh_plot(cmap=colormap, edgecolor=edgecolor, grid=plot_axes)\n",
    "# plotting.figure[dataExtractor_key].add_axes_label()\n",
    "# # plotting.figure[dataExtractor_key].add_color_bar(ticks=colorbar_ticks)\n",
    "# plotting.figure[dataExtractor_key].add_color_bar(ticks=np.linspace(0, 8, 5))\n",
    "# plotting.figure[dataExtractor_key].colorbar.ax.set_yticklabels(['{:.1f}'.format(x) for x in np.linspace(0, 8, 5)])\n",
    "\n",
    "# plotting.figure[dataExtractor_key].colorbar.ax.text(3.3, 7.69, '$\\\\times 10^{7}$', va='bottom', weight='bold', fontsize=20)\n",
    "\n",
    "# # plotting.figure[dataExtractor_key].set_color_bar_offset_position(0, 1.07)\n",
    "\n",
    "# # plotting.figure[dataExtractor_key].plot_axes.set_xlim([-30, -55])\n",
    "# # plotting.figure[dataExtractor_key].plot_axes.set_ylim([0.15, 0.65])\n",
    "\n",
    "# # plt.xticks([-40])\n",
    "\n",
    "# plotting.figure[dataExtractor_key].plot_axes.set_xlabel(\"3D-MOT detuning $\\delta_\\mathrm{3D} / \\Gamma_{626}$\", fontsize=20)\n",
    "# plotting.figure[dataExtractor_key].plot_axes.set_ylabel(\"3D-MOT gradient $b'_\\mathrm{3D}$ (G/cm)\", fontsize=20)\n",
    "# plotting.figure[dataExtractor_key].colorbar.set_label(\"Loaded atom number $N_\\mathrm{4s}$\", fontsize=20)\n",
    "\n",
    "# plot_axes = plt.subplot(grid[0, 0])\n",
    "# plot_axes.text(-0.17, 1.0, '(a)', va='bottom', weight='bold', fontsize=20)\n",
    "# plot_axes.set_axis_off()\n",
    "\n",
    "fig.savefig('figS3_v1.pdf', bbox_inches = \"tight\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "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.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
