{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Fitting with constraint\n",
    "result = []\n",
    "times = []\n",
    "x = np.linspace(0,shape[3],150)\n",
    "y = np.linspace(0,shape[2], 150)\n",
    "\n",
    "for i in range(0,shape[0]):\n",
    "    temp_res_arr = []\n",
    "    for j in range(0,shape[1]):\n",
    "        data = cropOD[i,j]\n",
    "        fitModel = lmfit.Model(density_profile_BEC_2d, independent_vars=['x','y'])\n",
    "        #fitModel.set_param_hint('deltax', value=5)\n",
    "\n",
    "        bval = result_x[i][j].best_values\n",
    "        S = np.max(blurred[i,j])/(bval['amp_bec'] + bval['amp_th'])\n",
    "\n",
    "        params = lmfit.Parameters()\n",
    "        #print(bval['sigma_th'])\n",
    "\n",
    "        params.add_many(\n",
    "            ('amp_bec', S * bval['amp_bec'], True, 0, 1.3 * np.max(data)),\n",
    "            ('amp_th',S * bval['amp_th'], True, 0, 1.3 * np.max(data)),\n",
    "            ('x0_bec',center[i,j,0], True, 0, 150),\n",
    "            ('y0_bec',center[i,j,1], True, 0, 150),\n",
    "            ('x0_th',center[i,j,0], True, 0, 150),\n",
    "            ('y0_th',center[i,j,1], True, 0, 150),\n",
    "            ('sigmax_bec',bval['sigma_bec'], True, 0, 50),\n",
    "            ('sigmay_bec',BEC_width_guess[i,j,1], True, 0, 50),\n",
    "            ('D_sigX', 1.93*bval['sigma_th'] - 1.22*bval['sigma_bec'], True, 0.1, 50),\n",
    "            ('D_sig_th', 0, True, -10, 10)\n",
    "        )\n",
    "\n",
    "        params.add('sigmax_th',bval['sigma_th'], min=0, expr=f'0.632*sigmax_bec + 0.518*D_sigX')\n",
    "        params.add('sigmay_th',bval['sigma_th'], min=0, expr=f'sigmax_th + D_sig_th')\n",
    "\n",
    "        X,Y = np.meshgrid(x, y)\n",
    "        X_1d = X.flatten()\n",
    "        Y_1d = Y.flatten()\n",
    "\n",
    "        data1d = data.flatten()\n",
    "        start = time.time()\n",
    "        res = fitModel.fit(data1d, x=X_1d, y=Y_1d, params=params)\n",
    "        stop = time.time()\n",
    "        print(stop-start)\n",
    "        times.append(stop-start)\n",
    "        temp_res_arr.append(res)\n",
    "    result.append(temp_res_arr)\n",
    "times = np.array(times)\n",
    "print(f\"fitting time = {np.mean(times)} +- {np.std(times, ddof=1)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# from opencv import moments\n",
    "shape = np.shape(cropOD)\n",
    "sigma = 0.4\n",
    "blurred = gaussian_filter(cropOD, sigma=sigma)\n",
    "\n",
    "thresh = np.zeros(shape)\n",
    "for i in range(0,shape[0]):\n",
    "    for j in range(0, shape[1]):\n",
    "        thresh[i,j] = np.where(blurred[i,j] < np.max(blurred[i,j])*0.5,0,1)\n",
    "\n",
    "center = calc_cen_bulk(thresh)\n",
    "\n",
    "BEC_width_guess = guess_BEC_width(thresh, center)\n",
    "\n",
    "X_guess_og = np.zeros((shape[0], shape[1], shape[3]))\n",
    "# Y_guess_og = np.zeros((shape[0], shape[1], shape[2]))\n",
    "\n",
    "for i in range(0, shape[0]):\n",
    "    for j in range(0, shape[1]):\n",
    "        X_guess_og[i,j] = np.sum(cropOD[i,j,round(center[i,j,1] - BEC_width_guess[i,j,1]/2) : round(center[i,j,1] + BEC_width_guess[i,j,1]/2) , :], 0) / len(cropOD[i,j,round(center[i,j,1] - BEC_width_guess[i,j,1]/2) : round(center[i,j,1] + BEC_width_guess[i,j,1]/2),0])\n",
    "\n",
    "        # Y_guess_og[i,j] = np.sum(cropOD[i,j, :, round(center[i,j,0] - BEC_width_guess[i,j,0]/2) : round(center[i,j,0] + BEC_width_guess[i,j,0]/2)], 1) / len(cropOD[i,j,0,round(center[i,j,0] - BEC_width_guess[i,j,0]/2) : round(center[i,j,0] + BEC_width_guess[i,j,0]/2)])\n",
    "\n",
    "#[nr images x, nr images y, X / Y, center / sigma]\n",
    "x = np.linspace(0,149,150)\n",
    "y = np.linspace(0,149, 200)\n",
    "\n",
    "popt = np.zeros((shape[0], shape[1], 6))\n",
    "\n",
    "p0 = np.ones((shape[0], shape[1], 6))\n",
    "\n",
    "max = np.zeros((shape[0], shape[1]))\n",
    "\n",
    "for i in range(0, shape[0]):\n",
    "    max[i] = np.ndarray.max(X_guess_og[i],axis=1)\n",
    "\n",
    "\n",
    "p0[:, :, 0] = center[:, :, 0]  # center BEC\n",
    "p0[:, :, 1] = center[:, :, 0]   # center th\n",
    "p0[:, :, 2] = 0.7 * max          # amp BEC\n",
    "p0[:, :, 3] = 0.3 * max         # amp th\n",
    "p0[:, :, 4] = BEC_width_guess[:, :, 0]   # sigma BEC\n",
    "p0[:, :, 5] = BEC_width_guess[:, :, 0] * 3          # sigma th\n",
    "\n",
    "start = time.time()\n",
    "for i in range(0, shape[0]):\n",
    "    for j in range(0, shape[1]):\n",
    "        popt[i,j], pcov = curve_fit(density_1d, x, X_guess_og[i,j] , p0[i, j], nan_policy='omit')\n",
    "stop = time.time()\n",
    "\n",
    "print(f'fitting time: {(stop-start)*1e3} ms')\n",
    "\n",
    "fsize=(15,10)\n",
    "vmax = 1.5\n",
    "fig, ax = plt.subplots(shape[0],shape[1],figsize=fsize)\n",
    "for i in range(0, shape[0]):\n",
    "    for j in range(0, shape[1]):\n",
    "        lab = f\"A$_{{BEC}}$ = {popt[i,j,0]:.1f} \\n A$_{{th}}$ = {popt[i,j,1]:.1f} \"\n",
    "        ax[i,j].plot(x, X_guess_og[i,j])\n",
    "        ax[i,j].plot(x, density_1d(x, *popt[i,j]), label = lab, zorder=3)\n",
    "        ax[i,j].plot(x, thermal(x, popt[i,j,1], popt[i,j, 3], popt[i,j, 5]))\n",
    "\n",
    "\n",
    "        #ax[i,j].legend(fontsize=10)\n",
    "        ax[i,j].set_facecolor('#FFFFFF')\n",
    "plt.show()\n"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# from opencv import moments\n",
    "start = time.time()\n",
    "shape = np.shape(cropOD)\n",
    "sigma = 0.4\n",
    "blurred = gaussian_filter(cropOD, sigma=sigma)\n",
    "\n",
    "thresh = np.zeros(shape)\n",
    "for i in range(0,shape[0]):\n",
    "    for j in range(0, shape[1]):\n",
    "        thresh[i,j] = np.where(blurred[i,j] < np.max(blurred[i,j])*0.5,0,1)\n",
    "\n",
    "center = calc_cen_bulk(thresh)\n",
    "\n",
    "BEC_width_guess = guess_BEC_width(thresh, center)\n",
    "\n",
    "X_guess_og = np.zeros((shape[0], shape[1], shape[3]))\n",
    "# Y_guess_og = np.zeros((shape[0], shape[1], shape[2]))\n",
    "\n",
    "for i in range(0, shape[0]):\n",
    "    for j in range(0, shape[1]):\n",
    "        X_guess_og[i,j] = np.sum(cropOD[i,j,round(center[i,j,1] - BEC_width_guess[i,j,1]/2) : round(center[i,j,1] + BEC_width_guess[i,j,1]/2) , :], 0) / len(cropOD[i,j,round(center[i,j,1] - BEC_width_guess[i,j,1]/2) : round(center[i,j,1] + BEC_width_guess[i,j,1]/2),0])\n",
    "\n",
    "        # Y_guess_og[i,j] = np.sum(cropOD[i,j, :, round(center[i,j,0] - BEC_width_guess[i,j,0]/2) : round(center[i,j,0] + BEC_width_guess[i,j,0]/2)], 1) / len(cropOD[i,j,0,round(center[i,j,0] - BEC_width_guess[i,j,0]/2) : round(center[i,j,0] + BEC_width_guess[i,j,0]/2)])\n",
    "\n",
    "#[nr images x, nr images y, X / Y, center / sigma]\n",
    "x = np.linspace(0,shape[3],150)\n",
    "y = np.linspace(0,shape[2], 150)\n",
    "\n",
    "popt = np.zeros((shape[0], shape[1], 6))\n",
    "\n",
    "p0 = np.ones((shape[0], shape[1], 6))\n",
    "\n",
    "max = np.zeros((shape[0], shape[1]))\n",
    "\n",
    "for i in range(0, shape[0]):\n",
    "    max[i] = np.ndarray.max(X_guess_og[i],axis=1)\n",
    "\n",
    "# Fitting x without math constr\n",
    "fitmodel = lmfit.Model(density_1d,independent_vars=['x'])\n",
    "\n",
    "result_x = []\n",
    "\n",
    "for i in range(0, shape[0]):\n",
    "    temp_res = []\n",
    "    for j in range(0, shape[1]):\n",
    "        t1 = time.time()\n",
    "        params = lmfit.Parameters()\n",
    "        params.add_many(\n",
    "            ('x0_bec', center[i,j,0], True,0, 200),\n",
    "            ('x0_th',center[i,j,0], True,0, 200),\n",
    "            ('amp_bec', 0.7 * max[i,j], True, 0, 1.3 * np.max(X_guess_og[i,j])),\n",
    "            ('amp_th', 0.3 * max[i,j], True, 0, 1.3 * np.max(X_guess_og[i,j])),\n",
    "            ('deltax', 0, True, 0,50),\n",
    "            ('sigma_bec',BEC_width_guess[i,j,0], True, 0, 50)\n",
    "        )\n",
    "        params.add('sigma_th', popt[i,j,5], min=0, expr=f'sigma_bec + deltax')\n",
    "\n",
    "        t2 = time.time()\n",
    "        res = fitmodel.fit(X_guess_og[i,j], x=x, params=params)\n",
    "        t3 = time.time()\n",
    "        temp_res.append(res)\n",
    "        t4 = time.time()\n",
    "        print(t2 - t1)\n",
    "        print(t3 - t2)\n",
    "        print(t4 - t3)\n",
    "        print(\"\")\n",
    "    result_x.append(temp_res)\n",
    "stop = time.time()\n",
    "\n",
    "print(f'total time: {(stop-start)*1e3} ms')3"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# from opencv import moments\n",
    "start = time.time()\n",
    "shape = np.shape(cropOD)\n",
    "sigma = 0.4\n",
    "blurred = gaussian_filter(cropOD, sigma=sigma)\n",
    "\n",
    "thresh = np.zeros(shape)\n",
    "for i in range(0,shape[0]):\n",
    "    for j in range(0, shape[1]):\n",
    "        thresh[i,j] = np.where(blurred[i,j] < np.max(blurred[i,j])*0.5,0,1)\n",
    "\n",
    "center = calc_cen_bulk(thresh)\n",
    "\n",
    "BEC_width_guess = guess_BEC_width(thresh, center)\n",
    "\n",
    "X_guess_og = np.zeros((shape[0], shape[1], shape[3]))\n",
    "# Y_guess_og = np.zeros((shape[0], shape[1], shape[2]))\n",
    "\n",
    "for i in range(0, shape[0]):\n",
    "    for j in range(0, shape[1]):\n",
    "        X_guess_og[i,j] = np.sum(cropOD[i,j,round(center[i,j,1] - BEC_width_guess[i,j,1]/2) : round(center[i,j,1] + BEC_width_guess[i,j,1]/2) , :], 0) / len(cropOD[i,j,round(center[i,j,1] - BEC_width_guess[i,j,1]/2) : round(center[i,j,1] + BEC_width_guess[i,j,1]/2),0])\n",
    "\n",
    "        # Y_guess_og[i,j] = np.sum(cropOD[i,j, :, round(center[i,j,0] - BEC_width_guess[i,j,0]/2) : round(center[i,j,0] + BEC_width_guess[i,j,0]/2)], 1) / len(cropOD[i,j,0,round(center[i,j,0] - BEC_width_guess[i,j,0]/2) : round(center[i,j,0] + BEC_width_guess[i,j,0]/2)])\n",
    "\n",
    "#[nr images x, nr images y, X / Y, center / sigma]\n",
    "x = np.linspace(0,shape[3],150)\n",
    "y = np.linspace(0,shape[2], 150)\n",
    "\n",
    "popt = np.zeros((shape[0], shape[1], 6))\n",
    "\n",
    "p0 = np.ones((shape[0], shape[1], 6))\n",
    "\n",
    "max = np.zeros((shape[0], shape[1]))\n",
    "\n",
    "for i in range(0, shape[0]):\n",
    "    max[i] = np.ndarray.max(X_guess_og[i],axis=1)\n",
    "\n",
    "# Fitting x without math constr\n",
    "fitmodel = lmfit.Model(density_1d,independent_vars=['x'])\n",
    "\n",
    "result_x = []\n",
    "\n",
    "for i in range(0, shape[0]):\n",
    "    temp_res = []\n",
    "    for j in range(0, shape[1]):\n",
    "        t1 = time.time()\n",
    "        params = lmfit.Parameters()\n",
    "        params.add_many(\n",
    "            ('x0_bec', center[i,j,0], True,0, 200),\n",
    "            ('x0_th',center[i,j,0], True,0, 200),\n",
    "            ('amp_bec', 0.7 * max[i,j], True, 0, 1.3 * np.max(X_guess_og[i,j])),\n",
    "            ('amp_th', 0.3 * max[i,j], True, 0, 1.3 * np.max(X_guess_og[i,j])),\n",
    "            ('sigma_bec',BEC_width_guess[i,j,0] / 1.22 , True, 0, 50),\n",
    "            ('sigma_th', 3 * BEC_width_guess[i,j,0], True, 0, 50)\n",
    "        )\n",
    "\n",
    "        t2 = time.time()\n",
    "        res = fitmodel.fit(X_guess_og[i,j], x=x, params=params)\n",
    "        t3 = time.time()\n",
    "        temp_res.append(res)\n",
    "        t4 = time.time()\n",
    "        # print(t2 - t1)\n",
    "        # print(t3 - t2)\n",
    "        # print(t4 - t3)\n",
    "        # print(\"\")\n",
    "    result_x.append(temp_res)\n",
    "stop = time.time()\n",
    "\n",
    "print(f'total time: {(stop-start)*1e3} ms')"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# get center of thresholded image\n",
    "\n",
    "def calc_thresh(data):\n",
    "    shape = np.shape(data)\n",
    "    thresh = np.zeros(shape)\n",
    "    sigma = 0.4\n",
    "\n",
    "    if len(shape) == 4:\n",
    "        blurred = gaussian_filter(data, sigma=sigma)\n",
    "        for i in range(0,shape[0]):\n",
    "            for j in range(0, shape[1]):\n",
    "                thresh[i,j] = np.where(blurred[i,j] < np.max(blurred[i,j])*0.5,0,1)\n",
    "\n",
    "    elif len(shape) == 3:\n",
    "        blurred = gaussian_filter(data, sigma=sigma)\n",
    "        for i in range(0,shape[0]):\n",
    "                thresh[i] = np.where(blurred[i] < np.max(blurred[i])*0.5,0,1)\n",
    "\n",
    "    else:\n",
    "        print(\"Shape of data is wrong, output is empty\")\n",
    "\n",
    "    return thresh\n",
    "\n",
    "def calc_cen(thresh1):\n",
    "    \"\"\"\n",
    "    returns array: [Y_center,X_center]\n",
    "    \"\"\"\n",
    "    cen = np.zeros(2)\n",
    "    (Y,X) = np.shape(thresh1)\n",
    "\n",
    "\n",
    "    thresh1 = thresh1 /np.sum(thresh1)\n",
    "\n",
    "    # marginal distributions\n",
    "    dx = np.sum(thresh1, 0)\n",
    "    dy = np.sum(thresh1, 1)\n",
    "\n",
    "    # expected values\n",
    "    cen[0] = np.sum(dx * np.arange(X))\n",
    "    cen[1] = np.sum(dy * np.arange(Y))\n",
    "    return cen\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], shape[1], 2))\n",
    "    for i in range(0, shape[0]):\n",
    "        for j in range(0, shape[1]):\n",
    "            cen[i,j] = calc_cen(thresh[i,j])\n",
    "    return cen\n",
    "\n",
    "\n",
    "def gaussian(x, x0, sigma, A):\n",
    "    return A * np.exp(-0.5 * (x-x0)**2 / sigma**2)"
   ],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}