mmuckley / torchkbnufft

A high-level, easy-to-deploy non-uniform Fast Fourier Transform in PyTorch.
https://torchkbnufft.readthedocs.io/
MIT License
209 stars 44 forks source link

single smaps for batches #48

Closed ckolbPTB closed 2 years ago

ckolbPTB commented 2 years ago

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 commented 2 years ago

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 commented 2 years ago

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 commented 2 years ago

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 commented 2 years ago

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

mmuckley commented 2 years ago

Thanks very much for these contributions @ckolbPTB!