diff --git a/backend/common/compute/estimate_distribution.py b/backend/common/compute/estimate_distribution.py index 55731838..85a66a76 100644 --- a/backend/common/compute/estimate_distribution.py +++ b/backend/common/compute/estimate_distribution.py @@ -5,27 +5,45 @@ from scipy import sparse from backend.common.constants import XApproximateDistribution -@numba.njit(fastmath=True, error_model="numpy", nogil=True) -def min_max(arr): +@numba.njit(error_model="numpy", nogil=True) +def min_max(arr: np.ndarray): """Return (min, max) values for the ndarray.""" - n = arr.size - odd = n % 2 - if not odd: - n -= 1 - max_val = min_val = arr[0] - i = 1 - while i < n: + + # initialize to first finite value in array. Normally, + # this will exit on the first value. + for i in range(arr.size): + min_val = max_val = arr[i] + if np.isfinite(min_val): + break + + # now find min/max, unrolled by two + odd = arr.size % 2 + unrolled_loop_limit = arr.size - 1 if odd else arr.size + i = 0 + while i < unrolled_loop_limit: x = arr[i] y = arr[i + 1] + + # ignore non-finites + x = x if np.isfinite(x) else min_val + y = y if np.isfinite(y) else min_val + if x > y: x, y = y, x min_val = min(x, min_val) max_val = max(y, max_val) i += 2 - if not odd: - x = arr[n] + + # handle the tail if any + if odd: + x = arr[arr.size - 1] + + # ignore non-finites + x = x if np.isfinite(x) else min_val + min_val = min(x, min_val) max_val = max(x, max_val) + return min_val, max_val @@ -38,6 +56,13 @@ def estimate_approximate_distribution(X) -> XApproximateDistribution: any (max-min) range in excess of 24 is implies tens of millions of observations of a single feature and so is extremely unlikely. """ + if X.dtype.kind not in ["i", "u", "f"]: + raise TypeError(f"Unsupported matrix dtype: {X.dtype.name}") + + if X.size == 0: + # default for empty array + return XApproximateDistribution.NORMAL + if sparse.isspmatrix_csc(X) or sparse.isspmatrix_csr(X): Xdata = X.data elif type(X) is np.ndarray: @@ -45,7 +70,7 @@ def estimate_approximate_distribution(X) -> XApproximateDistribution: X.size, ) else: - raise TypeError(f"Unsupported matrix type: {str(type(X))}") + raise TypeError(f"Unsupported matrix format: {str(type(X))}") CHUNKSIZE = 1 << 24 if Xdata.size > CHUNKSIZE: diff --git a/backend/test/test_server/unit/compute/test_est_dist.py b/backend/test/test_server/unit/compute/test_est_dist.py index b0457c02..2fb7bace 100644 --- a/backend/test/test_server/unit/compute/test_est_dist.py +++ b/backend/test/test_server/unit/compute/test_est_dist.py @@ -25,6 +25,9 @@ class EstDistTest(unittest.TestCase): def test_estimate_approximate_distribution(self): raw = np.random.exponential(scale=1000, size=(100, 40)) + # empty + self.assertEqual(estimate_approximate_distribution(np.zeros((0,))), XApproximateDistribution.NORMAL) + # ndarray self.assertEqual(estimate_approximate_distribution(raw), XApproximateDistribution.COUNT) self.assertEqual(estimate_approximate_distribution(np.log1p(raw)), XApproximateDistribution.NORMAL) @@ -35,7 +38,83 @@ class EstDistTest(unittest.TestCase): estimate_approximate_distribution(sparse.csr_matrix(np.log1p(raw))), XApproximateDistribution.NORMAL ) + # csc_matrix + self.assertEqual(estimate_approximate_distribution(sparse.csc_matrix(raw)), XApproximateDistribution.COUNT) + self.assertEqual( + estimate_approximate_distribution(sparse.csc_matrix(np.log1p(raw))), XApproximateDistribution.NORMAL + ) + # BIG (ie, trigger MT) big = np.random.exponential(scale=100, size=(1_000_000, 100)) self.assertEqual(estimate_approximate_distribution(big), XApproximateDistribution.COUNT) self.assertEqual(estimate_approximate_distribution(np.log1p(big)), XApproximateDistribution.NORMAL) + + def test_unsupported_throws(self): + # dtypes and matrix formats we do not support + with self.assertRaises(TypeError): + estimate_approximate_distribution(np.array(["a", "b"])) + with self.assertRaises(TypeError): + estimate_approximate_distribution(sparse.coo_matrix(np.array([[0, 1, 2], [3, 0, 2]]))) + + def test_nonfinites(self): + def put(arr, ind, vals): + # like np.put, but creates and returns a modified copy of original array + a = arr.copy() + np.put(a, ind, vals) + return a + + # non-finites + self.assertEqual(estimate_approximate_distribution(np.array([np.nan])), XApproximateDistribution.NORMAL) + self.assertEqual(estimate_approximate_distribution(np.array([np.PINF])), XApproximateDistribution.NORMAL) + self.assertEqual(estimate_approximate_distribution(np.array([np.NINF])), XApproximateDistribution.NORMAL) + self.assertEqual( + estimate_approximate_distribution(np.array([np.PINF, np.NINF, 0])), XApproximateDistribution.NORMAL + ) + self.assertEqual( + estimate_approximate_distribution(np.array([np.nan, np.PINF, np.NINF])), XApproximateDistribution.NORMAL + ) + + raw = np.random.exponential(scale=1000, size=(50, 3)) + logged = np.log1p(raw) + + self.assertEqual( + estimate_approximate_distribution(put(raw, [1], [np.nan])), + XApproximateDistribution.COUNT, + ) + self.assertEqual( + estimate_approximate_distribution(put(raw, [1], [np.PINF])), + XApproximateDistribution.COUNT, + ) + self.assertEqual( + estimate_approximate_distribution(put(raw, [1], [np.NINF])), + XApproximateDistribution.COUNT, + ) + self.assertEqual( + estimate_approximate_distribution(put(raw, [1, 3, 88], [np.nan, np.PINF, np.NINF])), + XApproximateDistribution.COUNT, + ) + self.assertEqual( + estimate_approximate_distribution(put(raw, [0, 1], [np.nan, np.nan])), + XApproximateDistribution.COUNT, + ) + + self.assertEqual( + estimate_approximate_distribution(put(logged, [1], [np.nan])), + XApproximateDistribution.NORMAL, + ) + self.assertEqual( + estimate_approximate_distribution(put(logged, [1], [np.PINF])), + XApproximateDistribution.NORMAL, + ) + self.assertEqual( + estimate_approximate_distribution(put(logged, [1], [np.NINF])), + XApproximateDistribution.NORMAL, + ) + self.assertEqual( + estimate_approximate_distribution(put(logged, [1, 3, 88], [np.nan, np.PINF, np.NINF])), + XApproximateDistribution.NORMAL, + ) + self.assertEqual( + estimate_approximate_distribution(put(logged, [0, 1], [np.nan, np.nan])), + XApproximateDistribution.NORMAL, + )