Problems with `find_registration`

[janden]

There appear to be two (unrelated) issues with Rotation.find_registration (which is used by Rotation.register). The first is that the way it chooses whether to apply J-conjugation doesn't seem to work. Indeed, if we compare the errors induced by the two rotation matrices Q1 and Q2 on the full loss function, these are not always in the same order as err1 and err2 (in other words, the proxy of looking at Q1 @ Q1.T and Q2 @ Q2.T looks incorrect). It works if there is an exact solution (that is, the matrices are rotations of one another and the loss function reaches zero), but not in the general case.

The second issue is that the Q matrix return is sometimes not a pure rotation (i.e., has determinant negative one).

The following script demonstrates both problems:

import numpy as np
from aspire.utils import Rotation

angs_true = np.load("angs_true.npy")
#angs_est01 = np.load("angs_est01.npy")
angs_est02 = np.load("angs_est02.npy")

rots_true = Rotation.from_euler(angs_true)
#rots_est01 = Rotation.from_euler(angs_est01)
rots_est02 = Rotation.from_euler(angs_est02)

#Q, flag = rots_est01.find_registration(rots_true)

#print(flag)                 # should be 0

Q, flag = rots_est02.find_registration(rots_true)

print(np.linalg.det(Q))     # should be >0

Data files are here: data.zip

EDIT: Test doesn't trigger issue with J-conjuacy. Updated angs_est02.npy.

[janden]

Turns out the first issue isn't triggered by the tester above. I'm still not convinced the J-conjugacy check is doing what we want it to. I haven't found any sources documenting why it should work, so I suggest we replace it with an actual calculation of the loss function (after finding the nearest rotation).

[garrettwrong]

Thanks for putting this patch together.

Yeah Im a little confused about that code actually. I gather that code checked Q is orthogonal by checking that is it normal. Sure. But I don't follow the Q2 situation in the simplest case. Take a single rotation that is exactly correct with reference (R == Rref).

Hopefully I don't embarrass myself by screwing this up too bad.

Q1 = R @ Rref.T = R @ R.T = I ; definitely normal, Q1 @ Q1.T will be pretty nice. Great, that should be hard to beat.

Q2 = (J@R@J) @ R.T
Then Q2.T = R @ (J@R@J).T
So Q2 @ Q2.T = (J@R@J) @ R.T @ R @ (J@R@J).T = (J@R@J) @(J@R@J).T = I. , uhhhh 😬

>>> A = Rotation.generate_random_rotations(1,seed=123).matrices[0]
>>> J = np.array([[1, 0, 0], [0, 1, 0], [0, 0, -1]])
>>> Q2 = (J @ A @ J) @ A.T
>>> Q2 @ Q2.T
array([[9.99999875e-01, 1.30374738e-08, 1.51338632e-08],
       [1.30374738e-08, 9.99999914e-01, 5.55418114e-08],
       [1.51338632e-08, 5.55418114e-08, 1.00000001e+00]])
>>> 
>>> A @ A.T 
array([[ 9.9999994e-01,  4.7545505e-09, -2.3542670e-09],
       [ 4.7545505e-09,  9.9999994e-01, -3.0722145e-08],
       [-2.3542670e-09, -3.0722145e-08,  1.0000000e+00]], dtype=float32)
>>> 

Sorry I used A instead of R, but anyway.... I played around with this in the interpreter a bit, and it would seem before introducing any actual Jconj that this method for a single rotation is questionable. I don't see why one would not be able to find some rotation where the numerical error in Q2 case wasn't smaller than Q1. When R is Rref, or very close, it just looks like floating point errors in both cases to me.... The last example I checked was close, but Q1 edged out, which would probably be the case most of the time due to accumulation of error, but, I'm not convinced always...

>>> np.linalg.norm(A @ A.T - np.eye(3), ord="fro")
np.float64(9.512846246542357e-08)
>>> np.linalg.norm(Q2 @ Q2.T - np.eye(3), ord="fro")
np.float64(1.7332543352097297e-07)

[j-c-c]

@garrettwrong I think one detail that makes a little bit of a difference is that in the perfect case where R == Rref, we have that Q1 = I, but Q2 does not. For just a single rotation our error function is no good since both Q1 @ Q1.T and Q2 @ Q2.T will be the identity (as you pointed out), but as soon as you add an additional rotation the accumulation in Q2, ie. Q2 = Q2 + (J @ R @ J) @ Rref.T, will almost certainly not produce an orthogonal matrix since we are not just accumulating Identity matrices as we would be for Q1.

Edit: I guess my point is moot, since we have the better option of using the actual loss function

[garrettwrong]

I don't think so. I think it will produce exactly this.... Rref is an array of reference rots, not the same rot for each iteration...

Assume we have a perfect set of R = Rref. R1 = Rref_1, R2=Rref_2 and so on. Each item accumulated would be as I described for the single case. I just avoided the loop to make it simpler.

[garrettwrong]

Unfortunately, if you agree with what I wrote for a single R, I'm becoming increasingly concerned this algorithm is incorrect 😬 .

If we had a perfect set of R, but it was just J conjugated, R = J@Rref@J, then Q1 and Q2 simply switch places. Both are Identity on paper. I might have made a different mistake though.

[j-c-c]

Assume we have a perfect set of R = Rref. R1 = Rref_1, R2=Rref_2 and so on. Each item accumulated would be as I described for the single case. I just avoided the loop to make it simpler.

Not sure I agree here. Let's consider the case of just two rotations that match the references exactly, R1 = Rref_1, R2=Rref_2. Then,

Q1 = (R1 @ R1.T + R2 @ R2.T) / 2 = (I + I) / 2 = I
and 
Q1 @ Q1.T = I @ I.T = I

but,
Q2 = ((J @ R1 @ J) @ R1.T + (J @ R2 @ J) @ R2.T) / 2
giving us
Q2 @ Q2.T = [(J @ R1 @ J) @ R1.T + (J @ R2 @ J) @ R2.T] / 2 @
                       [(J @ R1 @ J) @ R1.T + (J @ R2 @ J) @ R2.T].T / 2

                   = [(J @ R1 @ J @ R1.T + J @ R2 @ J @ R2.T) @ 
                      (R1 @ J @ R1.T @ J + R2 @ J @ R2.T @ J)] / 4 

                   = (J @ R1 @ J @ R1.T)(R1 @ J @ R1.T @ J ) + 
                      (J @ R2 @ J @ R2.T)(R2 @ J @ R2.T @ J) +
                      (J @ R1 @ J @ R1.T)(R2 @ J @ R2.T @ J) +
                      (J @ R2 @ J @ R2.T)(R1 @ J @ R1.T @ J )

                  = [I + I +  2 * Sym((J @ R1 @ J @ R1.T)(R2 @ J @ R2.T @ J)] / 4

The point is, we can't reduce the last two terms in that matrix multiplication distribution, so we will not get the identity (except for some edge cases like R1 = R2). It's totally possible I am messing something up here. Lot's of terms! lol

[garrettwrong]

Yeah I think you are correct. My mistake. The sum of the Q2 entries doesn't allow the situation I was describing. That is good.
Joakim's code should work then.

Okay, then the original code wasn't working because the the SVD always yielding the positive root convention or? What am I missing there.

[j-c-c]

Turns out the first issue isn't triggered by the tester above.

@janden It looks like angs_est01 and angs_est02 are the same arrays. Did you possibly save the wrong array for one of them?

[j-c-c]

Okay, then the original code wasn't working because the the SVD always yielding the positive root convention or? What am I missing there.

Now I'm getting a little suspicious about the use of J-conjugation here, ie. the term J @ R @ J that appears in Q2 = Q2 + (J @ R @ J) @ Rref.T. Most places in the code where there is J-conjugation it occurs on "relative rotations" which involve the product of two rotations. For example, J-conjugation on the relative rotation Rij = Ri @ Rj.T is really just J @ Rij @ J = J @ Ri @ Rj.T @ J = (J @ Ri) @ (J @ Rj).T.

In other words, when dealing with individual rotations, Ri, there should only be one J. I'm thinking this method should have Q2 = Q2 + (R @ J) @ Rref.T.

Looking at some of the common-line publications, it's stated that we are recovering the set {Ri} or {Ri @ J} up to some global alignment by orthogonal matrix O.

[janden]

@janden It looks like angs_est01 and angs_est02 are the same arrays. Did you possibly save the wrong array for one of them?

Fixed this. It was angs_est02 that was mistakenly set to angs_est01, so it doesn't affect the outcome.

Looking at some of the common-line publications, it's stated that we are recovering the set {Ri} or {Ri @ J} up to some global alignment by orthogonal matrix O.

Hmm. So then we wouldn't have proper (i.e. positive-determinant) rotations in that case. Is that acceptable?

[j-c-c]

Fixed this. It was angs_est02 that was mistakenly set to angs_est01, so it doesn't affect the outcome.

Great, I'll check it out.

Hmm. So then we wouldn't have proper (i.e. positive-determinant) rotations in that case. Is that acceptable?

The estimated rotations that we get from the common lines method will always be proper rotations, but they might be rotations that correspond to a volume of opposite chirality to the ground truth volume. In this case, we will have that Q @ R_est @ J = R_truth, where det(Q) = -1. In the case where the common lines method gives estimated rotations that corresponds to a volume of the correct chirality we will get that Q @ R_est = R_truth, where det(Q) = 1.

In terms of our find_registration method (and the other registration methods), the above means that if Q1 (ie. correct chirality) wins we should get det(Q_mat) = 1 with flag = 0. If Q2 wins we should get det(Q_mat) = -1 with flag = 1. This is not the behavior we are currently getting and I believe it is due to the extra J in the formulation.

The take away here is that I think it is correct that we will sometimes get a Q_mat with negative determinant.

[j-c-c]

To add a little more here, if we have the case where Q2 should win, the way the code is now we are accumulating terms that look like Q2 = (J @ R_est @ J) @ R_ref.T. If we have good estimates than R_est = Q @ R_ref @ J for some orthogonal Q that should have determinant = -1. Substituting that in we get Q2 = (J @ [Q @ R_ref @ J] @ J) @ R_ref.T = J @ Q. In other words, when Q2 wins, find_registration is actually computing Q_mat = J @ Q instead of just Q_mat = Q. This will always give us a determinant of 1, which is the behavior that I'm seeing.

Here's a script that illustrates this behavior. Switching between seeds 8 and 4 will give the two behaviors of Q1 (flag = 0) or Q2 (flag = 1) winning. You'll notice that we get positive determinant for Q in both cases.

import numpy as np

from aspire.downloader import emdb_2660
from aspire.source import Simulation, OrientedSource
from aspire.utils import Rotation, mean_aligned_angular_distance

original_vol = emdb_2660()

# Downsample the volume                                                                                                          
res = 41
vol = original_vol.downsample(res)

# seed = 4 results in wrong chirality rotations.                                                                                 
# seed = 8 results in correct chirality rotations.                                                                               
src = Simulation(
    n=50,  # number of projections                                                                                               
    vols=vol,  # volume source                                                                                                   
    offsets=0,  # Default: images are randomly shifted                                                                           
    seed=8,
)

oriented_src = OrientedSource(src)
rots_est = oriented_src.rotations
Q, flag = Rotation(src.rotations).find_registration(Rotation(rots_est))
print(Q, flag, np.linalg.det(Q))
print(mean_aligned_angular_distance(src.rotations, rots_est))

Note, these are good estimates, ie. mean_aligned_angular_distance < 1 degree. @janden with the example you provided the mean aligned angular distance is something like 106 degrees, so I think we'll get a bad Q_mat regardless of chirality.

[janden]

I'm not following exactly how this differs. From what I can tell, when we replace a volume with its flipped counterpart (i.e., we replace x(u) with x(Ju). This has the result of changing all rotation matrices R to J @ R @ J. If we add an arbitrary SO(3) rotation matrix Q to this, we get Q @ J @ R @ J, which in find_registration translates to

R_est = Q @ J @ R_ref @ J

Now substituting this into the formula for Q1, we get

Q1 = R_est @ R_ref.T = Q @ J @ R_ref @ J @ R_ref.T

which should have positive determinant, but will give us a lousy loss since Q1.T @ R_est = R_ref @ J @ R_ref.T @ J @ Q, which is not R_ref. Now Q2 on the other hand is

Q2 = (J @ R_est @ J) @ R_ref.T = J @ (Q @ J @ R_ref @ J) @ J @ R_ref.T = J @ Q @ J

which also has positive determinant (equal to one). Here the loss is zero since Q2.T @ (J @ R_est @ J) = J @ Q.T @ J @ J @ (Q @ J @ R_ref @ J) @ J = R_ref.

Now in your notation, you would have Qp = Q @ J, which in turns gives Q2 = J @ Qp, as you say, but as far as I can tell, both of these methods are equivalent.

[j-c-c]

From what I can tell, when we replace a volume with its flipped counterpart (i.e., we replace x(u) with x(Ju). This has the result of changing all rotation matrices R to J @ R @ J.

@janden, Ok, after convincing myself that this is true, I am in agreement with you. I think I confused myself because one of the paper uses J = diag(-1, -1, 1) and replaces volume x(u) with x(-u), which has the result of changing all rotations R to R @ J, but this is a different J.

In any case, you are correct that we should be using J @ R @ J. I guess this also means that Q should be in SO(3).

Sorry to add confusion to this issue!