feat: added optional tau/mu to Segment

Author: mrava87

pyproximal/optimization/segmentation.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 from pylops import BlockDiag, Gradient
+from pylops.utils.backend import get_array_module
 from pylops.utils.typing import NDArray
 
 from pyproximal import L21, Simplex, VStack
@@ -14,6 +15,8 @@ def Segment(
     cl: NDArray,
     sigma: float,
     alpha: float,
+    tau: float | None = 1.0,
+    mu: float | None = None,
     clsigmas: NDArray | None = None,
     z: NDArray | None = None,
     niter: int = 10,
@@ -37,6 +40,14 @@ def Segment(
         Positive scalar weight of the misfit term
     alpha : :obj:`float`
         Positive scalar weight of the regularization term
+    tau : :obj:`float`, optional
+        Stepsize of subgradient of segmentation term.
+        If ``None``, it is set from ``mu`` and the
+        Lipschitz constant of the gradient operator.
+    mu : :obj:`float`, optional
+        Stepsize of subgradient of regularization term.
+        If ``None``, it is set from ``tau`` and the
+        Lipschitz constant of the gradient operator.
     clsigmas : :obj:`numpy.ndarray`, optional
         Classes standard deviations
     z : :obj:`numpy.ndarray`, optional
@@ -64,6 +75,11 @@ def Segment(
         Estimated classes. This is a vector of the same size of the input data
         ``y`` with the selected classes at each pixel.
 
+    Raises
+    ------
+    ValueError
+        If both ``tau`` and ``mu`` are ``None``
+
     Notes
     -----
     This solver performs image segmentation over :math:`N_{cl}` classes solving
@@ -90,6 +106,11 @@ def Segment(
         Imaging and Vision, 40, 8pp. 120–145. 2011.
 
     """
+    if tau is None and mu is None:
+        msg = "Either tau or mu must be provided."
+        raise ValueError(msg)
+
+    ncp = get_array_module(y)
     kwargs_simplex = {} if kwargs_simplex is None else kwargs_simplex
 
     dims = y.shape
@@ -104,10 +125,7 @@ def Segment(
     g = g.ravel()
 
     # Gradient operator
-    sampling = 1.0
-    Gop = Gradient(
-        dims=dims, sampling=sampling, edge=False, kind="forward", dtype="float64"
-    )
+    Gop = Gradient(dims=dims, sampling=1.0, edge=False, kind="forward", dtype="float64")
     Gop = BlockDiag([Gop] * ncl)
 
     # Simplex and L21 proximal operators
@@ -119,9 +137,9 @@ def Segment(
     )
 
     # Steps
-    L = 8.0 / sampling**2
-    tau = 1.0
-    mu = 1.0 / (tau * L)
+    L = 4 * ndims
+    tau = 1.0 / (mu * L) if tau is None else tau
+    mu = 1.0 / (tau * L) if mu is None else mu
 
     # Inversion
     x: NDArray = PrimalDual(
@@ -132,14 +150,14 @@ def Segment(
         mu=mu,
         z=g if z is None else g + z,
         theta=1.0,
-        x0=np.zeros_like(g) if x0 is None else x0,
+        x0=ncp.zeros_like(g) if x0 is None else x0,
         niter=niter,
         callback=callback,
         show=show,
         returny=False,
     )
     x = x.reshape(ncl, dimsprod).T
-    cl = np.argmax(x, axis=1)
+    cl = ncp.argmax(x, axis=1)
     cl = cl.reshape(dims)
 
     return x, cl