{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "783efa55",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "e1d7df16",
   "metadata": {},
   "source": [
    "# Jupyter notebook for fitting omega scan profiles"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e80f9e9b-cd5a-4d20-967a-8cf73320199c",
   "metadata": {},
   "source": [
    "## import necessary modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bc31f46d-cc4f-4b51-b320-aa967f160908",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.cm as cm\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.cbook as cbook\n",
    "from matplotlib.path import Path\n",
    "from matplotlib.patches import PathPatch\n",
    "import nexusformat.nexus as nx\n",
    "import pandas as pd\n",
    "from PIL import ImageEnhance\n",
    "import os, fnmatch\n",
    "from scipy.integrate import quad\n",
    "import numpy as np\n",
    "from scipy import stats\n",
    "from scipy.optimize import curve_fit\n",
    "from pylab import *\n",
    "import math\n",
    "from scipy.fftpack import rfft, irfft, fftfreq\n",
    "from lmfit.models import PseudoVoigtModel, LinearModel, GaussianModel\n",
    "from lmfit import Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88478905-128f-474d-bdb6-aa1d89252b2e",
   "metadata": {},
   "source": [
    "## set paths\n",
    "\n",
    "below enter the paths to folders or the internal paths within the DAWN processed nxs file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "915a2b2f-6bc5-4ae1-a713-90ac7bf38673",
   "metadata": {},
   "outputs": [],
   "source": [
    "#path to where DAWN output has been saved\n",
    "DAWN_OUT_FOLDER='/dls/science/users/rpy65944/Demo_omegascans/'\n",
    "\n",
    "#path to where raw data has been saved\n",
    "ORIGINAL_DATA_FOLDER='/dls/staging/dls/i07/data/2023/si33555-1/Sample6/'\n",
    "\n",
    "#path from within the DAWN processed nxs file where the sum line profile is saved\n",
    "lineprofilepath='/processed/auxiliary/2-Sum/sum/data'\n",
    "crop_image_path='/processed/result/data'\n",
    "#change working directory to where the SXRDcorrections.py file is saved and import correction functions\n",
    "os.chdir('/dls/science/users/rpy65944/')    \n",
    "import SXRDcorrections as corrs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7698a8dd-6314-4104-8f6a-8afc47ffbde0",
   "metadata": {},
   "source": [
    "## define fitting choices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f4f33b10-ef7b-4536-b8e6-5171ff8da419",
   "metadata": {},
   "outputs": [],
   "source": [
    "pattern = \"*.nxs\"   # this is the file extension of the output from DAWN, will most likely be .nxs\n",
    "sample='sample6'    #name of sample used when creating subfolder to save data to\n",
    "fit_type='pvoigt'   #choose wither 'pvoigt' or 'gaussian', recomment to try pvoigt to start with"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e37fa95-4cdd-4903-869b-dee4645f25bf",
   "metadata": {},
   "source": [
    "## run fitting on all nxs output files in dawn out folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "b17e93b7-f424-477c-a574-7e937bb70711",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/dls/science/users/rpy65944/Demo_omegascans\n",
      "i07-475543_processed_250304_085816.nxs\n",
      "475543\n",
      "81 81\n",
      "1805149.00796 180514.90079600003 i07-475543_processed_250304_085816.nxs  \n",
      " corrections = cI 1.001 ,cA 0.173, cP 1.0307, cL 0.1725, delta 9.949582374, gamma 3.159939125, omega -153.76029166700002\n",
      "-0.9998 0.4999 0.5 1805149.00796 180514.90079600003\n",
      "\n",
      "i07-475544_processed_250304_130927.nxs\n",
      "475544\n",
      "81 81\n",
      "1563794.94016 156379.494016 i07-475544_processed_250304_130927.nxs  \n",
      " corrections = cI 1.002 ,cA 0.1731, cP 1.0306, cL 0.1723, delta 9.946380268, gamma 4.220432675, omega -153.419680556\n",
      "-0.9998 0.4999 0.666 1563794.94016 156379.494016\n",
      "\n",
      "i07-475545_processed_250304_130928.nxs\n",
      "475545\n",
      "81 81\n",
      "1762409.82892 176240.982892 i07-475545_processed_250304_130928.nxs  \n",
      " corrections = cI 1.001 ,cA 0.173, cP 1.0307, cL 0.1725, delta 9.950356853, gamma 3.1574823, omega -100.67338888900001\n",
      "-0.9998 -0.4999 0.5 1762409.82892 176240.982892\n",
      "\n",
      "i07-475546_processed_250304_130928.nxs\n",
      "475546\n",
      "29 29\n",
      "334909.60101 33490.960101 i07-475546_processed_250304_130928.nxs  \n",
      " corrections = cI 1.0001 ,cA 0.1723, cP 1.0306, cL 0.1723, delta 9.922676778131304, gamma 1.3023208105363422, omega 79.00097386099654\n",
      "1.0 0.5 0.2 334909.60101 33490.960101\n",
      "\n",
      "sample6_IAPL\n"
     ]
    }
   ],
   "source": [
    "\n",
    "%cd {DAWN_OUT_FOLDER}\n",
    "listOfFiles = os.listdir('.')\n",
    "listOfFiles.sort()\n",
    "\n",
    "def integrand(x1):\n",
    "    return result.eval_components(x=x1)['{}'.format(fit_type)]\n",
    "\n",
    "mypath=f\"{DAWN_OUT_FOLDER}/{sample}_IAPL/\"\n",
    "if not os.path.isdir(mypath):\n",
    "    os.makedirs(mypath)\n",
    "os.chdir(mypath)\n",
    "f = open(f'{sample}_omegas_datsIAPL.csv', \"w\")\n",
    "f.write('entry\\th\\tk\\tl\\tI\\tdI\\n')\n",
    "for entry in listOfFiles:  \n",
    "    print(entry)\n",
    "    if fnmatch.fnmatch(entry, pattern):\n",
    "        scan_number=entry.split('_')[0].replace('i07-','')\n",
    "        print(scan_number)\n",
    "        dawn_data=nx.nxload(f'{DAWN_OUT_FOLDER}/{entry}')\n",
    "\n",
    "        linedata=dawn_data[f'{lineprofilepath}'].nxdata\n",
    "        original_data=nx.nxload(f\"{ORIGINAL_DATA_FOLDER}/{entry.split('_')[0]}.nxs\")\n",
    "\n",
    "        omega=original_data['entry/instrument']['diff2omega/value_set'].nxdata\n",
    "        gamma=original_data['entry/instrument']['diff2gamma/value'].nxdata\n",
    "        alpha=original_data['entry/instrument']['diff2alpha/value'].nxdata\n",
    "        delta=original_data['entry/instrument']['diff2delta/value'].nxdata\n",
    "\n",
    "        checkangles=[alpha,delta,gamma]\n",
    "        alpha,delta,gamma=[np.mean(ang) if np.atleast_1d(ang).size > 1 else ang for ang in checkangles]\n",
    "        \n",
    "        #calculate corrections using equations from ANA software, in SXRDcorrections.py\n",
    "        corrI=corrs.intcorr(delta,gamma,alpha)\n",
    "        corrA=corrs.areacorr(delta,gamma,alpha)\n",
    "        corrP=corrs.polcorr(delta,gamma)\n",
    "        corrL=corrs.lorcorr(delta,gamma)\n",
    "        d5i=None #if d5i included in file put path here\n",
    "        \n",
    "        #set x and y values for fitting*(corrA)\n",
    "        if d5i != None:\n",
    "            y = 1000*linedata*(corrI)*(corrP)*(corrL)*(corrA)/d5i\n",
    "        else:\n",
    "            y= 1000*linedata*(corrI)*(corrP)*(corrL)*(corrA)\n",
    "        x = omega\n",
    "        print(len(x), len(y))\n",
    "        #set line model and calculate fit to data\n",
    "        line_mod = LinearModel(prefix='line_')\n",
    "   \n",
    "        if fit_type == 'gaussian':\n",
    "            fit_mod= GaussianModel()\n",
    "        elif fit_type =='pvoigt':\n",
    "            fit_mod=PseudoVoigtModel()\n",
    "        mod = line_mod + fit_mod \n",
    "        pars = line_mod.make_params(intercept=y.min(), slope=-700)\n",
    "        pars += fit_mod.guess(y, x=x)\n",
    "        result = mod.fit(y, pars, x=x)\n",
    "        xnew=arange(omega.min(),omega.max(),0.01)\n",
    "        comps = result.eval_components(x=xnew)\n",
    "                #dely = result.eval_uncertainty(sigma=1)\n",
    "\n",
    "        h=round(original_data['entry/instrument']['hkl/h'].nxdata.mean(),4)\n",
    "        k=round(original_data['entry/instrument']['hkl/k'].nxdata.mean(),4)\n",
    "        l=round(original_data['entry/instrument']['hkl/l'].nxdata.mean(),4)\n",
    "\n",
    "\n",
    "\n",
    "        I = round(quad(integrand, omega.min(), omega.max())[0],5)\n",
    "        dI=0.1*I\n",
    "\n",
    "        #write data to csv file\n",
    "        f.write(str(scan_number)+'\\t'+str(h)+'\\t'+str(k)+'\\t'+str(l)+'\\t'+str(I)+'\\t'+str(dI)+'\\n')\n",
    "        \n",
    "        #create and save figure\n",
    "        fig, (ax1, ax2,ax3) = plt.subplots(3, 1, figsize=(15,10),sharex=False,gridspec_kw={'height_ratios': [3, 1,1]})\n",
    "        ax1.plot(x,y)\n",
    "        ax1.plot(xnew,(comps['{}'.format(fit_type)]))\n",
    "        ax1.plot(xnew,(comps['line_']))\n",
    "        ax1.plot(xnew,(comps['{}'.format(fit_type)]+comps['line_']))\n",
    "        ax1.set_title('{}'.format(entry))\n",
    "        if d5i!=None:\n",
    "            ax2.plot(x,d5i,color='black')\n",
    "        else:\n",
    "            ax2.text(0.5,0.5,'no d5i data')\n",
    "\n",
    "        ax3.imshow(dawn_data[f'{crop_image_path}'][np.argmax(linedata)].nxdata)\n",
    "        ax1.axvline(x[np.argmax(linedata)],ls='--')\n",
    "        savefig(f'anaIAPL_{scan_number}.png')\n",
    "        plt.close(fig)\n",
    "\n",
    "        print(I,dI,entry,f\" \\n corrections = cI {round(corrI,4)} ,cA {round(corrA,4)}, cP {round(corrP,4)}, cL {round(corrL,4)}, delta {delta}, gamma {gamma}, omega {omega.mean()}\")\n",
    "        print(str(h)+' '+str(k)+' '+str(l)+' '+str(I)+' '+str(dI)+'\\n')\n",
    "                \n",
    "f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a312f5a-bf9b-427a-a41c-85495148031b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}