Joseba/multivariate alpha check

.gitignore

@@ -13,6 +13,7 @@
 # Virtual environment
 puncc-dev-env/
 puncc-user-env/
+.venv/
 
 # Local demo files
 **demo/

deel/puncc/api/utils.py

@@ -30,6 +30,7 @@
 from typing import Iterable
 from typing import Optional
 from typing import Tuple
+from typing import Union
 
 import numpy as np
 
@@ -144,7 +145,7 @@ def supported_types_check(*data: Iterable):
             )
 
 
-def alpha_calib_check(alpha: float, n: int, complement_check: bool = False):
+def alpha_calib_check(alpha: Union[float, np.ndarray], n: int, complement_check: bool = False):
     """Check if the value of alpha :math:`\\alpha` is consistent with the size
     of calibration set :math:`n`.
 
@@ -161,7 +162,7 @@ def alpha_calib_check(alpha: float, n: int, complement_check: bool = False):
 
         0 < \\alpha \cdot (1+1/n) < 1 \implies 0 < \\alpha < n/(n+1)
 
-    :param float alpha: target quantile order.
+    :param float or np.ndarray alpha: target quantile order.
     :param int n: size of the calibration dataset.
     :param bool complement_check: complementary check to compute the
         :math:`\\alpha \cdot (1+1/n)`-th quantile, required by some methods
@@ -171,26 +172,30 @@ def alpha_calib_check(alpha: float, n: int, complement_check: bool = False):
         calibration set.
     """
 
-    if alpha < 1 / (n + 1):
+    if np.any(alpha < 1 / (n + 1)):
         raise ValueError(
-            f"Alpha is too small: (alpha={alpha}, n={n}) {alpha} < {1/(n+1)}. "
+            f"Alpha is too small: (alpha={alpha}, n={n}) "
+            + f"{alpha} < {1/(n+1)} for at least one of its coordinates. "
             + "Increase alpha or the size of the calibration set."
         )
 
-    if alpha >= 1:
+    if np.any(alpha >= 1):
         raise ValueError(
-            f"Alpha={alpha} is too large. Decrease alpha such that alpha < 1."
+            f"Alpha={alpha} is too large. " 
+            + "Decrease alpha such that all of its coordinates are < 1."
         )
 
-    if complement_check and alpha > n / (n + 1):
+    if complement_check and np.any(alpha > n / (n + 1)):
         raise ValueError(
-            f"Alpha is too large: (alpha={alpha}, n={n}) {alpha} < {n/(n+1)}. "
+            f"Alpha is too large: (alpha={alpha}, n={n}) "
+            + f"{alpha} < {n/(n+1)} for at least on of its coordinates. "
             + "Decrease alpha or increase the size of the calibration set."
         )
 
-    if complement_check and alpha <= 0:
+    if complement_check and np.any(alpha <= 0):
         raise ValueError(
-            f"Alpha={alpha} is too small. Increase alpha such that alpha > 0."
+            f"Alpha={alpha} is too small. "
+            + "Increase alpha such that all of its coordinates > 0."
         )
 
 

tests/api/test_utils.py

@@ -133,6 +133,32 @@ def test_too_low_alpha(self):
         with self.assertRaises(ValueError):
             alpha_calib_check(alpha=0.01, n=10)
 
+    def test_proper_alpha_multidim(self):
+        max_n = 10000
+        n_vals = np.random.randint(10, high=max_n, size=100)
+        for n in n_vals:
+            lower_bound = 1 / (n + 1)
+            alpha_vals = (1 - lower_bound) * np.random.random_sample(
+                size=(100)
+            ) + lower_bound
+            alpha_calib_check(alpha_vals, n=n)
+
+    def test_out_lowerbound_alpha_multidim(self):
+        with self.assertRaises(ValueError):
+            alpha = np.array([0, 0.5, 0.5, 0.5])
+            alpha_calib_check(alpha, n=100000)
+
+    def test_out_upperbound_alpha_multidim(self):
+        with self.assertRaises(ValueError):
+            alpha = np.array([1, 0.5, 0.5, 0.5])
+            alpha_calib_check(alpha, n=100000)
+
+    def test_too_low_alpha_multidim(self):
+        with self.assertRaises(ValueError):
+            alpha = np.array([0.01, 0.5, 0.5, 0.5])
+            alpha_calib_check(alpha, n=10)
+
+
 
 class QuantileCheck(unittest.TestCase):
     def test_simple_quantile1d(self):