Alternative IPF plotting tool

[argerlt]

This is a functionality I have had for a few months and wanted to add to ORIX for some time, but have not had the time to do correctly. Leaving this here in case someone else would like to add it, and to get feedback from others.

ORIX currently has an pole density function plotting tool, but the S2 sampling it uses causes an extreme mismatch in the polar-to-azimuthal ratio around the poles. (See here for refernce). This picture is a good example of the end effect:

I've been using two rudimentary adaptions of this method that (i think) are more representative. 1) The first involves using a KDTree in S2 space to rapidly assign vectors to zones, then creates a map that is a matplotlib.tripcolor plot of the result. 2) The second uses the same trick as the current ORIX tool where binning is done in a 2D cartesian mapping so we can leverage np.histogram2d, but uses the equal area stereographic projection instead of polar-azimuthal, giving a much smoother plot.

Examples of both are shown in the attached code, and a picture of both techniques are shown here as well.

I will note though, these plots have some unsolved problems that need to be fixed before adding to ORIX. 1) I don't think the blurring I do is equal-angle, which the current ORIX PDF is. thus, the idea of a "broadened" PDF like the third picture above is not rigorously correct. 2) the plots themselves need some work, they are currently images as-is, so they don't play well with inserts or overlays. 3) These are scripts, not function classes, and I haven't considered outlier cases. 4) I'm playing fast and loose with PF versus IPF here. they are different and handle symmetries differently; this will need to be accounted for in anything that actually ends up in ORIX.

# -*- coding: utf-8 -*-
"""
Created on Fri Nov  8 14:36:52 2024

@author: agerlt
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.tri as tri
from scipy.spatial import KDTree
from orix.crystal_map import CrystalMap, Phase
import orix.io as io
from orix.plot import IPFColorKeyTSL
from orix.quaternion import Rotation, Orientation, Symmetry
import orix.quaternion.symmetry as symmetry
import orix.sampling as sampling
from orix.vector import Miller, Vector3d
from scipy.stats import norm
from scipy.ndimage import binary_dilation, gaussian_filter
from skimage import feature
from matplotlib import cm

# some constants to make life easer
hc_max = (3*np.pi/4)**(1/3)
r2d = 180/np.pi
d2r = np.pi/180

# equal area projection
# TODO: pretty sure this functionality is already in ORIX...
def vector3d2schmidt(v):
    r = 2*np.sin(v.polar/2)
    azimuth = v.azimuth
    x = r*np.cos(azimuth)
    y = r*np.sin(azimuth)
    return np.stack([x, y]).T

def schmidt2vector3d(x, y):
    r = np.sqrt(x**2 + y**2)
    polar = np.arcsin(0.5*r)*2
    # azimuth = np.zeros(x.shape)
    # mask = x**2 > 1E-6
    # azimuth[~mask] = np.pi
    # azimuth[mask] = np.arctan2(y[mask], x[mask])
    azimuth = np.arctan2(y, x)
    v = Vector3d.from_polar(azimuth, polar)
    return v

def pdf_from_kdtree(kdtree, vectors, kernel=False, hw=0.5):
    crds = vector3d2schmidt(vectors.flatten().unit)
    counts = np.zeros(kdtree.data.shape[0])
    if not kernel:
        a, b = np.unique(kdtree.query(crds)[1], False, False, True)
        counts[a] = b
        return counts
    # equal area is NOT equal distance projection, so calc spot contribution
    # SHOULD come from vector distance. however, it's slow, so fix later.
    r, idx = kdtree.query(crds, 50)
    # hw is in degrees, there is 90 degrees accross a 2^0.5 length.
    weights = norm.pdf((r/hw) * (90/np.sqrt(2)))
    for i in range(50):
        counts[idx[:, i]] += weights[:, i]
    return counts

# %% read in Greg's ebsd, make some poles
ebsd = io.load("GES_HEDM1_ebsdPCSHIFT_310_regi.ang")
ebsd.phases.point_group = symmetry.Oh
ph = ebsd.phases[0]
cs = symmetry.Oh
# gut check image
ori2rgb = IPFColorKeyTSL(cs)
raw_img = ori2rgb.orientation2color(ebsd.rotations).reshape(600, 600, 3)
ci = (ebsd.prop['unknown2']).reshape(600, 600)
img = raw_img*1
img[ci < 0.1] = 0
plt.imshow(img)

# %%
rots = ebsd.rotations.flatten()[ci.flatten() > 0.1]
poles = rots*Vector3d.zvector()
r = Miller(poles.data, phase=ph).in_fundamental_sector()
r_polar = np.stack(r.to_polar()[:2]).T

# %% using a KD tree in S2 space
# Do an icosohedron sampling of S2 space. throw out everything outside the FZ
s2_mesh = sampling.S2_sampling.sample_S2_icosahedral_mesh(0.15)
s2_fund_mesh = s2_mesh[np.min(
    s2_mesh.data == s2_mesh.in_fundamental_sector(symmetry.Oh).data, axis=1
    )]
# NOTE: to see why the plot edges look weird, do the following plot:
# xy = vector3d2schmidt(s2_fund_mesh)
# plt.scatter(*xy.T)
# basically, the seperation between the FZ bounds and the icosahedral mesh
# is not equal everywhere, so some faces in the triangulation are too big.
# this COULD be fixed with an area normalization, but I don't know how
# to do this efficiently.

# use a KD tree in polar/aximuth to figure out where to bin data
sp_fpm = np.stack(s2_fund_mesh.to_polar()[:2]).T
S2_f_kdtree = KDTree(sp_fpm)

# bin the poles
a, b = np.unique(S2_f_kdtree.query(r_polar)[1], False, False, True)
counts = np.zeros(s2_fund_mesh.shape)
counts[a] = b

# I think there is a better way to do this step where you make the bins equal
# to the icosohedron triangles we started with, but IDK how to do it, so this
# is good enough
# old equal-angle projection method:
# x = s2_fund_mesh.x/(1+s2_fund_mesh.z)
# y = s2_fund_mesh.y/(1+s2_fund_mesh.z)

xy = vector3d2schmidt(s2_fund_mesh)
triang = tri.Triangulation(*xy.T)

# %%% The (maybe?) better way, binning into a grid in xy-schmidt space

# find the min-max for the FZ in schmidt vector space
xy_line = np.arange(-2, 2.01, 0.01)
init_xy = [x.flatten() for x in np.meshgrid(xy_line, xy_line)]
iv = schmidt2vector3d(*init_xy)
init_mask = np.min(iv.in_fundamental_sector(cs).data == iv.data, axis=1)
min_max = [[x[init_mask].min()-0.02, x[init_mask].max()+0.02] for x in init_xy]

# make bins, calc FZ mask
x_bins = np.linspace(min_max[0][0], min_max[0][1], 781*2, True)
y_bins = np.linspace(min_max[1][0], min_max[1][1], 671*2, True)
bins = np.meshgrid(x_bins[:-1], y_bins[:-1])
pix_v = schmidt2vector3d(*bins)
fz_mask = np.min(pix_v.in_fundamental_sector(cs).data == pix_v.data, axis=2)

# bin the actual data
r_in_schmidt = vector3d2schmidt(r)
bin_counts = np.histogram2d(*r_in_schmidt.T, bins=[x_bins, y_bins])[0].T
border = binary_dilation(feature.canny(fz_mask), iterations=2)

# make the image and plot it
# img = cm.viridis(255*bin_counts/np.max(bin_counts))
img = cm.magma(255*bin_counts/np.max(bin_counts))
img[~fz_mask] = [1, 1, 1, 1]

# can also doa gaussian blur to make a pseudo-IPDF
# NOTE, this is NOT a "True" IPDF, which would require a distance function
# on the S2 sphere. probably the easiest way would be an RBF, which would
# be time consuming and give something very close to this anyway.
# Still, I would NOT publish this figure as-is in a paper, it lies about
# texture in an insidiously un-obvious way.
img2 = cm.magma(255*gaussian_filter(bin_counts, 2)/np.max(bin_counts))
img2[~fz_mask] = [1, 1, 1, 1]
img2[border] = [0, 0, 0, 1]

# %% plot it all

fig, ax = plt.subplots(1, 3)
ax = ax.flatten()

ax[0].tripcolor(triang, counts, shading='flat', vmax=30)
ax[0].set_xlim(-0.01, 0.78)
ax[0].set_ylim(-0.01, 0.67)

ax[1].imshow(img)
ax[1].set_xlim(0, 781*2)
ax[1].set_ylim(0, 671*2)

ax[2].imshow(img2)
ax[2].set_xlim(0, 781*2)
ax[2].set_ylim(0, 671*2)

for i in range(3):
    ax[i].set_aspect('equal')
    ax[i].set_xticks([])
    ax[i].set_yticks([])

ax[0].set_title("KD_tree and triangular mesh")
ax[1].set_title("2d histogram binning in schmidt space")
ax[2].set_title("2d_hist, plus gaussian blur for pseudo-IPDF")

plt.tight_layout()