Solver is not converging. Image deconvolution with spatially varying PSFs

[shnaqvi]

I'm trying to solve an implementation of the space-variant deconvolution algorithm presented in Flicker & Rigaut (2005) and reviewed here.

I have PSFs sampled over the field, which are centered and stacked in a matrix, of which we compute the SVD to get weights and kernels that form the low-rank approximation of the spatially varying blur forward model. (Image, im, weights, W, and kernels, U, are each vectors in R^n, n=1e7 (image of ~10million pixels)).

I construct my problem like shown below, but objective does not converge but increases without bounds. Also, FISTA seems to be taking really long >4min for just 2 iterations.

Can you please help check if I'm formulating the problem correctly, including forward, adjoint, TV prior, initialization and solver?

P.S. I'm running Python 3.10.1 on a macOS Ventura on an M1 Max 64GB RAM

import numpy as np
import pyunlocbox as plx
from scipy.signal import fftconvolve


# Sample data
im = np.random.randn(2748, 3840)
W = np.random.randn(20, np.prod(im.shape))
U = np.random.randn(np.prod(im.shape), 20)


def forward(im):
    im_sim = np.zeros_like(im)
    for i in range(15):
        weight = W[i,:].reshape(*im.shape)
        psf_mode = U[:,i].reshape(*im.shape)
        im_sim += fftconvolve(im* weight, psf_mode, mode='same')
    return im_sim

def forward_adj(im):
    im_sim = np.zeros_like(im)
    for i in range(15):
        weight = W[i,:].reshape(*im.shape)
        psf_mode = U[:,i].reshape(*im.shape)
        im_sim += fftconvolve(im, np.flipud(np.fliplr(psf_mode)), mode='same')* weight
    return im_sim

tau=10
f1 = plx.functions.norm_l2(y=im, A=forward, At=forward_adj, lambda_=tau)
f2 = plx.functions.norm_tv(maxit=50, dim=2)
solver = plx.solvers.forward_backward(step=0.5/tau) #, accel=plx.acceleration.fista()
ret = plx.solvers.solve([f1, f2], x0=im.copy(), solver=solver, maxit=10, verbosity='ALL')

    norm_l2 evaluation: 9.406386e+15
    norm_tv evaluation: 1.839802e+07
INFO: Forward-backward method
Iteration 1 of forward_backward:
Proximal TV Operator
Iter:  0  obj =  86027944765142.39  rel_obj =  1.0
Iter:  1  obj =  86027944708405.33  rel_obj =  6.59518981794953e-10
Prox_TV: obj = 86027944708405.33, rel_obj = 6.59518981794953e-10, TOL_EPS, iter = 1
exec_time =  1.062948226928711
    norm_l2 evaluation: 2.077274e+32
    norm_tv evaluation: 1.720559e+15
    objective = 2.08e+32
Iteration 2 of forward_backward:

[nperraud]

Hi @shnaqvi ,

I have not check the paper. Quickly here are two things...

Speed: How big is your image? One reason that it is slow is that the TV norm requires sub-iteration.

Convergence: How do you set tau? Tau is the Lipshitz constant of your gradient operator. Here it depends on weight and psf_mode.

The forward function is missing the weight term.

[shnaqvi]

Thanks @nperraud for commenting.

Image size: I've updated the code with Sample data matrices, im, U, W. My image flattened has ~10M pixels (3.8kx2.7k).

TV norm: Do you suggest that I omit the iterations from TV norm? What is the default or min. suggested iterations for it? Sorry I don't understand the role of the iterations.

Tau value: Currently, I'm just hand-waving the value of tau in the absence of logical intuition for its value. Thought it should be a small value, but once I start seeing convergence, I can play with it to fine tune. Do you suggest ways to come up with this value?

Weight terms: I do have the weight term in the forward function, defined by W. I couldn't find out how to pass in arguments to the forward and adjoint functions, so I assumed that these, U and W, are globally defined variables. Can you please comment on how to pass the arguments to the forward function?

[shnaqvi]

Hi @nperraud, following up on this. And I've also created a pull requested with my code in test_deconv_sv_psf.py under the examples folder. Would you please help me troubleshoot it?