scipp-imaging-2D.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to start\n",
    "\n",
    "Before starting you must:\n",
    "- Ensure that `scipp` and `mantid` are on your `PYTHONPATH`.\n",
    "- Generate the `scippconfig.py` file using `make_config.py` located in same directory as this notebook. Refer to the `README.md` or `python make_config.py --help` for information.\n",
    "- Install dependencies : `conda install fabio tifffile` (used for image handling)\n",
    "\n",
    "For `scipp` and `mantid` follow instructions at: https://scipp.readthedocs.io/en/latest/getting-started/installation.html.\n",
    "\n",
    "Converted to use scipp and notebook from [this repository](https://git.esss.dk/testbeamline/gp2/-/tree/stressexperiment).\n",
    "\n",
    "For Table of Contents install Jupyter extensions then reload the notebook:\n",
    "`conda install -c conda-forge jupyter_contrib_nbextensions`\n",
    "`jupyter contrib nbextension install --user`\n",
    "`jupyter nbextension enable toc2/main`\n",
    "\n",
    "# Experimental Summary\n",
    "\n",
    "This script has been developed to measure local strain ε defined as ε = ΔL/L0 in a FCC steel sample under elastic strain in a stress rig. Measured at V20, HZB, Berlin, September 2018 by Peter Kadletz.\n",
    "\n",
    "λ = 2dsinθ, where 2θ = π (transmission), edges characterise the Bragg condition and hence λ = 2d. Therefore strain is easily computed from the wavelength measurement of of a Bragg edge directly, using un-loaded vs loaded experimental runs (and reference mesurements). The known Miller indices of the crystal structure (FCC) are used to predict the wavelength where the Bragg edges should exist, which is bound by the reachable wavelength extents for the instrument. This provides an approximate region to apply a fit.  A complement error function is used to fit each Bragg edge, and a refined centre location (λ) for the edge is used in the strain measurement. Because each bragg edge can be identified individually, one can determine anisotropic strain across the unit cell in the reachable crystallographic directions. In addition the image processing allows for spacial grouping so localised effects, such as those on unconstrained edges of the sample or in necking regions of the sample can be treated seperately. The plotted outputs in the script aim to capture this.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to start\n",
    "\n",
    "Before starting you must:\n",
    "- Ensure that `scipp` and `mantid` are on your `PYTHONPATH`.\n",
    "- Generate the `scippconfig.py` file using `make_config.py` located in same directory as this notebook. Refer to the `README.md` or `python make_config.py --help` for information.\n",
    "- Install dependencies : `conda install fabio tifffile` (used for image handling)\n",
    "\n",
    "For `scipp` and `mantid` follow instructions at: https://scipp.readthedocs.io/en/latest/getting-started/installation.html.\n",
    "\n",
    "Converted to use scipp and notebook from [this repository](https://git.esss.dk/testbeamline/gp2/-/tree/stressexperiment).\n",
    "\n",
    "For Table of Contents install Jupyter extensions then reload the notebook:\n",
    "`conda install -c conda-forge jupyter_contrib_nbextensions`\n",
    "`jupyter contrib nbextension install --user`\n",
    "`jupyter nbextension enable toc2/main`\n",
    "\n",
    "# Experimental Summary\n",
    "\n",
    "This script has been developed to measure local strain ε defined as ε = ΔL/L0 in a FCC steel sample under elastic strain in a stress rig. Measured at V20, HZB, Berlin, September 2018 by Peter Kadletz.\n",
    "\n",
    "λ = 2dsinθ, where 2θ = π (transmission), edges characterise the Bragg condition and hence λ = 2d. Therefore strain is easily computed from the wavelength measurement of of a Bragg edge directly, using un-loaded vs loaded experimental runs (and reference mesurements). The known Miller indices of the crystal structure (FCC) are used to predict the wavelength where the Bragg edges should exist, which is bound by the reachable wavelength extents for the instrument. This provides an approximate region to apply a fit.  A complement error function is used to fit each Bragg edge, and a refined centre location (λ) for the edge is used in the strain measurement. Because each bragg edge can be identified individually, one can determine anisotropic strain across the unit cell in the reachable crystallographic directions. In addition the image processing allows for spacial grouping so localised effects, such as those on unconstrained edges of the sample or in necking regions of the sample can be treated seperately. The plotted outputs in the script aim to capture this.\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    },
    "scrolled": true
   },
   "source": [
    "# Script setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    import scipp\n",
    "except ImportError as e:\n",
    "    print(\"scipp is not available in the PYTHONPATH\")\n",
    "    raise e\n",
    "    \n",
    "try:\n",
    "    import mantid\n",
    "except ImportError as e:\n",
    "    print(\"mantid is not available in the PYTHONPATH\")\n",
    "    raise e\n",
    "    \n",
    "try:\n",
    "    import scippconfig\n",
    "except ImportError as e:\n",
    "    print(\"scippconfig is not available. Make sure you have generated it with `make_config.py`.\")\n",
    "    raise e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Set input and output dirs\n",
    "\n",
    "If your input directory has a different structure this is the cell to modify. \n",
    "Additionally the output directory can be renamed too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Lets get everything set up\n",
    "import os\n",
    "import sys\n",
    "\n",
    "import scipp as sc\n",
    "import numpy as np\n",
    "\n",
    "import imaging, operations\n",
    "\n",
    "from scipy import ndimage, signal\n",
    "\n",
    "DATA_DIR_NAME = \"data_GP2\"\n",
    "OUTPUT_DIR_NAME = \"output\"\n",
    "\n",
    "data_dir = os.path.join(scippconfig.script_root, DATA_DIR_NAME)\n",
    "output_dir = os.path.join(scippconfig.script_root, OUTPUT_DIR_NAME)\n",
    "instrument_file = os.path.join(scippconfig.script_root, 'IDF', 'V20_Definition_GP2.xml')\n",
    "\n",
    "tofs_path = os.path.join(data_dir, 'metadata', 'GP2_Stress_time_values.txt')\n",
    "raw_data_dir = os.path.join(data_dir, \"Stress_Experiments\")\n",
    "\n",
    "if not os.path.exists(data_dir):\n",
    "    raise FileNotFoundError(\"The following data directory does not exist,\"\n",
    "                            f\" check your make_config.py:\\n{data_dir}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    },
    "scrolled": true
   },
   "source": [
    "## Reduction Options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Customisable Options:\n",
    "\n",
    "# defining grouping of 2D detector pixels\n",
    "grouping_number = 27\n",
    "nx_target = grouping_number\n",
    "ny_target = grouping_number\n",
    "\n",
    "bin_width = (64 * 2.5)  # μS\n",
    "\n",
    "# Rebin regions for each of the 5 WFM frames. Values are in detector time : μS\n",
    "# in the format of [bin-start, bin-end, bin width].\n",
    "# used to crop each image, before stitching them together\n",
    "frame_parameters = [{\"start\": 15450, \"stop\":22942, \"step\":bin_width},\n",
    "                    {\"start\": 24800, \"stop\":32052, \"step\":bin_width},\n",
    "                    {\"start\": 33791, \"stop\":40084, \"step\":bin_width},\n",
    "                    {\"start\": 41763, \"stop\":47457, \"step\":bin_width},\n",
    "                    {\"start\": 49315, \"stop\":54500, \"step\":bin_width},\n",
    "                    {\"start\": 56500, \"stop\":58360, \"step\":bin_width}]\n",
    "\n",
    "# Used to rebin the summed frame in order to\n",
    "# cut off frames that contain no data\n",
    "rebin_parameters = {\"start\":8550, \"stop\":26000, \"step\":bin_width}\n",
    "\n",
    "# Used to shift the cropped frames so that their bins overlap\n",
    "# before summing them together into a single frame\n",
    "frame_shift_increments = [-6630, -2420, -2253, -2095, -1946, -1810]\n",
    "frame_shift_increments = [float(i) for i in frame_shift_increments]  # Work around #1114\n",
    "\n",
    "# Pulse references\n",
    "pulse_number_reference = 1.0 / 770956\n",
    "pulse_number_sample = 1.0 / 1280381\n",
    "pulse_number_sample_elastic = 1.0 / 2416839\n",
    "pulse_number_sample_plastic = 1.0 / 2614343\n",
    "\n",
    "# units of transmission, all pixels with transmission higher masking threshold are masked\n",
    "masking_threshold = 0.80\n",
    "\n",
    "# Toggles outputting masked and sliced tiff stacks\n",
    "output_tiff_stack = True\n",
    "\n",
    "# Experiment Metadata\n",
    "measurement_number = 11"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Reduction\n",
    "\n",
    "## Load and scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%% \n"
    }
   },
   "outputs": [],
   "source": [
    "# let's get the process started:\n",
    "\n",
    "# Load X values from the TOF file. Though raw contents are NOT TOF corrected yet.\n",
    "ds = sc.Dataset()\n",
    "ds.coords[\"t\"] = sc.Variable([\"t\"], unit=sc.units.us, values=imaging.read_x_values(tofs_path))\n",
    "ds.coords[\"t\"] *= 1e3\n",
    "\n",
    "def load_and_scale(folder_name, scale_factor):\n",
    "    to_load = os.path.join(raw_data_dir, folder_name)\n",
    "    variable = imaging.tiffs_to_variable(to_load)\n",
    "    variable *= scale_factor\n",
    "    return variable\n",
    "\n",
    "ds[\"reference\"] = load_and_scale(folder_name=\"1) R825 Open Beam\", scale_factor=pulse_number_reference)\n",
    "ds[\"sample\"] = load_and_scale(folder_name=\"2) R825\", scale_factor=pulse_number_sample)\n",
    "ds[\"sample_elastic\"] = load_and_scale(folder_name=\"3) R825 600 Mpa\", scale_factor=pulse_number_sample_elastic)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    },
    "scrolled": true
   },
   "source": [
    "## Add Spectra Coords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Adds a coordinate for the spectrum\n",
    "ds.coords[\"spectrum\"] = sc.Variable([\"spectrum\"], values=np.arange(ds[\"sample\"].shape[1]))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Frame shifts and stitching\n",
    "\n",
    "Pre-calculated frame shifts. This will be replaced with automatic WFM stitching in the future"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# TODO use Neil's wfmess, or similar to avoid hard-coding the edges. This approach is not very nice.\n",
    "frame_shifts = [sum(frame_shift_increments[:i + 1]) for i in range(len(frame_shift_increments))]\n",
    "\n",
    "def stitch(ds, frame_parameters, frame_shifts, rebin_parameters):\n",
    "    \"\"\"\n",
    "    Stitches the 5 wfm frames data.\n",
    "    It crops out each frame, then shifts it so that all frames align,\n",
    "    and then rebins to the operations bins used for all frames.\n",
    "    \"\"\"\n",
    "    frames = []\n",
    "\n",
    "    rebin_params = sc.Variable([\"t\"], unit=sc.units.us,\n",
    "                               values=np.arange(**rebin_parameters))\n",
    "\n",
    "    for i, (slice_bins, shift_parameter) in enumerate(\n",
    "            zip(frame_parameters, frame_shifts)):\n",
    "        bins = sc.Variable([\"t\"], values=np.arange(**slice_bins))\n",
    "        # Rebins the whole data to crop it to frame bins\n",
    "        rebinned = sc.rebin(ds, \"t\", bins)\n",
    "        # Shift each frame according to time shift (function of frequency and window opening)\n",
    "        rebinned.coords[\"t\"] += shift_parameter # Now in TOF\n",
    "        # Rebin to overarching coordinates so that the frame coordinates align\n",
    "        rebinned = sc.rebin(rebinned, \"t\", rebin_params)\n",
    "\n",
    "        frames.append(rebinned)\n",
    "    \n",
    "    # All frames now shifted and therefore represent actual TOF. Frames can be summed.\n",
    "    for f in frames[1:]:\n",
    "        frames[0] += f\n",
    "\n",
    "    frames[0].rename_dims({'t':'tof'})\n",
    "    return frames[0]\n",
    "\n",
    "# Comparing [\"sample\"][\"spectrum\", 37000] data stops matching, as Mantid\n",
    "# and skip appear to handle to partial bin at the end differently. This results\n",
    "# in different counts making data comparison difficult\n",
    "\n",
    "stitched = stitch(ds, frame_parameters=frame_parameters,frame_shifts=frame_shifts,rebin_parameters=rebin_parameters)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Masking\n",
    "### Integration step\n",
    "\n",
    "Integrate counts for each spectra by summing over TOF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "integrated = sc.sum(stitched, 'tof')\n",
    "\n",
    "# Pull out reference from out dataset to avoid checking for it in loops\n",
    "reference = stitched[\"reference\"].copy()\n",
    "integrated_reference = integrated[\"reference\"].copy()\n",
    "\n",
    "del stitched[\"reference\"]\n",
    "del integrated[\"reference\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Transmission Masking\n",
    "\n",
    "Divides the integrated sample counts with an open beam reference. Any values > masking threshold\n",
    "will be masked. The adj pixels step checks for random pixels which were left unmasked or masked\n",
    "with all their neighbours having the same mask value. These are forced to True or false depending on\n",
    "their neighbour value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "num_spectra = stitched.coords[\"spectrum\"].shape[0]\n",
    "bank_width = int(np.sqrt(num_spectra))\n",
    "\n",
    "def mask_non_sample_regions(integrated_dataset, stitched_dataset, var_name):\n",
    "    integrated_spectra = integrated_dataset[var_name]\n",
    "\n",
    "    # This should use sc.greater once #1178 is completed\n",
    "    spectra_masks = np.greater(integrated_spectra.values, masking_threshold)\n",
    "    spectra_masks = sc.Variable([\"spectrum\"], values=spectra_masks)\n",
    "\n",
    "    # Some pixels may be noisy / not strong enough signal so they erranously\n",
    "    # become unmasked / masked from the above operation. This step checks if *all*\n",
    "    # the surrounding pixels are True or False then sets the center pixel to\n",
    "    # that value if they are, removing \"salt and pepper\" noise\n",
    "    spectra_masks = sc.reshape(spectra_masks, dims=[\"y\", \"x\"], shape=(bank_width, bank_width))\n",
    "\n",
    "    final_mask = operations.mask_from_adj_pixels(spectra_masks)\n",
    "    return sc.Variable([\"spectrum\"], values=final_mask.values.ravel())\n",
    "\n",
    "# Calculate transmission value for all\n",
    "\n",
    "int_norm = integrated / integrated_reference\n",
    "sc.plot.plot(int_norm)\n",
    "\n",
    "# Store TOF coords back for the export to TIFF stack\n",
    "tof_coords = stitched.coords[\"tof\"]\n",
    "\n",
    "for k in integrated.keys():\n",
    "    print(f\"Masking non-sample regions in {k}\")\n",
    "\n",
    "    p = mask_non_sample_regions(int_norm, stitched, k)\n",
    "    print(k, sc.sum(p, dim='spectrum'))\n",
    "    stitched[k].masks['non-sample-region'] = mask_non_sample_regions(int_norm, stitched, k)\n",
    "\n",
    "    if output_tiff_stack:\n",
    "        print(f\"Exporting tiff stack for {k}\")\n",
    "        tiff_out_dir = os.path.join(output_dir, f\"tiffs_tof_sum_{k}\")\n",
    "        imaging.export_tiff_stack(dataset=int_norm, key=k,\n",
    "                                  base_name=f\"{k}_norm_sum\",\n",
    "                                  output_dir=tiff_out_dir,\n",
    "                                  x_len=324, y_len=324,\n",
    "                                  tof_values=tof_coords.values)\n",
    "        # Print a newline to prevent the saving messages overlapping\n",
    "        print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Normalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Normalize by open beam\n",
    "normalized = stitched / reference\n",
    "\n",
    "# Replace special values nan and inf\n",
    "replacement=sc.Variable(value=0.0, variance=0.0)\n",
    "kwargs = {\"nan\" : replacement, \"posinf\" : replacement, \"neginf\" : replacement}\n",
    "for k in normalized.keys():\n",
    "    sc.nan_to_num(normalized[k].data, out=normalized[k].data, **kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Output Intermediate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "if output_tiff_stack:\n",
    "    tof_values = normalized.coords[\"tof\"].values\n",
    "\n",
    "    imaging.export_tiff_stack(dataset=normalized, key=\"sample\", x_len=324, y_len=324,\n",
    "                              base_name=\"initial_tof\", output_dir=os.path.join(output_dir, \"tiffs_tof_initial\"),\n",
    "                              tof_values=tof_values)\n",
    "    imaging.export_tiff_stack(dataset=normalized, key=\"sample_elastic\", x_len=324, y_len=324,\n",
    "                              base_name=\"elastic_tof\", output_dir=os.path.join(output_dir, \"tiffs_tof_elastic\"),\n",
    "                              tof_values=tof_values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Convert to wavelength"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%% \n"
    }
   },
   "outputs": [],
   "source": [
    "# Load in IDF for positions of source/sample...etc\n",
    "sc.compat.mantid.load_component_info(normalized, instrument_file)\n",
    "wavelength = sc.neutron.convert(normalized, \"tof\", \"wavelength\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Apply filter function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "filter_func = operations.mean_from_adj_pixels # Mean filter\n",
    "#filter_func = operations.median_from_adj_pixels # Median filter\n",
    "for k in wavelength.keys():\n",
    "    print(f'{filter_func.__name__} filter over {k}')\n",
    "    yx = sc.reshape(wavelength[k].data, dims=[\"wavelength\", \"y\", \"x\"], shape=(wavelength[k].shape[0], bank_width, bank_width))\n",
    "    xy_filtered = filter_func(yx)\n",
    "    wavelength[k].data = sc.reshape(xy_filtered, dims=[\"wavelength\", \"spectrum\"], shape=(wavelength[k].shape[0],wavelength[k].shape[1]) )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Rebin to common bin width\n",
    "\n",
    "As each detector has a different relative position to the sample we\n",
    "need to rebin to a common set of wavelength bins for all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "wav_x_vals = wavelength[\"sample\"].coords[\"wavelength\"]\n",
    "start_wave = sc.min(wav_x_vals)\n",
    "stop_wave = sc.max(wav_x_vals)\n",
    "\n",
    "widths = wav_x_vals['wavelength', 1:] - wav_x_vals['wavelength', :-1]\n",
    "bin_width = sc.mean(sc.mean(widths, 'wavelength'), 'spectrum')\n",
    "# Arange does not include final bin so overshoot\n",
    "stop_wave += bin_width\n",
    "\n",
    "wavelength_coords = sc.Variable([\"wavelength\"], values=np.arange(start_wave.value, stop_wave.value,\n",
    "                                                                 bin_width.value))\n",
    "rebinned_wavelength = sc.rebin(wavelength, \"wavelength\", wavelength_coords)\n",
    "\n",
    "\n",
    "group_size = (324.0  / nx_target) * (324.0 / ny_target)\n",
    "factor = group_size\n",
    "rebinned_wavelength[\"sample\"] *= factor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Sample - Sum Spectra\n",
    "\n",
    "Performs a weighted sum for all spectra in the sample "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def weighted_sum(da : sc.DataArrayView, dim):\n",
    "    \n",
    "    variances = sc.variances(da.data)\n",
    "    data = da / variances\n",
    "    norm = sc.DataArray(data=sc.reciprocal(variances), masks=dict(da.masks.items()))\n",
    "    data.masks['zero_error'] = sc.equal(variances, 0.0 * variances.unit)\n",
    "    norm.masks['zero_error'] = sc.equal(variances, 0.0 * variances.unit)\n",
    "    return sc.sum(data, dim) / sc.sum(norm, dim)\n",
    "\n",
    "# If we are not interested in strain regions for the unloaded sample, we can just combine all spectra to improve statistics.\n",
    "# maybe another masking before summation, with tougher threshold to exclude boderline pixels.\n",
    "sample_final = weighted_sum(rebinned_wavelength[\"sample\"], \"spectrum\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Sample elastic - Group detectors\n",
    "\n",
    "Group detectors instead of summing all, since the elastic deformation could\n",
    "happen at any point we want a balance between signal-noise ratio and \n",
    "resolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "rebinned_wavelength.coords[\"spectrum_mapping\"] = imaging.make_detector_groups(bank_width, bank_width,\n",
    "                                                                              nx_target, ny_target)\n",
    "\n",
    "# Group detectors for the sample_elastic / sample_plastic\n",
    "grouped = sc.groupby(rebinned_wavelength, \"spectrum_mapping\").sum(\"spectrum\")\n",
    "grouped_sample_elastic = grouped[\"sample_elastic\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Pre-calculate detector positions\n",
    "\n",
    "Pre-calculate average detector positions for plotting later"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "histogram_matrices = [sample_final, grouped_sample_elastic]\n",
    "template_histogram = histogram_matrices[1]\n",
    "\n",
    "# Average out detector positions\n",
    "positions_dataset = sc.Dataset()\n",
    "positions_dataset[\"position\"] = sc.Variable(rebinned_wavelength.coords[\"position\"])\n",
    "positions_dataset.coords[\"spectrum_mapping\"] = rebinned_wavelength.coords[\"spectrum_mapping\"]\n",
    "grouped_positions = sc.groupby(positions_dataset, \"spectrum_mapping\").mean(\"spectrum\")\n",
    "\n",
    "plots = []\n",
    "for spec_idx in template_histogram.coords[\"spectrum_mapping\"].values:\n",
    "    pos = grouped_positions[\"position\"][\"spectrum_mapping\", int(spec_idx)]\n",
    "    # Append spectrum index and position\n",
    "    plots.append((spec_idx,pos))\n",
    "\n",
    "# We cant access position using (dim, I), since the vector 3d has not dim currently\n",
    "Y_key = 1\n",
    "X_key = 0\n",
    "\n",
    "# We must use round to prevent floating points errors during sorting like -0.05 > 0.05\n",
    "# Sort by y then by x across image plots contain tuples of (spec_index to position), sort by position y, x\n",
    "plots.sort(key=lambda t: (round(-t[1].value[Y_key], 10), round(t[1].value[X_key], 10)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Fitting\n",
    "## Calculate expected Bragg edge positions "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Bragg edge position roughly, in Angstrom\n",
    "FCC_a = 3.5 # Angstrom, taken from COD entry 9008469\n",
    "\n",
    "# These miller indices for the given unit cell yield bragg edges between the maximum and minimum wavelength range. \n",
    "#  TODO. should probably calculate which are visible from larger set.\n",
    "indices_FCC = [(1,1,1),(2,0,0),(2,2,0),(3,1,1)]\n",
    "\n",
    "def create_Braggedge_list(lattice_constant, miller_indices):\n",
    "    '''\n",
    "    :param miller-indices: like [(1,1,0),(2,0,0),...]\n",
    "    :type miller-indices: list of tuples\n",
    "    '''\n",
    "    coords = [str((h, k, l)) for h, k, l in miller_indices]    \n",
    "    interplanar_distances = [2.*lattice_constant/np.sqrt(h**2+k**2+l**2) for h, k, l in miller_indices]\n",
    "\n",
    "    d = sc.DataArray(sc.Variable(dims=[\"bragg-edge\"], values=np.array(interplanar_distances),\n",
    "                                  unit=sc.units.angstrom))\n",
    "    d.coords[\"miller-indices\"] = sc.Variable(dims=[\"bragg-edge\"], values=coords)\n",
    "    return d\n",
    "\n",
    "\n",
    "Bragg_edges_FCC = create_Braggedge_list(FCC_a, indices_FCC)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Bragg-edge fitting\n",
    "\n",
    "Fit Bragg-edge peaks for the sample initially. Then take this value to find the position within a window\n",
    "for the elastic sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\n",
    "def Bragg_edge_position(xpos_Bragg_edge,x_min_sides,x_max_sides,x_size,y_size,plots):\n",
    "    '''Takes limits around a roughly known Bragg edge and fits the Bragg edge to obtain\n",
    "    its exact position. Needs the size (x_size, y_size) and the workspaces containing histogram arrays/matrices.\n",
    "    x_min_side and x_max_side are arbitrary values and are calculated:\n",
    "    bragg_edge_pos * x_min_side or bragg_edge_pos * x_max_side.\n",
    "    '''\n",
    "\n",
    "    fit_list = []\n",
    "    spectrum_list_fitted_sample = []\n",
    "    spectrum_list_fitted_sample_elastic = []\n",
    "    elastic_ws = grouped_sample_elastic\n",
    "    sample_ws = sample_final\n",
    "\n",
    "    # loop for each Bragg edge in the list, find edge positions and save spectrums in lists\n",
    "    for edge_index in range(xpos_Bragg_edge.shape[0]):\n",
    "        fit_list.append([])\n",
    "        spectrum_list_fitted_sample.append([])\n",
    "        spectrum_list_fitted_sample_elastic.append([])\n",
    "\n",
    "        bragg_edge = xpos_Bragg_edge.coords[\"miller-indices\"][\"bragg-edge\", edge_index]\n",
    "\n",
    "        xpos_guess = xpos_Bragg_edge[\"bragg-edge\", edge_index].value\n",
    "        x_min_fit = xpos_guess - xpos_guess * abs(x_min_sides[edge_index])\n",
    "        x_max_fit = xpos_guess + xpos_guess * abs(x_max_sides[edge_index])\n",
    "\n",
    "        print(\"Now fitting Bragg edge {} at {:.3f} A (between {:.3f} A and {:.3f} A) across image groups\"\n",
    "              .format(bragg_edge, xpos_guess, x_min_fit, x_max_fit))\n",
    "\n",
    "        # if the full inital sample was taken and no grouping was done\n",
    "        #Fitting the masked sample\n",
    "        params_s, diff_s = sc.compat.mantid.fit(sample_ws, mantid_args={\n",
    "                                      'Function':f'name=LinearBackground,A0={230},A1={-4};name=UserFunction,Formula=h*erf(a*(x-x0)),h={16},a={-11},x0={xpos_guess}',\n",
    "                                      'StartX':x_min_fit, 'EndX':x_max_fit})\n",
    "\n",
    "        v_and_var_s=[params_s.data['parameter', i] for i in range(params_s['parameter',:].shape[0])] \n",
    "        params = dict(zip(params_s.coords[\"parameter\"].values, v_and_var_s))\n",
    "        d_sample = params[\"f1.x0\"] / 2.0 # See fit table definition for extract x0\n",
    "       \n",
    "        plot_cell_idx = 0 \n",
    "        for row in range(y_size):\n",
    "            for col in range(x_size):\n",
    "                spectrum_idx = int(plots[plot_cell_idx][0])\n",
    "                \n",
    "                # Fit Bragg edge using values from fit of unstrained sample\n",
    "                elastic_function = f\"name=LinearBackground,A0={params['f0.A0'].value},A1={params['f0.A1'].value};name=UserFunction,Formula=h*erf(a*(x-x0)),h={params['f1.h'].value},a={params['f1.a'].value},x0={params['f1.x0'].value}\"\n",
    "                params_e, diff_e = sc.compat.mantid.fit(elastic_ws['spectrum_mapping', spectrum_idx], mantid_args={'Function':elastic_function,\n",
    "                                                   'StartX':x_min_fit, 'EndX':x_max_fit})\n",
    "                \n",
    "                v_and_var_e=[params_e.data['parameter', i] for i in range(params_e['parameter',:].shape[0])] \n",
    "                elastic_params = dict(zip(params_e.coords[\"parameter\"].values, v_and_var_e))\n",
    "                d_sample_elastic = elastic_params[\"f1.x0\"] / 2.0 # See fit table definition for extract x0\n",
    "                lattice_strain = d_sample_elastic - d_sample\n",
    "\n",
    "                # define successful fitting\n",
    "                success = (params_s.coords[\"status\"].value == \"success\") and (params_e.coords[\"status\"].value == \"success\")\n",
    "\n",
    "                # fitted values STORED in list\n",
    "                fit_list[edge_index].append((spectrum_idx, (row, col), bragg_edge, d_sample, d_sample_elastic, lattice_strain, success))\n",
    "\n",
    "                # workspace created from sample fit, workspace index 1 for fitted spectrum (index 0 for data, index 2 for difference curve)\n",
    "                fitted_sample = diff_s['calculated'].values\n",
    "\n",
    "                if params_s.coords[\"status\"].value == \"success\":\n",
    "                    spectrum_list_fitted_sample[edge_index].append((diff_s.coords[\"wavelength\"].values, fitted_sample))\n",
    "                else:\n",
    "                    spectrum_list_fitted_sample[edge_index].append((diff_s.coords[\"wavelength\"].values,\n",
    "                                                           np.zeros_like(fitted_sample)))\n",
    "                    \n",
    "                fitted_elastic = diff_e['calculated'].values\n",
    "                # workspace created from fit of sample under elastic deformation\n",
    "                if params_e.coords[\"status\"].value == \"success\":\n",
    "                    spectrum_list_fitted_sample_elastic[edge_index].append((diff_e.coords[\"wavelength\"].values, fitted_elastic))\n",
    "                else:\n",
    "                    spectrum_list_fitted_sample_elastic[edge_index].append((diff_e.coords[\"wavelength\"].values,\n",
    "                                                                                               np.zeros_like(fitted_elastic)))\n",
    "                plot_cell_idx +=1\n",
    "\n",
    "    return fit_list, spectrum_list_fitted_sample, spectrum_list_fitted_sample_elastic\n",
    "\n",
    "fit_list, spectrum_list_fitted_sample, spectrum_list_fitted_sample_elastic = \\\n",
    "    Bragg_edge_position(xpos_Bragg_edge=Bragg_edges_FCC,\n",
    "                                      x_min_sides=[0.05, 0.1, 0.1, 0.05],#[0.1, 0.1, 0.1, 0.05]\n",
    "                                      x_max_sides=[0.1, 0.05, 0.1, 0.1],#[0.1, 0.1, 0.1, 0.1]\n",
    "                                      x_size=nx_target,\n",
    "                                      y_size=ny_target, \n",
    "                                      plots=plots)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Generate single tile plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def tileplot_colorcode(fit_list, nx_target, ny_target, outlier_threshold, output_filename_stem='tileplot_color'):\n",
    "    # for each part of the fit list containing the fit values of each Bragg edge, create a color plot.\n",
    "    output_filename_list = []\n",
    "    for fit_be_list in fit_list:\n",
    "        bragg_edge = fit_be_list[0][2].values  # the calculated Bragg edge TOF position\n",
    "        d_spacing = fit_be_list[0][3] # the calculated Bragg edge TOF position\n",
    "        output_filename = '{}_{:.3f}A_{}'.format(output_filename_stem, d_spacing.value, bragg_edge)\n",
    "        # Make a 2D image of the strain values\n",
    "        print(('Plotting color-coded tile plot of Bragg edge {}.'.format(bragg_edge)))\n",
    "        fig, ax = plt.subplots()\n",
    "        strains = np.zeros([ny_target, nx_target])\n",
    "        plots_counter = 0\n",
    "        for row in range(ny_target):\n",
    "            for col in range(nx_target):\n",
    "                if (fit_be_list[plots_counter][-1] and outlier_threshold[0] < fit_be_list[plots_counter][-2].value < outlier_threshold[1]):\n",
    "                    strains[row, col] = fit_be_list[plots_counter][-2].value\n",
    "                plots_counter += 1\n",
    "        import matplotlib.colors as colors\n",
    "        im = ax.imshow(strains, origin=\"upper\", norm=colors.SymLogNorm(linthresh=1.0e-3), cmap=\"RdBu\")\n",
    "        cb = plt.colorbar(im)\n",
    "        cb.set_label(\"${}$ Lattice strain $\\\\varepsilon$\".format(bragg_edge))\n",
    "        print(f'Saving {output_filename}.pdf and .png.')\n",
    "        fig.savefig(\"{}.pdf\".format(output_filename), bbox_inches=\"tight\")\n",
    "        fig.savefig(\"{}.png\".format(output_filename), bbox_inches=\"tight\")\n",
    "        output_filename_list.append(output_filename)\n",
    "    return output_filename_list\n",
    "\n",
    "\n",
    "\n",
    "# i.e. -1e-2, 1e-2 means only -1% to 1% values of lattice strain are shown\n",
    "outlier_threshold_color_plot = (-5e-2, 5e-2)\n",
    "\n",
    "output_filename_tileplot_color ='{:03d}_tileplot_color_{}xy_{}usbin_{}thresh_{:.3f}-{:.3f}plotthresh_initintegr'.format(\n",
    "                        measurement_number,grouping_number,bin_width,masking_threshold,\n",
    "                        outlier_threshold_color_plot[0], outlier_threshold_color_plot[1])\n",
    "\n",
    "fignames_col = tileplot_colorcode(fit_list=fit_list, nx_target=nx_target, ny_target=ny_target,\n",
    "                                  outlier_threshold=outlier_threshold_color_plot,\n",
    "                                  output_filename_stem=os.path.join(output_dir, output_filename_tileplot_color))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Generate tile plots for single fig."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def edges_to_centers(x):\n",
    "    \"\"\"\n",
    "    Convert array edges to centers\n",
    "    \"\"\"\n",
    "    return 0.5 * (x[1:] + x[:-1])\n",
    "\n",
    "def tileplot(x_min_plot,x_max_plot,y_min_plot,y_max_plot,x_size,y_size,\n",
    "             histogram_matrices,fitted_data, plots, fit_list, output_filename='tileplot_curves'):\n",
    "\n",
    "    fig, ax = plt.subplots(x_size, y_size, figsize=(12, 12))\n",
    "\n",
    "    labels = [\"Without load\", \"Elastic deformation\"]\n",
    "    markers = [\"o\", \"s\"]\n",
    "    # Plot some fake data so that we can have larger markers in the legend\n",
    "    ax[0][0].plot([0, 1], [0, 1], ls=\"None\", marker=markers[0], markersize=4, label=labels[0], color=\"C0\")\n",
    "    ax[0][0].plot([0, 1], [0, 1], ls=\"None\", marker=markers[1], markersize=4, label=labels[1], color=\"C1\")\n",
    "\n",
    "    # All bins are identical in a dataset\n",
    "    x0_nostrain = edges_to_centers(histogram_matrices[0].coords[\"wavelength\"].values)\n",
    "\n",
    "    # x and y values of inital state, samples without strain\n",
    "    y0_nostrain = histogram_matrices[0].values\n",
    "\n",
    "    plots_counter = 0\n",
    "    for row in range(y_size):\n",
    "        for col in range(x_size):\n",
    "            # for running workspace index\n",
    "            ws_index = plots[plots_counter][0]\n",
    "\n",
    "            # if the whole inital-sample-state is just one spectrum, always plot that.\n",
    "            # Plot no-strain curve, inital state of sample\n",
    "            ax[row][col].plot(x0_nostrain, y0_nostrain, ls=\"None\", marker= \"o\", markersize=1)\n",
    "\n",
    "            # Plot curves of other states given in histogram_matrices\n",
    "            for k in range(len(histogram_matrices)):\n",
    "                    x0 = edges_to_centers(histogram_matrices[1].coords[\"wavelength\"].values)\n",
    "                    y0 = histogram_matrices[1][\"spectrum_mapping\", int(ws_index)].values\n",
    "                    ax[row][col].plot(x0, y0, ls=\"None\", marker=markers[k], markersize=1)\n",
    "\n",
    "            # plotting the fitted spectra:\n",
    "            x1 = edges_to_centers(fitted_data[0][plots_counter][0])\n",
    "            x2 = edges_to_centers(fitted_data[1][plots_counter][0])\n",
    "            ax[row][col].plot(x1, fitted_data[0][plots_counter][1], label=\"{} fit\".format(labels[0]), lw=1, zorder=3)\n",
    "            ax[row][col].plot(x2, fitted_data[1][plots_counter][1], label=\"{} fit\".format(labels[1]), lw=1, zorder=3)\n",
    "\n",
    "            # Set axis limits\n",
    "            ax[row][col].set_xlim([x_min_plot, x_max_plot])\n",
    "            ax[row][col].set_ylim([y_min_plot, y_max_plot])\n",
    "            if col == x_size - 1:\n",
    "                ax[row][col].yaxis.tick_right()\n",
    "            else:\n",
    "                ax[row][col].set_yticklabels([])\n",
    "            if row == 0:\n",
    "                ax[row][col].xaxis.tick_top()\n",
    "            else:\n",
    "                ax[row][col].set_xticklabels([])\n",
    "            ax[row][col].tick_params(axis=\"x\", direction=\"in\", bottom=True, top=True)\n",
    "            ax[row][col].tick_params(axis=\"y\", direction=\"in\", left=True, right=True)\n",
    "\n",
    "            if row == y_size - 1:\n",
    "                ax[row][col].text(0.0, -0.1, \"{:.1f}\".format(0.1 * col),\n",
    "                                ha='center',va='top', transform=ax[row][col].transAxes)\n",
    "            if col == 0:\n",
    "                ax[row][col].text(-0.1, 0.0, \"{:.1f}\".format(0.1*(y_size - row - 1)),\n",
    "                                ha='right',va='center', transform=ax[row][col].transAxes)\n",
    "\n",
    "            if fit_list[plots_counter][-1]: # usual sign for lattice strain is \\\\varepsilon\n",
    "                ax[row][col].text(0.1, 0.8, \"${:.3f}$\".format(fit_list[plots_counter][-2].value),\n",
    "                                ha='left',va='top', transform=ax[row][col].transAxes, fontsize=4)\n",
    "            plots_counter += 1\n",
    "\n",
    "    # Last bin edge for xy coordinates\n",
    "    ax[-1][-1].text(1.0, -0.1, \"{:.1f}\".format(0.1 * x_size),\n",
    "                    ha='center',va='top', transform=ax[-1][-1].transAxes)\n",
    "    ax[0][0].text(-0.1, 1.0, \"{:.1f}\".format(0.1 * y_size),\n",
    "                  ha='right',va='center', transform=ax[0][0].transAxes)\n",
    "    # Display only one legend\n",
    "    ax[0][0].legend(loc=(0, 1.8), ncol=4)\n",
    "    fig.text(0.5, 0.905, \"Wavelength $[\\mathrm{\\AA}]$\")\n",
    "    fig.text(0.94, 0.5, \"Counts\", rotation=90, ha='center',va='center')\n",
    "    fig.text(0.5, 0.075, \"X position [m]\")\n",
    "    fig.text(0.09, 0.5, \"Y position [m]\", rotation=90, ha='center',va='center')\n",
    "    # Remove white space between subplots\n",
    "    fig.subplots_adjust(wspace=0.0, hspace=0.0)\n",
    "    print(('Saving {}.pdf.').format(output_filename))\n",
    "    fig.savefig(\"{}.pdf\".format(output_filename), bbox_inches=\"tight\")\n",
    "    fig.savefig(\"{}.png\".format(output_filename), bbox_inches=\"tight\")\n",
    "    return output_filename\n",
    "\n",
    "output_filename_tileplot = '{:03d}_tileplot_curves_{}xy_{}usbin_{}thresh_initintegr'\\\n",
    "    .format(measurement_number,grouping_number,bin_width,\n",
    "            masking_threshold)\n",
    "\n",
    "# Plotting: y_min_plot=200,y_max_plot=270 for 18x18, y_min_plot=90,y_max_plot=120 for 27x27.\n",
    "plotting_win = [(i*(1-.135), i*(1+.135)) for i in Bragg_edges_FCC.values]\n",
    "for i in range(Bragg_edges_FCC.values.shape[0]):\n",
    "    angstrom_val = Bragg_edges_FCC[\"bragg-edge\", i].value / 2.\n",
    "    lattice_indicies = Bragg_edges_FCC.coords[\"miller-indices\"][\"bragg-edge\", i].value\n",
    "    output_filename_tileplot_new = '{}_{:.3f}A_({})'.format(output_filename_tileplot, angstrom_val, lattice_indicies)\n",
    "\n",
    "    win = plotting_win[i]\n",
    "    fignames_curve = tileplot(\n",
    "        x_min_plot=win[0], x_max_plot=win[1],\n",
    "        y_min_plot=65610./(nx_target*ny_target), y_max_plot=87480./(nx_target*ny_target), # remove 0.9 later\n",
    "        x_size=nx_target, y_size=ny_target,\n",
    "        histogram_matrices=histogram_matrices,\n",
    "        fitted_data=[spectrum_list_fitted_sample[i], spectrum_list_fitted_sample_elastic[i]],\n",
    "        plots=plots, fit_list=fit_list[i], output_filename=os.path.join(output_dir, output_filename_tileplot_new))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}