nabu 2024.2.13__py3-none-any.whl → 2025.1.0__py3-none-any.whl

nabu/preproc/ctf.py

@@ -67,7 +67,7 @@ class GeoPars:
         else:
             self.M_vh = np.array([1, 1])
 
-        self.logger.debug("Magnification : h ({}) ;  v ({}) ".format(self.M_vh[1], self.M_vh[0]))
+        self.logger.debug(f"Magnification : h ({self.M_vh[1]}) ;  v ({self.M_vh[0]}) ")
 
         self.length_scale = length_scale
         self.wavelength = wavelength
@@ -80,17 +80,17 @@ class GeoPars:
 
         which_unit = int(np.sum(np.array([self.pix_size_rec > small for small in [1.0e-6, 1.0e-7]]).astype(np.int32)))
         self.pixelsize_string = [
-            "{:.1f} nm".format(self.pix_size_rec * 1e9),
-            "{:.3f} um".format(self.pix_size_rec * 1e6),
-            "{:.1f} um".format(self.pix_size_rec * 1e6),
+            f"{self.pix_size_rec * 1e9:.1f} nm",
+            f"{self.pix_size_rec * 1e6:.3f} um",
+            f"{self.pix_size_rec * 1e6:.1f} um",
         ][which_unit]
 
         if self.magnification:
             self.logger.debug(
-                "All images are resampled to smallest pixelsize: {}".format(self.pixelsize_string),
+                f"All images are resampled to smallest pixelsize: {self.pixelsize_string}",
             )
         else:
-            self.logger.debug("Pixelsize images:  {}".format(self.pixelsize_string))
+            self.logger.debug(f"Pixelsize images:  {self.pixelsize_string}")
 
 
 class CTFPhaseRetrieval:
@@ -157,7 +157,7 @@ class CTFPhaseRetrieval:
         """
         self.logger = LoggerOrPrint(logger)
         if not isinstance(geo_pars, GeoPars):
-            raise ValueError("Expected GeoPars instance for 'geo_pars' parameter")
+            raise TypeError("Expected GeoPars instance for 'geo_pars' parameter")
         self.geo_pars = geo_pars
         self._calc_shape(shape, padded_shape, padding_mode)
         self.delta_beta = delta_beta
@@ -292,7 +292,7 @@ class CTFPhaseRetrieval:
         self.cut_v = math.sqrt(1.0 / 2 / lambda_dist_vh[0]) / fsample_vh[0]
         self.cut_v = min(self.cut_v, 0.5)
 
-        self.logger.debug("Normalized cut-off = {:5.3f}".format(self.cut_v))
+        self.logger.debug(f"Normalized cut-off = {self.cut_v:5.3f}")
 
         self.r = fourier_filters.get_lowpass_filter(
             padded_img_shape,

nabu/preproc/ctf_cuda.py

@@ -48,7 +48,7 @@ class CudaCTFPhaseRetrieval(CTFPhaseRetrieval):
         padding_mode: str
             Padding mode. Default is "reflect".
 
-        Other parameters
+        Other Parameters
         -----------------
         Please refer to CTFPhaseRetrieval documentation.
         """

nabu/preproc/double_flatfield_cuda.py

@@ -7,7 +7,7 @@ from .ccd_cuda import CudaLog
 
 if __has_pycuda__:
     import pycuda.gpuarray as garray
-    import pycuda.cumath as cumath
+    from pycuda import cumath
 
 
 class CudaDoubleFlatField(DoubleFlatField):
@@ -109,7 +109,7 @@ class CudaDoubleFlatField(DoubleFlatField):
             Whether to recompute the double flatfield if already computed.
         """
         if not (isinstance(radios, garray.GPUArray)):  # pylint: disable=E0606
-            raise ValueError("Expected pycuda.gpuarray.GPUArray for radios")
+            raise TypeError("Expected pycuda.gpuarray.GPUArray for radios")
         if self._computed and not (recompute):
             return self.doubleflatfield
         acc = garray.zeros(radios[0].shape, "f")

nabu/preproc/double_flatfield_variable_region.py

@@ -1,5 +1,4 @@
 from .double_flatfield import (
-    DoubleFlatField,
     DoubleFlatField,
     check_shape,
     get_2D_3D_shape,

nabu/preproc/flatfield.py

@@ -1,8 +1,11 @@
+import os
 from multiprocessing.pool import ThreadPool
 from bisect import bisect_left
 import numpy as np
+from tomoscan.io import HDF5File
 from ..io.reader import load_images_from_dataurl_dict
 from ..utils import check_supported, deprecated_class, get_num_threads
+from .ccd import CCDFilter
 
 
 class FlatFieldArrays:
@@ -91,7 +94,7 @@ class FlatFieldArrays:
           flatfield_normalization(X') = (X' - D)/(F' - D) = (X - D) / (F - D) * sF/sX
         So current normalization boils down to a scalar multiplication after flat-field.
         """
-        if self._full_shape:
+        if self._full_shape:  # noqa: SIM102
             # this is never going to happen in this base class. But in the derived class for helical
             # which needs to keep the full shape
             if radios_indices is not None:
@@ -102,7 +105,7 @@ class FlatFieldArrays:
         self._precompute_flats_indices_weights()
         self._configure_srcurrent_normalization(radios_srcurrent, flats_srcurrent)
         self.distortion_correction = distortion_correction
-        self.n_threads = min(1, get_num_threads(n_threads) // 2)
+        self.n_threads = max(1, get_num_threads(n_threads) // 2)
 
     def _set_parameters(self, radios_shape, radios_indices, interpolation, nan_value):
         self._set_radios_shape(radios_shape)
@@ -450,3 +453,305 @@ class FlatFieldDataUrls(FlatField):
             radios_srcurrent=radios_srcurrent,
             flats_srcurrent=flats_srcurrent,
         )
+
+
+class PCAFlatsNormalizer:
+    """This class implement a flatfield normalization based on a PCA of a series of acquired flatfields.
+    The PCA decomposition is handled by a PCAFlatsDecomposer object.
+
+    This implementation was proposed by Jailin C. et al in https://journals.iucr.org/s/issues/2017/01/00/fv5055/
+
+    Code initially written by ID11 @ ESRF staff.
+    Jonathan Wright - Implementation based on research paper
+    Pedro D. Resende - Added saving and loading from file capabilities
+
+    Jerome Lesaint - Integrated the solution in Nabu.
+
+    """
+
+    def __init__(self, components, dark, mean):
+        """Initializes all variables needed to perform the flatfield normalization.
+
+        Parameters
+        -----------
+        components: ndarray
+            The components of the PCA decomposition.
+        dark: ndarray
+            The dark image. Should be one single 2D image.
+        mean: ndarray
+            The mean image of the series of flats.
+        """
+        ones = np.ones_like(components[0], dtype=np.float32)[np.newaxis]  # This comp will account for I0
+        self.components = np.concatenate([ones, components], axis=0)
+        self.dark = dark
+        self.mean = mean
+        self.n_threads = max(1, get_num_threads() // 2)
+        self._setmask()
+        self.ccdfilter = CCDFilter(mean.shape)
+
+    def _form_lsq_matrix(self):
+        """This function form the Least Square matrix, based on the flats components and the mask."""
+        self.Amat = np.stack([gg for gg in self.g], axis=1)  # JL: this is the matrix for the fit
+
+    def _setmask(self, prop=0.125):
+        """Sets the mask to select where the model is going to be fitted.
+
+        Parameters
+        ----------
+        prop: float, default: 0.125
+            The proportion of the image width to take on each side of the image as a mask.
+
+            By default it sets the strips on each side of the frame in the form:
+                mask[:, lim:] = True
+                mask[:, -lim:] = True
+            Where lim = prop * flat.shape[1]
+
+        If you need a custom mask, see update_mask() method.
+        """
+
+        lim = int(prop * self.mean.shape[1])
+        self.mask = np.zeros(self.mean.shape, dtype=bool)
+        self.mask[:, :lim] = True
+        self.mask[:, -lim:] = True
+        self.g = []
+        for component in self.components:
+            self.g.append(component[self.mask])
+
+        self._form_lsq_matrix()
+
+    def update_mask(self, mask: np.ndarray):
+        """Method to update the mask with a custom mask in the form of a boolean 2+D array.
+
+        Paramters
+        =========
+        mask: np.ndarray of Boolean;
+            The array of boolean allows the selection of the region of the image that will be used to fit against the components.
+
+        It will set the mask, replacing the standard mask created with setmask().
+        """
+        if mask.dtype == bool:
+            self.mask = mask
+            self.g = []
+            for component in self.components:
+                self.g.append(component[self.mask])
+
+            self._form_lsq_matrix()
+        else:
+            raise TypeError("Not a boolean array. Will keep the default mask")
+
+    def normalize_radios(self, projections, mask=None, prop=0.125):
+        """This is to keep the flatfield API in the pipeline."""
+        self.correct_stack(projections, mask=mask, prop=prop)
+
+    def correct_stack(
+        self,
+        projections: np.ndarray,
+        mask: np.ndarray = None,
+        prop: float = 0.125,
+    ):
+        """This functions normalizes the stack of projections.
+
+        Performs correction on a stack of projections based on the calculated decomposition.
+        The normalizations is done in-place. The previous projections before normalization are lost.
+
+        Parameters
+        ----------
+        projections: ndarray
+            Stack of projections to normalize.
+
+        prop: float (default: {0.125})
+            Fraction to mask on the horizontal field of view, assuming vertical rotation axis
+
+        mask: np.ndarray (default: None)
+            Custom mask if your data requires it.
+
+        Returns
+        -------
+        corrected projections: np.ndarray
+            Flat field corrected images. Note that the returned projections are exp-transformed to fit the pipeline (e.g. to allow for phase retrieval).
+        """
+        self.projections = self.ccdfilter.dezinger_correction(projections, self.dark)
+
+        if mask is not None:
+            self.update_mask(mask=mask)
+        else:
+            self._setmask(prop=prop)
+
+        with ThreadPool(self.n_threads) as tp:
+            for i, cor, sol in tp.map(self._readcorrect1, range(len(projections))):
+                # solution[i] = sol
+                projections[i] = np.exp(-cor)
+
+    def _readcorrect1(self, ii):
+        """Method to allow parallelization of the normalization."""
+        corr, s = self.correctproj(self.projections[ii])
+        return ii, corr, s
+
+    def correctproj(self, projection):
+        """Performs the correction on one projection of the stack.
+
+        Parameters
+        ----------
+        projection: np.ndarray, float
+            Radiograph from the acquisition stack.
+
+        Returns
+        -------
+        The fitted projection.
+        """
+        logp = np.log(projection.astype(np.float32) - self.dark)
+        corr = self.mean - logp
+        # model to be fitted !!
+        return self.fit(corr)
+
+    def fit(self, corr):
+        """Fit the (masked) projection to the (masked) components of the PCA decomposition.
+
+        This is for each projection, so worth optimising ...
+        """
+        y = corr[self.mask]
+        solution = np.linalg.lstsq(self.Amat, y, rcond=None)[0]
+        correction = np.einsum("ijk,i->jk", self.components, solution)
+        return corr - correction, solution
+
+
+class PCAFlatsDecomposer:
+    """This class implements a PCA decomposition of a serie of acquired flatfields.
+    The PCA decomposition is used to normalize the projections through a PCAFLatNormalizer object.
+
+    This implementation was proposed by Jailin C. et al in https://journals.iucr.org/s/issues/2017/01/00/fv5055/
+
+    Code initially written by ID11 @ ESRF staff.
+    Jonathan Wright - Implementation based on research paper
+    Pedro D. Resende - Added saving and loading from file capabilities
+
+    Jerome Lesaint - Integrated the solution in Nabu.
+    """
+
+    def __init__(self, flats: np.ndarray, darks: np.ndarray, nsigma=3):
+        """
+
+        Parameters
+        -----------
+        flats: np.ndarray
+            A stack of darks corrected flat field images
+        darks: np.ndarray
+            An image or stack of images of the dark current images of the camera.
+
+        Does the log scaling.
+        Subtracts mean and does eigenvector decomposition.
+        """
+        self.n_threads = max(1, get_num_threads() // 2)
+        self.flats = np.empty(flats.shape, dtype=np.float32)
+
+        darks = darks.astype(np.float32)
+        if darks.ndim == 3:
+            self.dark = np.median(darks, axis=0)
+        else:
+            self.dark = darks.copy()
+        del darks
+
+        self.nsigma = nsigma
+        self.ccdfilter = CCDFilter(self.dark.shape)
+        self._ccdfilter_and_log(flats)  # Log is taken here (after dezinger)
+
+        self.mean = np.mean(self.flats, axis=0)  # average
+
+        self.flats = self.flats - self.mean  # makes a copy
+        self.cov = self.compute_correlation_matrix()
+        self.compute_pca()
+        self.generate_pca_flats(nsigma=nsigma)  # Default nsigma=3
+
+    def __str__(self):
+
+        return f"PCA decomposition from flat images. \nThere are {self.components} components created at {self.sigma} level."
+
+    def _ccdfilter_and_log(self, flats: np.ndarray):
+        """Dezinger (substract dark, apply median filter) and takes log of the flat stack"""
+
+        self.ccdfilter.dezinger_correction(flats, self.dark)
+
+        with ThreadPool(self.n_threads) as tp:
+            for i, frame in enumerate(tp.map(np.log, flats)):
+                self.flats[i] = frame
+
+    @staticmethod
+    def ccij(args):
+        """Compute the covariance (img[i]*img[j]).sum() / npixels
+        It is a wrapper for threading.
+        args == i, j, npixels, imgs
+        """
+
+        i, j, NY, imgs = args
+        return i, j, np.einsum("ij,ij", imgs[i], imgs[j]) / NY
+
+    def compute_correlation_matrix(self):
+        """Computes an (nflats x nflats) correlation matrix"""
+
+        N = len(self.flats)
+        CC = np.zeros((N, N), float)
+        args = [(i, j, N, self.flats) for i in range(N) for j in range(i + 1)]
+        with ThreadPool(self.n_threads) as tp:
+            for i, j, result in tp.map(self.ccij, args):
+                CC[i, j] = CC[j, i] = result
+        return CC
+
+    def compute_pca(self):
+        """Gets eigenvectors and eigenvalues and sorts them into order"""
+
+        self.eigenvalues, self.eigenvectors = np.linalg.eigh(
+            self.cov
+        )  # Not sure why the eigh is needed. eig should be enough.
+        order = np.argsort(abs(self.eigenvalues))[::-1]  # high to low
+        self.eigenvalues = self.eigenvalues[order]
+        self.eigenvectors = self.eigenvectors[:, order]
+
+    def generate_pca_flats(self, nsigma=3):
+        """Projects the eigenvectors back into image space.
+
+        Parameters
+        ----------
+        nsigma: int (default: 3)
+        """
+
+        self.sigma = nsigma
+        av = abs(self.eigenvalues)
+        N = (av > (av[-2] * nsigma)).sum()  # Go for 3 sigma
+        self.components = [
+            None,
+        ] * N
+
+        def calculate(ii):
+            calc = np.einsum("i,ijk->jk", self.eigenvectors[:, ii], self.flats)
+            norm = (calc**2).sum()
+            return ii, calc / np.sqrt(norm)
+
+        with ThreadPool(self.n_threads) as tp:
+            for ii, result in tp.map(calculate, range(N)):
+                self.components[ii] = result
+
+        # simple gradients
+        r, c = self.components[0].shape
+        self.components.append(np.outer(np.ones(r), np.linspace(-1 / c, 1 / c, c)))
+        self.components.append(np.outer(np.linspace(-1 / r, 1 / r, r), np.ones(c)))
+
+    def save_decomposition(self, path="PCA_flats.h5", overwrite=True, entry="entry0000"):
+        """Saves the basic information of a PCA decomposition in view of the normalization of projections.
+
+        Parameters
+        ----------
+        path: str (default: "PCA_flats.h5")
+            Full path to the h5 file you want to save your results. It will overwrite!! Be careful.
+        """
+
+        file_exists = os.path.exists(path)
+        if overwrite or not file_exists:
+            with HDF5File(path, "w") as hout:
+                group = hout.create_group(entry)
+                group["eigenvalues"] = self.eigenvalues
+                group["dark"] = self.dark.astype(np.float32)
+                group["p_mean"] = self.mean.astype(np.float32)
+                group["p_components"] = np.array(self.components)
+                hout.flush()
+        else:
+            raise OSError(f"The file {path} already exists and you chose to NOT overwrite.")

nabu/preproc/flatfield_cuda.py

@@ -4,7 +4,6 @@ from nabu.cuda.processing import CudaProcessing
 from ..preproc.flatfield import FlatFieldArrays
 from ..utils import deprecated_class, get_cuda_srcfile
 from ..io.reader import load_images_from_dataurl_dict
-from ..cuda.utils import __has_pycuda__
 
 
 class CudaFlatFieldArrays(FlatFieldArrays):
@@ -90,7 +89,7 @@ class CudaFlatFieldArrays(FlatFieldArrays):
             Radios chunk.
         """
         if not (isinstance(radios, self.cuda_processing.array_class)):
-            raise ValueError("Expected a pycuda.gpuarray (got %s)" % str(type(radios)))
+            raise TypeError("Expected a pycuda.gpuarray (got %s)" % str(type(radios)))
         if radios.dtype != np.float32:
             raise ValueError("radios must be in float32 dtype (got %s)" % str(radios.dtype))
         if radios.shape != self.radios_shape:

nabu/preproc/flatfield_variable_region.py

@@ -1,5 +1,5 @@
 import numpy as np
-from .flatfield import FlatFieldArrays, load_images_from_dataurl_dict, check_supported
+from .flatfield import FlatFieldArrays, load_images_from_dataurl_dict
 
 
 class FlatFieldArraysVariableRegion(FlatFieldArrays):
@@ -27,9 +27,9 @@ class FlatFieldArraysVariableRegion(FlatFieldArrays):
             does not correspond to the length of sub_indexes which is {len(sub_indexes)}
             """
             raise ValueError(message)
-        do_flats_distortion_correction = self.distortion_correction is not None
+        # do_flats_distortion_correction = self.distortion_correction is not None
 
-        whole_dark = self.get_dark()
+        # whole_dark = self.get_dark()
         for i, (idx, sub_r) in enumerate(zip(sub_indexes, sub_regions_per_radio)):
             start_x, end_x, start_y, end_y = sub_r
             slice_x = slice(start_x, end_x)

nabu/preproc/phase.py

@@ -5,13 +5,12 @@ from scipy.fft import rfft2, irfft2, fft2, ifft2
 from ..utils import generate_powers, get_decay, check_supported, get_num_threads, deprecation_warning
 
 # COMPAT.
-from .ctf import CTFPhaseRetrieval
 
 #
 
 
 def lmicron_to_db(Lmicron, energy, distance):
-    """
+    r"""
     Utility to convert the "Lmicron" parameter of PyHST
     to a value of delta/beta.
 
@@ -55,7 +54,7 @@ class PaganinPhaseRetrieval:
         fftw_num_threads=None,
         fft_num_threads=None,
     ):
-        """
+        r"""
         Paganin Phase Retrieval for an infinitely distant point source.
         Formula (10) in [1].
 
@@ -98,7 +97,6 @@ class PaganinPhaseRetrieval:
 
         Notes
         ------
-
         **Padding methods**
 
         The phase retrieval is a convolution done in Fourier domain using FFT,

nabu/preproc/phase_cuda.py

@@ -107,13 +107,13 @@ class CudaPaganinPhaseRetrieval(PaganinPhaseRetrieval):
         assert data.dtype == np.float32
         # Rectangular memcopy
         # TODO profile, and if needed include this copy in the padding kernel
-        if isinstance(data, np.ndarray) or isinstance(data, self.cuda_processing.array_class):
+        if isinstance(data, np.ndarray) or isinstance(data, self.cuda_processing.array_class):  # noqa: SIM101
             self.d_radio_padded[: self.shape[0], : self.shape[1]] = data[:, :]
         elif isinstance(data, cuda.DeviceAllocation):
             # TODO manual memcpy2D
             raise NotImplementedError("pycuda buffers are not supported yet")
         else:
-            raise ValueError("Expected either numpy array, pycuda array or pycuda buffer")
+            raise TypeError("Expected either numpy array, pycuda array or pycuda buffer")
 
     def get_output(self, output):
         s0, s1 = self.shape

nabu/preproc/shift.py

@@ -42,7 +42,7 @@ class VerticalShift:
     def _init_interp_coefficients(self):
         self.interp_infos = []
         for s in self.shifts:
-            s0 = int(floor(s))
+            s0 = floor(s)
             f = s - s0
             self.interp_infos.append([s0, f])
 
@@ -51,7 +51,7 @@ class VerticalShift:
         assert np.max(iangles) < len(self.interp_infos)
         assert len(iangles) == radios.shape[0]
 
-    def apply_vertical_shifts(self, radios, iangles, output=None):
+    def apply_vertical_shifts(self, radios, iangles=None, output=None):
         """
         Parameters
         ----------
@@ -65,6 +65,8 @@ class VerticalShift:
             If given, it will be modified to contain the shifted radios.
             Must be of the same shape of `radios`.
         """
+        if iangles is None:
+            iangles = np.arange(radios.shape[0])
         self._check(radios, iangles)
 
         newradio = np.zeros_like(radios[0])

nabu/preproc/shift_cuda.py

@@ -1,5 +1,4 @@
 import numpy as np
-from ..cuda.utils import __has_pycuda__
 from ..cuda.processing import CudaProcessing
 from ..processing.muladd_cuda import CudaMulAdd
 from .shift import VerticalShift

nabu/preproc/tests/test_ctf.py

@@ -10,7 +10,7 @@ from nabu.preproc import ctf
 from nabu.estimation.distortion import estimate_flat_distortion
 from nabu.misc.filters import correct_spikes
 from nabu.preproc.distortion import DistortionCorrection
-from nabu.cuda.utils import __has_pycuda__, get_cuda_context
+from nabu.cuda.utils import __has_pycuda__
 
 __has_cufft__ = False
 if __has_pycuda__:
@@ -220,10 +220,10 @@ class TestCtf:
         # Test multi-core FFT
         ctf_fft = ctf.CtfFilter(*ctf_args, **ctf_kwargs, use_rfft=True, fft_num_threads=0)
         if ctf_fft.use_rfft:
-            phase_fft = ctf_fft.retrieve_phase(img)
+            # phase_fft = ctf_fft.retrieve_phase(img)
             self.check_result(phase_r2c, self.ref_plain, "Something wrong with CtfFilter-FFT")
 
-    @pytest.mark.skipif(not (__has_pycuda__ and __has_cufft__), reason="pycuda and (scikit-cuda or vkfft)")
+    @pytest.mark.skipif(not (__has_pycuda__ and __has_cufft__), reason="pycuda and (cupy? or vkfft)")
     def test_cuda_ctf(self):
         data = nabu_get_data("brain_phantom.npz")["data"]
         delta_beta = 50.0
@@ -232,7 +232,7 @@ class TestCtf:
         pix_size_m = 0.1e-6
 
         geo_pars = ctf.GeoPars(z2=distance_m, pix_size_det=pix_size_m, wavelength=1.23984199e-9 / energy_kev)
-        ctx = get_cuda_context()
+        # ctx = get_cuda_context()
 
         for normalize in [True, False]:
             ctf_filter = ctf.CTFPhaseRetrieval(

nabu/preproc/tests/test_double_flatfield.py

@@ -1,4 +1,4 @@
-import os.path as path
+from os import path
 from math import exp
 import tempfile
 import numpy as np

nabu/preproc/tests/test_flatfield.py

@@ -411,7 +411,7 @@ class FlatFieldTestDataset:
         self._generate_projections()
 
     def get_flat_idx(self, proj_idx):
-        flats_idx = sorted(list(self.flats.keys()))
+        flats_idx = sorted(self.flats.keys())
         if proj_idx <= flats_idx[0]:
             return (flats_idx[0],)
         elif proj_idx > flats_idx[0] and proj_idx < flats_idx[1]:

nabu/preproc/tests/test_paganin.py

@@ -77,9 +77,7 @@ class TestPaganin:
         errmax = np.max(np.abs(res - res_tomopy) / np.max(res_tomopy))
         assert errmax < self.rtol_pag, "Max error is too high"
 
-    @pytest.mark.skipif(
-        not (__has_pycuda__ and __has_cufft__), reason="Need pycuda and (scikit-cuda or vkfft) for this test"
-    )
+    @pytest.mark.skipif(not (__has_pycuda__ and __has_cufft__), reason="Need pycuda and (cupy? or vkfft) for this test")
     @pytest.mark.parametrize("config", scenarios)
     def test_gpu_paganin(self, config):
         paganin, data, pag_kwargs = self.get_paganin_instance_and_data(config, self.data)