ENH: spatial: Rotation: add split to break apart object with multiple values into list of rotations #21351
Description
Is your feature request related to a problem? Please describe.
This is essentially the opposite of Rotation.concatenate() and I think it should mirror the syntax of numpy.split() https://numpy.org/doc/stable/reference/generated/numpy.split.html
My current use case is making a pandas dataframe column with Rotations, which requires turning it into a list.
Should be pretty easy to implement with just calling numpy.split() under the hood, but want to get some feedback before opening an MR.
Describe the solution you'd like.
Current numpy.split() behavior:
import numpy as np
a = np.array([1, 2, 3, 4])
np.split(a, len(a))
# returns [array([1]), array([2]), array([3]), array([4])]
np.split(a, [1, 2])
# returns [array([1]), array([2]), array([3, 4])]New similar Rotation.split() behavior:
from scipy.spatial.transform import Rotation as R
r1 = R.from_rotvec([0, 0, 1])
r2 = R.from_rotvec([0, 0, 2])
r3 = R.from_rotvec([0, 0, 3])
r4 = R.from_rotvec([0, 0, 4])
r34 = R.concatenate([r3, r4])
rc = R.concatenate([r1, r2, r3, r4])
rc.split(len(rc))
# returns [r1, r2, r3, r4]
rc.split([1, 2])
# returns [r1, r2, r34]
# This equality would always hold, no matter what was in *
rc == R.concatenate(rc.split(*))Describe alternatives you've considered.
The following works, but is clunky and less flexible. This is pretty much exactly what would be happening under the hood.
[R.from_quat(np.squeeze(qs)) for qs in np.split(rc.as_quat(), len(rc))]
# returns [r1, r2, r3, r4]
[R.from_quat(np.squeeze(qs)) for qs in np.split(rc.as_quat(), [1, 2])]
# returns [r1, r2, r34]Additional context (e.g. screenshots, GIFs)
No response