Speed up rotations, closes #2097

[StFroese]

No description provided.

[codecov]

Codecov Report

Base: 92.51% // Head: 92.75% // Increases project coverage by +0.23% 🎉

Coverage data is based on head (948ecf8) compared to base (89add5c).
Patch coverage: 97.43% of modified lines in pull request are covered.

❗ Current head 948ecf8 differs from pull request most recent head fdd14a6. Consider uploading reports for the commit fdd14a6 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2098      +/-   ##
==========================================
+ Coverage   92.51%   92.75%   +0.23%     
==========================================
  Files         199      214      +15     
  Lines       16561    17876    +1315     
==========================================
+ Hits        15322    16580    +1258     
- Misses       1239     1296      +57     

Impacted Files	Coverage Δ
ctapipe/coordinates/ground_frames.py	98.50% <95.65%> (-1.50%)	⬇️
ctapipe/coordinates/tests/test_coordinates.py	100.00% <100.00%> (ø)
ctapipe/tools/process.py	87.87% <0.00%> (-4.35%)	⬇️
ctapipe/utils/datasets.py	87.90% <0.00%> (-1.82%)	⬇️
ctapipe/conftest.py	94.27% <0.00%> (-0.60%)	⬇️
ctapipe/core/tests/test_traits.py	99.45% <0.00%> (-0.55%)	⬇️
ctapipe/core/traits.py	93.37% <0.00%> (-0.20%)	⬇️
ctapipe/core/tool.py	93.51% <0.00%> (-0.06%)	⬇️
ctapipe/io/simteleventsource.py	96.05% <0.00%> (-0.02%)	⬇️
ctapipe/containers.py	100.00% <0.00%> (ø)
... and 43 more

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

[StFroese]

@maxnoe I also tried to implement the missing many_to_many transform. It's working mathematically but seems like astropy can't realize the frame because the pointing_direction has a shape of (3,2) and the new coordinates vectors have shape (6,3)

[StFroese]

So astropy blocks this because it thinks there to many coordinates (6) for only three transforms (3)

[StFroese]

Works now, but it's a little bit hacky...

[kosack]

Should we also replace other places where we use hand-coded rotation matrices? I see for example:

• ctapipe.utils.linalg.rotation_matrix_2d(), used in CameraGeometry.rotate() and toymodel
• ctapipe.reco.impact.ImPACTReconsturctor.rotate_translate()

Maybe there are others?

[StFroese]

Yes I think so. I also saw a few other by hand rotations in the code. But I think that I have to check if this is at least not slower right now with scipy

[kosack]

Yes I think so. I also saw a few other by hand rotations in the code. But I think that I have to check if this is at least not slower right now with scipy

I did a quick check with %timeit and a 2D rotation matrix using ctapipe.utils.linalg, and it was slightly faster to use scipy's method, at least there is not a huge difference (the small difference is likely due to units)

In [8]: %timeit rotation_matrix_2d(np.pi/2.0*u.rad)
38 µs ± 10.9 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

In [9]: %timeit spatial.transform.Rotation.from_euler("xy", [np.pi/2,0]).as_matr
   ...: ix()
29.3 µs ± 8 µs per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

[kosack]

However, I notice in #2052 that ImPACT now njit's the rotate_translate function, so maybe don't touch that one for now.

Certainly you can replace any usage of ctapipe.utils.linalg.rotation_matrix_2d with the scipy version

[StFroese]

The ground to tilted transform also seems to be faster.
On master:

In [1]: from ctapipe.coordinates import GroundFrame, TiltedGroundFrame
   ...: from astropy.coordinates import AltAz, SkyCoord
   ...: import astropy.units as u
   ...: import numpy as np
   ...: 
   ...: ground = GroundFrame(x=1000*[1] * u.m, y=1000 * [2] * u.m, z=1000* [3] * u.m)
   ...: 
   ...: pointing = SkyCoord(alt=[90] * u.deg, az=[180] * u.deg, frame=AltAz())
   ...: frame = TiltedGroundFrame(pointing_direction=pointing)
   ...: 
   ...: %timeit ground.transform_to(frame)v. of 7 runs, 1,000 loops each)
217 µs ± 649 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

On this branch:

170 µs ± 191 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

[StFroese]

But it's slower for the one_to_one case with multiple rotation matrices.
pointing = SkyCoord(alt=1000*[90] * u.deg, az=1000*[180] * u.deg, frame=AltAz())
master: 203 µs ± 624 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
here: 432 µs ± 846 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

[StFroese]

I will do some profiling to see if the generation of the matrices is the problem or the np.einsum

[kosack]

Could be something else, since when I use many phis (remember to construct them outside the timeit loop), I get this:

phis = np.asarray(1000*[np.pi])

In [12]: %timeit rotation_matrix_2d(phis*u.rad)
228 µs ± 5.35 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [13]: %timeit spatial.transform.Rotation.from_euler("x", phis).as_matrix()
55.5 µs ± 21.2 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

EDIT: I left the unit in the first one, which adds overhead. Removing it they are closer:

In [21]: phis = np.asarray(1000*[np.pi]) * u.rad

In [22]: %timeit rotation_matrix_2d(phis)
181 µs ± 24.6 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

In [23]: %timeit spatial.transform.Rotation.from_euler("x", phis.to_value(u.rad)).as_matrix()
98.9 µs ± 3.73 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

[StFroese]

So profiling says get_shower_trans_matrix is roughly 8 times slower on this branch. Looking at the SciPy Rotation code, it seem that they are using quaternions for their calculation. This requires more calculations than simply constructing the rotation matrix directly.
I think we can implement the matrix ourselves like it was before but maybe a little bit more readable and also put it in ctapipe.utils.linalg...

[StFroese]

I'm now using quaternions to produce the rotation matrices. It's faster on this branch now:

from ctapipe.coordinates import GroundFrame, TiltedGroundFrame
from astropy.coordinates import AltAz, SkyCoord
import astropy.units as u
import numpy as np
from cProfile import Profile

ground = GroundFrame(x=1000*[1] * u.m, y=1000*[2] * u.m, z=1000*[3] * u.m)

pointing = SkyCoord(alt=1000*[90] * u.deg, az=1000*[180] * u.deg, frame=AltAz())
frame = TiltedGroundFrame(pointing_direction=pointing)

this branch:

%timeit ground.transform_to(frame)
176 µs ± 714 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

main:

%timeit ground.transform_to(frame)
203 µs ± 607 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

[kosack]

Just to understand the difference: was it just that the code in scipy to turn Euler angles into a quaternion is slower than doing it manually? (i.e. Rotation.from_euler() is slower than constructing a quaternion yourself and passing it to Rotation(quat)?).

[StFroese]

Yes it is slower in this case because scipy constructs a quaternion for the first rotation axis and one for the second rotation axis. Those are multiplied via Hamilton product.
What I'm doing is to construct this product by hand (analytically) once for arbitrary rotations along y and z, so I don't have to do the whole cross product stuff and so on. This basically drops 12 of 16 calculations for this product.

[StFroese]

So 12 of 16 terms for this 2 axis rotation in scipy are always 0 but are computed everytime

[StFroese]

I think this is also the reason why your test for the 2d rotation is faster. This doesn't require the Hamilton product since it is only one axis. E.g. an xyz rotation in scipy computes 3 basic quaternions and multiplies them. Resulting in about 6 cos/sin calls and 32 multiplications

[maxnoe]

I am not sure these quaternion calculations are more readable than what we have before and njitting the Rotation matrix should be simple and probably even faster...

[StFroese]

Yes I think so too... it's not readable at all. It's even more difficult to understand this calculation than the matrix product used before

[StFroese]

I returned to the original implementation plus Max's faster _get_xyz, the many_to_many relation and some einsums.

on master:

In [2]: %timeit ground.transform_to(frame)
199 µs ± 574 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

this branch:

In [2]: %timeit ground.transform_to(frame)
124 µs ± 42.3 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

The title of this PR as well as the history is more or less a mess now. Should I open a new PR just for the current changes?

[maxnoe]

@StFroese Not needed, just update the title of the PR please. If you want you can rebase / squash into more meaningfull commits or I will just "squash and merge" at the end.

[StFroese]

@maxnoe best I can do for now. I let you decide if you want to squash it into one commit. I don't mind.