#!/usr/bin/env python
# coding: utf-8

# This notebook is part of the `kikuchipy` documentation https://kikuchipy.org.
# Links to the documentation won't work from the notebook.

# # Geometrical EBSD simulations
# 
# This section details how to inspect and visualize the results from pattern
# matching or Hough indexing **of cubic crystals** by plotting Kikuchi bands and
# zone axes onto an EBSD signal. We consider this a geometrical EBSD simulation,
# since it's only the Kikuchi band centres and zone axis positions that will be
# computed. These simulations are based on the work by Aimo Winkelmann in the
# supplementary material to
# <cite data-cite="britton2016tutorial">Britton et al. (2016)</cite>.
# 
# Let's import the necessary libraries and a small (3, 3) Nickel EBSD test data
# set

# In[ ]:


# Exchange inline for notebook or qt5 (from pyqt) for interactive plotting
get_ipython().run_line_magic('matplotlib', 'inline')

import tempfile
from diffsims.crystallography import ReciprocalLatticePoint
import hyperspy.api as hs
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from orix import crystal_map, quaternion
import kikuchipy as kp


s = kp.data.nickel_ebsd_small()  # Use kp.load("data.h5") to load your own data
s


# Let's enhance the Kikuchi bands by removing the static and dynamic backgrounds

# In[ ]:


s.remove_static_background()
s.remove_dynamic_background()


# In[ ]:


_ = hs.plot.plot_images(
    s, axes_decor=None, label=None, colorbar=False, tight_layout=True
)


# To project Kikuchi bands and zone axis positions onto our detector, we need
# 1. a description of the crystal phase, in the geometrical case only the space
#    group
# 2. the set of Kikuchi bands to consider, e.g. the {111}, {200}, {220}, and {311}
#    crystal plane families
# 3. the crystal orientations with respect to the reference frame
# 4. the position of the detector with respect to the sample reference frame,
#    in the form of a sample-detector model which includes the sample and detector
#    tilt and the projection center (shortes distance from the source point on the
#    sample to the detector), given here as (PC$_x$, PC$_y$, PC$_z$)
# 
# We'll store the crystal phase information in an
# [orix.crystal_map.Phase](https://orix.readthedocs.io/en/stable/reference.html#orix.crystal_map.Phase)
# instance

# In[ ]:


phase = crystal_map.Phase(name="ni", space_group=225)
phase


# We'll set up the relevant Kikuchi bands (and the zone axes from these) in a
# [diffsims.crystallography.ReciprocalLatticePoint](https://diffsims.readthedocs.io/en/latest/reference.html#diffsims.crystallography.ReciprocalLatticePoint)
# instance

# In[ ]:


rlp = ReciprocalLatticePoint(
    phase=phase, hkl=[[1, 1, 1], [2, 0, 0], [2, 2, 0], [3, 1, 1]]
)
rlp


# We can get a new instance with the symmetrically equivalent planes in each
# plane family using
# [ReciprocalLatticePoint.symmetrise()](https://diffsims.readthedocs.io/en/latest/reference.html#diffsims.crystallography.ReciprocalLatticePoint.symmetrise)

# In[ ]:


rlp2 = rlp.symmetrise()
rlp2


# We know from [pattern matching](pattern_matching.ipynb) of these nine patterns,
# to about 7 500 dynamically simulated patterns of orientations uniformly
# distributed in the orientation space of the point group $m\bar{3}m$, that they
# come from two grains with orientations of about $(\phi_1, \Phi, \phi_2) =
# (80^{\circ}, 34^{\circ}, -90^{\circ})$ and $(\phi_1, \Phi, \phi_2) =
# (115^{\circ}, 27^{\circ}, -95^{\circ})$. We store these orientations in an
# [orix.quaternion.Rotation](https://orix.readthedocs.io/en/stable/reference.html#orix.quaternion.Rotation)
# instance

# In[ ]:


grain1 = [80, 34, -90]
grain2 = [115, 27, -95]
r = quaternion.Rotation.from_euler(np.deg2rad([
        [grain1, grain2, grain2],
        [grain1, grain2, grain2],
        [grain1, grain2, grain2]
]))
r


# We describe the sample-detector model in an
# [kikuchipy.detectors.EBSDDetector](../reference.rst#kikuchipy.detectors.EBSDDetector)
# instance. From Hough indexing we know the projection center to be, in the EDAX
# TSL convention (see the [reference frame](reference_frames.rst) guide for the
# various conventions and more details on the use of the sample-detector model),
# $(x^{*}, y^{*}, z^{*}) = (0.421, 0.7794, 0.5049)$. The sample was tilted
# $70^{\circ}$ about the microscope X direction towards the detector, and the
# detector normal was orthogonal to the optical axis (beam direction)

# In[ ]:


detector = kp.detectors.EBSDDetector(
    shape=s.axes_manager.signal_shape[::-1],
    sample_tilt=70,
    pc=[0.421, 0.7794, 0.5049],
    convention="tsl"
)
detector


# Note that the projection center gets converted internally to the Bruker
# convention.
# 
# Let's create an
# [EBSDSimulationGenerator](../reference.rst#kikuchipy.generators.EBSDSimulationGenerator)
# instance

# In[ ]:


sim_gen = kp.generators.EBSDSimulationGenerator(
    detector=detector,
    phase=phase,
    rotations=r
)
sim_gen


# Now we're ready to simulate geometrical EBSD patterns from the generator and the
# sets of crystal plane families

# In[ ]:


sim = sim_gen.geometrical_simulation(reciprocal_lattice_point=rlp2)
sim


# We see that 27 of the 50 Kikuchi bands we're visible on the detector in the nine
# patterns, stored in an instance of the
# [kikuchipy.simulations.features.KikuchiBand](../reference.rst#kikuchipy.simulations.features.KikuchiBand)
# class, which is a class inheriting from the
# [ReciprocalLatticePoint](https://diffsims.readthedocs.io/en/latest/reference.html#diffsims.crystallography.ReciprocalLatticePoint)

# In[ ]:


sim.bands


# In[ ]:


sim.zone_axes.size


# We also see that there are 91 zone axes resulting from the 27 Kikuchi bands,
# stored in an instance of the
# [kikuchipy.simulations.features.ZoneAxis](../reference.rst#kikuchipy.simulations.features.ZoneAxis)
# class, also inheriting from `ReciprocalLatticePoint`.

# We can now add these simulations as markers to be displayed on top of our
# experimental EBSD signal when plotting

# In[ ]:


markers = sim.as_markers(pc=False)
s.add_marker(marker=markers, plot_marker=False, permanent=True)


# The markers update with the pattern when navigating, thus helping us determine
# whether an indexing was successful, and in labeling the bands and zone axes in
# the pattern

# In[ ]:


s.plot(navigator=None)


# Whether to plot only bands, zone axes, zone axes labels, projection center, or
# all of them, can be set in the
# [GeometricalEBSDSimulation.as_markers()](../reference.rst#kikuchipy.simulations.GeometricalEBSDSimulation.as_markers)
# method. Their appearance on the pattern can also be controlled to some extent.
# The above method itself calls
# [bands_as_markers()](../reference.rst#kikuchipy.simulations.GeometricalEBSDSimulation.bands_as_markers),
# [pc_as_markers()](../reference.rst#kikuchipy.simulations.GeometricalEBSDSimulation.pc_as_markers),
# [zone_axes_as_markers()](../reference.rst#kikuchipy.simulations.GeometricalEBSDSimulation.zone_axes_as_markers),
# and [zone_axes_labels_as_markers()](../reference.rst#kikuchipy.simulations.GeometricalEBSDSimulation.zone_axes_labels_as_markers).
# See their documentation for available modifications.
# 
# Let's first remove the markers from the signal, and add only the bands and zone
# axes

# In[ ]:


del s.metadata.Markers
s.add_marker(
    marker=sim.as_markers(
        pc=False,
        zone_axes_labels=False,
        bands_kwargs=dict(
            family_colors=["w", "magenta", "cyan", "lime"], linestyle="--",
        ),
        zone_axes_kwargs=dict(
            marker="s", size=150, facecolor="none", edgecolor="w",
        ),
    ),
    plot_marker=False,
    permanent=True,
)


# In[ ]:


s.plot(navigator=None)


# In[ ]:


del s.metadata.Markers
s.add_marker(marker=markers, plot_marker=False, permanent=True)


# We can write single EBSD patterns with the markers on top to file

# In[ ]:


fig = s._plot.signal_plot.figure
bbox = matplotlib.transforms.Bbox.from_extents(
    np.array(fig.axes[0].bbox.extents) / 72  # The denominator may vary
)


# In[ ]:


nav_x, nav_y = s.axes_manager.indices
temp_dir = tempfile.mkdtemp() + "/"
fname = temp_dir + f"geosim_y{nav_y}_x{nav_x}.png"
s._plot.signal_plot.figure.savefig(fname, bbox_inches=bbox, dpi=150)


# In[ ]:


s.axes_manager.indices = (2, 1)
s.plot(navigator=None, colorbar=False, axes_off=True, title="")


# In[ ]:


nav_x, nav_y = s.axes_manager.indices
fname = temp_dir + f"geosim_y{nav_y}_x{nav_x}.png"
s._plot.signal_plot.figure.savefig(fname, bbox_inches=bbox, dpi=150)


# The coordinates of these bands and zone axes are available as class attributes.
# For the bands, we can e.g. extract the plane trace coordinates (y0, x0, y1, x1)
# in either gnomonic or detector coordinates (taking into account the detector
# size in pixels) for all bands or per navigation position

# In[ ]:


sim.bands[0, 0].plane_trace_coordinates[:10]  # Gnomonic


# In[ ]:


sim.bands_detector_coordinates[0, 0, :10]  # Detector


# The NaN values signify that that particular band is not visible on the detector
# in that position. The crystal plane normal of each band, pointing from the
# source point to the detector, is also available

# In[ ]:


sim.bands[1, 1].hkl_detector[:10]


# And where the vector hits the detector, in either detector or gnomonic
# coordinates

# In[ ]:


sim.bands[0, 0].x_detector


# In[ ]:


sim.bands[0, 0].x_gnomonic


# The same information is available for the zone axes

# In[ ]:


sim.zone_axes_label_detector_coordinates[0, 0][20:30]  # Detector


# In[ ]:


sim.zone_axes[0, 1].uvw_detector[:10]


# In[ ]:


sim.zone_axes[0, 0].y_detector[:10]


# In[ ]:


sim.zone_axes[0, 0].y_gnomonic[:10]


# With this information, it should be straight forward to go around kikuchipy
# when plotting and only use matplotlib

# In[ ]:


nav_idx = (2, 1)[::-1]

fig, ax = plt.subplots(figsize=(5, 5))
ax.imshow(s.inav[nav_idx], cmap="gray")

print(sim.bands_detector_coordinates.shape)
for i in np.ndindex(sim.bands_detector_coordinates.shape[2]):
    sim_slice = nav_idx + (i,)
    coords = sim.bands_detector_coordinates[sim_slice][0]
    y0, x0, y1, x1 = coords
    ax.axline((y0, x0), (y1, x1), linestyle="--", color="w")

print(sim.zone_axes_detector_coordinates.shape)
for j in np.ndindex(sim.zone_axes_detector_coordinates.shape[2]):
    sim_slice = nav_idx + (j,)
    coords = sim.zone_axes_detector_coordinates[sim_slice][0]
    x, y = coords
    ax.scatter(x=x, y=y, zorder=5, s=50)

_ = ax.axis((0, 59, 59, 0))


# In[ ]:


# Remove files written to disk in this user guide
import os
for file in ["geosim_y0_x0.png", "geosim_y1_x2.png"]:
    os.remove(temp_dir + file)
os.rmdir(temp_dir)