vectorize all functions to support arbitrary dimensions

[ManuelPalermo]

Hi!

TLDR: It would be helpful if all functions supported vectorized operations (e.g for batch computation, fusing data from multiple sensors or in some cases across time).

Motivation

I've been using it for inertial pose estimation on wearable devices with multiple marg sensors (ex. 9, 17 marg sensors), and in multiple occasions have found it lacking in terms of support for vectorized operation, requiring extra lines of loop code to achieve the same results. I am aware that the library is intended for single AHRS data fusion, however, this change would not impact current functionality, while bringing big improvements in some cases.

Pitch

For applications were multiple ahrs sensors are used simultaneously (e.g. wearables), it would be beneficial to perform sensor fusion for all sensors in the same call. This can be easily achieved by using the vectorization provided by numpy, and perform operations only over the last axis in the arrays. This would decrease computation time when performing multiple sensor fusion and simplify the code in this cases. These benefits also extend to improved batch processing of multiple sequences of data, while maintaining current functionality.

For example, the TRIAD estimation can be vectorized as such:

    def estimate_vectorized(self, w1: np.ndarray, w2: np.ndarray, representation: str = 'rotmat') -> np.ndarray:
        w1, w2 = np.copy(w1), np.copy(w2)
        # Normalized Vectors
        w1 = w1 / np.linalg.norm(w1, axis=-1, keepdims=True)                                                             # (eq. 12-39a)
        w2 = w2 / np.linalg.norm(w2, axis=-1, keepdims=True)
        # First Triad
        w1xw2 = np.cross(w1, w2, axis=-1)
        s2 = w1xw2 / np.linalg.norm(w1xw2, axis=-1, keepdims=True)                                                # (eq. 12-39b)
        s3 = np.cross(w1, w1xw2, axis=-1) / np.linalg.norm(w1xw2, axis=-1, keepdims=True)          # (eq. 12-39c)
        # Second Triad
        v1xv2 = np.cross(self.v1, self.v2, axis=-1)
        r2 = v1xv2 / np.linalg.norm(v1xv2, axis=-1, keepdims=True)
        r3 = np.cross(self.v1, v1xv2, axis=-1) / np.linalg.norm(v1xv2, axis=-1, keepdims=True)
        # Solve TRIAD
        Mb = np.stack([w1, s2, s3], axis=-1)                            # (eq. 12-41)
        Mr = np.stack([self.v1, r2, r3], axis=-1)                        # (eq. 12-42)
        A = np.matmul(Mb,Mr.transpose((-1,-2)))                    # (eq. 12-45)

        # Return according to desired representation
        #if representation.lower() == 'quaternion':
        #    return chiaverini(A)                                        # would also need vectorizing
        return A

Some example results:

import time
import numpy as np

from ahrs.filters.triad import TRIAD

triad = TRIAD()
triad.v1 = np.array([0.0, 0.0, 1.0])                          # Normalized reference gravity vector (g)
triad.v2 = np.array([0.4328755, 0.02747412, 0.90103495])      # Normalized reference geomagnetic field (h)
a = np.array([4.098297, 8.663757, 2.1355896])                 # Measured acceleration
m = np.array([-28715.50512, -25927.43566, 4756.83931])        # Measured magnetic field

# repeated values over multiple dimensions to check vectorized version
repeat_a = np.tile(a, (1, 16, 7, 9, 1))
repeat_m = np.tile(m, (1, 16, 7, 9, 1))

# example in TRIAD docs
stime = time.perf_counter()
o1 = triad.estimate(w1=a, w2=m)
etime = time.perf_counter() - stime

# perform same computation over multiple dimensions simultaneously
stime_vec = time.perf_counter()
o2 = triad.estimate_vectorized(w1=repeat_a, w2=repeat_m)
etime_vec = time.perf_counter() - stime_vec

print(a.shape, m.shape, repeat_a.shape, repeat_m.shape)    # [out]: (3,) (3,) (1, 16, 7, 9, 1) (1, 16, 7, 9, 1)
print(o1.shape, o2.shape)                                  # [out]: (3, 3) (1, 16, 7, 9, 3, 3)

# results are still the same
print(np.all(o1 == o2[0,0,0,0]))                           # [out]: True

# way faster computation than having to loop through data in all dimensions (although might have slight overhead for a single sample)
print(np.round(etime, 5), np.round(etime_vec, 5))          # [out]: 0.00033 0.00082

Discussion

This changes would take some time to implement and test since almost all functions would need to be refactored, however, when creating new code, it can be implemented from the beginning to support vectorization, which does not take much more time. It would also not require any additional dependency, current code usage is still preserved and the overhead for processing a single sample is minimal. This also removes dimension checks in many functions.

Thank you a lot for this great library! Let me know what you think of these changes!

[Mayitzin]

Excellent recommendation! I will finish the unit tests for these estimators, so that they can be easily tested.