Fix density update for gpuNUFFT (#190)

* Fix gpuDens

* Added tests

* Fix the style

* Fix test

* Fix tests

* sa

* Update the density stuff

* fix the style

---------

Co-authored-by: chaithyagr <[email protected]>

src/mrinufft/operators/base.py

@@ -442,10 +442,10 @@ def compute_density(self, method=None):
                 or `torchkbnufft-gpu`.
         """
         if isinstance(method, np.ndarray):
-            self.density = method
+            self._density = method
             return None
         if not method:
-            self.density = None
+            self._density = None
             return None
         if method is True:
             method = "pipe"
@@ -466,7 +466,7 @@ def compute_density(self, method=None):
                 shape,
                 **kwargs,
             )
-        self.density = method(self.samples, self.shape, **kwargs)
+        self._density = method(self.samples, self.shape, **kwargs)
 
     def get_lipschitz_cst(self, max_iter=10, **kwargs):
         """Return the Lipschitz constant of the operator.

src/mrinufft/operators/interfaces/cufinufft.py

@@ -279,6 +279,15 @@ def samples(self, samples):
             self.raw_op._set_pts(typ, samples)
         self.compute_density(self._density_method)
 
+    @FourierOperatorBase.density.setter
+    def density(self, new_density):
+        """Update the density compensation."""
+        xp = get_array_module(new_density)
+        if xp.__name__ == "numpy":
+            self._density = cp.array(new_density)
+        elif xp.__name__ == "cupy":
+            self._density = new_density
+
     @with_numpy_cupy
     @nvtx_mark()
     def op(self, data, ksp_d=None):

src/mrinufft/operators/interfaces/gpunufft.py

@@ -551,6 +551,21 @@ def samples(self, samples):
             density=self.density,
         )
 
+    @FourierOperatorBase.density.setter
+    def density(self, density):
+        """Set the density for the Fourier Operator.
+
+        Parameters
+        ----------
+        density: np.ndarray
+            The density for the Fourier Operator.
+        """
+        self._density = density
+        self.raw_op.set_pts(
+            self._samples,
+            density=density,
+        )
+
     @classmethod
     def pipe(
         cls,

tests/operators/test_update.py

@@ -128,7 +128,7 @@ def new_smaps(operator):
 
 
 @param_array_interface
-def test_op(
+def test_op_samples(
     operator,
     array_interface,
     image_data,
@@ -145,7 +145,7 @@ def test_op(
 
 
 @param_array_interface
-def test_adj_op(
+def test_adj_op_samples(
     operator,
     array_interface,
     kspace_data,
@@ -162,6 +162,36 @@ def test_adj_op(
     npt.assert_allclose(image_changed, image_true, atol=1e-3, rtol=1e-3)
 
 
+@param_array_interface
+def test_adj_op_density(
+    operator,
+    array_interface,
+    kspace_data,
+):
+    """Test the batch type 1 (adjoint)."""
+    kspace_data = to_interface(kspace_data, array_interface)
+    jitter = np.random.rand(operator.samples.shape[0]).astype(np.float32)
+    # Add very little noise to the trajectory, variance of 1e-3
+    if operator.uses_density:
+        # Test density can be updated
+        operator.density += jitter / 100
+    else:
+        # Test that operator can handle density added later
+        operator.density = 1e-2 + jitter / 100
+    new_operator = update_operator(operator)
+    image_changed = from_interface(operator.adj_op(kspace_data), array_interface)
+    image_true = from_interface(new_operator.adj_op(kspace_data), array_interface)
+    # Reduced accuracy for the GPU cases...
+    npt.assert_allclose(image_changed, image_true, atol=1e-3, rtol=1e-3)
+    if operator.uses_density:
+        # Check if the operator can handle removing density compensation
+        operator.density = None
+        new_operator = update_operator(operator)
+        image_changed = from_interface(operator.adj_op(kspace_data), array_interface)
+        image_true = from_interface(new_operator.adj_op(kspace_data), array_interface)
+        npt.assert_allclose(image_changed, image_true, atol=1e-3, rtol=1e-3)
+
+
 @param_array_interface
 def test_op_smaps_update(
     operator,