Add support for density compensation estimation with cufinufft #195
Conversation
paquiteau
changed the title
|
cufinufft 2.3.1 is here!! Ill update the PR to pull it. |
.github/workflows/test-ci.yml
Outdated
| python -m pip install finufft pooch brainweb-dl torch fastmri | ||
|
|
||
| - name: Install GPU related interfaces | ||
| - name: Point to CUDA 12.1 #TODO: This can be combined from other jobs |
There was a problem hiding this comment.
I don't get your comment, what combinaison do you have in mind ?
There was a problem hiding this comment.
We do this even in other places in the CI, we not really re-do it again and again
examples/GPU/example_density.py
Outdated
There was a problem hiding this comment.
Since there is two backend, you could plot the different density compensation vectors to show the differences (as cufinufft and gpunufft does not use the same interpolation kernel)
| @@ -54,7 +54,7 @@ def __init__(self, inital_trajectory): | |||
| data=torch.Tensor(inital_trajectory), | |||
| requires_grad=True, | |||
| ) | |||
| self.operator = get_operator("gpunufft", wrt_data=True, wrt_traj=True)( | |||
There was a problem hiding this comment.
why do you prefer cufinufft ?
There was a problem hiding this comment.
The idea is to run cufinufft with density compensation, basically increase the coverage
pyproject.toml
Outdated
| torchkbnufft-cpu = ["torchkbnufft", "cupy-cuda12x"] | ||
| torchkbnufft-gpu = ["torchkbnufft", "cupy-cuda12x"] |
There was a problem hiding this comment.
those "CI backends" could be put all together and with a comment above for explaining.
you could also do something like torchkbnufft-cpu = ["mri-nufft[torchkbnufft]"] to avoid repeating the dependency
| cp.array(samples[:, 1], copy=False), | ||
| cp.array(samples[:, 2], copy=False) if self.ndim == 3 else None, | ||
| ) | ||
| plan.setpts(self._kx, self._ky, self._kz) |
There was a problem hiding this comment.
just to be sure, this does not create copies of the arrays ( maybe related to #147 as well)
| if normalize: | ||
| test_op = MRICufiNUFFT(samples=kspace_loc, shape=original_shape, **kwargs) | ||
| test_im = cp.ones(original_shape, dtype=test_op.cpx_dtype) | ||
| test_im_recon = test_op.adj_op(density_comp * test_op.op(test_im)) | ||
| density_comp /= cp.mean(cp.abs(test_im_recon)) |
There was a problem hiding this comment.
this could be refactored into a def _normalize_density(backend, samples, shape, density_comp) so that it could be used for other density compensation methods (e.g. voronoi)
This is crucial support for many use cases. I have some temporary things added to ensure tests still run good.