From 098c018f38c9e7f825814221448252fca3ca6bef Mon Sep 17 00:00:00 2001
From: EStorvik <erlend.storvik@hvl.no>
Date: Fri, 14 Jul 2023 10:59:40 +0200
Subject: [PATCH 1/3] Tried to implement with voxel size, but it works poorly

---
 examples/regularization.py | 42 ++++++++++++++++++++++++++++++--------
 1 file changed, 33 insertions(+), 9 deletions(-)

diff --git a/examples/regularization.py b/examples/regularization.py
index ae8513d4..313ab315 100644
--- a/examples/regularization.py
+++ b/examples/regularization.py
@@ -6,27 +6,49 @@
 import darsia
 
 # Load 3D image
-image_folder = f"{os.path.dirname(__file__)}/images/"
+# image_folder = f"{os.path.dirname(__file__)}/images/"
+image_folder = "C:/Users/erst/src/image_analysis/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)
+
+print("Image shape: ", img.img.shape)
 
 # Make regularization parameter heterogenous (for illustration purposes)
+mu_physical = 4.7e-4 * np.ones_like(img.img)
+mu_physical[:, 0:100, :] = 0.5 * 4.7e-4
+
+mu_physical = 4.7e-4
+
 mu = np.ones_like(img.img)
 mu[:, 0:100, :] = 0.5
 
 # Regularize image using anisotropic tvd
-img_regularized_tvd = darsia.tvd(
+img_regularized_physical_tvd = darsia.tvd(
     img=img,
     method="heterogeneous bregman",
     isotropic=False,
-    weight=mu,
+    weight=mu_physical,
     omega=0.1,
     max_num_iter=30,
-    eps=1e-6,
+    eps=1e-4,
     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),
+)
+
+img_regularized_tvd = darsia.tvd(
+    img=img.img,
+    method="heterogeneous bregman",
+    isotropic=False,
+    weight=mu,
+    omega=0.1,
+    max_num_iter=30,
+    eps=1e-4,
+    dim=3,
+    verbose=True,
+    # solver=darsia.Jacobi(maxiter=20),
+    solver=darsia.CG(maxiter=20, tol=1e-3),
 )
 
 # Regularize image using H1 regularization
@@ -35,14 +57,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 physical tvd slice")
+plt.imshow(img_regularized_physical_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()

From 478f7111855699eb8818f89eae28529c262ce0ad Mon Sep 17 00:00:00 2001
From: EStorvik <erlend.storvik@hvl.no>
Date: Fri, 14 Jul 2023 10:59:57 +0200
Subject: [PATCH 2/3] Tried to implement with voxel size but it works poorly

---
 src/darsia/restoration/split_bregman_tvd.py | 65 ++++++++++++++++-----
 src/darsia/restoration/tvd.py               | 13 ++++-
 src/darsia/utils/derivatives.py             | 16 +++--
 src/darsia/utils/linear_solvers/cg.py       |  4 +-
 src/darsia/utils/linear_solvers/jacobi.py   | 64 ++++++++++++++------
 src/darsia/utils/linear_solvers/solver.py   |  5 +-
 6 files changed, 127 insertions(+), 40 deletions(-)

diff --git a/src/darsia/restoration/split_bregman_tvd.py b/src/darsia/restoration/split_bregman_tvd.py
index 24286373..22307f68 100644
--- a/src/darsia/restoration/split_bregman_tvd.py
+++ b/src/darsia/restoration/split_bregman_tvd.py
@@ -11,7 +11,7 @@
 
 
 def split_bregman_tvd(
-    img: np.ndarray,
+    img: Union[np.ndarray, da.Image],
     mu: Union[float, np.ndarray] = 1.0,
     omega: Union[float, np.ndarray] = 1.0,
     ell: Optional[Union[float, np.ndarray]] = None,
@@ -31,7 +31,7 @@ def split_bregman_tvd(
     shrinkage step. The regularization term is weighted by the parameter ell.
 
     Args:
-        img (array): image array
+        img (array, da.Image): image
         mu (float or array): TV penalization parameter
         omega (float or array): mass penalization parameter
         ell (float or array): regularization parameter
@@ -48,16 +48,33 @@ def split_bregman_tvd(
         array: denoised image
 
     """
+
+    # Check input image and extract voxel size (h)
+    if isinstance(img, da.Image):
+        if img.series:
+            voxel_size = [img.dimensions[i] / img.img.shape[i] for i in range(dim)]
+        else:
+            voxel_size = img.voxel_size
+        img = img.img
+    else:
+        voxel_size = np.ones(dim).tolist()
+
     # 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)
 
     # 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,
@@ -68,7 +85,11 @@ def split_bregman_tvd(
     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)))
+                np.sum(
+                    np.abs(
+                        mu * da.backward_diff(img=x, axis=j, dim=dim, h=voxel_size[j])
+                    )
+                )
                 for j in range(dim)
             ]
         )
@@ -81,12 +102,14 @@ def _shrink(x: np.ndarray, k: Union[float, np.ndarray]) -> np.ndarray:
     def _rhs_function(dt: np.ndarray, bt: np.ndarray) -> np.ndarray:
         return omega * img + ell * sum(
             [
-                da.forward_diff(img=bt[..., i] - dt[..., i], axis=i, dim=dim)
+                da.forward_diff(
+                    img=bt[..., i] - dt[..., i], axis=i, dim=dim, h=voxel_size[i]
+                )
                 for i in range(dim)
             ]
         )
 
-    # 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
@@ -104,35 +127,49 @@ def _rhs_function(dt: np.ndarray, bt: np.ndarray) -> np.ndarray:
     # Bregman iterations
     for iter in range(max_num_iter):
         # First step - solve the stabilized diffusion system.
-        img_new = solver(x0=img_iter, rhs=_rhs_function(d, b))
+        img_new = solver(x0=img_iter, rhs=_rhs_function(d, b), h=voxel_size)
 
         # Second step - shrinkage.
         if isotropic:
             dub = b.copy()
             for j in range(dim):
-                dub[..., j] += da.backward_diff(img=img_new, axis=j, dim=dim)
+                dub[..., j] += da.backward_diff(
+                    img=img_new, axis=j, dim=dim, voxel_size=voxel_size[j]
+                )
             s = np.linalg.norm(dub, 2, axis=-1)
             shrinkage_factor = np.maximum(s - mu / ell, 0) / (s + 1e-18)
             d = dub * shrinkage_factor[..., None]
             b = dub - d
+        elif isinstance(mu, np.ndarray) and (mu.shape != img_new.shape):
+            for j in range(dim):
+                dub = (
+                    da.backward_diff(img=img_new, axis=j, dim=dim, h=voxel_size[j])
+                    + b[..., j]
+                )
+                d[..., j] = _shrink(dub, mu[j, ...] / ell[j, ...])
+                b[..., j] = dub - d[..., j]
         else:
             for j in range(dim):
-                dub = da.backward_diff(img=img_new, axis=j, dim=dim) + b[..., j]
+                dub = (
+                    da.backward_diff(img=img_new, axis=j, dim=dim, h=voxel_size[j])
+                    + b[..., j]
+                )
                 d[..., j] = _shrink(dub, mu / ell)
                 b[..., j] = dub - d[..., j]
 
         # 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:
diff --git a/src/darsia/restoration/tvd.py b/src/darsia/restoration/tvd.py
index 3bf517f7..f074a881 100644
--- a/src/darsia/restoration/tvd.py
+++ b/src/darsia/restoration/tvd.py
@@ -125,7 +125,18 @@ def _tvd_image(self, img: darsia.Image) -> darsia.Image:
 
         """
         img_copy = img.copy()
-        img_copy.img = self._tvd_array(img.img)
+        if self.method == "heterogeneous bregman":
+            img_copy.img = darsia.split_bregman_tvd(
+                img,
+                mu=self.weight,
+                omega=self.omega,
+                ell=self.regularization,
+                max_num_iter=self.max_num_iter,
+                eps=self.eps,
+                **self.kwargs,
+            )
+        else:
+            img_copy.img = self._tvd_array(img.img)
         return img_copy
 
 
diff --git a/src/darsia/utils/derivatives.py b/src/darsia/utils/derivatives.py
index 66a2feb7..334f398e 100644
--- a/src/darsia/utils/derivatives.py
+++ b/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,13 +31,16 @@ 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.
 
@@ -45,7 +48,7 @@ def forward_diff(
         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,6 +58,9 @@ 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
 
 
@@ -62,7 +68,7 @@ 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:
diff --git a/src/darsia/utils/linear_solvers/cg.py b/src/darsia/utils/linear_solvers/cg.py
index faa27609..e6079552 100644
--- a/src/darsia/utils/linear_solvers/cg.py
+++ b/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.
 
diff --git a/src/darsia/utils/linear_solvers/jacobi.py b/src/darsia/utils/linear_solvers/jacobi.py
index d8bcdcef..4df1cf12 100644
--- a/src/darsia/utils/linear_solvers/jacobi.py
+++ b/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:
diff --git a/src/darsia/utils/linear_solvers/solver.py b/src/darsia/utils/linear_solvers/solver.py
index 2c6b5531..dfcc4e55 100644
--- a/src/darsia/utils/linear_solvers/solver.py
+++ b/src/darsia/utils/linear_solvers/solver.py
@@ -58,12 +58,15 @@ def update_params(
         )
 
     @abc.abstractmethod
-    def __call__(self, x0: np.ndarray, rhs: np.ndarray) -> np.ndarray:
+    def __call__(
+        self, x0: np.ndarray, rhs: np.ndarray, h: Optional[Union[float, list]] = None
+    ) -> np.ndarray:
         """Main method of the solver - run the solver.
 
         Args:
             x0 (np.ndarray): initial guess
             rhs (np.ndarray): right hand side
+            h (float or list, optional): voxel size
 
         Returns:
             np.ndarray: solution

From 9252aa19be2e000a54c3b6fceb9f4d7ec652334f Mon Sep 17 00:00:00 2001
From: EStorvik <erlend.storvik@hvl.no>
Date: Fri, 14 Jul 2023 15:58:19 +0200
Subject: [PATCH 3/3] Removed physical derivatives from tvd

---
 examples/regularization.py                  | 28 ++++-----
 src/darsia/restoration/split_bregman_tvd.py | 69 +++++++++------------
 src/darsia/restoration/tvd.py               | 13 +---
 3 files changed, 42 insertions(+), 68 deletions(-)

diff --git a/examples/regularization.py b/examples/regularization.py
index 313ab315..b125e473 100644
--- a/examples/regularization.py
+++ b/examples/regularization.py
@@ -6,31 +6,27 @@
 import darsia
 
 # Load 3D image
-# image_folder = f"{os.path.dirname(__file__)}/images/"
-image_folder = "C:/Users/erst/src/image_analysis/images/"
+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], dim=3)
 
 print("Image shape: ", img.img.shape)
 
-# Make regularization parameter heterogenous (for illustration purposes)
-mu_physical = 4.7e-4 * np.ones_like(img.img)
-mu_physical[:, 0:100, :] = 0.5 * 4.7e-4
-
-mu_physical = 4.7e-4
-
+# 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_physical_tvd = darsia.tvd(
+img_regularized_anisotropic_tvd = darsia.tvd(
     img=img,
     method="heterogeneous bregman",
     isotropic=False,
-    weight=mu_physical,
+    weight=mu_anisotropic,
     omega=0.1,
     max_num_iter=30,
-    eps=1e-4,
+    eps=1e-6,
     dim=3,
     verbose=True,
     solver=darsia.Jacobi(maxiter=20),
@@ -44,11 +40,11 @@
     weight=mu,
     omega=0.1,
     max_num_iter=30,
-    eps=1e-4,
+    eps=1e-6,
     dim=3,
     verbose=True,
-    # solver=darsia.Jacobi(maxiter=20),
-    solver=darsia.CG(maxiter=20, tol=1e-3),
+    solver=darsia.Jacobi(maxiter=20),
+    # solver=darsia.CG(maxiter=20, tol=1e-3),
 )
 
 # Regularize image using H1 regularization
@@ -63,8 +59,8 @@
 
 plt.figure("img slice")
 plt.imshow(img.img[100, :, :])
-plt.figure("regularized physical tvd slice")
-plt.imshow(img_regularized_physical_tvd.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[100, :, :])
 plt.figure("regularized H1 slice")
diff --git a/src/darsia/restoration/split_bregman_tvd.py b/src/darsia/restoration/split_bregman_tvd.py
index 22307f68..023f0f49 100644
--- a/src/darsia/restoration/split_bregman_tvd.py
+++ b/src/darsia/restoration/split_bregman_tvd.py
@@ -12,7 +12,7 @@
 
 def split_bregman_tvd(
     img: Union[np.ndarray, da.Image],
-    mu: Union[float, np.ndarray] = 1.0,
+    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,
@@ -32,7 +32,9 @@ def split_bregman_tvd(
 
     Args:
         img (array, da.Image): image
-        mu (float or array): TV penalization parameter
+        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
@@ -49,23 +51,16 @@ def split_bregman_tvd(
 
     """
 
-    # Check input image and extract voxel size (h)
-    if isinstance(img, da.Image):
-        if img.series:
-            voxel_size = [img.dimensions[i] / img.img.shape[i] for i in range(dim)]
-        else:
-            voxel_size = img.voxel_size
-        img = img.img
-    else:
-        voxel_size = np.ones(dim).tolist()
-
     # 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 mass coefficient if no value is provided.
     if ell is None:
@@ -83,16 +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, h=voxel_size[j])
-                    )
-                )
-                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:
@@ -102,9 +101,7 @@ def _shrink(x: np.ndarray, k: Union[float, np.ndarray]) -> np.ndarray:
     def _rhs_function(dt: np.ndarray, bt: np.ndarray) -> np.ndarray:
         return omega * img + ell * sum(
             [
-                da.forward_diff(
-                    img=bt[..., i] - dt[..., i], axis=i, dim=dim, h=voxel_size[i]
-                )
+                da.forward_diff(img=bt[..., i] - dt[..., i], axis=i, dim=dim)
                 for i in range(dim)
             ]
         )
@@ -127,33 +124,25 @@ def _rhs_function(dt: np.ndarray, bt: np.ndarray) -> np.ndarray:
     # Bregman iterations
     for iter in range(max_num_iter):
         # First step - solve the stabilized diffusion system.
-        img_new = solver(x0=img_iter, rhs=_rhs_function(d, b), h=voxel_size)
+        img_new = solver(x0=img_iter, rhs=_rhs_function(d, b))
 
         # Second step - shrinkage.
         if isotropic:
             dub = b.copy()
             for j in range(dim):
-                dub[..., j] += da.backward_diff(
-                    img=img_new, axis=j, dim=dim, voxel_size=voxel_size[j]
-                )
+                dub[..., j] += da.backward_diff(img=img_new, axis=j, dim=dim)
             s = np.linalg.norm(dub, 2, axis=-1)
             shrinkage_factor = np.maximum(s - mu / ell, 0) / (s + 1e-18)
             d = dub * shrinkage_factor[..., None]
             b = dub - d
-        elif isinstance(mu, np.ndarray) and (mu.shape != img_new.shape):
+        elif isinstance(mu, list):
             for j in range(dim):
-                dub = (
-                    da.backward_diff(img=img_new, axis=j, dim=dim, h=voxel_size[j])
-                    + b[..., j]
-                )
-                d[..., j] = _shrink(dub, mu[j, ...] / ell[j, ...])
+                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, h=voxel_size[j])
-                    + b[..., j]
-                )
+                dub = da.backward_diff(img=img_new, axis=j, dim=dim) + b[..., j]
                 d[..., j] = _shrink(dub, mu / ell)
                 b[..., j] = dub - d[..., j]
 
diff --git a/src/darsia/restoration/tvd.py b/src/darsia/restoration/tvd.py
index f074a881..3bf517f7 100644
--- a/src/darsia/restoration/tvd.py
+++ b/src/darsia/restoration/tvd.py
@@ -125,18 +125,7 @@ def _tvd_image(self, img: darsia.Image) -> darsia.Image:
 
         """
         img_copy = img.copy()
-        if self.method == "heterogeneous bregman":
-            img_copy.img = darsia.split_bregman_tvd(
-                img,
-                mu=self.weight,
-                omega=self.omega,
-                ell=self.regularization,
-                max_num_iter=self.max_num_iter,
-                eps=self.eps,
-                **self.kwargs,
-            )
-        else:
-            img_copy.img = self._tvd_array(img.img)
+        img_copy.img = self._tvd_array(img.img)
         return img_copy