Weighted mean quaternion/orientation

[DorianDepriester]

Hello,
I was looking for a way to compute the weighted mean orientation of crystals. Starting from the discussion in issue #434, and mixing it with the current implementation of the mean method(), together with the the full Eq. (13) in the given reference, I ended with this:

def weighted_mean(o, weights):
  o2 = o.map_into_symmetry_reduced_zone()
  q = o2.data.T
  qq = (weights*q).dot(q.T)
  w, v = np.linalg.eig(qq)
  w_max = np.argmax(w)
  qmean = v[:, w_max]
  return Orientation(Quaternion(qmean), o.symmetry)

I am probably wrong (that's why I haven't made a pull request). Anyway, I think this feature could be useful (at least for me 😉).

[hakonanes]

Thanks for this Dorian, this would be very useful indeed!

Started testing this a bit with the following procedure (in an IPython console), and it looks promising to me. To probe the weighting, I generated a grid of six orientations, +/- 30 degrees about each sample direction x, y and z, and then plotted mean orientations with different weights, like so:

%matplotlib qt5
import numpy as np
from orix.quaternion import Orientation, Quaternion, symmetry


def weighted_mean(o, weights):
    o2 = o.map_into_symmetry_reduced_zone()
    q = o2.data.T
    qq = (weights*q).dot(q.T)
    w, v = np.linalg.eig(qq)
    w_max = np.argmax(w)
    qmean = v[:, w_max]
    return Orientation(Quaternion(qmean), o.symmetry)

ori = Orientation.from_axes_angles(
    [[1, 0, 0], [1, 0, 0], [0, 1, 0], [0, 1, 0], [0, 0, 1], [0, 0, 1]],
    [30, -30, 30, -30, 30, -30],
    symmetry=symmetry.Oh,
    degrees=True
)

fig = ori.scatter(return_figure=True, c=np.arange(ori.size))
ax = fig.axes[0]
ax.set(xlabel="x [rad]", ylabel="y [rad]", zlabel="z [rad]")
ax.axis(True)
ax.grid(False)

# Same as mean: red
ori_w1 = weighted_mean(ori, [1, 1, 1, 1, 1, 1])
ax.scatter(ori_w1, c="r")

# Weight 30deg [100] 10 times more than the others: blue
ori_w2 = weighted_mean(ori, [10, 1, 1, 1, 1, 1])
ax.scatter(ori_w2, c="b")

# Weight 30deg about [100] and [001] 10 times more than the others: cyan
ori_w3 = weighted_mean(ori, [10, 1, 1, 1, 10, 1])
ax.scatter(ori_w3, c="c")

# Save different views of axis-angle space (using fig.savefig(), not shown)
ax.view_init(90, -90)  # XY plane
ax.view_init(0, -90)  # -XZ plane

We must supplement this visual test with some quantitative ones... Once those are found, we should add this feature (with a PR from you if you're willing).

[hakonanes]

Actually, I guess a valid test is just to compare the mean when having two of the same orientations in a set, or just having one but weighting it twice. If such a test passes, I'm happy with including this feature.

[DorianDepriester]

I agree. Your previous test actually seems robust to me!

[DorianDepriester]

Yet, as another step toward complexity, let's assume that one wants to use a matrix of weights. The following function seems to work without any loop:

def weighted_mean_matrix(o, mat):
    o2 = o.map_into_symmetry_reduced_zone()
    q = o2.data
    qq = np.einsum('pi,ij,ik->pjk', mat, q, q)
    w, v = np.linalg.eig(qq)
    w_max = np.argmax(w, axis=1)
    q_mean= v[np.arange(mat.shape[0]), :, w_max]
    return Orientation(Quaternion(q_mean), o.symmetry)

Again, this should be tested (as made above)

[hakonanes]

By matrix, do you mean that your orientations have ndim > 1? The einsum you use works for 2D, but how about 3D? Another solution is to reshape q to shape (n, 4) and ravel weights, and then reshape q_mean again before returning, right?

[DorianDepriester]

Consider n orientations (o.shape=(n,)). If one wants to compute m weighted mean orientations, the matrix will be of size (m,n). It basically does the same as the an iteration of the function defined above.

[hakonanes]

Ah, I see, sorry for my misunderstanding. Your generalization seems to cover all cases, then! In the implementation we should flatten the orientations before applying the weights. After that I think this should be robust.