Fix: use alpha + tau as the branch threshold in Huber.prox

Author: raashish1601

pyproximal/proximal/Huber.py

@@ -36,10 +36,14 @@ class Huber(ProxOperator):
 
         \prox_{\tau H_\alpha(\cdot)}(\mathbf{x}) =
         \begin{cases}
-        \prox_{\frac{\tau}{2 \alpha} |x_i|^2}(x_i), & |x_i| \leq \alpha \\
-        \prox_{\tau |x_i|}(x_i), & |x_i| > \alpha
+        \prox_{\frac{\tau}{2 \alpha} |x_i|^2}(x_i), & |x_i| \leq \alpha + \tau \\
+        \prox_{\tau |x_i|}(x_i), & |x_i| > \alpha + \tau
         \end{cases}
 
+    The branches switch at :math:`|x_i| = \alpha + \tau` rather than at
+    :math:`|x_i| = \alpha`, as that is where the minimizer moves from the
+    quadratic to the linear part of the Huber function.
+
     """
 
     def __init__(self, alpha: float) -> None:
@@ -60,12 +64,11 @@ def __call__(self, x: NDArray) -> float:
     def prox(self, x: NDArray, tau: float) -> NDArray:
         y = np.zeros_like(x)
         xabs = np.abs(x)
-        mask = xabs > self.alpha
+        mask = xabs > self.alpha + tau
         y[~mask] = self.l2.prox(x[~mask], tau)
         y[mask] = self.l1.prox(x[mask], tau)
-        # alternative from https://math.stackexchange.com/questions/1650411/
-        # proximal-operator-of-the-huber-loss-function... currently commented
-        # as it does not provide the same result
+        # equivalent closed form, from https://math.stackexchange.com/questions/1650411/
+        # proximal-operator-of-the-huber-loss-function
         # y = (1. - tau / np.maximum(np.abs(x), tau + self.alpha)) * x
 
         return y

pytests/test_norms.py

@@ -478,6 +478,21 @@ def test_Huber(par):
     assert moreau(hub, x, tau)
 
 
+@pytest.mark.parametrize("par", [(par1), (par2)])
+def test_Huber_prox(par):
+    """Huber proximal operator matches a brute-force minimizer, including for
+    alpha < |x| <= alpha + tau where the quadratic branch still applies
+    """
+    alpha, tau = par["sigma"], 1.5
+    hub = Huber(alpha=alpha)
+
+    x = np.linspace(-2 * (alpha + tau), 2 * (alpha + tau), 41)
+    ys = np.linspace(-3 * (alpha + tau), 3 * (alpha + tau), 120001)
+    hys = np.where(np.abs(ys) <= alpha, ys**2 / (2 * alpha), np.abs(ys) - alpha / 2)
+    expected = np.array([ys[np.argmin(tau * hys + (ys - xi) ** 2 / 2)] for xi in x])
+    assert_array_almost_equal(hub.prox(x, tau), expected, decimal=3)
+
+
 @pytest.mark.parametrize("par", [(par1), (par2)])
 def test_HuberCircular(par):
     """Circular Huber norm and proximal/dual proximal"""