get_beta_softlabels now returns a numpy array

ayrna · Apr 16, 2024 · e2b0c83 · e2b0c83
1 parent 5a5243d
commit e2b0c83
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/dlordinal/soft_labelling/beta_distribution.py b/dlordinal/soft_labelling/beta_distribution.py
@@ -1,3 +1,4 @@
+import numpy as np
 from scipy.special import gamma, hyp2f1
 
 from .utils import get_intervals
@@ -285,4 +286,4 @@ def get_beta_softlabels(J, params_set="standard"):
         raise ValueError(f"Invalid params_set: {params_set}")
 
     params = _beta_params_sets[params_set]
-    return [_get_beta_softlabel(J, p, q, a) for (p, q, a) in params[J]]
+    return np.array([_get_beta_softlabel(J, p, q, a) for (p, q, a) in params[J]])
diff --git a/dlordinal/soft_labelling/tests/test_beta_distribution.py b/dlordinal/soft_labelling/tests/test_beta_distribution.py
@@ -95,8 +95,9 @@ def test_beta_softlabels():
     n = 5
     result = get_beta_softlabels(n)
 
-    assert len(result) == n
-    assert len(result[0]) == n
+    assert len(result.shape) == 2
+    assert result.shape[0] == n
+    assert result.shape[1] == n
 
     expected_result = [
         [