single smaps for batches

[ckolbPTB]

Running this code

import torch
import torchkbnufft as tkbn

image = torch.randn(11,1,8,8) + 1j * torch.randn(11,1,8,8)
smaps = torch.randn(11,1,8,8) + 1j * torch.randn(11,1,8,8)
omega = torch.rand(11,2,12)

toep_ob = tkbn.ToepNufft()
kernel = tkbn.calc_toeplitz_kernel(omega, im_size=[8,8])
image = toep_ob(image, kernel, smaps=smaps)

will carry out the Forward/backward NUFFT with Toeplitz embedding for all 11 batches. Nevertheless, often batches correspond to different dynamics and for those smaps is always the same (as it depends on the hardware rather than the underlying object to be imaged). In order to avoid having to make nbatch copies of smaps this PR allows for the following:

image = torch.randn(11,1,8,8) + 1j * torch.randn(11,1,8,8)
smaps = torch.randn(1,1,8,8) + 1j * torch.randn(1,1,8,8)
omega = torch.rand(11,2,12)

toep_ob = tkbn.ToepNufft()
kernel = tkbn.calc_toeplitz_kernel(omega, im_size=[8,8])
image = toep_ob(image, kernel, smaps=smaps)

Remark: The above code with smaps.shape = [1,1,8,8] will also run with the current main-branch but it will yield an image of dimension [1,1,8,8] rather than [11,1,8,8].

[mmuckley]

Hello @ckolbPTB this looks like a great contribution! Could you make sure that your code passes the formatting lints (particularly, black) so that I can merge it?

[ckolbPTB]

Hi, sure no problem. I have never used this automatic format checking. Can I see somewhere which code lines cause the problem, or is it simply a matter of pip install black and this will automatically fix the problem? Thanks for your help!

[mmuckley]

Hello @ckolbPTB,

For torchkbnufft, you can see the exact version that CI uses in dev-requirements.txt

The following two lines should fix it if run from the root directory:

pip install black==21.10b0
black torchkbnufft

You should get a message that reformats the offending file.

[ckolbPTB]

I think it passed all the checks now. Thanks for reviewing the PR so quickly!

[mmuckley]

LGTM.

[mmuckley]

Thanks very much for these contributions @ckolbPTB!