lanl / scico

Scientific Computational Imaging COde
BSD 3-Clause "New" or "Revised" License
105 stars 17 forks source link

Use of TV prior in Accelerated PGM solver #446

Closed shnaqvi closed 1 year ago

shnaqvi commented 1 year ago

I've been enforcing TV prior (1-norm of the gradient) in case of ADMM by supplying L21Norm as a functional in g_list and FiniteDifference as a math operator in C_list. However, since ADMM doesn't allow for using my custom forward operator be passed in SquaredL2Loss, I have to resort to using another solver and, in particular, AcceleratedPGM is showing lightning fast convergence compared to ProximalADMM and PDHG. However, PGM doesn't seem to allow for passing in both the FiniteDifference and L21Norm to it.

Can you please tell me how can I may be compose a new functional by combining L21Norm and FiniteDifference that effectively works as a TV prior, so I can enforce blockiness in the reconstructed image that's essential for solving the deblur problem?

PGM code:

im_jx = jax.device_put(im_s) 
psf_jx = jax.device_put(psf2_cropped)  

C = linop.CircularConvolve(h=psf_jx, input_shape=im_jx.shape, h_center=[psf_jx.shape[0] // 2, psf_jx.shape[1] // 2])
Cx = C(im_jx)

# APGM
f = loss.SquaredL2Loss(y=Cx, A=C)
lbd = 5e-2  # L1 norm regularization parameter
g = lbd * functional.L21Norm()
D = linop.FiniteDifference(input_shape=im_jx.shape, circular=True)

maxiter = 20

L0=5e8
solver_apgm = AcceleratedPGM(
    f=f,
    g=g,
    L0=L0, 
    x0=Cx, 
    maxiter=maxiter, 
    itstat_options={"display": True, "period": 10}
)
print("\nProximal Gradient solver")
x = solver_apgm.solve()
hist_apgm = solver_apgm.itstat_object.history(transpose=True)

ADMM code for reference:

im_jx = jax.device_put(im_s)  
psf_jx = jax.device_put(psf2_cropped)  

C = linop.CircularConvolve(h=psf_jx, input_shape=im_jx.shape, h_center=[psf_jx.shape[0] // 2, psf_jx.shape[1] // 2])
Cx = C(im_jx)

# ADMM
f = loss.SquaredL2Loss(y=Cx, A=C) 
lbd = 5e-1  # L1 norm regularization parameter
g = lbd * functional.L21Norm()
D = linop.FiniteDifference(input_shape=im_jx.shape, circular=True)

rho = 5e0#5e2  # ADMM penalty parameter
maxiter = 20  # number of ADMM iterations

solver = ADMM(
    f=f,
    g_list=[g],
    C_list=[D],
    rho_list=[rho],
    x0=Cx,
    maxiter=maxiter,
    subproblem_solver=CircularConvolveSolver(),
    itstat_options={"display": True, "period": 10},
)

print(f"Solving on {device_info()}\n")
x = solver.solve()
hist = solver.itstat_object.history(transpose=True)
bwohlberg commented 1 year ago

TV denoising can be solved via APGM using a dual formulation of the problem (see this example, but there isn't any straightforward way of solving a problem with a forward operator and TV regularization via APGM. There is a possible approach via approximation of the TV norm (see doi:10.1109/ICASSP.2016.7472568), but it isn't implemented in SCICO.

shnaqvi commented 1 year ago

Got it. Can this be requested as an enhancement, i.e. adding TV prior to the SCICO API? Several of the computational imaging problems use this prior, e.g. this approach that I'm trying to use for space-variant deblurring.

bwohlberg commented 1 year ago

We can add it to the list, but, realistically, given the resources we have at the moment, it's unlikely to be done anytime soon, particularly since the TV prior is supported via other optimization algorithms in SCICO.

shnaqvi commented 1 year ago

Oh ok. I've created an official enhancement request issue, that you can label as such. I've tried to build my case there, so I hope that it can be somewhat prioritized. Thank you!