Plug and play regularisers in pfb

[landmanbester]

Chatting to @TariqBlecher about incorporating some learnt regularisers into one of the minor cycles in pfb. Suggest we start with the approach in https://arxiv.org/pdf/2312.00125.pdf. I have created a new taivas worker (feel free to change the name @TariqBlecher) to get the ball rolling. After running the init and grid steps using eg.

$ pfb init --ms path/to/MS -o out/test -cpi -1 -ipi -1 -dc DATA -fov 2.0 --scheduler distributed -nw 4 -ntd 2 -nvt 2
$ pfb grid -o out/test -nb 1 -fov 1.5 -srf 2 --robustness -1.5 --scheduler distributed -nw 4 -ntd 2 -nvt 2

You should be able to run the taivas worker with

$ pfb taivas -o out/test -nb 1 --niter 5 --scheduler threads -nw 1 -ntd 1 -nvt 8

This will drop you into ipdb (you may have to pip install ipdb) at the point where the new minor cycle can be inserted. Don't worry about the parallelisation settings (i.e. the stuff after --scheduler for the time being, the above should ok to start with).

My suggestion is that you start by:

1. Defining and testing a psf convolution function that can run on pytorch. You can test it against the cpu version psf_concolve(.) that is available at the debug point.
2. Create the likelihood function vdot(x, psf_convolve(x) - 2*dirty) for pytorch. I assume this will have to support autodiff (at least grad). Again you can test it against the numpy implementation (the grad is simply 2*psf_convolve(x) - 2*dirty so you can test that too).
3. Load a pretrained model and run crrfista from the repo to get a new model.

I have set -nb 1 above assuming you are going to start with a single band. If you want to slice out a subset of the frequency axis you can just use --freq-range fmin:fmax with fmin and fmax in Hz when you run the init worker. Going to multiple bands, I think a sensible strategy is to run them separately in parallel for now. We can think about imposing regularisation along the frequency axis at a later stage. My thinking is that we can impose smoothness along frequency during the forward step, thus leaving the rest of the algo essentially unchanged.

Let's discuss progress/problems etc. here