Don't know what's going on behind the scenes but skimage TV is a fair bit quicker than your implementation. Could be looped in with something like the following:
from pyproximal.ProxOperator import _check_tau
from skimage.restoration import denoise_tv_chambolle, denoise_tv_bregman
class TV:
def __init__(self, sigma, dims, isotropic=True, max_iter=1000):
self.sigma = sigma
self.isotropic = isotropic
self.max_iter = max_iter
self.count = 0
self.dims = dims
def _increment_count(func):
"""Increment counter"""
def wrapped(self, *args, **kwargs):
self.count += 1
return func(self, *args, **kwargs)
return wrapped
@_increment_count
@_check_tau
def prox(self, x, tau):
x = x.reshape(self.dims)
if self.isotropic:
return denoise_tv_chambolle(
x, weight=tau * self.sigma, max_num_iter=self.max_iter
).ravel()
else:
return denoise_tv_bregman(
x,
weight=tau * self.sigma,
isotropic=self.isotropic,
max_num_iter=self.max_iter,
).ravel()
Don't know what's going on behind the scenes but skimage TV is a fair bit quicker than your implementation. Could be looped in with something like the following: