Anisotropic regularization parameter

examples/regularization.py

@@ -8,25 +8,43 @@
 # Load 3D image
 image_folder = f"{os.path.dirname(__file__)}/images/"
 img_path = image_folder + "dicom_3d.npy"
-img = darsia.imread(img_path, dimensions=[0.1, 0.1, 0.1])
+img = darsia.imread(img_path, dimensions=[0.1, 0.1, 0.1], dim=3)
 
-# Make regularization parameter heterogenous (for illustration purposes)
+print("Image shape: ", img.img.shape)
+
+# Make regularization parameter heterogenous and anisotropic (for illustration purposes)
 mu = np.ones_like(img.img)
 mu[:, 0:100, :] = 0.5
 
+mu_anisotropic = [10 * mu, mu, mu]
+
 # Regularize image using anisotropic tvd
-img_regularized_tvd = darsia.tvd(
+img_regularized_anisotropic_tvd = darsia.tvd(
     img=img,
     method="heterogeneous bregman",
     isotropic=False,
+    weight=mu_anisotropic,
+    omega=0.1,
+    max_num_iter=30,
+    eps=1e-6,
+    dim=3,
+    verbose=True,
+    solver=darsia.Jacobi(maxiter=20),
+    # solver=darsia.CG(maxiter=20, tol=1e-3),
+)
+
+img_regularized_tvd = darsia.tvd(
+    img=img.img,
+    method="heterogeneous bregman",
+    isotropic=False,
     weight=mu,
     omega=0.1,
     max_num_iter=30,
     eps=1e-6,
     dim=3,
     verbose=True,
     solver=darsia.Jacobi(maxiter=20),
-    # solver = darsia.CG(maxiter = 20, tol = 1e-3)
+    # solver=darsia.CG(maxiter=20, tol=1e-3),
 )
 
 # Regularize image using H1 regularization
@@ -35,14 +53,16 @@
     mu=1,
     omega=1,
     dim=3,
-    solver=darsia.Jacobi(maxiter=30, tol=1e-6, verbose=True),
+    solver=darsia.Jacobi(maxiter=30, tol=1e-4, verbose=True),
     # solver = darsia.MG(3)
 )
 
 plt.figure("img slice")
 plt.imshow(img.img[100, :, :])
+plt.figure("regularized anisotropic weight tvd slice")
+plt.imshow(img_regularized_anisotropic_tvd.img[100, :, :])
 plt.figure("regularized tvd slice")
-plt.imshow(img_regularized_tvd.img[100, :, :])
+plt.imshow(img_regularized_tvd[100, :, :])
 plt.figure("regularized H1 slice")
-plt.imshow(img_regularized_H1.img[100, :, :])
+plt.imshow(img_regularized_h1.img[100, :, :])
 plt.show()

src/darsia/restoration/split_bregman_tvd.py

@@ -11,8 +11,8 @@
 
 
 def split_bregman_tvd(
-    img: np.ndarray,
-    mu: Union[float, np.ndarray] = 1.0,
+    img: Union[np.ndarray, da.Image],
+    mu: Union[float, np.ndarray, list] = 1.0,
     omega: Union[float, np.ndarray] = 1.0,
     ell: Optional[Union[float, np.ndarray]] = None,
     dim: int = 2,
@@ -31,8 +31,10 @@ def split_bregman_tvd(
     shrinkage step. The regularization term is weighted by the parameter ell.
 
     Args:
-        img (array): image array
-        mu (float or array): TV penalization parameter
+        img (array, da.Image): image
+        mu (float or array or list): inverse TV penalization parameter, if list
+            it must have the lenght of the dimension and each entry corresponds
+            to the penalization in the respective direction.
         omega (float or array): mass penalization parameter
         ell (float or array): regularization parameter
         dim (int): spatial dimension of the image
@@ -48,16 +50,26 @@ def split_bregman_tvd(
         array: denoised image
 
     """
+
     # Keep track of input type and convert input image to float for further calculations
     img_dtype = img.dtype
 
     # Store input image and its norm for convergence check
     img_nrm = np.linalg.norm(img)
 
+    # Controll that mu has correct length if its a list
+    if isinstance(mu, list):
+        assert len(mu) == dim, "mu must be a list of length dim"
+
     # Feed the solver with parameters, follow the suggestion to use double the weight
-    # for the diffusion coefficient if no value is provided.
+    # for the mass coefficient if no value is provided.
     if ell is None:
-        ell = 2 * mu
+        if isinstance(omega, float):
+            ell = 2 * omega
+        elif isinstance(omega, np.ndarray):
+            ell = 2 * omega.copy()
+            ell[ell == 0] = 1
+
     solver.update_params(
         mass_coeff=omega,
         diffusion_coeff=ell,
@@ -66,12 +78,20 @@ def split_bregman_tvd(
 
     # Define energy functional for verbose
     def _functional(x: np.ndarray) -> float:
-        return 0.5 * np.linalg.norm(omega * (x - img)) ** 2 + sum(
-            [
-                np.sum(np.abs(mu * da.backward_diff(img=x, axis=j, dim=dim)))
-                for j in range(dim)
-            ]
-        )
+        if isinstance(mu, list):
+            return 0.5 * np.linalg.norm(omega * (x - img)) ** 2 + sum(
+                [
+                    np.sum(np.abs(mu[j] * da.backward_diff(img=x, axis=j, dim=dim)))
+                    for j in range(dim)
+                ]
+            )
+        else:
+            return 0.5 * np.linalg.norm(omega * (x - img)) ** 2 + sum(
+                [
+                    np.sum(np.abs(mu * da.backward_diff(img=x, axis=j, dim=dim)))
+                    for j in range(dim)
+                ]
+            )
 
     # Define shrinkage operator (shrinks element-wise by k)
     def _shrink(x: np.ndarray, k: Union[float, np.ndarray]) -> np.ndarray:
@@ -86,7 +106,7 @@ def _rhs_function(dt: np.ndarray, bt: np.ndarray) -> np.ndarray:
             ]
         )
 
-    # Define initial guess if provided, otherwise start with input image and allovate
+    # Define initial guess if provided, otherwise start with input image and allocate
     # zero arrays for the split Bregman variables.
     if x0 is not None:
         img0, d0, b0 = x0
@@ -115,6 +135,11 @@ def _rhs_function(dt: np.ndarray, bt: np.ndarray) -> np.ndarray:
             shrinkage_factor = np.maximum(s - mu / ell, 0) / (s + 1e-18)
             d = dub * shrinkage_factor[..., None]
             b = dub - d
+        elif isinstance(mu, list):
+            for j in range(dim):
+                dub = da.backward_diff(img=img_new, axis=j, dim=dim) + b[..., j]
+                d[..., j] = _shrink(dub, mu[j] / ell)
+                b[..., j] = dub - d[..., j]
         else:
             for j in range(dim):
                 dub = da.backward_diff(img=img_new, axis=j, dim=dim) + b[..., j]
@@ -123,16 +148,17 @@ def _rhs_function(dt: np.ndarray, bt: np.ndarray) -> np.ndarray:
 
         # Monitor performance
         relative_increment = np.linalg.norm(img_new - img_iter) / img_nrm
+
+        # Update of result
+        img_iter = img_new.copy()
+
         if verbose if isinstance(verbose, bool) else iter % verbose == 0:
             print(
                 f"""Split Bregman iteration {iter} - """
-                f"""relative increment: {relative_increment}, """
+                f"""relative increment: {round(relative_increment,5)}, """
                 f"""energy functional: {_functional(img_iter)}"""
             )
 
-        # Update of result
-        img_iter = img_new.copy()
-
         # Convergence check
         if eps is not None:
             if relative_increment < eps:

src/darsia/utils/derivatives.py

@@ -13,7 +13,7 @@
 
 
 def backward_diff(
-    img: np.ndarray, axis: int, dim: int = 2, h: Optional[float] = None
+    img: np.ndarray, axis: int, dim: int = 2, h: Optional[Union[float, list]] = None
 ) -> np.ndarray:
     """Backward difference of image matrix in direction of axis.
 
@@ -31,21 +31,24 @@ def backward_diff(
     if h is None:
         return np.diff(img, axis=axis, append=da.array_slice(img, axis, -1, None, 1))
     else:
+        if isinstance(h, list):
+            assert len(h) == dim, "h must be a list of length dim"
+            h = h[axis]
         return (
             np.diff(img, axis=axis, append=da.array_slice(img, axis, -1, None, 1)) / h
         )
 
 
 def forward_diff(
-    img: np.ndarray, axis: int, dim: int = 2, h: Optional[float] = None
+    img: np.ndarray, axis: int, dim: int = 2, h: Optional[Union[float, list]] = None
 ) -> np.ndarray:
     """Forward difference of image matrix in direction of axis.
 
     Args:
         img (np.ndarray): image array
         axis (int): axis along which the difference is taken
         dim (int): dimension of image array
-        h (Optional[float]): grid spacing
+        h (Optional[float, list]): grid spacing
 
     Returns:
         np.ndarray: forward difference image matrix
@@ -55,14 +58,17 @@ def forward_diff(
     if h is None:
         return np.diff(img, axis=axis, prepend=da.array_slice(img, axis, 0, 1, 1))
     else:
+        if isinstance(h, list):
+            assert len(h) == dim, "h must be a list of length dim"
+            h = h[axis]
         return np.diff(img, axis=axis, prepend=da.array_slice(img, axis, 0, 1, 1)) / h
 
 
 def laplace(
     img: np.ndarray,
     axis: Optional[int] = None,
     dim: int = 2,
-    h: Optional[float] = None,
+    h: Optional[Union[float, list]] = None,
     diffusion_coeff: Union[np.ndarray, float] = 1,
 ) -> np.ndarray:
     """Laplace operator with homogeneous boundary conditions.
@@ -73,7 +79,7 @@ def laplace(
         img (np.ndarray): image array
         axis (int): axis along which the difference is taken
         dim (int): dimension of image array
-        h (Optional[float]): grid spacing
+        h (Optional[float, list]): grid spacing
         diffision_coeff (Optional[np.ndarray]): diffusion coefficient
 
     Returns:

src/darsia/utils/linear_solvers/cg.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Union
 
 import numpy as np
 import scipy.sparse as sps
@@ -14,7 +14,7 @@ class CG(da.Solver):
     """
 
     def __call__(
-        self, x0: np.ndarray, rhs: np.ndarray, h: Optional[float] = None
+        self, x0: np.ndarray, rhs: np.ndarray, h: Optional[Union[float, list]] = None
     ) -> np.ndarray:
         """Solve the problem.
 

src/darsia/utils/linear_solvers/jacobi.py

@@ -3,7 +3,7 @@
 """
 from __future__ import annotations
 
-from typing import Union
+from typing import Optional, Union
 
 import numpy as np
 
@@ -21,7 +21,9 @@ class Jacobi(da.Solver):
 
     """
 
-    def _neighbor_accumulation(self, im: np.ndarray) -> np.ndarray:
+    def _neighbor_accumulation(
+        self, im: np.ndarray, h: Optional[Union[float, list]] = None
+    ) -> np.ndarray:
         """Accumulation of neighbor pixels.
 
         Accumulates for each entry, the entries of neighbors, regardless of dimension.
@@ -37,44 +39,72 @@ def _neighbor_accumulation(self, im: np.ndarray) -> np.ndarray:
 
         """
         im_av: np.ndarray = np.zeros_like(im)
-        for ax in range(im.ndim):
-            im_av += np.concatenate(
-                (da.array_slice(im, ax, 0, 1), da.array_slice(im, ax, 0, -1)), axis=ax
-            ) + np.concatenate(
-                (da.array_slice(im, ax, 1, None), da.array_slice(im, ax, -1, None)),
-                axis=ax,
-            )
+        if h is None:
+            for ax in range(self.dim):
+                im_av += np.concatenate(
+                    (da.array_slice(im, ax, 0, 1), da.array_slice(im, ax, 0, -1)),
+                    axis=ax,
+                ) + np.concatenate(
+                    (da.array_slice(im, ax, 1, None), da.array_slice(im, ax, -1, None)),
+                    axis=ax,
+                )
+        else:
+            if isinstance(h, float):
+                h = [h] * self.dim
+            for ax in range(self.dim):
+                im_av += (
+                    np.concatenate(
+                        (da.array_slice(im, ax, 0, 1), da.array_slice(im, ax, 0, -1)),
+                        axis=ax,
+                    )
+                    + np.concatenate(
+                        (
+                            da.array_slice(im, ax, 1, None),
+                            da.array_slice(im, ax, -1, None),
+                        ),
+                        axis=ax,
+                    )
+                ) / h[ax] ** 2
         return im_av
 
-    def _diag(self, h: float = 1) -> Union[float, np.ndarray]:
+    def _diag(self, h: Optional[Union[float, list]] = None) -> Union[float, np.ndarray]:
         """
         Compute diagonal of the stiffness matrix.
 
         The stiffness matrix is understood in finite difference sense, with the
-        possibility to identify pixel sizes with phyiscal mesh sizes. This is a helper
+        possibility to identify pixel sizes with phyiscal voxel sizes. This is a helper
         function for the main Jacobi solver.
 
         Args:
-            h (float): mesh diameter
+            h (float): voxel size
 
         Returns:
             Union[float, np.ndarray]: diagonal of the stiffness matrix
 
         """
-        return self.mass_coeff + self.diffusion_coeff * 2 * self.dim / h**2
+        if h is None:
+            voxel_weights = self.dim
+        else:
+            if isinstance(h, float):
+                voxel_weights = self.dim / (h**2)
+            else:
+                assert len(h) == self.dim, "h must be a float or a list of length dim"
+                voxel_weights = sum([1 / h[i] ** 2 for i in range(self.dim)])
+
+        return self.mass_coeff + 2 * self.diffusion_coeff * voxel_weights
 
     def __call__(
         self,
         x0: np.ndarray,
         rhs: np.ndarray,
-        h: float = 1.0,
+        h: Optional[Union[float, list]] = None,
     ) -> np.ndarray:
         """One iteration of a Jacobi solver for linear systems.
 
         Args:
             x0 (np.ndarray): initial guess
             rhs (np.ndarray): right hand side of the linear system
-            h (float): mesh diameter
+            h (optional[float, list]): voxel size
 
         Returns:
             np.ndarray: solution to the linear system
@@ -88,14 +118,14 @@ def __call__(
         if self.tol is None:
             for _ in range(self.maxiter):
                 x = (
-                    rhs + self.diffusion_coeff * self._neighbor_accumulation(x) / h**2
+                    rhs + self.diffusion_coeff * self._neighbor_accumulation(x, h)
                 ) / diag
                 if self.verbose:
                     print(f"Jacobi iteration {_} of {self.maxiter} completed.")
         else:
             for _ in range(self.maxiter):
                 x_new = (
-                    rhs + self.diffusion_coeff * self._neighbor_accumulation(x) / h**2
+                    rhs + self.diffusion_coeff * self._neighbor_accumulation(x, h)
                 ) / diag
                 err = np.linalg.norm(x_new - x) / np.linalg.norm(x0)
                 if err < self.tol: