Take NonUniformFastFourierOp out of FourierOp

[ckolbPTB]

• NonUniformFastFourierOp separate to FourierOp
• nufft_typ1/2-functions created in init and used in forward/adjoint (easy to switch to finufft)
• reshaping of trajectory done in init
• data is reshaped to (dim with separate trajectory)(dim with same trajectory)(nufft dims) (possible speed up)
• multiple other dimensions are possible in forward/adjoint
• kzyx and k210 do not have to aligned but if they are not FFTOp and NUFFTOp cannot be used together in FourierOp

Still to do:

• make sure operator can be moved between cpu and gpu

[github-actions]

Coverage Report

File	Stmts	Miss	Cover	Missing
src/mrpro/algorithms/csm
inati.py	24	1	96%	44
walsh.py	16	1	94%	34
src/mrpro/algorithms/dcf
dcf_voronoi.py	53	4	92%	15, 48–49, 76
src/mrpro/algorithms/optimizers
adam.py	20	1	95%	69
src/mrpro/algorithms/reconstruction
DirectReconstruction.py	28	16	43%	51–71, 85
IterativeSENSEReconstruction.py	13	1	92%	76
Reconstruction.py	50	22	56%	42, 54–56, 80–87, 104–113
RegularizedIterativeSENSEReconstruction.py	41	17	59%	96–100, 114–139
src/mrpro/data
AcqInfo.py	128	3	98%	26, 169, 207
CsmData.py	29	3	90%	15, 82–84
DcfData.py	45	8	82%	18, 66, 78–83
IData.py	67	9	87%	119, 125, 129, 159–167
IHeader.py	75	7	91%	75, 109, 127–131
KData.py	214	25	88%	111–112, 127, 134, 144, 152, 206–207, 245, 250–251, 270–281, 435, 437, 500, 515, 552, 583, 592
KHeader.py	153	17	89%	25, 119–123, 150, 199, 210, 217–218, 221, 228, 260–271
KNoise.py	31	15	52%	39–52, 56–61
KTrajectory.py	81	12	85%	108–113, 116–118, 203–207
MoveDataMixin.py	140	18	87%	15, 113, 129, 143–145, 207, 323–325, 338, 417, 437–438, 440, 455–456, 458
QData.py	39	7	82%	42, 65–73
Rotation.py	674	35	95%	100, 198, 335, 433, 477, 495, 581, 583, 592, 626, 628, 691, 768, 773, 776, 791, 808, 813, 889, 1077, 1082, 1085, 1109, 1113, 1240, 1242, 1250–1251, 1315, 1397, 1690, 1846, 1881, 1885, 1996
SpatialDimension.py	232	21	91%	34, 104, 141, 148, 154, 274–276, 289–291, 325, 343, 356, 369, 382, 395, 404–405, 420, 429
acq_filters.py	12	1	92%	47
src/mrpro/data/traj_calculators
KTrajectoryCalculator.py	25	2	92%	23, 45
KTrajectoryIsmrmrd.py	13	2	85%	41, 50
KTrajectoryPulseq.py	23	1	96%	55
src/mrpro/operators
CartesianSamplingOp.py	89	3	97%	118, 157, 280
ConstraintsOp.py	60	2	97%	46, 48
EndomorphOperator.py	65	2	97%	228, 234
FiniteDifferenceOp.py	27	2	93%	40, 105
FourierOp.py	73	2	97%	230, 235
Functional.py	71	5	93%	20–22, 117, 119
GridSamplingOp.py	136	9	93%	72–73, 82–83, 90–91, 94, 96, 98
LinearOperator.py	168	10	94%	55, 91, 190, 220, 261, 270, 278, 287, 295, 320
LinearOperatorMatrix.py	158	16	90%	82, 119, 152, 161, 166, 175–178, 191–194, 203, 215, 304, 331, 359
MultiIdentityOp.py	13	2	85%	43, 48
NonUniformFastFourierOp.py	169	4	98%	89, 384, 407, 412
Operator.py	78	2	97%	25, 74
ProximableFunctionalSeparableSum.py	39	3	92%	50, 103, 110
SliceProjectionOp.py	173	8	95%	44, 61, 63, 69, 206, 227, 260, 300
WaveletOp.py	120	5	96%	152, 170, 205, 210, 233
ZeroPadOp.py	16	1	94%	30
src/mrpro/utils
filters.py	62	2	97%	44, 49
reshape.py	60	1	98%	191
slice_profiles.py	46	6	87%	20, 36, 113–116, 149
sliding_window.py	34	1	97%	34
split_idx.py	10	2	80%	43, 47
summarize_tensorvalues.py	11	9	18%	20–29
typing.py	18	11	39%	8–23
zero_pad_or_crop.py	31	6	81%	26, 30, 54, 57, 60, 63
TOTAL	4976	363	93%

Tests	Skipped	Failures	Errors	Time
2315	0 💤	0 ❌	0 🔥	1m 17s ⏱️

[github-actions]

📚 Documentation

📁 Download as zip
🔍 View online

[fzimmermann89]

Mid-term, I would still really like to decouple us from torchkbnufft and have a NufftOperator option in the FourierOperator,
where you can either specify give a TorchKBNufft class or a Finufft class which all implement a common interface.

Also, long-term, I would like to fix the axes sorting code, such that for example a RPE with fft direction along k1 and kz works, i.e. a mismatch in the order of the axes.

[fzimmermann89]

Not sure about if we should pass on traj as KTrajectory

If not to difficult, yes. Or at least allow both, KTrajectory and tensor inputs.

and if it should be the original one or already matching recon_ and encoding_matrix.

In case cartesian fft removes stuff? Does it matter?
Also, does our current code even work in that case?
Shouldnt for that to work the order of NUFFT and FFT be swapped?

[ckolbPTB]

In case cartesian fft removes stuff? Does it matter?
Also, does our current code even work in that case?
Shouldnt for that to work the order of NUFFT and FFT be swapped?

FFT and NUFFT always act on separate dimensions so it does not matter if FFT changes stuff. Only case where it matters would be if the non-uniform trajectory changes along a uniform dimension. But the NUFFT acts along the original k-space uniform dimension so this should be fine.

[fzimmermann89]

For this or another PR:
Now that we have both operators separated, I think it would be good to have a test that tests nufft with cartesian coordinates vs an FFT.

[ckolbPTB]

@fzimmermann89 I added a test to verify against the analytic solution. Is this what you had in mind?

[fzimmermann89]

Testing with analytical solution is even better!

Final remarks:

I would remove .forward from the operator calls.
It should generally be call that is called on pytorch modules (which all our operator are) as you can in general set global debugging hooks that only work in that case.
I think we should not give a bad example here.

We could either set fast_fourier_op and nufft op always, just with empty dims, or set them to identity op (which we didn't have when we wrote this)
I am wondering why mypy is not complaining about them possibly being None (or I missed something).

[fzimmermann89]

now that the NufftOp can be created separately, there are some cases not handled:
What if dim is empty?
What if dim contains invalid axes?
What if dim contains positive axes?

[fzimmermann89]

I would say dim is empty, the operator should do nothing.
We should add that dim talks about kx ky kz and not k0 k1 k2 here.
So maybe not dim bit direction?

[fzimmermann89]

also, what about fastfourierop with empty dim?

[ckolbPTB]

I would say dim is empty, the operator should do nothing.

added

also, what about fastfourierop with empty dim?

should be fixed in separate PR

[ckolbPTB]

What if dim contains invalid axes?
What if dim contains positive axes?

We could transform them to negative axes based on the dimensions of traj and then verify the result is within (-3,-2,-1)?

[ckolbPTB]

I guess we have a similar problem with your FFTOp where we also assume that kz ky kx and k2 k1 k0 align, i.e. if I set dim = -1 then I expect to carry out the FFT along k0 with the parameters of (k)x.

We could either say that FFTOp and NUFFTOp work in k2 k1 k0 space and we have to provide matrix sizes and trajectory and so on in this space. Then we could do the magic of figuring all this out in the FourierOp. This would leave the FFTOp to be more intuitive but would mean that for the NUFFTOp users would have to figure out the traj by themselves.

Other option would be to provide dim_kzkykx and dim_k2k1k0 to both operators which would make FFTOp less intuitive but mean that NUFFTOp could get the original trajectory as input.

[fzimmermann89]

I see basically the same two options here :

• nuffOp gets data and omega in shape (batch_different_traj, batch_same_traj, samples), all reshaping is done in FourierOp.

• NufftOp to gets 'x','y','z' as a direction argument. We transform it any ways always to x,y,z where we use it.
  This would mean a minimal change to the code. Nufft does not need dim_k2k1k0 at all.
  (Midterm, it would be nice to support in FourierOp also RPE with read-out in z direction, but saved in k0. For this, FFTOp would need both dim and direction. But as I said, this is IMHO a separate piece of work for the future)

Not really sure which I prefer, I think a bit the second for know, as it is easier to implement and switch to the first one if it required later on without losing work. Am I missing substantial downsides of the second option?

[ckolbPTB]

I think Nufft and Fft need k2 k1 k0 because we need to know which dimensions of kdata to combine to samples and which to use as batch. currently we assume this is the same as kz ky kx and hence do the rearrange of kdata in forward/adjoint using nufft_dims which is in kz ky kx space although we are actually rearranging k2 k1 k0 dimensions.

I would prefer it if FftOp and NufftOp used a similar interface. This I think would be much easier to understand by users.

[fzimmermann89]

Ok, so option two would be to provide direction Literal['x', 'y', 'z',-3,-2,-1] and dim Literal[-3,-2,-1,'k0,'k1',k2'] to nufft.
FFTOp would still only get dim:int. Maye at some point we can add direction here as well to support the mixed case.

Option one is still all reshaping in FourierOp

[ckolbPTB]

A bit more work on this topic:
Now nufft_directions:Sequence[Literal['x', 'y', 'z',-3,-2,-1]] is provided by the user to define along which directions the nufft should be performed. Based on this ky, kx and kz are selected for Nufft. In a second step it is checked which k0, k1, k2 dimensions ky, kx and kz define (nufft_dims). These are then automatically selected.

In the forward and adjoint this means that the image is permuted/reshaped based on nufft_directions, the k-space trajectory and k-space data based on nufft_dims.

This means users don't have to worry about k2-k1-k0-space but can define everything in x-y-z/kx-ky-kz-space.

It also means that

support in FourierOp also RPE with read-out in z direction, but saved in k0

also works now.

The interface is also now more similar to FFT where we either allow SpatialDimension or Sequence[int] for the recon and encoding matrix.

[fzimmermann89]

is the misalignment k210 and kzyx still a problem ?

regarding review, let me know if i should have another look at this.
(i will not get to reviewing this today, maybe for next meeting, sorry.)