{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lmfit\n",
    "import xarray as xr\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import copy\n",
    "\n",
    "import glob\n",
    "\n",
    "import xrft\n",
    "import finufft\n",
    "\n",
    "from uncertainties import ufloat\n",
    "from uncertainties import unumpy as unp\n",
    "from uncertainties import umath\n",
    "\n",
    "from datetime import datetime\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "#test\n",
    "plt.rcParams['font.size'] = 18\n",
    "\n",
    "from scipy.ndimage import gaussian_filter, rotate\n",
    "import matplotlib as mpl\n",
    "from scipy.interpolate import CubicSpline\n",
    "from scipy.optimize import curve_fit\n",
    "mpl.rc('xtick', labelsize=8)\n",
    "mpl.rc('ytick', labelsize=8)\n",
    "\n",
    "import sys\n",
    "sys.path.append(\"C:/Users/Jianshun Gao/VisualCodeProjects/analysescript\")\n",
    "sys.path.append(\"C:/Users/Jianshun Gao/PycharmProjects/Foglight4/Data\")\n",
    "\n",
    "from DataContainer.ReadData import read_hdf5_file, read_hdf5_global, read_hdf5_run_time, read_csv_file\n",
    "from Analyser.ImagingAnalyser import ImageAnalyser\n",
    "from Analyser.FitAnalyser import FitAnalyser\n",
    "from Analyser.FitAnalyser import ThomasFermi2dModel, DensityProfileBEC2dModel, Polylog22dModel\n",
    "from Analyser.FFTAnalyser import fft, ifft, fft_nutou\n",
    "from ToolFunction.ToolFunction import *\n",
    "\n",
    "import time\n",
    "import logging as log\n",
    "import scipy.constants as const\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()\n",
    "\n",
    "\n",
    "import Analyser.FitAnalyser as dp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "C:/Users/Jianshun Gao/PycharmProjects/Foglight4/Data/josch/high_cond_1.nc\n"
     ]
    }
   ],
   "source": [
    "path = \"C:/Users/Jianshun Gao/PycharmProjects/Foglight4/Data\"\n",
    "file = 'high_cond_1'\n",
    "file_path = f'{path}/josch/{file}.nc'\n",
    "print(file_path)\n",
    "with xr.open_dataarray(file_path) as data_x:\n",
    "    data_x\n",
    "\n",
    "data = data_x.to_numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "is_debug = True\n",
    "fitm = DensityProfileBEC2dModel(is_debug=is_debug, atom_n_conv=144/2)\n",
    "\n",
    "def calc_cen_bulk(thresh):\n",
    "    \"\"\"\n",
    "    returns array in shape of input, containing array with [Y_center,X_center] for each image\n",
    "    \"\"\"\n",
    "    shape = np.shape(thresh)\n",
    "    cen = np.zeros((shape[0], 2))\n",
    "    for i in range(0, shape[0]):\n",
    "        cen[i] = fitm.calc_cen_pix(thresh[i])\n",
    "    return cen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "shape = np.shape(data)\n",
    "data_rot = np.empty(shape)\n",
    "\n",
    "for i in range(0,shape[0]):\n",
    "    data_rot[i] = rotate(data[i], 36, reshape=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[599.42377261 959.62015504]\n",
      " [599.46599496 959.50629723]\n",
      " [599.47858942 959.49118388]\n",
      " [599.47505938 959.41330166]\n",
      " [599.68149883 959.26697892]\n",
      " [599.4627907  959.58372093]\n",
      " [599.46412556 959.46188341]]\n"
     ]
    }
   ],
   "source": [
    "# cropping\n",
    "pr_data = data_rot\n",
    "\n",
    "cut_width_x = 200\n",
    "cut_width_y = 150\n",
    "thresh = fitm.calc_thresh(pr_data)\n",
    "\n",
    "center = calc_cen_bulk(thresh)\n",
    "\n",
    "print(center)\n",
    "\n",
    "\n",
    "if len(shape) == 3:\n",
    "    cropOD = np.zeros((shape[0], cut_width_y, cut_width_x))\n",
    "    for i in range(0,shape[0]):\n",
    "        for j in range(0, shape[1]):\n",
    "            cropOD[i] = pr_data[i, round(center[i,0]-cut_width_y/2):round(center[i,0]+cut_width_y/2), round(center[i,1]-cut_width_x/2):round(center[i,1]+cut_width_x/2)]\n",
    "\n",
    "shape = np.shape(cropOD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 196620 characters]",
      "text/plain": [
       "<Figure size 1750x250 with 7 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(1, shape[0], figsize=(2.5 *shape[0],2.5))\n",
    "for i in range(0, shape[0]):\n",
    "    ax[i].pcolormesh(cropOD[i], cmap='jet')\n",
    "\n",
    "    ax[i].set_aspect('equal')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " \n",
      "file high_cond_1 \n",
      " image: 0\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 78324 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: (200, 150)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 44772 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y smaller x, 1d fit along y\n",
      "\n",
      "1d fit initialization\n",
      "center = [99.49748744 75.5033557 ]\n",
      "BEC widths: [23 13]\n",
      "\n",
      "1d init fit values\n",
      "Name          Value      Min      Max   Stderr     Vary     Expr Brute_Step\n",
      "amp_bec       1.821        0    4.734     None     True     None     None\n",
      "amp_th        1.821        0    4.734     None     True     None     None\n",
      "deltax           39        0      200     None     True     None     None\n",
      "sigma_bec     10.66        0       26     None     True     None     None\n",
      "sigma_th      26.94        0      inf     None    False 0.632*sigma_bec + 0.518*deltax     None\n",
      "x0_bec         75.5     65.5     85.5     None     True     None     None\n",
      "x0_th          75.5     65.5     85.5     None     True     None     None\n",
      "1d fitted values\n",
      "x0_bec: 75.720, (init = 75.503), bounds = [65.50 : 85.50] \n",
      "x0_th: 75.824, (init = 75.503), bounds = [65.50 : 85.50] \n",
      "amp_bec: 2.943, (init = 1.821), bounds = [0.00 : 4.73] \n",
      "amp_th: 0.742, (init = 1.821), bounds = [0.00 : 4.73] \n",
      "sigma_bec: 10.919, (init = 10.656), bounds = [0.00 : 26.00] \n",
      "sigma_th: 9.846, (init = 26.936), bounds = [0.00 : inf] \n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 47360 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Init Params\n",
      "Name                    Value      Min      Max   Stderr     Vary     Expr Brute_Step\n",
      "amp_bec                   3.5        0    6.577     None     True     None     None\n",
      "amp_th                 0.8824        0    6.577     None     True     None     None\n",
      "atom_number_bec         905.4     -inf      inf     None    False amp_bec / 5 * 2 * 3.14159265359 * sigmax_bec * sigmay_bec     None\n",
      "atom_number_th          393.2     -inf      inf     None    False amp_th * 2 * 3.14159265359 * 1.20206 / 1.643 * sigma_th * sigma_th     None\n",
      "condensate_fraction    0.6972     -inf      inf     None    False atom_number_bec / (atom_number_bec + atom_number_th)     None\n",
      "rot_angle                  40       10       70     None     True     None     None\n",
      "sigma_th                9.846        0      200     None     True     None     None\n",
      "sigmax_bec              18.85        0       46     None     True     None     None\n",
      "sigmay_bec              10.92        0       26     None     True     None     None\n",
      "x0_bec                   99.5     89.5    109.5     None     True     None     None\n",
      "x0_th                    99.5     89.5    109.5     None     True     None     None\n",
      "y0_bec                   75.5     65.5     85.5     None     True     None     None\n",
      "y0_th                    75.5     65.5     85.5     None     True     None     None\n",
      "\n",
      "bval first fit\n",
      "amp_bec: 3.361, (init = 3.500), bounds = [0.00 : 6.58] \n",
      "amp_th: 1.033, (init = 0.882), bounds = [0.00 : 6.58] \n",
      "x0_bec: 100.041, (init = 99.497), bounds = [89.50 : 109.50] \n",
      "y0_bec: 76.005, (init = 75.503), bounds = [65.50 : 85.50] \n",
      "x0_th: 99.794, (init = 99.497), bounds = [89.50 : 109.50] \n",
      "y0_th: 75.935, (init = 75.503), bounds = [65.50 : 85.50] \n",
      "sigmax_bec: 23.817, (init = 18.852), bounds = [0.00 : 46.00] \n",
      "sigmay_bec: 11.257, (init = 10.919), bounds = [0.00 : 26.00] \n",
      "sigma_th: 9.384, (init = 9.846), bounds = [0.00 : 200.00] \n",
      "rot_angle: 36.040, (init = 40.000), bounds = [10.00 : 70.00] \n",
      "\n",
      "Name                    Value      Min      Max   Stderr     Vary     Expr Brute_Step\n",
      "amp_bec                   3.5        0    6.577     None     True     None     None\n",
      "amp_th                 0.8824        0    6.577     None     True     None     None\n",
      "atom_number_bec         905.4     -inf      inf     None    False amp_bec / 5 * 2 * 3.14159265359 * sigmax_bec * sigmay_bec     None\n",
      "atom_number_th          393.2     -inf      inf     None    False amp_th * 2 * 3.14159265359 * 1.20206 / 1.643 * sigma_th * sigma_th     None\n",
      "condensate_fraction    0.6972     -inf      inf     None    False atom_number_bec / (atom_number_bec + atom_number_th)     None\n",
      "rot_angle                  40       10       70     None     True     None     None\n",
      "sigma_th                9.846        0      200     None     True     None     None\n",
      "sigmax_bec              18.85        0       46     None     True     None     None\n",
      "sigmay_bec              10.92        0       26     None     True     None     None\n",
      "x0_bec                   99.5     89.5    109.5     None     True     None     None\n",
      "x0_th                    99.5     89.5    109.5     None     True     None     None\n",
      "y0_bec                   75.5     65.5     85.5     None     True     None     None\n",
      "y0_th                    75.5     65.5     85.5     None     True     None     None\n",
      "amp_bec: 3.361, (init = 3.500), bounds = [0.00 : 6.58] \n",
      "amp_th: 1.033, (init = 0.882), bounds = [0.00 : 6.58] \n",
      "x0_bec: 100.041, (init = 99.497), bounds = [89.50 : 109.50] \n",
      "y0_bec: 76.005, (init = 75.503), bounds = [65.50 : 85.50] \n",
      "x0_th: 99.794, (init = 99.497), bounds = [89.50 : 109.50] \n",
      "y0_th: 75.935, (init = 75.503), bounds = [65.50 : 85.50] \n",
      "sigmax_bec: 23.817, (init = 18.852), bounds = [0.00 : 46.00] \n",
      "sigmay_bec: 11.257, (init = 10.919), bounds = [0.00 : 26.00] \n",
      "sigma_th: 9.384, (init = 9.846), bounds = [0.00 : 200.00] \n",
      "rot_angle: 36.040, (init = 40.000), bounds = [10.00 : 70.00] \n",
      "\n",
      "\n",
      "Atom numbers:\n",
      "   N_bec: 80595\n",
      "   N_th: 29742\n",
      "   N_ges: 110338\n",
      "   Cond. frac: 73.04 %\n",
      "\n",
      "Temperature: 31.40 nK\n"
     ]
    }
   ],
   "source": [
    "times = []\n",
    "x = np.linspace(0,shape[2],shape[2])\n",
    "y = np.linspace(0,shape[1], shape[1])\n",
    "\n",
    "X,Y = np.meshgrid(x, y)\n",
    "X_1d = X.flatten()\n",
    "Y_1d = Y.flatten()\n",
    "result = []\n",
    "for i in range(0, shape[0]):\n",
    "# for i in [0]:\n",
    "    print('\\n ')\n",
    "    print(f'file {file} \\n image: {i}')\n",
    "    start = time.time()\n",
    "    \n",
    "    init = fitm.guess(cropOD[i].flatten(), X_1d, Y_1d, rot_angle=40, vary_rot=True, is_debug=True)\n",
    "\n",
    "\n",
    "    res = fitm.fit(cropOD[i].flatten(), x=X_1d, y=Y_1d, params=init)\n",
    "\n",
    "    init.pretty_print()\n",
    "    fitm.print_bval(res)\n",
    "    stop = time.time()\n",
    "    fitm.return_atom_number(res, X_1d, Y_1d)\n",
    "    fitm.return_temperature(res, omg=1036.5, tof=26e-3, is_print=False)\n",
    "    fitm.return_temperature(res, tof=26e-3)\n",
    "\n",
    "    # print(f'   time = {stop-start:.3f} s')\n",
    "\n",
    "    result.append(res)\n",
    "    times.append(stop-start)\n",
    "    break\n",
    "\n",
    "times = np.array(times)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1561024 characters]",
      "text/plain": [
       "<Figure size 2000x2800 with 42 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# fig, axs = plt.subplots(shape[0] , 5, figsize=(14, 4 * shape[0] ),dpi = 100)\n",
    "\n",
    "fig, axs = plt.subplots(shape[0] , 6, figsize=(20,4 * shape[0] ),dpi = 100)\n",
    "fig.suptitle(f'{file}: ν = ({data_x.trap_f[0]}, {data_x.trap_f[1]}, {data_x.trap_f[2]}) Hz, T = {data_x.temp[0]*1e9:.0f}nK, tof = {data_x.tof * 1e3:.0f}ms, a = {data_x.a_fac :.0f}a$_0$, N = {data_x.N:.0e} ', y=-0.005*shape[0]+0.965, fontsize=24)\n",
    "# fig.suptitle('test', va='bottom')\n",
    "\n",
    "for i in range(0,shape[0]):\n",
    "\n",
    "    axs[i,0].set_title(f'image {i}, cond. frac = {fitm.cond_frac(result[i], X, Y) *1e2:.2f} %', loc='left', fontsize=18)\n",
    "\n",
    "    lmfit.fit_report(result[i])\n",
    "    bval = result[i].best_values\n",
    "    fit = dp.density_profile_BEC_2d(X,Y, **bval)\n",
    "    vmax = np.max(fit)\n",
    "\n",
    "    if bval['amp_bec'] > bval['amp_th']:\n",
    "        cen_str = 'bec'\n",
    "    else:\n",
    "        cen_str = 'th'\n",
    "    cen_x = round(bval[f'x0_{cen_str}'])\n",
    "    cen_y = round(bval[f'y0_{cen_str}'])\n",
    "\n",
    "    ax = axs[i,0]\n",
    "    ax.pcolormesh(X, Y, cropOD[i], vmin=0, vmax=vmax, cmap='jet', shading='auto')\n",
    "    #plt.colorbar(art, ax=ax, label='z')\n",
    "\n",
    "\n",
    "    # Plot gaussian 2d Fit + legend including Width parameters\n",
    "    ax = axs[i,1]\n",
    "\n",
    "    ax.pcolormesh(X, Y, fit, vmin=0, vmax=vmax, cmap='jet', shading='auto')\n",
    "    #plt.colorbar(art, ax=ax, label='z')\n",
    "\n",
    "    ax = axs[i,2]\n",
    "\n",
    "    ax.pcolormesh(X, Y, fit-cropOD[i], vmin=0, vmax=0.2, cmap='jet', shading='auto')\n",
    "\n",
    "\n",
    "    ax = axs[i,3]\n",
    "\n",
    "    ax.plot(x, cropOD[i, cen_y, :])\n",
    "    ax.plot(x, fit[cen_y, :], label='bimodal')\n",
    "    ax.plot(x, dp.thermal(x, bval['x0_th'], bval['amp_th'], bval['sigma_th']), label='thermal')\n",
    "    ax.legend(fontsize=8)\n",
    "\n",
    "    ax = axs[i,4]\n",
    "\n",
    "    ax.plot(y, cropOD[i, :, cen_x])\n",
    "    ax.plot(y, fit[:, cen_x], label='bimodal')\n",
    "    ax.plot(y, dp.thermal(y, bval['y0_th'], bval['amp_th'], bval['sigma_th']), label='thermal')\n",
    "    ax.legend(fontsize=8)\n",
    "\n",
    "    ax = axs[i,5]\n",
    "    Na = fitm.return_atom_number(result[i], X, Y, is_print=False)\n",
    "    ax.axis('off')\n",
    "    col_labels = ['Sim', 'Fit']\n",
    "    row_labels = ['N', f'N$_{{bec}}$', f'N$_{{th}}$','f', 'T (nk)']\n",
    "    table_vals = [[f'{data_x.N:.2e}',f\"{Na['N']:.2e}\"],\n",
    "                  [f'{data_x.cond_frac.data[i] * data_x.N:.2e}',f\"{Na['N_bec']:.2e}\"],\n",
    "                  [f'{(1-data_x.cond_frac.data[i]) * data_x.N:.2e}',f\"{Na['N_th']:.2e}\"],\n",
    "                  [f'{data_x.cond_frac.data[i] *1e2:.1f} %', f\"{Na['cond_f'] *1e2:.2f} %\"],\n",
    "                  [f'{data_x.temp[i] *1e9:.0f} nK', f'{fitm.return_temperature(result[i], tof=data_x.tof, is_print=False)*1e9:.1f} nK']\n",
    "                  ]\n",
    "\n",
    "\n",
    "    my_table = ax.table(cellText=table_vals,\n",
    "                        colWidths=[0.4]*2,\n",
    "                        rowLabels=row_labels,\n",
    "                        colLabels=col_labels,\n",
    "                        loc='upper right')\n",
    "    my_table.scale(1,2)\n",
    "    ax.text(0.1,0.1,f'fitting time = {times[i]:.2f} s', fontsize=14)\n",
    "    ax.text(0.1,0,f'rot angle = {bval[\"rot_angle\"]}°', fontsize=14)\n",
    "\n",
    "t_fonts=20\n",
    "axs[0,0].set_title(f'Data \\n ', fontsize=t_fonts)\n",
    "axs[0,1].set_title('Fit \\n ', fontsize=t_fonts)\n",
    "axs[0,2].set_title('Data - Fit \\n ', fontsize=t_fonts)\n",
    "axs[0,3].set_title('cut along x \\n ', fontsize=t_fonts)\n",
    "axs[0,4].set_title('cut along y \\n ', fontsize=t_fonts)\n",
    "\n",
    "\n",
    "# plt.tight_layout()\n",
    "# plt.savefig(f'fit_output/final/{file}_{saveprefix}.png')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([3.e-08, 3.e-08, 3.e-08, 3.e-08, 3.e-08, 3.e-08, 3.e-08])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_x.temp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.9.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
