giwaxs_plotting.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b3c13da0-ce45-4653-9fb2-5f71078d5a36",
   "metadata": {},
   "source": [
    "# GIWAXS plotting notebook - plotting single images from loaded zarr datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96625ca6-7ec2-4690-bf01-72b422801f76",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "dd667c0e-baba-4a5d-857a-ca8bd5ce1407",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports:\n",
    "import pathlib\n",
    "import json\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LogNorm\n",
    "import xarray as xr\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "# For lmfit later\n",
    "from scipy import optimize, signal\n",
    "from lmfit import models\n",
    "\n",
    "# Choose a colormap:\n",
    "cmap = plt.cm.turbo\n",
    "cmap.set_bad('black')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dffa6de-0360-4fcb-b0bf-f320927837d0",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Define & check paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e63ad6f-1b77-4925-9708-38f4fd8ccefe",
   "metadata": {},
   "outputs": [],
   "source": [
    "def select_attrs(data_arrays_iterable, selected_attrs_dict):\n",
    "    \"\"\"\n",
    "    Selects data arrays whose attributes match the specified values.\n",
    "\n",
    "    Parameters:\n",
    "    data_arrays_iterable: Iterable of xarray.DataArray objects.\n",
    "    selected_attrs_dict: Dictionary where keys are attribute names and \n",
    "                         values are the attributes' desired values.\n",
    "\n",
    "    Returns:\n",
    "    List of xarray.DataArray objects that match the specified attributes.\n",
    "    \"\"\"    \n",
    "    sublist = list(data_arrays_iterable)\n",
    "    \n",
    "    for attr_name, attr_values in selected_attrs_dict.items():\n",
    "        sublist = [da.copy() for da in sublist if da.attrs[attr_name] in attr_values]\n",
    "                \n",
    "    return sublist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8db0fc93-6739-457a-a7fe-ba695bb41716",
   "metadata": {},
   "outputs": [],
   "source": [
    "# I like pathlib for its readability & checkability, it's also necessary for the loadSeries function later on\n",
    "# Replace the paths with the ones relevant to your data, you can use the \".exists()\" method to make sure you defined a path correctly\n",
    "suitePath = pathlib.Path('/Users/andrew/Library/CloudStorage/OneDrive-UCB-O365/research/data_analysis/giwaxs_suite')\n",
    "outPath = suitePath.joinpath('processed_data/xenocs/AX_films_01')\n",
    "zarrsPath = outPath.joinpath('zarrs')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4fdac03-46b0-4814-8fed-6b0b3de72404",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# List zarrs\n",
    "sorted([f.name for f in zarrsPath.iterdir()])  # a way to list just the filenames and not the whole path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41265fdd-51bc-4d12-ba8c-2cc1febef9bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# make quick dictionary to from sample id to sample name for plotting titles:\n",
    "sn = {\n",
    "    'A1s':'A1 Static',\n",
    "    'A1d':'A1 Dynamic',\n",
    "    'A2s':'A2 Static',\n",
    "    'A2d':'A2 Dynamic',\n",
    "    'A3s':'A3 Static',\n",
    "    'A3d':'A3 Dynamic'\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "783090af-9f5f-49e7-b846-84ba5540dd18",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Cartesian plots"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57330ec8-0edf-4856-ba7d-2d6fb191893f",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Raw data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca531215-d669-4eea-8ea8-30c93d835465",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "filename = 'raw.zarr'\n",
    "raw_DS = xr.open_zarr(zarrsPath.joinpath(filename)).compute()\n",
    "raw_DS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2267272e-fc88-462a-8560-a4f851e75ada",
   "metadata": {},
   "outputs": [],
   "source": [
    "for DA in raw_DS.data_vars.values():\n",
    "    cmin = DA.quantile(0.01)\n",
    "    cmax = DA.quantile(0.99)\n",
    "    ax = DA.sel(pix_y=slice(None,900)).plot.imshow(\n",
    "        cmap=cmap, norm=plt.Normalize(cmin,cmax), origin='upper', figsize=(5,4), interpolation='antialiased')\n",
    "    ax.axes.set(title=f'{sn[DA.film]}:\\n Stitched & averaged raw image')\n",
    "    ax.colorbar.set_label('Intensity [arb. units]', rotation=270, labelpad=12)\n",
    "\n",
    "    # savePath = outPath.joinpath('raw_pix_plots_v1')\n",
    "    # savePath.mkdir(exist_ok=True)\n",
    "    # ax.figure.savefig(savePath.joinpath(f'{DA.film}_{DA.film_number}_th{DA.incident_angle}.png'), dpi=150)\n",
    "    \n",
    "    # plt.show()\n",
    "    plt.close('all')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0aeb3817-807a-4b79-be63-1b81884a7b25",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Recip data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36f06119-6393-42a9-ad7d-f2f0fd5371d3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "filename = 'recip.zarr'\n",
    "recip_DS = xr.open_zarr(zarrsPath.joinpath(filename)).compute()\n",
    "recip_DS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c48c13a-44f4-4fdc-8e8f-4665a05ffac1",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 2D reciprocal space images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de3e97bd-7b57-4d22-8af4-851c3a0b7965",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%matplotlib widget\n",
    "plt.close('all')\n",
    "\n",
    "# 2D reciprocal space cartesian plots\n",
    "qxy_min = -1.1\n",
    "qxy_max = 2.1\n",
    "qz_min = 0\n",
    "qz_max = 2.2\n",
    "\n",
    "# Make & save or show plots\n",
    "selected_attrs_dict = {}\n",
    "selected_DAs = select_attrs(recip_DS.data_vars.values(), selected_attrs_dict)\n",
    "        \n",
    "for DA in tqdm(selected_DAs):\n",
    "    # Slice data for selected q ranges (will need to rename q_xy if dimensions are differently named)\n",
    "    sliced_DA = DA.sel(q_xy=slice(qxy_min, qxy_max), q_z=slice(qz_min, qz_max))\n",
    "\n",
    "    # Feel free to mess around with the quantile selection...\n",
    "    cmin = sliced_DA.quantile(0.3)\n",
    "    cmax = sliced_DA.quantile(0.994) \n",
    "    \n",
    "    # Plot\n",
    "    ax = sliced_DA.plot.imshow(cmap=cmap, norm=plt.Normalize(cmin, cmax), interpolation='antialiased', figsize=(5.5,3.3))\n",
    "    ax.colorbar.set_label('Intensity [arb. units]', rotation=270, labelpad=15)\n",
    "    ax.axes.set(aspect='equal', title=f'Cartesian Plot: {sn[DA.film]}',\n",
    "                xlabel='q$_{xy}$ [Å$^{-1}$]', ylabel='q$_z$ [Å$^{-1}$]')\n",
    "    ax.axes.xaxis.set_minor_locator(plt.MultipleLocator(0.1))\n",
    "    ax.axes.yaxis.set_minor_locator(plt.MultipleLocator(0.1))\n",
    "    ax.figure.set(tight_layout=True, dpi=130)\n",
    "    \n",
    "    # # Save\n",
    "    # savePath = outPath.joinpath('recip_plots_v1')\n",
    "    # savePath.mkdir(exist_ok=True)\n",
    "    # ax.figure.savefig(savePath.joinpath(f'{DA.film}_{DA.film_number}_qxy{qxy_min}to{qxy_max}_qz{qz_min}to{qz_max}_th{DA.incident_angle}.png'), dpi=150)\n",
    "\n",
    "    plt.show()  # or show\n",
    "    # plt.close('all')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9fd26a5-3875-47bb-a594-b1af830fa824",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Polar plots"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc1b333a-53cc-4ace-8769-f61a5a91f60c",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fbc28bc-15c1-4a00-8c65-cfcaa32e35d1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "filename = 'caked.zarr'\n",
    "raw_DS = xr.open_zarr(zarrsPath.joinpath(filename)).compute()\n",
    "raw_DS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ec68108-69d5-4e80-8238-7f278a7b3d69",
   "metadata": {},
   "source": [
    "### Apply sin(chi), thickness, and other scalar corrections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "392d0fa7-6f28-4f31-b219-a735ed59de90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# +/- 10 nm error\n",
    "thickness_error = 10\n",
    "thicknesses = {\n",
    "    'A1s':40,\n",
    "    'A1d':40,\n",
    "    'A2s':55,\n",
    "    'A2d':55,\n",
    "    'A3s':55,\n",
    "    'A3d':55\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e4ea1de-89fc-4c16-9002-46e8a3b05bf4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Apply a sin chi correction\n",
    "# Also add in thickness and/or exposure time corrections if applicable\n",
    "\n",
    "sin_chi_DA = np.sin(np.radians(np.abs(raw_DS.chi)))\n",
    "\n",
    "corr_DS = raw_DS.copy()\n",
    "# corr_DS = corr_DS * sin_chi_DA  # This works mathematically, but does not preserve attributes\n",
    "for i, var in enumerate(corr_DS.data_vars):\n",
    "    corrected = corr_DS[var] * sin_chi_DA\n",
    "    corrected = corrected / thicknesses[var]  # how you could implement thickness correction with an iterable, dict would probably be better\n",
    "    # corrected = corrected / (thicknesses[var] - thickness_error)  # how you could implement thickness correction with an iterable, dict would probably be better\n",
    "    # corrected = corrected / (thicknesses[var] + thickness_error)  # how you could implement thickness correction with an iterable, dict would probably be better\n",
    "\n",
    "    corr_DS[var].values = corrected.values\n",
    "    \n",
    "corr_DS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57eef254-2427-4373-9cf6-ac17e09ee7d3",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 2D caked images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d662afe1-739a-4c98-a9a7-0866b85df471",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### raw(chi) intensity polar plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ba48bbe-2db8-45db-b8f5-f4ea1b0d6542",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Polar plots, for raw(chi) intensities\n",
    "DS = raw_DS.copy()\n",
    "\n",
    "# Set chi range: Full range\n",
    "chi_min = 0\n",
    "chi_max = 90\n",
    "q_min = 0.1\n",
    "q_max = 2.2\n",
    "\n",
    "# Select attribute\n",
    "selected_attrs_dict = {}\n",
    "selected_DAs = select_attrs(DS.data_vars.values(), selected_attrs_dict)\n",
    "for DA in tqdm(selected_DAs):\n",
    "    # Slice dataarray to select plotting region \n",
    "    sliced_DA = DA.sel(chi=slice(chi_min,chi_max), qr=slice(q_min, q_max))\n",
    "    \n",
    "    # Set color limits\n",
    "    cmin = sliced_DA.quantile(0.01)\n",
    "    cmax = sliced_DA.quantile(0.995)\n",
    "    \n",
    "    # Plot sliced dataarray\n",
    "    ax = sliced_DA.plot.imshow(cmap=cmap, norm=plt.Normalize(cmin, cmax), figsize=(5,4), interpolation='antialiased')  # plot, optional parameter interpolation='antialiased' for image smoothing\n",
    "\n",
    "    # ax.axes.set(title=f'Polar plot: {DA.material} {DA.solvent} {DA.rpm}, $\\\\alpha_i$ = {float(DA.incident_angle[2:])}°, raw intensity')\n",
    "    ax.axes.set(title=f'Polar plot: {sn[DA.film]}, $\\\\alpha_i$ = {DA.incident_angle}°, raw intensity')\n",
    "\n",
    "    ax.colorbar.set_label('Raw intensity [arb. units]', rotation=270, labelpad=15)  # set colorbar label & parameters \n",
    "    ax.axes.set(xlabel='q$_r$ [Å$^{-1}$]', ylabel='$\\\\chi$ [°]')  # set title, axis labels, misc\n",
    "    ax.figure.set(tight_layout=True, dpi=130)  # Adjust figure dpi & plotting style\n",
    "    \n",
    "    \n",
    "    # # Save\n",
    "    # savePath = outPath.joinpath('raw_caked_plots_v1')\n",
    "    # savePath.mkdir(exist_ok=True)\n",
    "    # ax.figure.savefig(savePath.joinpath(f'{DA.film}_{DA.film_number}_qr{q_min}to{q_max}_{DA.incident_angle}.png'), dpi=150)\n",
    "\n",
    "    plt.show()  # Comment to mute plotting output\n",
    "    plt.close('all')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb2b88a0-1c2d-424b-a05c-bf78015285e7",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### sin(chi) intensity polar plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97b0b8ba-8464-4a94-8a5c-ff5ef645b147",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd66083e-9491-48f3-b691-0404de41a4df",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.close('all')\n",
    "# Polar plots, for sin(chi) intensities\n",
    "DS = corr_DS.copy()\n",
    "\n",
    "# Set chi range: Full range\n",
    "chi_min = 0\n",
    "chi_max = 90\n",
    "q_min = 0.1\n",
    "q_max = 2.2\n",
    "\n",
    "# Select attribute\n",
    "# selected_attrs_dict = {'film': ['A3d']}\n",
    "selected_attrs_dict = {}\n",
    "selected_DAs = select_attrs(DS.data_vars.values(), selected_attrs_dict)\n",
    "for DA in tqdm(selected_DAs):\n",
    "    # Slice dataarray to select plotting region \n",
    "    sliced_DA = DA.sel(chi=slice(chi_min,chi_max), qr=slice(q_min, q_max))\n",
    "    \n",
    "    # Set color limits\n",
    "    cmin = sliced_DA.quantile(0.01)\n",
    "    cmax = sliced_DA.quantile(0.995)  \n",
    "    \n",
    "    # Plot sliced dataarray\n",
    "    ax = sliced_DA.plot.imshow(origin='upper', cmap=cmap, norm=plt.Normalize(cmin, cmax), figsize=(5,4), interpolation='antialiased')  # plot, optional parameter interpolation='antialiased' for image smoothing\n",
    "    ax.axes.set(title=f'Polar Plot: {sn[DA.film]}, sin($\\\\chi$) corr.')\n",
    "    ax.colorbar.set_label('Intensity * sin($\\\\chi$) [arb. units]', rotation=270, labelpad=15)  # set colorbar label & parameters \n",
    "    ax.axes.set(xlabel='q$_r$ [Å$^{-1}$]', ylabel='$\\\\chi$ [°]')  # set title, axis labels, misc\n",
    "    ax.figure.set(tight_layout=True, dpi=130)  # Adjust figure dpi & plotting style\n",
    "        \n",
    "    # # Save\n",
    "    # savePath = outPath.joinpath('corr_caked_plots_v1')\n",
    "    # savePath.mkdir(exist_ok=True)\n",
    "    # ax.figure.savefig(savePath.joinpath(f'{DA.film}_{DA.film_number}_qr{q_min}to{q_max}_{DA.incident_angle}.png'), \n",
    "    #                   dpi=150)\n",
    "    # # ax.figure.savefig(savePath.joinpath(f'pipi_{DA.film}_{DA.film_number}_qr{q_min}to{q_max}_{DA.incident_angle}.png'), \n",
    "    # #                   dpi=150)\n",
    "    \n",
    "    plt.show()\n",
    "    \n",
    "    # plt.close('all')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b47c5689-4a4e-4bb0-8c4e-48cbf0a8d252",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### 1D linecuts along chi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c61bdaab-f028-458e-9c2e-e57766111885",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Polar plots, for sin(chi) intensities\n",
    "DS = corr_DS.copy()\n",
    "\n",
    "# Set chi bounds & bins, q bounds, etc.\n",
    "# chi_min, chi_width, chi_bins = [14, 9.25, 8]  # for full chi cut\n",
    "chi_min, chi_width, chi_bins = [14, 4, 18]  # for full chi cut, precise wedges\n",
    "# chi_min, chi_width, chi_bins = [14, 4, 6]  # for π-π chi cut\n",
    "chi_max = chi_min + (chi_width * chi_bins)\n",
    "colors = plt.cm.viridis_r(np.linspace(0.15,1,chi_bins))\n",
    "\n",
    "# q_min = 0.1  # for full qr cut\n",
    "q_min = 0.9  # for π-π qr cut\n",
    "q_max = 2.07\n",
    "\n",
    "# Select attribute\n",
    "# selected_attrs_dict = {'film': ['A3s']}\n",
    "selected_attrs_dict = {}\n",
    "selected_DAs = select_attrs(DS.data_vars.values(), selected_attrs_dict)\n",
    "\n",
    "for DA in tqdm(selected_DAs):\n",
    "    sliced_DA = DA.sel(chi=slice(chi_min,chi_max), qr=slice(q_min,q_max))\n",
    "    binned_DA = sliced_DA.groupby_bins('chi', chi_bins).sum('chi')\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(5.5,3), dpi=150, tight_layout=True)\n",
    "    for i, chi_bin in enumerate(binned_DA.chi_bins.data[::-1]):\n",
    "        binned_DA.sel(chi_bins=chi_bin).plot.line(ax=ax, color=colors[i], label=chi_bin)\n",
    "\n",
    "    ax.set(title=f'{sn[DA.film]}: $\\\\chi$-binned-summed linecuts:\\n'+ \n",
    "                 f'{chi_bins}, {chi_width}° $\\\\chi$ bins: from {int(chi_min)} to {int(chi_max)}°', \n",
    "           xlabel='$q_r$ $[Å^{-1}]$', \n",
    "           ylabel='$sin(\\\\chi) * Intensity$ [arb. units]')\n",
    "    ax.axes.xaxis.set_minor_locator(plt.MultipleLocator(0.1))\n",
    "    # ax.legend(title='Chi Bins')\n",
    "\n",
    "    # Create a colormap and normalizer to match the plotted data\n",
    "    sm = plt.cm.ScalarMappable(cmap=plt.cm.viridis, norm=plt.Normalize(vmin=chi_min, vmax=chi_max))\n",
    "    sm.set_array([])  # You need to set the array for the scalar mappable\n",
    "\n",
    "    # Create the colorbar\n",
    "    cbar = plt.colorbar(sm, ax=ax, orientation='vertical')\n",
    "    cbar.set_label('Azimuthal Angle [°]', rotation=270, labelpad=15)\n",
    "    \n",
    "    # Define the ticks for the colorbar based on the chi bins\n",
    "    chi_values = np.linspace(chi_min, chi_max, int(chi_bins/2), endpoint=True)\n",
    "    cbar.set_ticks(chi_values)\n",
    "    cbar.set_ticklabels([f\"{value:.1f}\" for value in chi_values])\n",
    "\n",
    "    # # Save\n",
    "    # outPath.joinpath('chi-binned_linecuts_v1').mkdir(exist_ok=True)\n",
    "    # # savePath = outPath.joinpath('chi-binned_linecuts_v1', f'full_chiWidth-{chi_width}_chiBins-{chi_bins}_chiRange{chi_min}-{chi_max}')\n",
    "    # savePath = outPath.joinpath('chi-binned_linecuts_v1', f'pipi_chiWidth-{chi_width}_chiBins-{chi_bins}_chiRange{chi_min}-{chi_max}')\n",
    "    # savePath.mkdir(exist_ok=True)\n",
    "    # ax.figure.savefig(savePath.joinpath(f'{DA.film}_{DA.film_number}_qr{q_min}to{q_max}_{DA.incident_angle}.png'), \n",
    "    #                   dpi=150)\n",
    "    \n",
    "plt.show()\n",
    "plt.close('all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c9ccaee-abde-4247-a135-19d19e42d61e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# %matplotlib widget"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a62b90b-919c-4dfe-91fe-4408eaf40ab2",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 1D line fitting*\n",
    "Use lmfit to perform the linefits"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a32d682-3aca-47d9-a062-f73f9a3f72aa",
   "metadata": {},
   "source": [
    "#### π-π lmfit functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06a8fdc7-3746-4e5a-b652-389126325f57",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def pipi_lmfit(sliced_DA):\n",
    "    \"\"\"\n",
    "    Function currently utilizes global variables, define in notebook!\n",
    "    \n",
    "    Inputs:\n",
    "    - sliced_DA: 1D data for a selected chi bin, already meaned over chi\n",
    "    \"\"\"    \n",
    "    point_y = float(sliced_DA.sel(qr=slice(0.9,1.1)).mean('qr'))\n",
    "    point_y = point_y - point_y*(0.1)\n",
    "\n",
    "    x = sliced_DA.qr.data\n",
    "    y = sliced_DA.data\n",
    "\n",
    "    # Define all models to include in fitting\n",
    "    bkg_mod = models.LinearModel(prefix='bkg_')\n",
    "    pars = bkg_mod.make_params(intercept=point_y, slope=0)\n",
    "    pars['bkg_intercept'].set(vary=False)\n",
    "    pars['bkg_slope'].set(vary=False)\n",
    "\n",
    "    pipi_mod = models.PseudoVoigtModel(prefix='pipi_')  # pi-pi peak\n",
    "    pars += pipi_mod.guess(y, x, center=1.63)\n",
    "    pars['pipi_amplitude'].set(min=0)\n",
    "    pars['pipi_center'].set(min=1.4, max=1.78)\n",
    "\n",
    "    # Combine into full model\n",
    "    mod = bkg_mod + pipi_mod \n",
    "\n",
    "    # Run fit and store all info in a ModelResult object\n",
    "    out = mod.fit(y, pars, x=x)\n",
    "    return out\n",
    "\n",
    "def pipi_fit(sliced_DA):   \n",
    "    # Run lmfit\n",
    "    out = pipi_lmfit(sliced_DA)\n",
    "    FWHM = np.round(float(out.params['pipi_fwhm']), 2)\n",
    "    Lc = np.round((2*np.pi*0.9)/float(out.params['pipi_fwhm']), 2)\n",
    "    center = np.round(float(out.params['pipi_center']), 2)\n",
    "    dspacing = np.round((2*np.pi)/float(out.params['pipi_center']), 2)\n",
    "    \n",
    "    amplitude = np.round(float(out.params['pipi_amplitude']), 3)\n",
    "\n",
    "    chi_range = f'{str(round(sliced_DA.chi_bin.left))}-{str(round(sliced_DA.chi_bin.right))}'\n",
    "    \n",
    "    # rsquared = np.round(out.rsquared, 3)\n",
    "    redchi = np.round(out.redchi, 5)\n",
    "\n",
    "    AA1 = '$\\\\AA^{-1}$'\n",
    "    \n",
    "    # Plot\n",
    "    q = sliced_DA.qr.data\n",
    "    I = sliced_DA.data\n",
    "    \n",
    "    fig, ax = plt.subplots()\n",
    "    fig.set(size_inches=(6.5,3.5), dpi=120, tight_layout=True)\n",
    "       \n",
    "    ax.plot(q, I, label='data', linewidth=2.5)\n",
    "    ax.plot(q, out.best_fit, '--', label='full_fit')\n",
    "    for key in out.eval_components():\n",
    "        ax.plot(q, out.eval_components()[key], label=f'{key}')\n",
    "    ax.set(xlabel=f'Q [{AA1}]', ylabel='Intensity [arb. units]', yscale='linear')\n",
    "    ax.set_title(\n",
    "        f'π-π peak fit: {sn[sliced_DA.film]}, {chi_range}° chi; peak area = {amplitude}\\n' + \n",
    "        f'center = {center}, FWHM = {FWHM} {AA1} (d-spacing = {dspacing}, $L_c$ = {Lc} $\\\\AA$)', \n",
    "        x=0.6)\n",
    "    ax.legend(title='$\\\\chi_{red}$ '+f'= {redchi}', loc='upper left', bbox_to_anchor=(1,1))\n",
    "    ax.grid(visible=True, axis='x')\n",
    "    ax.xaxis.set_minor_locator(plt.MultipleLocator(0.1))\n",
    "    \n",
    "    return out, fig, ax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16a77e44-f0ed-4ca2-85f6-8492814701e3",
   "metadata": {},
   "source": [
    "#### Run π-π fits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18d891a2-fe25-45e5-b7e6-a9262773500f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Peak fitting for each chi bin\n",
    "DS = corr_DS.copy()\n",
    "\n",
    "# Set chi bounds & bins, q bounds, etc.\n",
    "# chi_min, chi_width, chi_bins = [14, 9.25, 8]  # for full chi cut\n",
    "# chi_min, chi_width, chi_bins = [14, 4, 18]  # for full chi cut, precise wedges\n",
    "# chi_min, chi_width, chi_bins = [14, 4, 6]  # for π-π chi cut -> 38\n",
    "chi_min, chi_width, chi_bins = [14, 4, 7]  # for π-π chi cut -> 42\n",
    "chi_max = chi_min + (chi_width * chi_bins)\n",
    "colors = plt.cm.viridis_r(np.linspace(0.15,1,chi_bins))\n",
    "\n",
    "# q_min = 0.1  # for full qr cut\n",
    "q_min = 0.9  # for π-π qr cut\n",
    "q_max = 2.07\n",
    "\n",
    "# Select attributes\n",
    "# selected_attrs_dict = {'film': ['A3s']}\n",
    "selected_attrs_dict = {}\n",
    "selected_DAs = select_attrs(DS.data_vars.values(), selected_attrs_dict)\n",
    "\n",
    "# For each selected DA, fit the π-π peak\n",
    "all_outs = {}\n",
    "all_params = {}\n",
    "for DA in tqdm(selected_DAs):\n",
    "    sliced_DA = DA.sel(chi=slice(chi_min,chi_max), qr=slice(q_min,q_max))\n",
    "    binned_DA = sliced_DA.groupby_bins('chi', chi_bins).sum('chi')\n",
    "    outs = {}\n",
    "    params = {}\n",
    "    for i, chi_bin in enumerate(binned_DA.chi_bins.data[::-1]):\n",
    "        key = f'{str(round(chi_bin.left))}-{str(round(chi_bin.right))}'\n",
    "        sel_DA = binned_DA.sel(chi_bins=chi_bin)\n",
    "        sel_DA.attrs['chi_bin'] = chi_bin\n",
    "        outs[key], fig, ax = pipi_fit(sel_DA)\n",
    "        params[key] = outs[key].params.valuesdict()\n",
    "\n",
    "        # # Save\n",
    "        # outPath.joinpath('pi-pi_linefits_v3').mkdir(exist_ok=True)\n",
    "        # savePath = outPath.joinpath('pi-pi_linefits_v3', f'chiWidth-{chi_width}_chiBins-{chi_bins}_chiRange{chi_min}-{chi_max}')\n",
    "        # savePath.mkdir(exist_ok=True)\n",
    "        # fig.savefig(savePath.joinpath(f'{DA.film}_chi{key}_pi-pi_fit.png'), dpi=150)\n",
    "\n",
    "        # plt.show()\n",
    "        plt.close('all')\n",
    "        \n",
    "    all_outs[DA.film] = outs\n",
    "    all_params[DA.film] = params\n",
    "\n",
    "# # Save model_params.json in same save folder\n",
    "# json_data = json.dumps(all_params)\n",
    "# with open(str(savePath.joinpath('model_params.json')), 'w') as f:\n",
    "#     f.write(json_data)\n",
    "\n",
    "# # Create full fit reports folder and save entire report .txt files inside:\n",
    "# savePath.joinpath('fit_reports').mkdir(exist_ok=True)\n",
    "# for film, outs in all_outs.items():\n",
    "#     for chi_bin, out in outs.items():\n",
    "#         with savePath.joinpath('fit_reports', f'{film}_{chi_bin}_fit_result.txt').open(mode='w') as f:\n",
    "#             f.write(out.fit_report())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea41cc9c-7322-4dde-9836-766c1fb81034",
   "metadata": {},
   "source": [
    "#### Plot fit results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68542d5e-a2d3-4abf-afdc-7b28323b0e0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load params from json file\n",
    "# chi_width = 4\n",
    "# chi_bins = 6\n",
    "# chi_range = '14-38'\n",
    "# chi_width = 4\n",
    "# chi_bins = 7\n",
    "# chi_range = '14-42'\n",
    "chi_width = 4\n",
    "chi_bins = 18\n",
    "chi_range = '14-86'\n",
    "\n",
    "with open(str(outPath.joinpath('pi-pi_linefits_v3', f'chiWidth-{chi_width}_chiBins-{chi_bins}_chiRange{chi_range}', 'model_params.json')), 'r') as f:\n",
    "    all_params = json.load(f)\n",
    "\n",
    "len(all_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26bf1a5c-5ff1-410d-9cba-44d9e6714151",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot pseudovoigt fractions, d-spacings, scherrer coherence lengths, and normalized peak areas vs chi\n",
    "\n",
    "# Initialize figures to be populated later\n",
    "fig_psvoigt, ax_psvoigt = plt.subplots(figsize=((5.5,2.5)), dpi=150, tight_layout=True)\n",
    "fig_dspacing, ax_dspacing = plt.subplots(figsize=((5.5,2.5)), dpi=150, tight_layout=True)\n",
    "fig_ccl, ax_ccl = plt.subplots(figsize=((5.5,2.5)), dpi=150, tight_layout=True)\n",
    "fig_pole, ax_pole = plt.subplots(figsize=((5.5,2.5)), dpi=150, tight_layout=True)\n",
    "\n",
    "# Define custom colors for each film, here I've chosen just two shades along the same sequention colorbar for each sample group\n",
    "colors_dict = {\n",
    "    'A1s':plt.cm.Blues(np.linspace(0.6,0.8,2))[0],\n",
    "    'A1d':plt.cm.Blues(np.linspace(0.6,0.8,2))[1],\n",
    "    'A2s':plt.cm.Greens(np.linspace(0.6,0.8,2))[0],\n",
    "    'A2d':plt.cm.Greens(np.linspace(0.6,0.8,2))[1],\n",
    "    'A3s':plt.cm.Oranges(np.linspace(0.6,0.8,2))[0],\n",
    "    'A3d':plt.cm.Oranges(np.linspace(0.6,0.8,2))[1]\n",
    "}\n",
    "\n",
    "### Extracting info from all_params & populating ___ vs chi plots:\n",
    "summed_peak_areas = {}\n",
    "avg_dspacings = {}\n",
    "avg_coherence_lengths = {}\n",
    "for film in all_params:\n",
    "    peak_areas = []\n",
    "    pvoigt_fracs = []\n",
    "    peak_centers = []\n",
    "    peak_fwhms = []\n",
    "    chis = []\n",
    "    for chi_bin, params in all_params[film].items():\n",
    "        # Get starting value of chi bin\n",
    "        chi = int(chi_bin.split('-')[0])\n",
    "        chis.append(chi)\n",
    "        # Get peak areas\n",
    "        peak_area = params['pipi_amplitude']\n",
    "        peak_areas.append(peak_area)\n",
    "        # Get pseudovoigt fraction\n",
    "        pvoigt_frac = params['pipi_fraction']\n",
    "        pvoigt_fracs.append(pvoigt_frac)\n",
    "        # Get peak centers\n",
    "        peak_center = params['pipi_center']\n",
    "        peak_centers.append(peak_center)\n",
    "        # Get peak fwhms\n",
    "        peak_fwhm = params['pipi_fwhm']\n",
    "        peak_fwhms.append(peak_fwhm)\n",
    "        \n",
    "    # Get normalized peak areas (divide by sum of all peak areas; adds to 1) and write area sum to dict\n",
    "    peak_areas = np.array(peak_areas)\n",
    "    summed_peak_areas[film] = np.sum(peak_areas)  # Write sum of peak area out, this corresponds to relative extent of crystallinity if thickness normalized\n",
    "    normed_peak_areas = peak_areas / np.sum(peak_areas)\n",
    "\n",
    "    # Calculate d-spacings from peak centers, calculate peak-area-weighted average of d-spacing\n",
    "    dspacings = (2*np.pi) / np.array(peak_centers) \n",
    "    avg_dspacing = np.round(np.sum(dspacings*normed_peak_areas), 2)\n",
    "    avg_dspacings[film] = avg_dspacing  \n",
    "\n",
    "    # Calculate coherence length from peak fwhm, calculate peak-area-weighted average of coherence length    \n",
    "    coherence_lengths = (2*np.pi*0.9) / np.array(peak_fwhms) \n",
    "    avg_coherence_length = np.round(np.sum(coherence_lengths*normed_peak_areas), 2)\n",
    "    avg_coherence_lengths[film] = avg_coherence_length\n",
    "\n",
    "    ### Plotting\n",
    "    # Plot pseudovoigt fraction vs chi for each film\n",
    "    ax_psvoigt.plot(chis, pvoigt_fracs, label=film, color=colors_dict[film], marker='o')\n",
    "    ax_psvoigt.set(title='Pseudo-voigt lorentzian-to-gaussian ratio versus $\\\\chi$',\n",
    "                   ylabel='Lorentzian peak fraction',\n",
    "                   xlabel='Binned $\\\\chi$ value [°]')\n",
    "    ax_psvoigt.legend(title='Film', loc='upper left', bbox_to_anchor=(1,1))\n",
    "\n",
    "    # Plot d-spacing vs chi for each film\n",
    "    ax_dspacing.plot(chis, dspacings, label=film, color=colors_dict[film], marker='o')\n",
    "    ax_dspacing.set(title='D-spacing versus $\\\\chi$',\n",
    "                    ylabel='d-spacing [Å]',\n",
    "                    xlabel='Binned $\\\\chi$ value [°]')\n",
    "    ax_dspacing.legend(title='Film', loc='upper left', bbox_to_anchor=(1,1))\n",
    "\n",
    "    # Plot coherence length vs chi for each film\n",
    "    ax_ccl.plot(chis, coherence_lengths, label=film, color=colors_dict[film], marker='o')\n",
    "    ax_ccl.set(title='Crystalline coherence length (CCL) versus $\\\\chi$',\n",
    "               ylabel='CCL [Å]',\n",
    "               xlabel='Binned $\\\\chi$ value [°]')\n",
    "    ax_ccl.legend(title='Film', loc='upper left', bbox_to_anchor=(1,1))\n",
    "    \n",
    "    # Plot 'pole figure' for each film\n",
    "    ax_pole.plot(chis, normed_peak_areas, label=film, color=colors_dict[film], marker='o')\n",
    "    ax_pole.set(title='Normalized π-π peak areas versus $\\\\chi$:\\n'+\n",
    "                 'peak area divided by sum of all peak areas ',\n",
    "                ylabel='Peak area [arb. units]',\n",
    "                xlabel='Binned $\\\\chi$ value [°]')\n",
    "    ax_pole.legend(title='Film', loc='upper left', bbox_to_anchor=(1,1.1))    \n",
    "\n",
    "    # Save\n",
    "    chi_width = chis[0]-chis[1]\n",
    "    chi_bins = len(chis)\n",
    "    chi_min = chis[-1]\n",
    "    chi_max = chis[0]+chi_width\n",
    "\n",
    "    parent_folder = 'pi-pi_fit_results_plots_v2'\n",
    "    outPath.joinpath(parent_folder).mkdir(exist_ok=True)\n",
    "    savePath = outPath.joinpath(parent_folder, f'chiWidth-{chi_width}_chiBins-{chi_bins}_chiRange{chi_min}-{chi_max}')\n",
    "    savePath.mkdir(exist_ok=True)\n",
    "    fig_psvoigt.savefig(savePath.joinpath(f'pseudovoigt_ratios_pi-pi_fit.png'), dpi=150)\n",
    "    fig_dspacing.savefig(savePath.joinpath(f'dspacings_pi-pi_fit.png'), dpi=150)\n",
    "    fig_ccl.savefig(savePath.joinpath(f'coherence_lengths_pi-pi_fit.png'), dpi=150)\n",
    "    fig_pole.savefig(savePath.joinpath(f'peak_areas_pi-pi_fit.png'), dpi=150)\n",
    "    \n",
    "plt.show()\n",
    "plt.close('all')  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11309ced-2bc7-4d4e-8cfe-b8633282c2f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot average d-spacing, scherrer coherence length, summed peak area vs film\n",
    "fig_dspacing, ax_dspacing = plt.subplots(figsize=((5.5,2.5)), dpi=150, tight_layout=True)\n",
    "fig_ccl, ax_ccl = plt.subplots(figsize=((5.5,2.5)), dpi=150, tight_layout=True)\n",
    "fig_area, ax_area = plt.subplots(figsize=((5.5,2.5)), dpi=150, tight_layout=True)\n",
    "\n",
    "ax_dspacing.scatter(list(avg_dspacings.keys()), list(avg_dspacings.values()))\n",
    "ax_ccl.scatter(list(avg_coherence_lengths.keys()), list(avg_coherence_lengths.values()))\n",
    "ax_area.scatter(list(summed_peak_areas.keys()), list(summed_peak_areas.values()))\n",
    "\n",
    "ax_dspacing.set(title='Peak-area-weighted average d-spacings', xlabel='Film', ylabel='D-spacing [Å]')\n",
    "ax_ccl.set(title='Peak-area-weighted average coherence lengths', xlabel='Film', ylabel='CCL [Å]')\n",
    "ax_area.set(title='Summed peak areas', xlabel='Film', ylabel='Intensity [arb. units]')\n",
    "\n",
    "plt.show()\n",
    "plt.close('all')\n",
    "display(avg_dspacings)\n",
    "display(avg_coherence_lengths)\n",
    "display(summed_peak_areas)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "312c3464-7a2d-4371-8169-97b5098caa26",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Pole Figures"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cfbb2a1-f0c6-4d9e-b38c-83c19a8ea09a",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Pi-pi "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e75594d9-cfc1-422b-b884-c2a565ad0996",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "### Set limits\n",
    "q_min = 1.68\n",
    "q_max = 1.95\n",
    "chi_min = 12\n",
    "chi_max = 80  # to ~45 least until I mask out the silicon(?) background scattering that interferes\n",
    "\n",
    "\n",
    "# selected_attrs_dict = {'material':['PM6'], 'incident_angle':['th0.120']}\n",
    "# # selected_attrs_dict = {}\n",
    "# selected_DAs = select_attrs(folded_corr_DAs, selected_attrs_dict)\n",
    "good_scans_list = ['sam1_10s', 'sam2_10s-2', 'sam3_30s', 'sam4_30s', 'sam5_30s', 'sam6_2s']\n",
    "selected_DAs = []\n",
    "for DA in DS.data_vars.values():\n",
    "    sample_name = DA.attrs['sample_id']\n",
    "    exp_time = DA.attrs['exposure_time']\n",
    "    if f'{sample_name}_{exp_time}' in good_scans_list:\n",
    "        selected_DAs.append(DA)\n",
    "        \n",
    "peak_areas_dict = {}    \n",
    "for DA in tqdm(selected_DAs):\n",
    "    chis = DA.chi.sel(chi=slice(chi_min,chi_max)).data\n",
    "    \n",
    "    peak_areas = np.array([])\n",
    "    peak_centers = np.array([])\n",
    "    \n",
    "    for chi in chis:\n",
    "        sliced_DA = DA.sel(qr=slice(q_min-0.1, q_max+0.1), chi=float(chi))\n",
    "\n",
    "        points_x = [q_min, q_max]\n",
    "        points_y = [float(sliced_DA.sel(qr=slice(points_x[0]-0.05, points_x[0]+0.05)).mean('qr')), \n",
    "                    float(sliced_DA.sel(qr=slice(points_x[1]-0.05, points_x[1]+0.05)).mean('qr'))]\n",
    "        m = (points_y[1]-points_y[0])/(points_x[1]-points_x[0])\n",
    "        b = points_y[1] - (m*points_x[1])\n",
    "        y_fit = np.polyval([m, b], sliced_DA.qr)\n",
    "        \n",
    "        sub_sliced_DA = sliced_DA-y_fit\n",
    "\n",
    "        peak_area = sub_sliced_DA.sel(qr=slice(q_min,q_max)).integrate('qr')\n",
    "        peak_area = np.nan if peak_area<0 else peak_area  # remove negatives as they are nonphysical\n",
    "        \n",
    "        peak_areas = np.append(peak_areas, float(peak_area))\n",
    "                        \n",
    "#         # Plot linecuts for each chi\n",
    "#         fig, ax = plt.subplots()\n",
    "#         sliced_DA.plot(ax=ax)\n",
    "#         ax.plot(points_x, points_y)\n",
    "#         plt.show()\n",
    "\n",
    "#         plt.close('all')\n",
    "\n",
    "    normed_peak_areas = peak_areas / np.nansum(peak_areas)\n",
    "    data = np.vstack((chis, normed_peak_areas)).T\n",
    "    peak_areas_dict[DA.attrs['sample_id']+'_'+DA.attrs['exposure_time']] = data\n",
    "\n",
    "    # Individual measurement pole figure plot\n",
    "    fig, ax = plt.subplots(figsize=(5,3.3))\n",
    "    fig.set(dpi=120)\n",
    "    \n",
    "    ax.plot(chis, normed_peak_areas)\n",
    "    fig.suptitle(f'{DA.sample_id} {DA.exposure_time} Pole Figure; $\\\\alpha_i = $ {DA.incident_angle[2:-1]}°', x=0.54)\n",
    "    ax.set(title=f'{q_min} to {q_max} 1/Å Q Bounds; {chi_min}° to {chi_max}° $\\chi$')\n",
    "    ax.set(xlabel='$\\chi$ [°]', ylabel='Integrated Pi-Pi Peak Area [arb. units]')\n",
    "    plt.subplots_adjust(top=0.85, bottom=0.2, left=0.2)\n",
    "    \n",
    "    # fig.savefig(savePath.joinpath(f'{DA.material}-{DA.solvent}-{DA.rpm}_chi{chi_min}to{chi_max}_q{q_min}to{q_max}_{DA.incident_angle}.png'), dpi=150)\n",
    "    # fig.savefig(savePath.joinpath(f'{DA.material}-{DA.solvent}_chi{chi_min}to{chi_max}_q{q_min}to{q_max}_{DA.incident_angle}.png'), dpi=150)\n",
    "\n",
    "    plt.show()\n",
    "    \n",
    "    plt.close('all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e42f2d83-b3f4-4e3b-b47b-82a36ce5a18b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "### Set savePath\n",
    "# savePath = outPath.joinpath('var/pipi_pole_figures_v1')\n",
    "savePath = outPath.joinpath('fix/pipi_pole_figures_v1')\n",
    "\n",
    "# for th in tqdm(['th0.080', 'th0.100', 'th0.120', 'th0.140']):\n",
    "for th in tqdm(['th0.120']):\n",
    "    selected_dict_items = [(name, data) for (name, data) in peak_areas_dict.items()]\n",
    "    # selected_dict_items = [(name, data) for (name, data) in peak_areas_dict.items() if \n",
    "    #                        # # 'Y7BO-' in name and\n",
    "    #                        # 'CF' in name and '3000' in name and \n",
    "    #                        th in name]    \n",
    "    #                        # 'CFCN-' in name]    \n",
    "\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(5,3.3))\n",
    "    fig.set(dpi=120)\n",
    "    \n",
    "    colors = plt.cm.Dark2(np.linspace(0,1,len(selected_dict_items)))\n",
    "\n",
    "    for i, (measurement, pole_fig_arr_data) in enumerate(selected_dict_items):\n",
    "        chis = pole_fig_arr_data[:,0]\n",
    "        normed_peak_areas = pole_fig_arr_data[:,1]\n",
    "        ax.plot(chis, normed_peak_areas, color=colors[i], label=measurement)\n",
    "        \n",
    "    fig.suptitle('Pole Figure', x=0.55)\n",
    "    ax.set(title=f'{q_min} to {q_max} 1/Å Q Fit Region; {chi_min}° to {chi_max}° $\\chi$')\n",
    "    ax.set(xlabel='$\\chi$ [°]', ylabel='Normalized Peak Area')\n",
    "    ax.set_xlim(None,40)\n",
    "    ax.set_ylim(None,0.14)\n",
    "\n",
    "    \n",
    "    plt.subplots_adjust(top=0.85, bottom=0.2, left=0.2)\n",
    "    ax.legend()\n",
    "    \n",
    "    savePath = outPath.joinpath('pipi_pole_figures_v1')\n",
    "    savePath.mkdir(exist_ok=True)\n",
    "    fig.savefig(savePath.joinpath(f'all_chi{chi_min}to{chi_max}_q{q_min}to{q_max}_{th}.png'), dpi=150)\n",
    "    \n",
    "    plt.show()\n",
    "    plt.close('all')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abfe059c-0c3b-4383-84ed-b838175123b9",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Lamella"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9371635b-56ba-4d10-b9fb-919e476cd5d8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "### Set limits\n",
    "q_min = 0.2\n",
    "q_max = 0.45\n",
    "chi_min = 10\n",
    "chi_max = 82  # to ~45 least until I mask out the silicon(?) background scattering that interferes\n",
    "\n",
    "# Set savePath\n",
    "savePath = outPath.joinpath('PM6-Y6series/lamella_pole_figures_v1')\n",
    "\n",
    "selected_attrs_dict = {'material':['PM6'], 'incident_angle': ['th0.120']}\n",
    "selected_DAs = select_attrs(folded_corr_DAs, selected_attrs_dict)\n",
    " \n",
    "peak_areas_dict = {}    \n",
    "for DA in tqdm(selected_DAs):\n",
    "    chis = DA.chi.sel(chi=slice(chi_min,chi_max)).data\n",
    "    \n",
    "    peak_areas = np.array([])\n",
    "    peak_centers = np.array([])\n",
    "    \n",
    "    for chi in chis:\n",
    "        sliced_DA = DA.sel(qr=slice(q_min-0.05, q_max+0.05), chi=float(chi))\n",
    "\n",
    "        points_x = [q_min, q_max]\n",
    "        points_y = [float(sliced_DA.sel(qr=slice(points_x[0]-0.05, points_x[0]+0.05)).mean('qr')), \n",
    "                    float(sliced_DA.sel(qr=slice(points_x[1]-0.05, points_x[1]+0.05)).mean('qr'))]\n",
    "        # points_y = [float(sliced_DA.sel(qr=q_min, method='nearest')), \n",
    "        #             float(sliced_DA.sel(qr=q_max, method='nearest'))]\n",
    "        m = (points_y[1]-points_y[0])/(points_x[1]-points_x[0])\n",
    "        b = points_y[1] - (m*points_x[1])\n",
    "        y_fit = np.polyval([m, b], sliced_DA.qr)\n",
    "        \n",
    "        sub_sliced_DA = sliced_DA-y_fit\n",
    "\n",
    "        peak_area = sub_sliced_DA.sel(qr=slice(q_min,q_max)).integrate('qr')\n",
    "\n",
    "        # peak_area = sliced_DA.sel(qr=slice(None,0.45)).integrate('qr')\n",
    "        peak_area = np.nan if peak_area<0 else peak_area  # remove negatives as they are nonphysical\n",
    "        \n",
    "        peak_areas = np.append(peak_areas, float(peak_area))\n",
    "                        \n",
    "#         # Plot linecuts for each chi\n",
    "#         fig, ax = plt.subplots()\n",
    "#         sliced_DA.plot(ax=ax)\n",
    "#         ax.plot(points_x, points_y)\n",
    "#         plt.show()\n",
    "\n",
    "#         plt.close('all')\n",
    "\n",
    "    normed_peak_areas = peak_areas / np.nansum(peak_areas)\n",
    "    data = np.vstack((chis, normed_peak_areas)).T\n",
    "    # peak_areas_dict[DA.attrs['material']+'-'+DA.attrs['solvent']+'-'+DA.attrs['rpm']+'-'+DA.attrs['incident_angle']] = data\n",
    "    peak_areas_dict[DA.attrs['material']+'-'+DA.attrs['solvent']+'-'+DA.attrs['incident_angle']] = data\n",
    "\n",
    "\n",
    "    # Individual measurement pole figure plot\n",
    "    fig, ax = plt.subplots(figsize=(5,3.3))\n",
    "    fig.set(dpi=120)\n",
    "    \n",
    "    ax.plot(chis, normed_peak_areas)\n",
    "    # fig.suptitle(f'{DA.material} {DA.solvent} {DA.rpm} Pole Figure; $\\\\alpha_i = $ {DA.incident_angle[2:-1]}°', x=0.54)\n",
    "    fig.suptitle(f'{DA.material} {DA.solvent} Pole Figure; $\\\\alpha_i = $ {DA.incident_angle[2:-1]}°', x=0.54)\n",
    "\n",
    "    ax.set(title=f'{q_min} to {q_max} 1/Å Q Bounds; {chi_min}° to {chi_max}° $\\chi$')\n",
    "    ax.set(xlabel='$\\chi$ [°]', ylabel='Integrated Pi-Pi Peak Area [arb. units]')\n",
    "    plt.subplots_adjust(top=0.85, bottom=0.2, left=0.2)\n",
    "    \n",
    "    # fig.savefig(savePath.joinpath(f'{DA.material}-{DA.solvent}-{DA.rpm}_chi{chi_min}to{chi_max}_q{q_min}to{q_max}_{DA.incident_angle}.png'), dpi=150)\n",
    "    # fig.savefig(savePath.joinpath(f'{DA.material}-{DA.solvent}_chi{chi_min}to{chi_max}_q{q_min}to{q_max}_{DA.incident_angle}.png'), dpi=150)\n",
    "    \n",
    "    plt.show()\n",
    "    \n",
    "    plt.close('all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efa555b4-63dc-4804-940a-1e42217ea9ad",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "### Set savePath\n",
    "# savePath = outPath.joinpath('PM6-Y6series/var/lamella_pole_figures_v1')\n",
    "savePath = outPath.joinpath('fix/lamella_pole_figures_v1')\n",
    "\n",
    "# for th in tqdm(['th0.080', 'th0.100', 'th0.120', 'th0.140']):\n",
    "for th in tqdm(['th0.120']):\n",
    "    selected_dict_items = [(name, data) for (name, data) in peak_areas_dict.items() if \n",
    "                           # # 'Y7BO-' in name and\n",
    "                           # 'CF' in name and '3000' in name and \n",
    "                           # 'CF-' in name]     \n",
    "                           th in name] \n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(5,3.3))\n",
    "    fig.set(dpi=120)\n",
    "    \n",
    "    colors = plt.cm.Dark2(np.linspace(0,1,len(selected_dict_items)))\n",
    "\n",
    "    for i, (measurement, pole_fig_arr_data) in enumerate(selected_dict_items):\n",
    "        chis = pole_fig_arr_data[:,0]\n",
    "        normed_peak_areas = pole_fig_arr_data[:,1]\n",
    "        ax.plot(chis, normed_peak_areas, color=colors[i], label=measurement)\n",
    "        \n",
    "    fig.suptitle('Pole Figure', x=0.55)\n",
    "    ax.set(title=f'{q_min} to {q_max} 1/Å Q Fit Region ; {chi_min}° to {chi_max}° $\\chi$')\n",
    "    ax.set(xlabel='$\\chi$ [°]', ylabel='Normalized Peak Area')\n",
    "    \n",
    "    plt.subplots_adjust(top=0.85, bottom=0.2, left=0.2)\n",
    "    ax.legend()\n",
    "\n",
    "    # fig.savefig(savePath.joinpath(f'multiple-PM6_chi{chi_min}to{chi_max}_q{q_min}to{q_max}_{th}.png'), dpi=150)\n",
    "\n",
    "\n",
    "    plt.show()\n",
    "    plt.close('all')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dccce40-450c-4d67-aa6b-573a50e67ef4",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### Backbone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d593ac2-9cbd-4570-a7fb-6c5b829833cc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "### Set limits\n",
    "q_min = 0.55\n",
    "q_max = 0.77\n",
    "chi_min = 10\n",
    "chi_max = 82  # to ~45 least until I mask out the silicon(?) background scattering that interferes\n",
    "\n",
    "# Set savePath\n",
    "savePath = outPath.joinpath('PM6-Y6series/lamella_pole_figures_v1')\n",
    "\n",
    "selected_attrs_dict = {'material':['PM6'], 'incident_angle': ['th0.120']}\n",
    "selected_DAs = select_attrs(folded_corr_DAs, selected_attrs_dict)\n",
    " \n",
    "peak_areas_dict = {}    \n",
    "for DA in tqdm(selected_DAs):\n",
    "    chis = DA.chi.sel(chi=slice(chi_min,chi_max)).data\n",
    "    \n",
    "    peak_areas = np.array([])\n",
    "    peak_centers = np.array([])\n",
    "    \n",
    "    for chi in chis:\n",
    "        sliced_DA = DA.sel(qr=slice(q_min, q_max), chi=float(chi))\n",
    "\n",
    "        points_x = [q_min, q_max]\n",
    "        points_y = [float(sliced_DA.sel(qr=slice(points_x[0]-0.05, points_x[0]+0.05)).mean('qr')), \n",
    "                    float(sliced_DA.sel(qr=slice(points_x[1]-0.05, points_x[1]+0.05)).mean('qr'))]\n",
    "        # points_y = [float(sliced_DA.sel(qr=q_min, method='nearest')), \n",
    "        #             float(sliced_DA.sel(qr=q_max, method='nearest'))]\n",
    "        m = (points_y[1]-points_y[0])/(points_x[1]-points_x[0])\n",
    "        b = points_y[1] - (m*points_x[1])\n",
    "        y_fit = np.polyval([m, b], sliced_DA.qr)\n",
    "        \n",
    "        sub_sliced_DA = sliced_DA-y_fit\n",
    "\n",
    "        peak_area = sub_sliced_DA.sel(qr=slice(q_min,q_max)).integrate('qr')\n",
    "\n",
    "        # peak_area = sliced_DA.sel(qr=slice(None,0.45)).integrate('qr')\n",
    "        peak_area = np.nan if peak_area<0 else peak_area  # remove negatives as they are nonphysical\n",
    "        \n",
    "        peak_areas = np.append(peak_areas, float(peak_area))\n",
    "                        \n",
    "#         # Plot linecuts for each chi\n",
    "#         fig, ax = plt.subplots()\n",
    "#         sliced_DA.plot(ax=ax)\n",
    "#         ax.plot(points_x, points_y)\n",
    "#         plt.show()\n",
    "\n",
    "#         plt.close('all')\n",
    "\n",
    "    normed_peak_areas = peak_areas / np.nansum(peak_areas)\n",
    "    data = np.vstack((chis, normed_peak_areas)).T\n",
    "    # peak_areas_dict[DA.attrs['material']+'-'+DA.attrs['solvent']+'-'+DA.attrs['rpm']+'-'+DA.attrs['incident_angle']] = data\n",
    "    peak_areas_dict[DA.attrs['material']+'-'+DA.attrs['solvent']+'-'+DA.attrs['incident_angle']] = data\n",
    "\n",
    "\n",
    "    # Individual measurement pole figure plot\n",
    "    fig, ax = plt.subplots(figsize=(5,3.3))\n",
    "    fig.set(dpi=120)\n",
    "    \n",
    "    ax.plot(chis, normed_peak_areas)\n",
    "    # fig.suptitle(f'{DA.material} {DA.solvent} {DA.rpm} Pole Figure; $\\\\alpha_i = $ {DA.incident_angle[2:-1]}°', x=0.54)\n",
    "    fig.suptitle(f'{DA.material} {DA.solvent} Pole Figure; $\\\\alpha_i = $ {DA.incident_angle[2:-1]}°', x=0.54)\n",
    "\n",
    "    ax.set(title=f'{q_min} to {q_max} 1/Å Q Bounds; {chi_min}° to {chi_max}° $\\chi$')\n",
    "    ax.set(xlabel='$\\chi$ [°]', ylabel='Integrated Pi-Pi Peak Area [arb. units]')\n",
    "    plt.subplots_adjust(top=0.85, bottom=0.2, left=0.2)\n",
    "    \n",
    "    # fig.savefig(savePath.joinpath(f'{DA.material}-{DA.solvent}-{DA.rpm}_chi{chi_min}to{chi_max}_q{q_min}to{q_max}_{DA.incident_angle}.png'), dpi=150)\n",
    "    # fig.savefig(savePath.joinpath(f'{DA.material}-{DA.solvent}_chi{chi_min}to{chi_max}_q{q_min}to{q_max}_{DA.incident_angle}.png'), dpi=150)\n",
    "    \n",
    "    plt.show()\n",
    "    \n",
    "    plt.close('all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2afef5f2-58d2-45f2-be7e-198831a11e04",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Set parameters:\n",
    "q_min = 0.55\n",
    "q_max = 0.77\n",
    "chi_min = 45\n",
    "chi_max = 82\n",
    "\n",
    "# Set savePath\n",
    "# savePath = outPath.joinpath('PM6-Y6series/var/backbone_pole_figures_v1')\n",
    "savePath = outPath.joinpath('fix/backbone_pole_figures_v1')\n",
    "\n",
    "# for th in tqdm(['th0.080', 'th0.100', 'th0.120', 'th0.140']):\n",
    "for th in tqdm(['th0.120']):\n",
    "    selected_dict_items = [(name, data) for (name, data) in peak_areas_dict.items() if \n",
    "                           # # 'Y7BO-' in name and\n",
    "                           # 'CF' in name and '3000' in name and \n",
    "                           # 'CF-' in name]     \n",
    "                           th in name] \n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(5,3.3))\n",
    "    fig.set(dpi=120)\n",
    "    \n",
    "    colors = plt.cm.Dark2(np.linspace(0,1,len(selected_dict_items)))\n",
    "\n",
    "    for i, (measurement, pole_fig_arr_data) in enumerate(selected_dict_items):\n",
    "        chis = pole_fig_arr_data[:,0]\n",
    "        normed_peak_areas = pole_fig_arr_data[:,1]\n",
    "        ax.plot(chis, normed_peak_areas, color=colors[i], label=measurement)\n",
    "        \n",
    "    fig.suptitle('Pole Figure', x=0.55)\n",
    "    ax.set(title=f'{q_min} to {q_max} 1/Å Q Fit Region ; {chi_min}° to {chi_max}° $\\chi$')\n",
    "    ax.set(xlabel='$\\chi$ [°]', ylabel='Normalized Peak Area')\n",
    "    ax.set_xlim(45,None)\n",
    "    \n",
    "    plt.subplots_adjust(top=0.85, bottom=0.2, left=0.2)\n",
    "    ax.legend()\n",
    "    \n",
    "    \n",
    "    fig.savefig(savePath.joinpath(f'multiple-PM6_chi{chi_min}to{chi_max}_q{q_min}to{q_max}_{th}.png'), dpi=150)\n",
    "\n",
    "\n",
    "    # plt.show()\n",
    "    plt.close('all')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f5306ce-5a7f-4a83-a7c2-842fd2168e8f",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "# Original notebook code"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbfcb7d6-878f-4fc8-a624-36ed8d3022a3",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 2D plots"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e4fa209-1b91-44f5-990e-8643c21f79b9",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### Caked Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96d06371-37ed-4b6b-947e-e955785b9cda",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "filename = 'fix_caked_stitched.zarr'\n",
    "DS = xr.open_zarr(zarrsPath.joinpath(filename))\n",
    "DS = DS.where(DS>1e-5)\n",
    "DS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "519afdb6-8f5b-4e6f-aec1-9372e1c8b6bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # How one could apply a sin chi correction\n",
    "# sin_chi_DA = np.sin(np.radians(np.abs(DA.chi)))\n",
    "# # sin_chi_DA\n",
    "\n",
    "# corr_DA = DA * sin_chi_DA\n",
    "# # corr_DA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93ef1d8d-30f2-4eb3-a72c-e437484bc2b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# A way to select dataarrays based on attribute values:\n",
    "selected_DAs = [da for da in DS.data_vars.values() if \n",
    "                da.attrs['material']]\n",
    "                # da.attrs['material'] in ['Y6', 'Y7'] and \n",
    "                # da.attrs['incident_angle'] in ['th0.120', 'th0.140']]\n",
    "\n",
    "len(selected_DAs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b58f83a8-a030-4947-be20-e14ccea6107a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Plot and optionally save selected dataarrays:\n",
    "# Set chi range: Full range\n",
    "chi_min = -90\n",
    "chi_max = 90\n",
    "\n",
    "for DA in tqdm(selected_DAs):\n",
    "    # Slice dataarray to select plotting region \n",
    "    sliced_DA = DA.sel(chi=slice(chi_min,chi_max))\n",
    "    \n",
    "    # real_min = float(DA.sel(q_xy=slice(-0.5, -0.1), q_z=slice(0.1, 0.4)).compute().quantile(1e-3))\n",
    "    real_min = float(DA.compute().quantile(0.01))\n",
    "    cmin = 1 if real_min < 1 else real_min\n",
    "    \n",
    "    # cmax = float(DA.sel(q_xy=slice(-0.5, -0.1), q_z=slice(0.1, 2)).compute().quantile(1))   \n",
    "    cmax = float(DA.compute().quantile(0.99))   \n",
    "    \n",
    "    # Plot sliced dataarray\n",
    "    ax = sliced_DA.plot.imshow(cmap=cmap, norm=plt.Normalize(cmin, cmax), figsize=(5,4), interpolation='antialiased')  # plot, optional parameter interpolation='antialiased' for image smoothing\n",
    "    ax.colorbar.set_label('Intensity [arb. units]', rotation=270, labelpad=15)  # set colorbar label & parameters \n",
    "    ax.axes.set(title=f'Polar Plot: {DA.material} {DA.solvent}, {float(DA.incident_angle[2:])}° Incidence',\n",
    "                xlabel='q$_r$ [Å$^{-1}$]', ylabel='$\\chi$ [°]')  # set title, axis labels, misc\n",
    "    ax.figure.set(tight_layout=True, dpi=130)  # Adjust figure dpi & plotting style\n",
    "    \n",
    "    plt.show()  # Comment to mute plotting output\n",
    "    \n",
    "    # Uncomment below line and set savepath/savename for saving plots, I usually like to check \n",
    "    # ax.figure.savefig(outPath.joinpath('PM6-Y6set_waxs', f'polar-2D_{DA.sample_id}_{chi_min}to{chi_max}chi_{DA.incident_angle}.png'), dpi=150)\n",
    "    plt.close('all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f1c0569-17bb-4237-a711-72cb191351f8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Plot and optionally save selected dataarrays:\n",
    "# Set chi range: Full range\n",
    "chi_min = -90\n",
    "chi_max = 90\n",
    "\n",
    "for DA in tqdm(selected_DAs):\n",
    "    # Slice dataarray to select plotting region \n",
    "    sliced_DA = DA.sel(chi=slice(chi_min,chi_max))\n",
    "    \n",
    "    # real_min = float(DA.sel(q_xy=slice(-0.5, -0.1), q_z=slice(0.1, 0.4)).compute().quantile(1e-3))\n",
    "    real_min = float(DA.compute().quantile(0.01))\n",
    "    cmin = 1 if real_min < 1 else real_min\n",
    "    \n",
    "    # cmax = float(DA.sel(q_xy=slice(-0.5, -0.1), q_z=slice(0.1, 2)).compute().quantile(1))   \n",
    "    cmax = float(DA.compute().quantile(0.999))   \n",
    "    \n",
    "    # Plot sliced dataarray\n",
    "    ax = sliced_DA.plot.imshow(cmap=cmap, norm=plt.Normalize(cmin, cmax), figsize=(5,4), interpolation='antialiased')  # plot, optional parameter interpolation='antialiased' for image smoothing\n",
    "    ax.colorbar.set_label('Intensity [arb. units]', rotation=270, labelpad=15)  # set colorbar label & parameters \n",
    "    ax.axes.set(title=f'Polar Plot: {DA.material} {DA.solvent} {DA.rpm}, {float(DA.incident_angle[2:])}° Incidence',\n",
    "                xlabel='q$_r$ [Å$^{-1}$]', ylabel='$\\chi$ [°]')  # set title, axis labels, misc\n",
    "    ax.figure.set(tight_layout=True, dpi=130)  # Adjust figure dpi & plotting style\n",
    "    \n",
    "    plt.show()  # Comment to mute plotting output\n",
    "    \n",
    "    # Uncomment below line and set savepath/savename for saving plots, I usually like to check \n",
    "    # ax.figure.savefig(outPath.joinpath('PM6-Y6set_waxs', f'polar-2D_{DA.sample_id}_{chi_min}to{chi_max}chi_{DA.incident_angle}.png'), dpi=150)\n",
    "    plt.close('all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b89ad7b1-406b-4c59-a9c4-70cb8b3d05b1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Plot and optionally save selected dataarrays:\n",
    "# Set chi range: In plane slice, choose a smooth section without detector gap/edge effects\n",
    "chi_min = 72\n",
    "chi_max = 82\n",
    "\n",
    "for DA in tqdm(selected_DAs):\n",
    "    # Slice dataarray to select plotting region \n",
    "    sliced_DA = DA.sel(chi=slice(chi_min,chi_max), qr=slice(0.23,2.05))\n",
    "    cmin = float(sliced_DA.compute().quantile(1e-2))  # Set color minimum value, based on quantile \n",
    "    cmax = float(sliced_DA.compute().quantile(1-1e-6))  # Set color maximum value, based on quantile\n",
    "    \n",
    "    # Plot sliced dataarray\n",
    "    ax = sliced_DA.plot.imshow(cmap=cmap, norm=LogNorm(cmin, cmax), figsize=(5,4))  # plot\n",
    "    ax.colorbar.set_label('Intensity [arb. units]', rotation=270, labelpad=15)  # set colorbar label & parameters \n",
    "    ax.axes.set(title=f'Polar Plot: {DA.sample_id}, {float(DA.incident_angle[2:])}° Incidence',\n",
    "                xlabel='q$_r$ [Å$^{-1}$]', ylabel='$\\chi$ [°]')  # set title, axis labels, misc\n",
    "    ax.figure.set(tight_layout=True, dpi=130)  # Adjust figure dpi & plotting style\n",
    "    \n",
    "    plt.show()  # Comment to mute plotting output\n",
    "    \n",
    "    # Uncomment below line and set savepath/savename for saving plots, I usually like to check \n",
    "    # ax.figure.savefig(outPath.joinpath('PM6-Y6set_waxs', f'polar-2D_{DA.sample_id}_{chi_min}to{chi_max}chi_{DA.incident_angle}.png'), dpi=150)\n",
    "    plt.close('all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2326db38-33f2-4b64-9449-e3bdf2b30a2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # A way to save data as csv files  \n",
    "# for DA in DS.data_vars.values():\n",
    "#     # qr columns, chi rows\n",
    "#     DA.to_pandas().to_csv(outPath.joinpath('PM6-Y6_waxs', f'polar-2D_{DA.polymer}-{DA.weight_percent}_{DA.incident_angle}_{DA.scan_id}.csv'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fdc502d-edae-4eb6-88b4-2f6eeecd9241",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5892802-e160-40b0-89d6-0b985500860f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "filename = 'fix_raw_stitched.zarr'\n",
    "DS = xr.open_zarr(zarrsPath.joinpath(filename))\n",
    "DS = DS.where(DS>1e-6)\n",
    "DS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41578001-68a8-4003-8893-2a0e202be4bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "selected_DAs = [da for da in DS.data_vars.values() if \n",
    "                da.attrs['incident_angle']]\n",
    "len(selected_DAs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5d0bd7f-58b9-4a52-ba7c-a7181a304cca",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Plot & optionally save each selected polymer: fix data set\n",
    "for DA in tqdm(selected_DAs):\n",
    "    # Slice data for selected q ranges (will need to rename q_xy if dimensions are differently named)\n",
    "    sliced_DA = DA.copy()\n",
    "    sliced_DA = DA.sel(pix_y=slice(None, 680), pix_x=slice(150,None))\n",
    "\n",
    "    # real_min = float(DA.sel(q_xy=slice(-0.5, -0.1), q_z=slice(0.1, 0.4)).compute().quantile(1e-3))\n",
    "    real_min = float(DA.compute().quantile(0.01))\n",
    "    cmin = 1 if real_min < 1 else real_min\n",
    "    \n",
    "    # cmax = float(DA.sel(q_xy=slice(-0.5, -0.1), q_z=slice(0.1, 2)).compute().quantile(1))   \n",
    "    cmax = float(DA.compute().quantile(0.999))   \n",
    "    \n",
    "    # Same plotting procedure as above\n",
    "    # ax = sliced_DA.plot.imshow(cmap=cmap, norm=LogNorm(cmin, cmax), interpolation='antialiased', figsize=(5.5,3.3))\n",
    "    ax = sliced_DA.plot.imshow(cmap=cmap, norm=plt.Normalize(cmin, cmax), interpolation='antialiased', figsize=(5.5,3.3),\n",
    "                               origin='upper')\n",
    "\n",
    "    ax.colorbar.set_label('Intensity [arb. units]', rotation=270, labelpad=15)\n",
    "    ax.axes.set(title=f'Raw Detector: {DA.material} {DA.solvent}, {float(DA.incident_angle[2:])}° Incidence',\n",
    "                aspect='equal')\n",
    "    ax.figure.set(tight_layout=True, dpi=130)\n",
    "    \n",
    "    # ax.figure.savefig(outPath.joinpath('recip_plots/stitched_v2', f'{DA.material}_{DA.solvent}_{DA.incident_angle}.png'), dpi=120)\n",
    "    plt.show()\n",
    "    plt.close('all')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a8a902e-8ce1-4b80-9760-975d907eb354",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Reciprocal Space Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d456826-6758-4de8-a0ca-5765ef10b813",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "filename = 'var_recip_stitched.zarr'\n",
    "DS = xr.open_zarr(zarrsPath.joinpath(filename))\n",
    "DS = DS.where(DS>1e-6)\n",
    "DS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b992c21-2ee9-466c-90d8-737a000efdfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "selected_DAs = [da for da in DS.data_vars.values() if \n",
    "                da.attrs['incident_angle']]\n",
    "len(selected_DAs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e56947e3-ec74-4df2-905c-11c4e61a1ea9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Plot & optionally save each selected polymer: fix data set\n",
    "for DA in tqdm(selected_DAs):\n",
    "    # Slice data for selected q ranges (will need to rename q_xy if dimensions are differently named)\n",
    "    sliced_DA = DA.sel(q_xy=slice(-1.1, 2.1), q_z=slice(-0.01, 2.2))\n",
    "    # sliced_DA = DA.sel(q_xy=slice(-0.5, -0.25), q_z=slice(1.5, 1.75))\n",
    "\n",
    "    # real_min = float(DA.sel(q_xy=slice(-0.5, -0.1), q_z=slice(0.1, 0.4)).compute().quantile(1e-3))\n",
    "    real_min = float(DA.compute().quantile(0.01))\n",
    "    cmin = 1 if real_min < 1 else real_min\n",
    "    \n",
    "    # cmax = float(DA.sel(q_xy=slice(-0.5, -0.1), q_z=slice(0.1, 2)).compute().quantile(1))   \n",
    "    cmax = float(DA.compute().quantile(0.999))   \n",
    "    \n",
    "    # Same plotting procedure as above\n",
    "    # ax = sliced_DA.plot.imshow(cmap=cmap, norm=LogNorm(cmin, cmax), interpolation='antialiased', figsize=(5.5,3.3))\n",
    "    ax = sliced_DA.plot.imshow(cmap=cmap, norm=plt.Normalize(cmin, cmax), interpolation='antialiased', figsize=(5.5,3.3))\n",
    "\n",
    "    ax.colorbar.set_label('Intensity [arb. units]', rotation=270, labelpad=15)\n",
    "    ax.axes.set(title=f'Cartesian Plot: {DA.material} {DA.solvent}, {float(DA.incident_angle[2:])}° Incidence',\n",
    "                aspect='equal', xlabel='q$_{xy}$ [Å$^{-1}$]', ylabel='q$_z$ [Å$^{-1}$]')\n",
    "    ax.figure.set(tight_layout=True, dpi=130)\n",
    "    \n",
    "    ax.figure.savefig(outPath.joinpath('recip_plots/stitched_v2', f'{DA.material}_{DA.solvent}_{DA.incident_angle}.png'), dpi=120)\n",
    "    # plt.show()\n",
    "    plt.close('all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de525325-773d-4ce5-908c-aa6facf88dc7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Plot & optionally save each selected polymer: var data set\n",
    "for DA in tqdm(selected_DAs):\n",
    "    # Slice data for selected q ranges (will need to rename q_xy if dimensions are differently named)\n",
    "    sliced_DA = DA.sel(q_xy=slice(-1.1, 2.1), q_z=slice(-0.01, 2.2))\n",
    "    # sliced_DA = DA.sel(q_xy=slice(-0.5, -0.25), q_z=slice(1.5, 1.75))\n",
    "\n",
    "    # real_min = float(DA.sel(q_xy=slice(-0.5, -0.1), q_z=slice(0.1, 0.4)).compute().quantile(1e-3))\n",
    "    real_min = float(DA.compute().quantile(0.01))\n",
    "    cmin = 1 if real_min < 1 else real_min\n",
    "    \n",
    "    # cmax = float(DA.sel(q_xy=slice(-0.5, -0.25), q_z=slice(0.1, 1.75)).compute().quantile(1-1e-10))   \n",
    "    cmax = float(DA.compute().quantile(0.999))   \n",
    "    \n",
    "    # Same plotting procedure as above\n",
    "    ax = sliced_DA.plot.imshow(cmap=cmap, norm=plt.Normalize(cmin, cmax), interpolation='antialiased', figsize=(5.5,3.3))\n",
    "    ax.colorbar.set_label('Intensity [arb. units]', rotation=270, labelpad=15)\n",
    "    ax.axes.set(title=f'Cartesian Plot: {DA.material} {DA.solvent} {DA.rpm}, {float(DA.incident_angle[2:])}° Incidence',\n",
    "                aspect='equal', xlabel='q$_{xy}$ [Å$^{-1}$]', ylabel='q$_z$ [Å$^{-1}$]')\n",
    "    ax.figure.set(tight_layout=True, dpi=130)\n",
    "    \n",
    "    ax.figure.savefig(outPath.joinpath('recip_plots/stitched_v2', f'{DA.material}_{DA.solvent}_{DA.rpm}_{DA.incident_angle}.png'), dpi=120)\n",
    "    # plt.show()\n",
    "    plt.close('all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12a0b2cb-f757-446d-85e4-8a130b147e44",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # A way to save data as csv files\n",
    "# for DA in tqdm(DS.data_vars.values()):\n",
    "#     # qxy columns, qz rows\n",
    "#     DA.to_pandas().to_csv(outPath.joinpath('PM6-Y6_waxs', f'cartesian-2D_{DA.polymer}-{DA.weight_percent}_{DA.incident_angle}_{DA.scan_id}.csv'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f164c8c0-98bb-43b3-8d63-9d6de075640e",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 1D Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38841f0e-dcb4-4c6b-abb4-cd70d77f7e8f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "filename = 'caked_PM6-Y6_waxs_stitched.zarr'\n",
    "DS = xr.open_zarr(samplesPath.joinpath(filename))\n",
    "DS = DS.where(DS>1e-5)\n",
    "DS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3a4cd9d-e268-4e87-897e-9102f1980336",
   "metadata": {},
   "outputs": [],
   "source": [
    "selected_DAs = [da for da in DS.data_vars.values() if \n",
    "                da.attrs['incident_angle'] == 'th0.120']\n",
    "len(selected_DAs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "954d1667-14e5-425f-9bb5-b74154e72d3f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fpath = outPath.joinpath('PM6-Y6set_waxs', f'linecut_IP_{DA.polymer}-{DA.weight_percent}_{chi_min}to{chi_max}chi_{DA.incident_angle}.csv')\n",
    "OOP_linecut_DA.to_dataframe('OOP Intensity').to_csv(fpath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d3f6381-4597-41f5-916d-cd805f1ba00b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Plot linecuts for selected chi ranges, here I've put both in plane and out of plane selections into the loop\n",
    "\n",
    "for DA in tqdm(DS.data_vars.values()):\n",
    "    # OOP\n",
    "    chi_min = -18\n",
    "    chi_max = -8\n",
    "    OOP_linecut_DA = DA.sel(chi=slice(chi_min, chi_max), qr=slice(0.14,2.01)).sum('chi')\n",
    "    OOP_linecut_DA.plot.line(figsize=(6,4))\n",
    "\n",
    "    # A plot.line xarray plot does not return an AxesImage object like imshow does, so I use plt.gca() and plt.gcf() to access the axes & figure parameters\n",
    "    ax = plt.gca()\n",
    "    fig = plt.gcf()\n",
    "    \n",
    "    ax.set(title=f'OOP Linecut, {chi_min}° to {chi_max}° $\\chi$: {DA.polymer}-{DA.weight_percent}, {float(DA.incident_angle[2:])}° Incidence',\n",
    "           yscale='log', ylabel='Intensity [arb. units]', xlabel='q$_r$ [Å$^{-1}$]')\n",
    "    ax.grid(visible=True, which='major', axis='x')\n",
    "    fig.set(tight_layout=True, dpi=130)\n",
    "    \n",
    "    plt.show()\n",
    "    fpath = outPath.joinpath('PM6-Y6set_waxs', f'linecut_OOP_{DA.polymer}-{DA.weight_percent}_{chi_min}to{chi_max}chi_{DA.incident_angle}.csv')\n",
    "    print('Saving csv data...')\n",
    "    OOP_linecut_DA.to_dataframe('OOP_Intensity').to_csv(fpath)\n",
    "    print('Saved!')\n",
    "    # fig.savefig(outPath.joinpath('bladecoated-set_waxs', f'linecut_OOP_{DA.polymer}-{DA.weight_percent}_{chi_min}to{chi_max}chi_{DA.incident_angle}.png'), dpi=150)\n",
    "    plt.close('all')\n",
    "    \n",
    "    # IP\n",
    "    chi_min = 72\n",
    "    chi_max = 82\n",
    "    IP_linecut_DA = DA.sel(chi=slice(chi_min, chi_max), qr=slice(0.23,2.01)).sum('chi')\n",
    "    IP_linecut_DA.plot.line(figsize=(6,4))  \n",
    "    \n",
    "    ax = plt.gca()\n",
    "    fig = plt.gcf()\n",
    "    \n",
    "    ax.set(title=f'IP Linecut, {chi_min}° to {chi_max}° $\\chi$: {DA.polymer}-{DA.weight_percent}, {float(DA.incident_angle[2:])}° Incidence',\n",
    "           yscale='log', ylabel='Intensity [arb. units]', xlabel='q$_r$ [Å$^{-1}$]')\n",
    "    ax.grid(visible=True, which='major', axis='x')\n",
    "    fig.set(tight_layout=True, dpi=130)\n",
    "    \n",
    "    plt.show()\n",
    "    fpath = outPath.joinpath('PM6-Y6set_waxs', f'linecut_IP_{DA.polymer}-{DA.weight_percent}_{chi_min}to{chi_max}chi_{DA.incident_angle}.csv')\n",
    "    print('Saving csv data...')\n",
    "    IP_linecut_DA.to_dataframe('IP_Intensity').to_csv(fpath)  \n",
    "    print('Saved!')\n",
    "    # fig.savefig(outPath.joinpath('PM6-Y6set_waxs', f'linecut_IP_{DA.polymer}-{DA.weight_percent}_{chi_min}to{chi_max}chi_{DA.incident_angle}.png'), dpi=150)\n",
    "    plt.close('all')\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d21e65b-229b-4cba-9cd1-c80ab7a271e2",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}