#!/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.

# # Kikuchi pattern simulations
# 
# This section explains how to inspect and visualize the results from EBSD
# indexing by plotting Kikuchi lines and zone axes onto an EBSD signal. We
# consider this a *geometrical* EBSD simulation, since it's only positions of
# Kikuchi lines and zone axes that are 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>.
# 
# We'll also show how to perform kinematical Kikuchi pattern simulations.
# 
# 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 diffpy.structure import Atom, Lattice, Structure
from diffsims.crystallography import ReciprocalLatticeVector
import hyperspy.api as hs
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from orix.crystal_map import Phase
from orix.quaternion import Rotation
import kikuchipy as kp
import pyvista


# Plotting parameters
plt.rcParams.update(
    {"figure.figsize": (10, 10), "font.size": 20, "lines.markersize": 10}
)
pyvista.global_theme.window_size = [700, 700]
pyvista.set_jupyter_backend("pythreejs")

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


# ## Geometrical simulations

# 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 lines and zone axes onto our detector, we need
# 
# 1. a description of the crystal phase
# 
# 2. the set of Kikuchi bands to consider, e.g. the sets of planes {111}, {200},
#    {220}, and {311}
# 
# 3. the crystal orientations with respect to the reference frame
# 
# 4. the position of the detector with respect to the sample, 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 = Phase(
    space_group=225,
    structure=Structure(
        atoms=[Atom("Ni", [0, 0, 0])], lattice=Lattice(3.52, 3.52, 3.52, 90, 90, 90)
    ),
)

print(phase)
print(phase.structure)


# We'll build up the reflector list using
# [diffsims.crystallography.ReciprocalLatticeVector](https://diffsims.readthedocs.io/en/latest/reference.html#diffsims.crystallography.ReciprocalLatticeVector)

# In[ ]:


ref = ReciprocalLatticeVector(
    phase=phase, hkl=[[1, 1, 1], [2, 0, 0], [2, 2, 0], [3, 1, 1]]
)
ref


# We'll obtain the symmetrically equivalent vectors and plot each family of
# vectors in a distinct colour in the stereographic projection

# In[ ]:


ref = ref.symmetrise().unique()
ref.size


# In[ ]:


ref.print_table()


# In[ ]:


# Dictionary with {hkl} as key and indices into `ref` as values
hkl_sets = ref.get_hkl_sets()
hkl_sets


# In[ ]:


hkl_colors = np.zeros((ref.size, 3))
for idx, color in zip(
    hkl_sets.values(),
    [[1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 1, 0]],  # Red, green, blue, yellow
):
    hkl_colors[idx] = color


# In[ ]:


hkl_labels = []
for hkl in ref.hkl.round(0).astype(int):
    hkl_labels.append(str(hkl).replace("[", "(").replace("]", ")"))


# In[ ]:


ref.scatter(c=hkl_colors, grid=True, ec="k", vector_labels=hkl_labels)


# We can also plot the plane traces, i.e. the Kikuchi lines, in both hemispheres
# (they are identical for Ni)

# In[ ]:


ref.draw_circle(
    color=hkl_colors, hemisphere="both", figure_kwargs=dict(figsize=(15, 10))
)


# We know from [pattern matching](pattern_matching.ipynb) of these nine patterns
# to dynamically simulated patterns of orientations uniformly
# distributed in the orientation space of the proper point group $432$, 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 = np.deg2rad((80, 34, -90))
grain2 = np.deg2rad((115, 27, -95))
rot = Rotation.from_euler(
    [[grain1, grain2, grain2], [grain1, grain2, grain2], [grain1, grain2, grain2]]
)
rot


# 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="edax",
)
detector


# Note that the projection center gets converted internally to the Bruker
# convention.
# 
# Now we're ready to create geometrical simulations. We create simulations using
# the
# [kikuchipy.simulations.KikuchiPatternSimulator](../reference.rst#kikuchipy.simulations.KikuchiPatternSimulator),
# which takes the reflectors as input

# In[ ]:


simulator = kp.simulations.KikuchiPatternSimulator(ref)


# In[ ]:


sim = simulator.on_detector(detector, rot)


# By passing the detector and crystal orientations to
# [KikuchiPatternSimulator.on_detector()](../reference.rst#kikuchipy.simulations.KikuchiPatternSimulator.on_detector),
# we've obtained a
# [kikuchipy.simulations.GeometricalKikuchiPatternSimulation](../reference.rst#kikuchipy.simulations.GeometricalKikuchiPatternSimulation),
# which stores the detector and gnomonic coordinates of the Kikuchi lines and
# zone axes for each crystal orientation

# In[ ]:


sim


# We see that not all 50 of the reflectors in the reflector list are present in
# some pattern.
# 
# These geometrical simulations can be plotted one-by-one by themselves

# In[ ]:


sim.plot()


# Or, they be plotted on top of patterns in three ways: passing a pattern to
# [GeometricalKikuchiPatternSimulation.plot()](../reference.rst#kikuchipy.simulations.GeometricalKikuchiPatternSimulation.plot)

# In[ ]:


sim.plot(index=(1, 2), pattern=s.inav[2, 1].data)


# Or obtaining a collection of lines, zone axes and zone axes labels as Matplotlib
# objects via
# [GeometricalKikuchiPatternSimulation.as_collections()](../reference.rst#kikuchipy.simulations.GeometricalKikuchiPatternSimulation.as_collections)
# and adding them to an existing Matplotlib axis

# In[ ]:


fig, ax = plt.subplots(ncols=3, nrows=3, figsize=(15, 15))

for idx in np.ndindex(s.axes_manager.navigation_shape[::-1]):
    ax[idx].imshow(s.data[idx], cmap="gray")
    ax[idx].axis("off")

    lines, zone_axes, zone_axes_labels = sim.as_collections(
        idx,
        zone_axes=True,
        zone_axes_labels=True,
        zone_axes_labels_kwargs=dict(fontsize=12),
    )
    ax[idx].add_collection(lines)
    ax[idx].add_collection(zone_axes)
    for label in zone_axes_labels:
        ax[idx].add_artist(label)

fig.tight_layout()


# Or obtaining the lines, zone axes, zone axes labels and PCs as HyperSpy markers
# via
# [GeometricalKikuchiPatternSimulation.as_markers()](../reference.rst#kikuchipy.simulations.GeometricalKikuchiPatternSimulation.as_markers)
# and adding them to a signal of the same navigation shape as the simulation
# instance. This enables navigating the patterns *with* the geometrical simulations

# In[ ]:


markers = sim.as_markers()

# To delete previously added permanent markers, do
# del s.metadata.Markers

s.add_marker(markers, plot_marker=False, permanent=True)


# In[ ]:


s.plot()


# ## Kinematical simulations
# 
# We can obtain kinematical master patterns using
# [KikuchiPatternSimulator.calculate_master_pattern()](../reference.rst#kikuchipy.simulations.KikuchiPatternSimulator.calculate_master_pattern),
# provided that the simulator is created from a
# [ReciprocalLatticeVector](https://diffsims.readthedocs.io/en/stable/reference.html#diffsims.crystallography.ReciprocalLatticeVector)
# instance that satisfy these conditions:
# 
# 1. The unit cell, i.e. the `structure` used to create the `phase` used in
#    `ReciprocalLatticeVector`, must have all asymmetric atom positions filled,
#    which can either be done by creating a `Phase` instance from a valid CIF file
#    with
#    [Phase.from_cif()](https://orix.readthedocs.io/en/stable/reference.html#orix.crystal_map.Phase.from_cif)
#    or calling 
#    [ReciprocalLatticeVector.sanitise_phase()](https://diffsims.readthedocs.io/en/stable/reference.html#diffsims.crystallography.ReciprocalLatticeVector.sanitise_phase)
# 
# 2. The atoms in the `structure` have their elements described by the symbol
#    (Ni), not by the atomic number (28)
# 
# 3. The lattice parameters are in given in Ångström.
# 
# 4. Kinematical structure factors $F_{hkl}$ have been calculated with
#    [ReciprocalLatticeVector.calculate_structure_factor()](https://diffsims.readthedocs.io/en/stable/reference.html#diffsims.crystallography.ReciprocalLatticeVector.calculate_structure_factor)
# 
# 5. Bragg angles $\theta_B$ have been calculated with
#    [ReciprocalLatticeVector.calculate_theta()](https://diffsims.readthedocs.io/en/stable/reference.html#diffsims.crystallography.ReciprocalLatticeVector.calculate_theta)
# 
# Let's simulate three master patterns:
# 
# * nickel
# 
# * variant of the $\sigma$-phase (Fe, Cr) in steels
# 
# * silicon carbide 6H.

# ### Nickel
# 
# We'll compare our kinematical simulations to dynamical simulations performed
# with EMsoft (see
# <cite data-cite="callahan2013dynamical">Callahan and De Graef (2013)</cite>),
# since we have a Ni master pattern available in the `kikuchipy.data` module

# In[ ]:


mp_ni_dyn = kp.data.nickel_ebsd_master_pattern_small(projection="stereographic")


# Inspect phase

# In[ ]:


phase_ni = mp_ni_dyn.phase.deepcopy()

print(phase_ni)
print(phase_ni.structure.lattice)


# Change lattice parameters from nm to Ångström

# In[ ]:


lat_ni = phase_ni.structure.lattice  # Shallow copy
lat_ni.setLatPar(lat_ni.a * 10, lat_ni.b * 10, lat_ni.c * 10)

print(phase_ni.structure.lattice)


# We'll build up the reflector list by:
# 
# 1. Finding all reflectors with a minimal interplanar spacing $d$
# 
# 2. Keeping those that have a structure factor above 0.5% of the reflector with
#    the highest structure factor

# In[ ]:


ref_ni = ReciprocalLatticeVector.from_min_dspacing(phase_ni, 0.5)

ref_ni = ref_ni[
    ref_ni.allowed
]  # Exclude non-allowed reflectors (not available for hexagonal or trigonal phases!)
ref_ni = ref_ni.unique(use_symmetry=True).symmetrise()


# Sanitise `phase`

# In[ ]:


ref_ni.phase.structure


# In[ ]:


ref_ni.sanitise_phase()
ref_ni.phase.structure


# We can now calculate the structure factors. Two parametrizations are available,
# from <cite data-cite="kirkland1998advanced">Kirkland (1998)</cite> (`"xtables"`,
# the default) or
# <cite data-cite="lobato2014accurate">Lobato and Van Dyck (2014)</cite> (`"lobato"`)

# In[ ]:


ref_ni.calculate_structure_factor()


# In[ ]:


structure_factor_ni = abs(ref_ni.structure_factor)
ref_ni = ref_ni[structure_factor_ni > 0.05 * structure_factor_ni.max()]

ref_ni.print_table()


# In[ ]:


ref_ni.calculate_theta(20e3)


# We can now create our simulator and plot the simulation

# In[ ]:


simulator_ni = kp.simulations.KikuchiPatternSimulator(ref_ni)
simulator_ni.reflectors.size


# Plotting the band centers with intensities scaled by the structure factor

# In[ ]:


simulator_ni.plot()


# Or no scaling (`scaling="square"` for the structure factor squared)

# In[ ]:


simulator_ni.plot(scaling=None)


# We can also plot the Kikuchi bands, showing both hemispheres, also adding the
# crystal axes alignment

# In[ ]:


fig = simulator_ni.plot(hemisphere="both", mode="bands", return_figure=True)

ax = fig.axes[0]
ax.scatter(simulator_ni.phase.a_axis, c="r")
ax.scatter(simulator_ni.phase.b_axis, c="g")
ax.scatter(simulator_ni.phase.c_axis, c="b")


# The simulation can be plotted in spherical projection as well using
# `Matplotlib`, or `PyVista` provided that it is [installed](../installation.rst#with-pip)

# In[ ]:


simulator_ni.plot("spherical", mode="bands")


# In[ ]:


# Intensity scaling is not available when plotting in a notebook
simulator_ni.plot("spherical", mode="bands", backend="pyvista")  # Interactive!


# When we're happy with the reflector list in the simulator, we can generate our
# kinematical master pattern

# In[ ]:


mp_ni_kin = simulator_ni.calculate_master_pattern(half_size=200)


# The returned master pattern is an instance of
# [EBSDMasterPattern](../reference.rst#kikuchipy.signals.EBSDMasterPattern) in the
# stereographic projection.

# In[ ]:


mp_ni_kin


# In[ ]:


mp_ni_kin.plot_spherical(style="points")  # Interactive!


# Comparing kinematical and dynamical simulations

# In[ ]:


# Exclude outside equator
ni_dyn_data = mp_ni_dyn.data.astype(np.float32)
ni_kin_data = mp_ni_kin.data.astype(np.float32)
mask = ni_dyn_data == 0
ni_dyn_data[mask] = np.nan
ni_kin_data[mask] = np.nan

fig, ax = plt.subplots(ncols=2)
ax[0].imshow(ni_kin_data, cmap="gray")
ax[1].imshow(ni_dyn_data, cmap="gray")
ax[0].axis("off")
ax[1].axis("off")
ax[0].set_title("Ni kinematical 20 kV")
ax[1].set_title("Ni dynamical 20 kV")
fig.tight_layout()


# <div class="alert alert-warning">
# 
# Warning
#     
# Use dynamical simulations when performing pattern matching, not kinematical
# simulations. The latter intensities are not realistic, as demonstrated in the
# above comparison.
# 
# </div>
# 
# Finally, we can transform the master pattern in the stereographic projection
# to one in the Lambert projection

# In[ ]:


mp_ni_kin_lp = mp_ni_kin.as_lambert()


# In[ ]:


mp_ni_kin_lp.plot()


# We can then project parts of this pattern onto our EBSD detector using
# [get_patterns()](../reference.rst#kikuchipy.signals.EBSDMasterPattern.get_patterns)

# In[ ]:


s_kin = mp_ni_kin_lp.get_patterns(rot, detector, energy=20, compute=True)


# In[ ]:


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


# Compare these to the ones the first plot!

# ### $\sigma$-phase

# In[ ]:


phase_sigma = Phase(
    name="sigma",
    space_group=136,
    structure=Structure(
        atoms=[
            Atom("Cr", [0, 0, 0], 0.5),
            Atom("Fe", [0, 0, 0], 0.5),
            Atom("Cr", [0.31773, 0.31773, 0], 0.5),
            Atom("Fe", [0.31773, 0.31773, 0], 0.5),
            Atom("Cr", [0.06609, 0.26067, 0], 0.5),
            Atom("Fe", [0.06609, 0.26067, 0], 0.5),
            Atom("Cr", [0.13122, 0.53651, 0], 0.5),
            Atom("Fe", [0.13122, 0.53651, 0], 0.5),
        ],
        lattice=Lattice(8.802, 8.802, 4.548, 90, 90, 90),
    ),
)
phase_sigma


# In[ ]:


ref_sigma = ReciprocalLatticeVector.from_min_dspacing(phase_sigma, 1)

ref_sigma.sanitise_phase()

ref_sigma.calculate_structure_factor("lobato")

structure_factor = abs(ref_sigma.structure_factor)
ref_sigma = ref_sigma[structure_factor > 0.05 * structure_factor.max()]

ref_sigma.calculate_theta(20e3)

ref_sigma.print_table()


# In[ ]:


simulator_sigma = kp.simulations.KikuchiPatternSimulator(ref_sigma)
simulator_sigma


# In[ ]:


fig = simulator_sigma.plot(hemisphere="both", mode="bands", return_figure=True)

ax = fig.axes[0]
ax.scatter(simulator_sigma.phase.a_axis, c="r")
ax.scatter(simulator_sigma.phase.b_axis, c="g")
ax.scatter(simulator_sigma.phase.c_axis, c="b")
fig.tight_layout()


# In[ ]:


simulator_sigma.plot("spherical", mode="bands", backend="pyvista")


# In[ ]:


mp_sigma = simulator_sigma.calculate_master_pattern()


# In[ ]:


mp_sigma.plot()


# In[ ]:


mp_sigma.plot_spherical(style="points")  # Interactive!


# ### Silicon carbide 6H

# In[ ]:


phase_sic = Phase(
    name="sic_6h",
    space_group=186,
    structure=Structure(
        atoms=[
            Atom("Si", [1 / 3, 2 / 3, 0.20778]),
            Atom("C", [1 / 3, 2 / 3, 0.33298]),
            Atom("Si", [1 / 3, 2 / 3, 0.54134]),
            Atom("C", [1 / 3, 2 / 3, 0.66647]),
            Atom("C", [0, 0, 0]),
            Atom("Si", [0, 0, 0.37461]),
        ],
        lattice=Lattice(3.081, 3.081, 15.2101, 90, 90, 120),
    ),
)
phase_sic


# In[ ]:


ref_sic = ReciprocalLatticeVector.from_min_dspacing(phase_sic)  # 0.7 Å, default
ref_sic.sanitise_phase()

ref_sic.calculate_structure_factor()

structure_factor = abs(ref_sic.structure_factor)
ref_sic = ref_sic[structure_factor > 0.05 * structure_factor.max()]

ref_sic.calculate_theta(20e3)

ref_sic.print_table()


# In[ ]:


simulator_sic = kp.simulations.KikuchiPatternSimulator(ref_sic)
simulator_sic


# In[ ]:


fig = simulator_sic.plot(hemisphere="both", mode="bands", return_figure=True)

ax = fig.axes[0]
ax.scatter(simulator_sic.phase.a_axis, c="r")
ax.scatter(simulator_sic.phase.b_axis, c="g")
ax.scatter(simulator_sic.phase.c_axis, c="b")


# In[ ]:


simulator_sic.plot("spherical", mode="bands", backend="pyvista")


# In[ ]:


mp_sic = simulator_sic.calculate_master_pattern(hemisphere="both", half_size=200)


# In[ ]:


mp_sic


# In[ ]:


mp_sic.plot(navigator=None)


# In[ ]:


mp_sic.plot_spherical(style="points")  # Interactive!