06 Nanocrystal segmentation in SPED data - Demonstration on partly overlapping MgO cubes.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook demonstrates two approaches to nanocrystal segmentation:\n",
    "1. Virtual dark-field (VDF) imaging-based segmentation\n",
    "2. Non-negative matrix factorisation (NMF)-based segmentation\n",
    "\n",
    "The segmentation is demonstrated on a SPED dataset of partly overlapping MgO nanoparticles, where some of the particles share the same orientation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This functionality has been checked to run with pyxem-0.11.0 (May 2020). Bugs are always possible, do not trust the code blindly, and if you experience any issues please report them here: https://github.com/pyxem/pyxem-demos/issues"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Contents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. <a href='#gen'> Setting up, Loading Data, Pre-processing</a>\n",
    "2. <a href='#vdf'> Virtual Image Based Segmentation</a>\n",
    "3. <a href='#nmf'> NMF Based Segmentation</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='gen'></a> 1. Setting up, Loading Data, Pre-processing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import pyxem and other required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib qt\n",
    "import numpy as np\n",
    "import hyperspy.api as hs\n",
    "import matplotlib.pyplot as plt\n",
    "import pyxem as pxm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load demonstration data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dp = hs.load('./data/06/mgo_nanoparticles.hdf5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot data to inspect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dp.plot(cmap='magma_r')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remove the background"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma_min = 1.7\n",
    "sigma_max = 13.2\n",
    "\n",
    "dp_rb = dp.subtract_diffraction_background('difference of gaussians', \n",
    "                             min_sigma=sigma_min, \n",
    "                             max_sigma=sigma_max)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the background subtracted data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dp_rb.plot(cmap='magma_r')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Find the position of the direct beam in the background subtracted data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shifts = dp_rb.center_direct_beam(method='cross_correlate',\n",
    "                                  half_square_width=15,\n",
    "                                  return_shifts=True,\n",
    "                                  radius_start=2,\n",
    "                                  radius_finish=6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this step is likely to be quiet slow\n",
    "shifts.compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Apply the same shifts to the raw data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dp.align2D(shifts=shifts, crop=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set calibrations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scale = 0.03246 \n",
    "scale_real = 3.03\n",
    "dp.set_diffraction_calibration(scale)\n",
    "dp.set_scan_calibration(scale_real)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# see https://github.com/pyxem/pyxem/issues/717 for why have to do this\n",
    "dp_rb.compute()\n",
    "dp_rb = pxm.signals.ElectronDiffraction2D(dp_rb)\n",
    "\n",
    "dp_rb.set_diffraction_calibration(scale)\n",
    "dp_rb.set_scan_calibration(scale_real)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='vdf'></a> 2. Virtual Image Based Segmentation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.1. Peak Finding & Refinement"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Find all diffraction peaks for all PED patterns. \n",
    "The parameters were found by interactive peak finding:\n",
    "\n",
    "`peaks = dp_rb.find_peaks_interactive(imshow_kwargs={'cmap': 'magma_r'})`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#This is broken, see https://github.com/hyperspy/hyperspy/issues/2622\n",
    "peaks = dp_rb.find_peaks(method='laplacian_of_gaussian', \n",
    "                         min_sigma=0.7,\n",
    "                         max_sigma=10,\n",
    "                         num_sigma=30, \n",
    "                         threshold=0.046, \n",
    "                         overlap=0.5, \n",
    "                         log_scale=False,\n",
    "                         exclude_border=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualise the number of diffraction peaks found at each probe position"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diff_map = peaks.get_diffracting_pixels_map()\n",
    "diff_map.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Exclude peaks too close to the detector edge for sub-pixel refinement. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "peaks_filtered = peaks.filter_detector_edge(exclude_width=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Refine the peak positions using center of mass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from pyxem.generators.subpixelrefinement_generator import SubpixelrefinementGenerator\n",
    "from pyxem.signals.diffraction_vectors import DiffractionVectors\n",
    "\n",
    "\n",
    "refine_gen = SubpixelrefinementGenerator(dp_rb, peaks_filtered)\n",
    "\n",
    "peaks_refined = DiffractionVectors(refine_gen.center_of_mass_method(square_size=4))\n",
    "\n",
    "peaks_refined.axes_manager.set_signal_dimension(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.2. Determine Unique Peaks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Determine the unique diffraction peaks by clustering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "distance_threshold = scale*0.89\n",
    "min_samples = 10\n",
    "\n",
    "unique_peaks = peaks_refined.get_unique_vectors(method='DBSCAN',\n",
    "                                                distance_threshold=distance_threshold,\n",
    "                                                min_samples=min_samples)\n",
    "print(np.shape(unique_peaks.data)[0], ' unique vectors were found.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualise the detected unique peaks by plotting them on the maximum of the signal. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "radius_px = dp_rb.axes_manager.signal_shape[0]/2\n",
    "reciprocal_radius = radius_px * scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "unique_peaks.plot_diffraction_vectors(\n",
    "    method='DBSCAN',\n",
    "    unique_vectors=unique_peaks,\n",
    "    distance_threshold=distance_threshold,\n",
    "    xlim=reciprocal_radius,\n",
    "    ylim=reciprocal_radius,\n",
    "    min_samples=min_samples,\n",
    "    image_to_plot_on=dp_rb.max(),\n",
    "    image_cmap='magma_r',\n",
    "    plot_label_colors=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualise both the clusters and the unique peaks obtained after DBSCAN clustering. \n",
    "\n",
    "*NB The cluster colors are randomly generated, so run it again if it is hard to discern two close clusters.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "peaks_refined.plot_diffraction_vectors(\n",
    "    method='DBSCAN',\n",
    "    xlim=reciprocal_radius, \n",
    "    ylim=reciprocal_radius,\n",
    "    unique_vectors=unique_peaks, \n",
    "    distance_threshold=distance_threshold,\n",
    "    min_samples=min_samples, \n",
    "    image_to_plot_on=dp_rb.max(), \n",
    "    image_cmap='gray_r',\n",
    "    plot_label_colors=True, \n",
    "    distance_threshold_all=scale*0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Filter the unique vectors by magnitude in order to exclude the direct beam from the following analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Gs = unique_peaks.filter_magnitude(min_magnitude=10*scale,\n",
    "                                   max_magnitude=np.inf)\n",
    "print(np.shape(Gs)[0], ' unique vectors.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the unique vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Gs.plot_diffraction_vectors(unique_vectors=Gs,\n",
    "                            distance_threshold=distance_threshold,\n",
    "                            xlim=reciprocal_radius,\n",
    "                            ylim=reciprocal_radius,\n",
    "                            min_samples=min_samples,\n",
    "                            image_to_plot_on=dp_rb.max(),\n",
    "                            image_cmap='magma',\n",
    "                            plot_label_colors=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Optionally save and load the unique peaks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save('peaks.npy', Gs.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Gs = np.load('peaks.npy', allow_pickle=True)\n",
    "Gs = pxm.DiffractionVectors(Gs)\n",
    "Gs.axes_manager.set_signal_dimension(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.3. Virtual Imaging & Segmentation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate VDF images for all unique peaks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyxem.generators.vdf_generator import VDFGenerator\n",
    "\n",
    "radius=scale*2\n",
    "\n",
    "vdfgen = VDFGenerator(dp_rb, Gs)\n",
    "VDFs = vdfgen.get_vector_vdf_images(radius=radius)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the VDF images for inspection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "VDFs.plot(cmap='magma', scalebar=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First find adequate parameters by looking at watershed segmentation of a single VDF image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyxem.utils.segment_utils import separate_watershed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "min_distance = 5.5\n",
    "min_size = 10\n",
    "max_size = 1000\n",
    "max_number_of_grains = 1000\n",
    "marker_radius = 2\n",
    "exclude_border = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 25\n",
    "sep_i = separate_watershed(\n",
    "    VDFs.inav[i].data, min_distance=min_distance, min_size=min_size,\n",
    "    max_size=max_size, max_number_of_grains=max_number_of_grains,\n",
    "    exclude_border=exclude_border, marker_radius=marker_radius,\n",
    "    threshold=True, plot_on=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform segmentation on all the VDF images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "segs = VDFs.get_vdf_segments(min_distance=min_distance,\n",
    "                             min_size=min_size,\n",
    "                             max_size = max_size,\n",
    "                             max_number_of_grains = max_number_of_grains,\n",
    "                             exclude_border=exclude_border,\n",
    "                             marker_radius=marker_radius,\n",
    "                             threshold=True)\n",
    "\n",
    "print(np.shape(segs.segments)[0],' segments were found.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the segments for inspection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "segs.segments.plot(cmap='magma_r')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate normalised cross-correlations between all VDF image segments to identify those that are related to the same crystal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ncc_vdf = segs.get_ncc_matrix()\n",
    "ncc_vdf.plot(scalebar=False, cmap='RdBu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the correlation value exceeds *corr_threshold* for certain segments, those segments are summed. These segments are discarded if the number of these segments are below *vector_threshold*, as this number corresponds to the number of detected diffraction peaks associated with the single crystal. The *vector_threshold* criteria is included to avoid including segment images resulting from noise or incorrect segmentation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "corr_threshold=0.7\n",
    "vector_threshold=5\n",
    "segment_threshold=4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "corrsegs = segs.correlate_vdf_segments(\n",
    "    corr_threshold=corr_threshold, vector_threshold=vector_threshold,\n",
    "    segment_threshold=segment_threshold)\n",
    "print(np.shape(corrsegs.segments)[0],' correlated segments were found.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Simulate virtual diffraction patterns for each summed segment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma = scale*1.5\n",
    "\n",
    "virtual_sig = corrsegs.get_virtual_electron_diffraction(\n",
    "    calibration=scale, shape=(int(radius_px*2), int(radius_px*2)), sigma=sigma)\n",
    "virtual_sig.set_diffraction_calibration(scale)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the final results from the VDF image-based segmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "hs.plot.plot_images(corrsegs.segments, cmap='magma_r', axes_decor='off',\n",
    "                    per_row=np.shape(corrsegs.segments)[0],\n",
    "                    suptitle='', scalebar=False, scalebar_color='white',\n",
    "                    colorbar=False,\n",
    "                    padding={'top': 0.95, 'bottom': 0.05,\n",
    "                             'left': 0.05, 'right':0.78})\n",
    "hs.plot.plot_images(virtual_sig, cmap='magma_r', axes_decor='off',\n",
    "                    per_row=np.shape(corrsegs.segments)[0],\n",
    "                    suptitle='', scalebar=False, scalebar_color='white',\n",
    "                    colorbar=False,\n",
    "                    padding={'top': 0.95, 'bottom': 0.05,\n",
    "                             'left': 0.05, 'right': 0.78})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='nmf'></a> 3. NMF Based Segmentation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3.1. NMF Decomposition"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a signal mask so that the region in the centre of each PED pattern, including the direct beam, can be excluded in the machine learning. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dpm = pxm.Diffraction2D(dp.inav[0,0])\n",
    "signal_mask = dpm.get_direct_beam_mask(radius=10)\n",
    "signal_mask.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform single value decomposition (SVD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dp.change_dtype('float32')\n",
    "dp.decomposition(algorithm='svd',\n",
    "                 normalize_poissonian_noise=True,\n",
    "                 centre='variables',\n",
    "                 signal_mask=signal_mask.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dp.plot_decomposition_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Investigate the scree plot and use it as a guide to determine the number of components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_comp=11\n",
    "\n",
    "ax = dp.plot_explained_variance_ratio(n=200, threshold=num_comp,\n",
    "                                      hline=True, xaxis_labeling='ordinal',\n",
    "                                      signal_fmt={'color':'k', 'marker':'.'}, \n",
    "                                      noise_fmt={'color':'gray', 'marker':'.'})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform NMF decomposition with specified number of components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dp.decomposition(normalize_poissonian_noise=True,\n",
    "                 algorithm='nmf',\n",
    "                 output_dimension=num_comp,\n",
    "                 signal_mask=signal_mask.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dp_nmf = dp.get_decomposition_model(components=np.arange(num_comp))\n",
    "factors = dp_nmf.get_decomposition_factors()\n",
    "loadings = dp_nmf.get_decomposition_loadings()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the NMF results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hs.plot.plot_images(loadings, cmap='magma_r', axes_decor='off', per_row=11,\n",
    "             suptitle='', scalebar=False, scalebar_color='white', colorbar=False,\n",
    "             padding={'top': 0.95, 'bottom': 0.05,\n",
    "                      'left': 0.05, 'right':0.78})\n",
    "hs.plot.plot_images(factors, cmap='magma_r', axes_decor='off', per_row=11,\n",
    "             suptitle='', scalebar=False, scalebar_color='white', colorbar=False,\n",
    "             padding={'top': 0.95, 'bottom': 0.05,\n",
    "                      'left': 0.05, 'right':0.78})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Discard the components related to background (\\#0) and to the carbon film (\\#4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hyperspy.signals import Signal2D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "factors = Signal2D(np.delete(factors.data, [0, 4], axis = 0))\n",
    "loadings = Signal2D(np.delete(loadings.data, [0, 4], axis = 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hs.plot.plot_images(factors, cmap='magma_r', axes_decor='off',\n",
    "                    per_row=9, suptitle='', scalebar=False,\n",
    "                    scalebar_color='white', colorbar=False,\n",
    "                    padding={'top': 0.95, 'bottom': 0.05,\n",
    "                             'left': 0.05, 'right':0.78})\n",
    "\n",
    "hs.plot.plot_images(loadings, cmap='magma_r', axes_decor='off',\n",
    "                    per_row=9, suptitle='', scalebar=False,\n",
    "                    scalebar_color='white', colorbar=False,\n",
    "                    padding={'top': 0.95, 'bottom': 0.05,\n",
    "                             'left': 0.05, 'right':0.78})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.2. Correlate NMF Loading Maps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NMF often leads to splitting of some crystals into several components. Therefore the correlation between loadings and between component patterns are calculated, and if both the correlation values for loadings and factors exceed threshold values, those loadings and factors are summed. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate the matrix of normalised cross-correlation for both the loadings and patterns first, to find suitable correlation threshold values. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyxem.signals.segments import LearningSegment\n",
    "learn = LearningSegment(factors=factors, loadings=loadings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ncc_nmf = learn.get_ncc_matrix()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ncc_nmf.plot(scalebar=False, cmap='RdBu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "corr_th_factors = 0.45\n",
    "corr_th_loadings = 0.3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform correlation and summation of the factors and loadings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn_corr = learn.correlate_learning_segments(corr_th_factors=corr_th_factors,\n",
    "                                               corr_th_loadings=corr_th_loadings)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the NMF reuslts after correlation and summation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hs.plot.plot_images(learn_corr.loadings, cmap='magma_r', axes_decor='off',\n",
    "                    per_row=7, suptitle='', scalebar=False,\n",
    "                    scalebar_color='white', colorbar=False,\n",
    "                    padding={'top': 0.95, 'bottom': 0.05,\n",
    "                             'left': 0.05, 'right':0.78})\n",
    "hs.plot.plot_images(learn_corr.factors, cmap='magma_r', axes_decor='off',\n",
    "                    per_row=7, suptitle='', scalebar=False,\n",
    "                    scalebar_color='white', colorbar=False,\n",
    "                    padding={'top': 0.95, 'bottom': 0.05,\n",
    "                             'left': 0.05, 'right':0.78})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First investigate how the parameters influence the segmentation on\n",
    "one single loading map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyxem.utils.segment_utils import separate_watershed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "min_distance = 10\n",
    "min_size = 50\n",
    "max_size = 100000\n",
    "max_number_of_grains = 100000\n",
    "marker_radius = 2\n",
    "exclude_border = 1\n",
    "threshold = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i =1\n",
    "sep_i = separate_watershed(\n",
    "    learn_corr.loadings.data[i], min_distance=min_distance,\n",
    "    min_size=min_size, max_size=max_size, \n",
    "    max_number_of_grains=max_number_of_grains,\n",
    "    exclude_border=exclude_border, \n",
    "    marker_radius=marker_radius, threshold=True, plot_on=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set a threshold for the minimum intensity value that a loading segment must contain in order to be kept. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "min_intensity_threshold = 10000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn_corr_seg = learn_corr.separate_learning_segments(\n",
    "    min_intensity_threshold=min_intensity_threshold,\n",
    "    min_distance = min_distance, min_size = min_size,\n",
    "    max_size = max_size, \n",
    "    max_number_of_grains = max_number_of_grains,\n",
    "    exclude_border = exclude_border,\n",
    "    marker_radius = marker_radius, threshold = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the final results from the NMF-based segmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hs.plot.plot_images(learn_corr_seg.loadings, \n",
    "                    cmap='magma_r', axes_decor='off',\n",
    "                    per_row=10, suptitle='', scalebar=False,\n",
    "                    scalebar_color='white', colorbar=False,\n",
    "                    padding={'top': 0.95, 'bottom': 0.05,\n",
    "                             'left': 0.05, 'right':0.78})\n",
    "\n",
    "hs.plot.plot_images(learn_corr_seg.factors, \n",
    "                    cmap='magma_r', axes_decor='off',\n",
    "                    per_row=10, suptitle='', scalebar=False,\n",
    "                    scalebar_color='white', colorbar=False,\n",
    "                    padding={'top': 0.95, 'bottom': 0.05,\n",
    "                             'left': 0.05, 'right':0.78})"
   ]
  },
  {
   "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}