Add array normalization in CudaSinoNormalization

nabu/reconstruction/sinogram_cuda.py

@@ -155,9 +155,9 @@ CudaSinoProcessing = deprecated_class(
 
 
 class CudaSinoNormalization(SinoNormalization):
-    def __init__(self, kind="chebyshev", sinos_shape=None, radios_shape=None, cuda_options=None):
+    def __init__(self, kind="chebyshev", sinos_shape=None, radios_shape=None, normalization_array=None, cuda_options=None):
         super().__init__(
-            kind=kind, sinos_shape=sinos_shape, radios_shape=radios_shape,
+            kind=kind, sinos_shape=sinos_shape, radios_shape=radios_shape, normalization_array=normalization_array
         )
         self._get_shapes(sinos_shape, radios_shape)
         self.cuda_processing = CudaProcessing(**(cuda_options or {}))
@@ -166,41 +166,89 @@ class CudaSinoNormalization(SinoNormalization):
 
     _get_shapes = SinoBuilder._get_shapes
 
+    #
+    # Chebyshev normalization
+    #
+
     def _init_cuda_normalization(self):
-        self._cuda_kernel = CudaKernel(
-            self._get_norm_func_name(),
-            filename=get_cuda_srcfile("normalization.cu"),
-            signature="Piii",
-        )
-        self._cuda_kernel_args = [
-            np.int32(self.n_x), np.int32(self.n_angles), np.int32(self.n_z)
-        ]
-        blk = (1, 64, 16) # TODO tune ?
-        self._cuda_kernel_kwargs = {
-            "block": blk,
-            "grid": (1, int(updiv(self.n_angles, blk[1])), int(updiv(self.n_z, blk[2]))),
-        }
+        self._d_tmp = garray.zeros(self.sinos_shape[-2:], "f")
+        if self.normalization_kind == "chebyshev":
+            self._chebyshev_kernel = CudaKernel(
+                self._get_norm_func_name(),
+                filename=get_cuda_srcfile("normalization.cu"),
+                signature="Piii",
+            )
+            self._chebyshev_kernel_args = [
+                np.int32(self.n_x), np.int32(self.n_angles), np.int32(self.n_z)
+            ]
+            blk = (1, 64, 16) # TODO tune ?
+            self._chebyshev_kernel_kwargs = {
+                "block": blk,
+                "grid": (1, int(updiv(self.n_angles, blk[1])), int(updiv(self.n_z, blk[2]))),
+            }
+            return
+        elif self.normalization_array is not None:
+            normalization_array = self.normalization_array
+            n_angles, n_x = self.sinos_shape[-2:]
+            # If normalization_array is 1D, make a 2D array by repeating the line
+            if normalization_array.ndim == 1:
+                normalization_array = np.tile(normalization_array, (n_angles, 1))
+            if normalization_array.shape != (n_angles, n_x):
+                raise ValueError(
+                    "Expected normalization array to be of shape (n_angles, n_x) = (%d, %d)"
+                    % (n_angles, n_x)
+                )
+            self._d_normalization_array = garray.to_gpu(normalization_array.astype("f"))
 
 
-    def normalize(self, sinos):
+    def _normalize_chebyshev(self, sinos):
         if sinos.flags.c_contiguous:
-            self._cuda_kernel(
-                sinos, *self._cuda_kernel_args, **self._cuda_kernel_kwargs
+            self._chebyshev_kernel(
+                sinos, *self._chebyshev_kernel_args, **self._chebyshev_kernel_kwargs
             )
         else:
-            d_tmp = garray.zeros(sinos.shape[1:], "f")
             # This kernel seems to have an issue on arrays that are not C-contiguous.
             # We have to process image per image.
             nz = np.int32(1)
             nthreadsperblock = (1, 32, 1) # TODO tune
             nblocks = (1, int(updiv(self.n_angles, nthreadsperblock[1])), 1)
             for i in range(sinos.shape[0]):
-                d_tmp[:] = sinos[i][:]
-                self._cuda_kernel(
-                    d_tmp,
+                self._d_tmp[:] = sinos[i][:]
+                self._chebyshev_kernel(
+                    self._d_tmp,
                     np.int32(self.n_x), np.int32(self.n_angles), np.int32(1),
                     grid=nblocks,
                     block=nthreadsperblock
                 )
-                sinos[i][:] = d_tmp[:]
+                sinos[i][:] = self._d_tmp[:]
         return sinos
+
+
+    #
+    # Array subtraction
+    #
+
+    def _normalize_subtract_array(self, sino):
+        if sino.ndim == 2:
+            # Things can go wrong if "sino" is a non-contiguous 2D array
+            # But this should be handled outside this function, as the processing is in-place
+            print(type(sino))
+            print(sino.shape)
+            print(sino.dtype)
+            print(type(self.normalization_array))
+            print(self.normalization_array.shape)
+            print(self.normalization_array.dtype)
+            sino -= self._d_normalization_array
+        else:
+            if sino.flags.forc:
+                # Contiguous 3D array. But pycuda wants the same shape for both operands.
+                for i in range(sino.shape[0]):
+                    sino[i] -= self._d_normalization_array
+            else:
+                # Non-contiguous 2D array. Make a temp. copy
+                for i in range(sino.shape[0]):
+                    self._d_tmp[:] = sino[i][:]
+                    self._d_tmp -= self._d_normalization_array
+                    sino[i][:] = self._d_tmp[:]
+        return sino
+