Optimize multiannotator.py for performance

[gogetron]

Summary

This PR partially addresses #862

🎯 Purpose: Improve performance of get_label_quality_multiannotator with default arguments.

[ ✏️ Write your summary here. ]
After profiling this function, the parts that took most time were a few loops that were reimplemented using numpy operations and the apply_along_axis methods which were also reimplemented with faster numpy operations.

I noticed there was a small issue with the warning in assert_valid_inputs_multiannotator. The warning was only raised when the labels were sorted. This can be reproduced with the code in the added test:

import numpy as np

from cleanlab.internal.multiannotator_utils import assert_valid_inputs_multiannotator

not_agree_labels = np.array([[2, 1, np.nan, 0], [1, 0, 2, np.nan]])
assert_valid_inputs_multiannotator(not_agree_labels)
# Does not raise any warning

The function will raise the warning correctly now and a new test was added to assert this behavior.

For memory I used the memory-profiler library. The code I used for benchmarking is copied below. In addition I sorted the imports in the modified files.

Code Setup

import numpy as np

from cleanlab.multiannotator import get_label_quality_multiannotator

np.random.seed(0)
%load_ext memory_profiler


M = 12
N = 500_000
K = 20
labels = np.random.randint(K, size=(N, M))
pred_probs = np.random.random((N, K))
pred_probs /= pred_probs.sum(axis=1, keepdims=True)

# Execute once to avoid the tensorflow import time in the benchmark.
_ = get_label_quality_multiannotator(labels[:5, :], pred_probs[:5, :])

labels_with_some_nans = labels.astype(np.float64)
labels_with_some_nans[: labels_with_some_nans.shape[0] // 2, np.random.randint(M, size=(M))] = (
    np.nan
)

Current version

%%timeit
%memit get_label_quality_multiannotator(labels_with_some_nans, pred_probs)
# peak memory: 1125.32 MiB, increment: 319.85 MiB
# peak memory: 1178.05 MiB, increment: 316.38 MiB
# peak memory: 1125.34 MiB, increment: 237.05 MiB
# peak memory: 1125.34 MiB, increment: 261.64 MiB
# peak memory: 1147.72 MiB, increment: 284.02 MiB
# peak memory: 1147.72 MiB, increment: 283.83 MiB
# peak memory: 1147.36 MiB, increment: 281.45 MiB
# peak memory: 1147.72 MiB, increment: 281.81 MiB
# 3min 24s ± 1.14 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

This PR

%%timeit
%memit get_label_quality_multiannotator(labels_with_some_nans, pred_probs)
# peak memory: 1131.42 MiB, increment: 298.31 MiB
# peak memory: 1131.44 MiB, increment: 268.32 MiB
# peak memory: 1152.99 MiB, increment: 289.86 MiB
# peak memory: 1131.69 MiB, increment: 268.56 MiB
# peak memory: 1153.24 MiB, increment: 289.86 MiB
# peak memory: 1153.25 MiB, increment: 289.86 MiB
# peak memory: 1119.70 MiB, increment: 256.31 MiB
# peak memory: 1153.25 MiB, increment: 289.86 MiB
# 12.5 s ± 95.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Testing

🔍 Testing Done: Existing test suite and a new added test.

References

Reviewer Notes

💡 Include any specific points for the reviewer to consider during their review.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 96.21%. Comparing base (abd0924) to head (5596137).
Report is 31 commits behind head on master.

❗ Current head 5596137 differs from pull request most recent head 3182b56. Consider uploading reports for the commit 3182b56 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1077      +/-   ##
==========================================
+ Coverage   96.15%   96.21%   +0.05%     
==========================================
  Files          74       74              
  Lines        5850     5861      +11     
  Branches     1044     1044              
==========================================
+ Hits         5625     5639      +14     
+ Misses        134      132       -2     
+ Partials       91       90       -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[huiwengoh]

Thank you so much for the amazing contribution and speedup on the multiannotator module!

I made some comments, mostly about code readability (many of these methods are doing many non-trivial operations, so it would be super helpful to write some comments documenting the steps).

Besides that, the only real issue I found is in the assert_valid_inputs_multiannotator method, made a comment there with more details.

Please don't hesitate to ask for more clarification if anything isn't clear!

[huiwengoh]

I made a small edit to make sure that the mask variable names are clear, but otherwise the changes look good to me!

Awesome work on the speedup @gogetron!