nabu 2023.2.1__py3-none-any.whl → 2024.1.0rc3__py3-none-any.whl

nabu/pipeline/estimators.py

@@ -8,12 +8,17 @@ import scipy.fft  #  pylint: disable=E0611
 from silx.io import get_data
 from typing import Union, Optional
 import math
+from numbers import Real
+from scipy import ndimage as nd
+
 from ..preproc.flatfield import FlatFieldDataUrls
 from ..estimation.cor import (
     CenterOfRotation,
     CenterOfRotationAdaptiveSearch,
     CenterOfRotationSlidingWindow,
     CenterOfRotationGrowingWindow,
+    CenterOfRotationFourierAngles,
+    CenterOfRotationOctaveAccurate,
 )
 from ..estimation.cor_sino import SinoCorInterface
 from ..estimation.tilt import CameraTilt
@@ -23,7 +28,7 @@ from ..resources.utils import extract_parameters
 from ..utils import check_supported, is_int
 from .params import tilt_methods
 from ..resources.dataset_analyzer import get_0_180_radios
-from ..misc.rotation import Rotation
+from ..processing.rotation import Rotation
 from ..io.reader import ChunkReader
 from ..preproc.ccd import Log, CCDFilter
 from ..misc import fourier_filters
@@ -50,9 +55,15 @@ def estimate_cor(method, dataset_info, do_flatfield=True, cor_options: Optional[
     else:
         raise TypeError(f"cor_options_str is expected to be a dict or a str. {type(cor_options)} provided")
 
-    # Dispatch
+    # Dispatch. COR estimation is always expressed in absolute number of pixels (i.e. from the center of the first pixel column)
     if method in CORFinder.search_methods:
-        cor_finder = CORFinder(method, dataset_info, do_flatfield=do_flatfield, cor_options=cor_options, logger=logger)
+        cor_finder = CORFinder(
+            method,
+            dataset_info,
+            do_flatfield=do_flatfield,
+            cor_options=cor_options,
+            logger=logger,
+        )
         estimated_cor = cor_finder.find_cor()
     elif method in SinoCORFinder.search_methods:
         cor_finder = SinoCORFinder(
@@ -109,21 +120,102 @@ class CORFinderBase:
             raise TypeError(
                 f"cor_options is expected to be an optional instance of dict. Get {cor_options} ({type(cor_options)}) instead"
             )
-        self.cor_options = cor_options or {}
+        self.cor_options = {}
+        if isinstance(cor_options, dict):
+            self.cor_options.update(cor_options)
 
-        cor_class = self.search_methods[method]["class"]
-        self.cor_finder = cor_class(logger=self.logger)
+        # tomotools internal meeting 07 feb 2024: Merge of options 'near_pos' and 'side'.
+        # See [minutes](https://gitlab.esrf.fr/tomotools/minutes/-/blob/master/minutes-20240207.md?ref_type=heads)
 
+        detector_width = self.dataset_info.radio_dims[0]
         default_lookup_side = "right" if self.dataset_info.is_halftomo else "center"
+        near_init = self.cor_options.get("side", None)
+
+        if near_init is None:
+            near_init = default_lookup_side
+
+        if near_init == "from_file":
+            try:
+                near_pos = self.dataset_info.dataset_scanner.estimated_cor_frm_motor  # relative pos in pixels
+                if isinstance(near_pos, Real):
+                    # near_pos += detector_width // 2  # Field in NX is relative.
+                    self.cor_options.update({"near_pos": int(near_pos)})
+                else:
+                    near_init = default_lookup_side
+            except:
+                self.logger.warning(
+                    "COR estimation from motor position absent from NX file. Global search is performed."
+                )
+                near_init = default_lookup_side
+        elif isinstance(near_init, Real):
+            self.cor_options.update({"near_pos": near_init})
+            near_init = "near"  # ???
+        elif near_init == "near":  # Legacy
+            if not isinstance(self.cor_options["near_pos"], Real):
+                self.logger.warning("Side option set to 'near' but no 'near_pos' option set.")
+                self.logger.warning("Set side to right if HA, center otherwise.")
+                near_init = default_lookup_side
+        elif near_init in ("left", "right", "center", "all"):
+            pass
+        else:
+            self.logger.warning(
+                f"COR option 'side' received {near_init} and should be either 'from_file' (default), 'left', 'right', 'center', 'near' or a number."
+            )
+
+        if isinstance(self.cor_options.get("near_pos", None), Real):
+            # Check validity of near_pos
+            if np.abs(self.cor_options["near_pos"]) > detector_width / 2:
+                self.logger.warning(
+                    f"Relative COR passed is greater than half the size of the detector. Did you enter a absolute COR position?"
+                )
+                self.logger.warning("Instead, the center of the detector is used.")
+                self.cor_options["near_pos"] = 0
+
+            # Set side from near_pos if passed.
+            if self.cor_options["near_pos"] < 0.0:
+                self.cor_options.update({"side": "left"})
+                near_init = "left"
+            else:
+                self.cor_options.update({"side": "right"})
+                near_init = "right"
+
+        self.cor_options.update({"side": near_init})
+
+        # At this stage : side is set to one of left, right, center near.
+        # and near_pos to a numeric value.
+
+        # if isinstance(self.cor_options["near_pos"], Real):
+        #    # estimated_cor_frm_motor value is supposed to be relative. Since the config documentation expects the "near_pos" options
+        #    # to be given as an absolute COR estimate, a conversion is needed.
+        #    self.cor_options["near_pos"] += detector_width // 2  # converted in absolute nb of pixels.
+        # if not (isinstance(self.cor_options["near_pos"], Real) or self.cor_options["near_pos"] == "ignore"):
+        #    self.cor_options.update({"near_pos": "ignore"})
+
+        # At this stage, cor_options["near_pos"] is either
+        #   - 'ignore':
+        #   - an (absolute) integer value (either the user-provided one if present or the NX one).
+
+        cor_class = self.search_methods[method]["class"]
+        self.cor_finder = cor_class(logger=self.logger, cor_options=self.cor_options)
+
         lookup_side = self.cor_options.get("side", default_lookup_side)
 
+        # OctaveAccurate
+        if cor_class == CenterOfRotationOctaveAccurate:
+            lookup_side = "center"
+        angles = self.dataset_info.rotation_angles
+
         self.cor_exec_args = []
         self.cor_exec_args.extend(self.search_methods[method].get("default_args", []))
 
         # CenterOfRotationSlidingWindow is the only class to have a mandatory argument ("side")
         # TODO - it would be more elegant to have it as a kwarg...
-        if len(self.cor_exec_args) > 0 and cor_class is CenterOfRotationSlidingWindow:
-            self.cor_exec_args[0] = lookup_side
+        if len(self.cor_exec_args) > 0:
+            if cor_class in (CenterOfRotationSlidingWindow, CenterOfRotationOctaveAccurate):
+                self.cor_exec_args[0] = lookup_side
+            elif cor_class in (CenterOfRotationFourierAngles,):
+                self.cor_exec_args[0] = angles
+                self.cor_exec_args[1] = lookup_side
         #
         self.cor_exec_kwargs = update_func_kwargs(self.cor_finder.find_shift, self.cor_options)
 
@@ -148,6 +240,10 @@ class CORFinder(CORFinderBase):
         "growing-window": {
             "class": CenterOfRotationGrowingWindow,
         },
+        "octave-accurate": {
+            "class": CenterOfRotationOctaveAccurate,
+            "default_args": ["center"],
+        },
     }
 
     def __init__(self, method, dataset_info, do_flatfield=True, cor_options=None, logger=None):
@@ -225,13 +321,14 @@ class SinoCORFinder(CORFinderBase):
         "sino-coarse-to-fine": {
             "class": SinoCorInterface,
         },
-        "sliding-window": {
+        "sino-sliding-window": {
             "class": CenterOfRotationSlidingWindow,
             "default_args": ["right"],
         },
-        "growing-window": {
+        "sino-growing-window": {
             "class": CenterOfRotationGrowingWindow,
         },
+        "fourier-angles": {"class": CenterOfRotationFourierAngles, "default_args": [None, "center"]},
     }
 
     def __init__(
@@ -286,6 +383,7 @@ class SinoCORFinder(CORFinderBase):
                 self.projs_indices = np.round(indices_float).astype(np.int32).tolist()
             else:  # Subsampling step
                 self.projs_indices = projs_idx[::subsampling]
+                self.angles = self.dataset_info.rotation_angles[::subsampling]
         else:  # Angular step
             raise NotImplementedError()
 
@@ -338,7 +436,7 @@ class SinoCORFinder(CORFinderBase):
         self.logger.info("Estimating center of rotation")
         self.logger.debug("%s.find_shift(%s)" % (self.cor_finder.__class__.__name__, str(self.cor_exec_kwargs)))
         img_1, img_2 = self._split_sinogram(self.sinogram)
-        shift = self.cor_finder.find_shift(img_1, img_2, *self.cor_exec_args, **self.cor_exec_kwargs)
+        shift = self.cor_finder.find_shift(img_1, np.fliplr(img_2), *self.cor_exec_args, **self.cor_exec_kwargs)
         return self.shape[1] / 2 + shift
 
 
@@ -346,7 +444,7 @@ class SinoCORFinder(CORFinderBase):
 SinoCOREstimator = SinoCORFinder
 
 
-class CompositeCORFinder:
+class CompositeCORFinder(CORFinderBase):
     """
     Class and method to prepare sinogram and calculate COR
     The pseudo sinogram is built with shrinked radios taken every theta_interval degres
@@ -363,6 +461,11 @@ class CompositeCORFinder:
            by several order of magnitude without modifing the final result
     """
 
+    search_methods = {
+        "composite-coarse-to-fine": {
+            "class": CenterOfRotation,  # Hack. Not used. Everything is done in the find_cor() func.
+        }
+    }
     _default_cor_options = {"low_pass": 0.4, "high_pass": 10, "side": "center", "near_pos": 0, "near_width": 20}
 
     def __init__(
@@ -377,6 +480,9 @@ class CompositeCORFinder:
         logger=None,
         norm_order=1,
     ):
+        super().__init__(
+            "composite-coarse-to-fine", dataset_info, do_flatfield=True, cor_options=cor_options, logger=logger
+        )
         if norm_order not in [1, 2]:
             raise ValueError(
                 f""" the norm order (nom_order parameter) must be either 1 or 2. You passed {norm_order}
@@ -388,6 +494,12 @@ class CompositeCORFinder:
         self.dataset_info = dataset_info
         self.logger = LoggerOrPrint(logger)
 
+        self.sx, self.sy = self.dataset_info.radio_dims
+
+        default_cor_options = self._default_cor_options.copy()
+        default_cor_options.update(self.cor_options)
+        self.cor_options = default_cor_options
+
         # the algorithm can work for angular ranges larger than 1.2*pi
         # up to an arbitrarily number of turns as it is the case in helical scans
         self.spike_threshold = spike_threshold
@@ -398,8 +510,6 @@ class CompositeCORFinder:
         self.angle_min = self.unwrapped_rotation_angles.min()
         self.angle_max = self.unwrapped_rotation_angles.max()
 
-        self.sx, self.sy = self.dataset_info.radio_dims
-
         if (self.angle_max - self.angle_min) < 1.2 * np.pi:
             useful_span = None
             raise ValueError(
@@ -414,7 +524,7 @@ class CompositeCORFinder:
             if useful_span < np.pi:
                 theta_interval = theta_interval * useful_span / np.pi
 
-        self._get_cor_options(cor_options)
+        # self._get_cor_options(cor_options)
 
         self.take_log = take_log
         self.ovs = oversampling
@@ -474,7 +584,6 @@ class CompositeCORFinder:
         # initialize sinograms and radios arrays
         self.sino = np.zeros([2 * self.nprobed * n_subsampling_y, (self.sx - 1) * self.ovs + 1], "f")
         self._loaded = False
-
         self.high_pass = self.cor_options["high_pass"]
         img_filter = fourier_filters.get_bandpass_filter(
             (self.sino.shape[0] // 2, self.sino.shape[1]),
@@ -492,16 +601,9 @@ class CompositeCORFinder:
         """oversampling in the horizontal direction"""
         if self.ovs == 1:
             return radio
-        result = np.zeros([radio.shape[0], (radio.shape[1] - 1) * self.ovs + 1], "f")
-        result[:, :: self.ovs] = radio
-
-        for i in range(1, self.ovs):
-            f = i / self.ovs
-            result[:, i :: self.ovs] = (1 - f) * result[:, 0 : -self.ovs : self.ovs] + f * result[
-                :, self.ovs :: self.ovs
-            ]
-
-        return result
+        else:
+            ovs_2D = [1, self.ovs]
+        return oversample(radio, ovs_2D)
 
     def _get_cor_options(self, cor_options):
         default_dict = self._default_cor_options.copy()
@@ -644,45 +746,59 @@ class CompositeCORFinder:
 
         best_overlap = overlap_min
         best_error = np.inf
+
+        blurred_radio1 = nd.gaussian_filter(abs(radio1), [0, self.high_pass])
+        blurred_radio2 = nd.gaussian_filter(abs(radio2), [0, self.high_pass])
+
         for z in range(int(overlap_min), int(overlap_max) + 1):
             if z <= ovsd_sx:
                 my_z = z
                 my_radio1 = radio1
                 my_radio2 = radio2
+                my_blurred_radio1 = blurred_radio1
+                my_blurred_radio2 = blurred_radio2
             else:
                 my_z = ovsd_sx - (z - ovsd_sx)
                 my_radio1 = np.fliplr(radio1)
                 my_radio2 = np.fliplr(radio2)
+                my_blurred_radio1 = np.fliplr(blurred_radio1)
+                my_blurred_radio2 = np.fliplr(blurred_radio2)
 
             common_left = np.fliplr(my_radio1[:, ovsd_sx - my_z :])[:, : -int(math.ceil(self.ovs * self.high_pass * 2))]
             # adopt a 'safe' margin considering high_pass value (possibly float)
             common_right = my_radio2[:, ovsd_sx - my_z : -int(math.ceil(self.ovs * self.high_pass * 2))]
 
+            common_blurred_left = np.fliplr(my_blurred_radio1[:, ovsd_sx - my_z :])[
+                :, : -int(math.ceil(self.ovs * self.high_pass * 2))
+            ]
+            # adopt a 'safe' margin considering high_pass value (possibly float)
+            common_blurred_right = my_blurred_radio2[:, ovsd_sx - my_z : -int(math.ceil(self.ovs * self.high_pass * 2))]
+
             if common_right.size == 0:
                 continue
 
-            error = self.error_metric(common_right, common_left)
+            error = self.error_metric(common_right, common_left, common_blurred_right, common_blurred_left)
 
             min_error = min(best_error, error)
 
             if min_error == error:
                 best_overlap = z
                 best_error = min_error
-            self.logger.debug(
-                "testing an overlap of %.2f pixels, actual best overlap is %.2f pixels over %d\r"
-                % (z / self.ovs, best_overlap / self.ovs, ovsd_sx / self.ovs),
-            )
+            # self.logger.debug(
+            #     "testing an overlap of %.2f pixels, actual best overlap is %.2f pixels over %d\r"
+            #     % (z / self.ovs, best_overlap / self.ovs, ovsd_sx / self.ovs),
+            # )
 
         offset = (ovsd_sx - best_overlap) / self.ovs / 2
         cor_abs = (self.sx - 1) / 2 + offset
 
         return cor_abs
 
-    def error_metric(self, common_right, common_left):
+    def error_metric(self, common_right, common_left, common_blurred_right, common_blurred_left):
         if self.norm_order == 2:
             return self.error_metric_l2(common_right, common_left)
         elif self.norm_order == 1:
-            return self.error_metric_l1(common_right, common_left)
+            return self.error_metric_l1(common_right, common_left, common_blurred_right, common_blurred_left)
         else:
             assert False, "this cannot happen"
 
@@ -699,14 +815,56 @@ class CompositeCORFinder:
 
         return res
 
-    def error_metric_l1(self, common_right, common_left):
-        common = common_right - common_left
+    def error_metric_l1(self, common_right, common_left, common_blurred_right, common_blurred_left):
+        common = (common_right - common_left) / (common_blurred_right + common_blurred_left)
 
         res = abs(common).mean()
 
         return res
 
 
+def oversample(radio, ovs_s):
+    """oversampling an image in arbitrary directions.
+    The first and last point of each axis will still remain as  extremal points of the new axis.
+    """
+    result = np.zeros([(radio.shape[0] - 1) * ovs_s[0] + 1, (radio.shape[1] - 1) * ovs_s[1] + 1], "f")
+
+    # Pre-initialisation: The original data falls exactly on the following strided positions in the new data array.
+    result[:: ovs_s[0], :: ovs_s[1]] = radio
+
+    for k in range(0, ovs_s[0]):
+        # interpolation coefficient for axis 0
+        g = k / ovs_s[0]
+        for i in range(0, ovs_s[1]):
+            if i == 0 and k == 0:
+                # this case subset was already exactly matched from before the present double loop,
+                # in the pre-initialisation line.
+                continue
+            # interpolation coefficent for axis 1
+            f = i / ovs_s[1]
+
+            # stop just a bit before cause we are not extending beyond the limits.
+            # If we are exacly on a vertical or horizontal original line, then no shift will be applied,
+            # and we will exploit the equality f+(1-f)=g+(1-g)=1 adding twice the same contribution with
+            # interpolation factors which become dummies pour le coup.
+            stop0 = -ovs_s[0] if k else None
+            stop1 = -ovs_s[1] if i else None
+
+            # Once again, we exploit the  g+(1-g)=1 equality
+            start0 = ovs_s[0] if k else 0
+            start1 = ovs_s[1] if i else 0
+
+            # and what is done below makes clear the corundum above.
+            result[k :: ovs_s[0], i :: ovs_s[1]] = (1 - g) * (
+                (1 - f) * result[0 : stop0 : ovs_s[0], 0 : stop1 : ovs_s[1]]
+                + f * result[0 : stop0 : ovs_s[0], start1 :: ovs_s[1]]
+            ) + g * (
+                (1 - f) * result[start0 :: ovs_s[0], 0 : stop1 : ovs_s[1]]
+                + f * result[start0 :: ovs_s[0], start1 :: ovs_s[1]]
+            )
+    return result
+
+
 # alias
 CompositeCOREstimator = CompositeCORFinder
 

nabu/pipeline/fullfield/chunked.py

@@ -16,10 +16,10 @@ from ...preproc.phase import PaganinPhaseRetrieval
 from ...preproc.ctf import CTFPhaseRetrieval, GeoPars
 from ...reconstruction.sinogram import SinoNormalization
 from ...reconstruction.filtering import SinoFilter
-from ...misc.rotation import Rotation
-from ...reconstruction.rings import MunchDeringer
-from ...misc.unsharp import UnsharpMask
-from ...misc.histogram import PartialHistogram, hist_as_2Darray
+from ...processing.rotation import Rotation
+from ...reconstruction.rings import MunchDeringer, SinoMeanDeringer, VoDeringer
+from ...processing.unsharp import UnsharpMask
+from ...processing.histogram import PartialHistogram, hist_as_2Darray
 from ..utils import use_options, pipeline_step, get_subregion
 from ..datadump import DataDumpManager
 from ..writer import WriterManager
@@ -49,7 +49,9 @@ class ChunkedPipeline:
     UnsharpMaskClass = UnsharpMask
     ImageRotationClass = Rotation
     VerticalShiftClass = VerticalShift
-    SinoDeringerClass = MunchDeringer
+    MunchDeringerClass = MunchDeringer
+    SinoMeanDeringerClass = SinoMeanDeringer
+    VoDeringerClass = VoDeringer
     MLogClass = Log
     SinoNormalizationClass = SinoNormalization
     SinoFilterClass = SinoFilter
@@ -425,6 +427,8 @@ class ChunkedPipeline:
             output_is_mlog=False,
             average_is_on_log=avg_is_on_log,
             sigma_filter=options["sigma"],
+            log_clip_min=options["log_min_clip"],
+            log_clip_max=options["log_max_clip"],
         )
 
     @use_options("ccd_correction", "ccd_correction")
@@ -468,7 +472,7 @@ class ChunkedPipeline:
                 lim1=options["ctf_lim1"],
                 lim2=options["ctf_lim2"],
                 logger=self.logger,
-                fftw_num_threads=None,  # TODO tune in advanced params of nabu config file
+                fft_num_threads=None,  # TODO tune in advanced params of nabu config file
                 use_rfft=True,
                 normalize_by_mean=options["ctf_normalize_by_mean"],
                 translation_vh=translations_vh,
@@ -482,12 +486,7 @@ class ChunkedPipeline:
                 pixel_size=options["pixel_size_m"],
                 padding=options["padding_type"],
                 # TODO tune in advanced params of nabu config file
-                fftw_num_threads=None,
-            )
-        if self.phase_retrieval.use_fftw:
-            self.logger.debug(
-                "%s using FFTW with %d threads"
-                % (self.phase_retrieval.__class__.__name__, self.phase_retrieval.fftw.num_threads)
+                fft_num_threads=None,
             )
 
     @use_options("unsharp_mask", "unsharp_mask")
@@ -519,11 +518,25 @@ class ChunkedPipeline:
 
     @use_options("sino_rings_correction", "sino_deringer")
     def _init_sino_rings_correction(self):
-        options = self.processing_options["sino_rings_correction"]
-        fw_params = extract_parameters(options["user_options"])
-        fw_sigma = fw_params.pop("sigma", 1.0)
         n_a, n_z, n_x = self.radios_cropped_shape
-        self.sino_deringer = self.SinoDeringerClass(fw_sigma, (n_z, n_a, n_x), **fw_params)
+        sinos_shape = (n_z, n_a, n_x)
+        options = self.processing_options["sino_rings_correction"]
+
+        destriper_params = extract_parameters(options["user_options"])
+        if options["method"] == "munch":
+            # TODO MunchDeringer does not have an API consistent with the other deringers
+            fw_sigma = destriper_params.pop("sigma", 1.0)
+            self.sino_deringer = self.MunchDeringerClass(fw_sigma, sinos_shape, **destriper_params)
+        elif options["method"] == "vo":
+            self.sino_deringer = self.VoDeringerClass(sinos_shape, **destriper_params)
+        elif options["method"] == "mean-subtraction":
+            self.sino_deringer = self.SinoMeanDeringerClass(
+                sinos_shape, mode="subtract", fft_num_threads=None, **destriper_params
+            )
+        elif options["method"] == "mean-division":
+            self.sino_deringer = self.SinoMeanDeringerClass(
+                sinos_shape, mode="divide", fft_num_threads=None, **destriper_params
+            )
 
     @use_options("reconstruction", "reconstruction")
     def _init_reconstruction(self):
@@ -598,6 +611,7 @@ class ChunkedPipeline:
             "float_clip_values": options["float_clip_values"],
             "tiff_single_file": options.get("tiff_single_file", False),
             "single_output_file_initialized": getattr(self.process_config, "single_output_file_initialized", False),
+            "raw_vol_metadata": {"voxelSize": self.dataset_info.pixel_size},  # legacy...
         }
         writer_extra_options.update(extra_options)
         self.writer = WriterManager(
@@ -738,8 +752,11 @@ class ChunkedPipeline:
 
     @pipeline_step("writer", "Saving data")
     def _write_data(self, data=None):
-        if data is None:
+        if data is None and self.reconstruction is not None:
             data = self.recs
+        if data is None:
+            self.logger.info("No data to write")
+            return
         self.writer.write_data(data)
         self.logger.info("Wrote %s" % self.writer.fname)
         self._write_histogram()

nabu/pipeline/fullfield/chunked_cuda.py

@@ -6,10 +6,10 @@ from ...preproc.phase_cuda import CudaPaganinPhaseRetrieval
 from ...preproc.ctf_cuda import CudaCTFPhaseRetrieval
 from ...reconstruction.sinogram_cuda import CudaSinoBuilder, CudaSinoNormalization
 from ...reconstruction.filtering_cuda import CudaSinoFilter
-from ...reconstruction.rings_cuda import CudaMunchDeringer
-from ...misc.unsharp_cuda import CudaUnsharpMask
-from ...misc.rotation_cuda import CudaRotation
-from ...misc.histogram_cuda import CudaPartialHistogram
+from ...reconstruction.rings_cuda import CudaMunchDeringer, CudaSinoMeanDeringer, CudaVoDeringer
+from ...processing.unsharp_cuda import CudaUnsharpMask
+from ...processing.rotation_cuda import CudaRotation
+from ...processing.histogram_cuda import CudaPartialHistogram
 from ...reconstruction.fbp import Backprojector
 from ...reconstruction.cone import __have_astra__, ConebeamReconstructor
 from ...cuda.utils import get_cuda_context, __has_pycuda__, __pycuda_error_msg__
@@ -36,7 +36,9 @@ class CudaChunkedPipeline(ChunkedPipeline):
     UnsharpMaskClass = CudaUnsharpMask
     ImageRotationClass = CudaRotation
     VerticalShiftClass = CudaVerticalShift
-    SinoDeringerClass = CudaMunchDeringer
+    MunchDeringerClass = CudaMunchDeringer
+    SinoMeanDeringerClass = CudaSinoMeanDeringer
+    VoDeringerClass = CudaVoDeringer
     MLogClass = CudaLog
     SinoBuilderClass = CudaSinoBuilder
     SinoNormalizationClass = CudaSinoNormalization

nabu/pipeline/fullfield/computations.py

@@ -3,9 +3,11 @@ from silx.image.tomography import get_next_power
 from ...utils import check_supported
 
 
-def estimate_required_memory(process_config, delta_z=None, delta_a=None, max_mem_allocation_GB=None, debug=False):
+def estimate_required_memory(
+    process_config, delta_z=None, delta_a=None, max_mem_allocation_GB=None, fft_plans=True, debug=False
+):
     """
-    Estimate the memory (RAM) needed for a reconstruction.
+    Estimate the memory (RAM) in Bytes needed for a reconstruction.
 
     Parameters
     -----------
@@ -87,20 +89,31 @@ def estimate_required_memory(process_config, delta_z=None, delta_a=None, max_mem
     # Phase retrieval
     # ---------------
     if "phase" in processing_steps:
-        # Phase retrieval is done image-wise, so near in-place, but needs to
-        # allocate some images, fft plans, and so on
+        # Phase retrieval is done image-wise, so near in-place, but needs to allocate some memory:
+        # filter with padded shape, radio_padded, radio_padded_fourier, and possibly FFT plan.
+        # CTF phase retrieval uses "2 filters" (num and denom) but let's neglect this.
         Nx_p = get_next_power(2 * Nx)
         Nz_p = get_next_power(2 * Nz)
-        img_size_real = 2 * 4 * Nx_p * Nz_p
-        img_size_cplx = 2 * 8 * ((Nx_p * Nz_p) // 2 + 1)
-        total_memory_needed += 2 * img_size_real + 3 * img_size_cplx
+        img_size_real = Nx_p * Nz_p * 4
+        img_size_cplx = ((Nx_p * Nz_p) // 2 + 1) * 8  # assuming RFFT
+        factor = 1
+        if fft_plans:
+            factor = 2
+        total_memory_needed += (2 * img_size_real + img_size_cplx) * factor
 
     # Sinogram de-ringing
     # -------------------
     if "sino_rings_correction" in processing_steps:
-        # Process is done image-wise.
-        # Needs one Discrete Wavelets transform and one FFT/IFFT plan for each scale
-        total_memory_needed += (Nx * Na * 4) * 5.5  # approx.
+        method = process_config.processing_options["sino_rings_correction"]["method"]
+        if method == "munch":
+            # Process is done image-wise.
+            # Needs one Discrete Wavelets transform and one FFT/IFFT plan for each scale
+            factor = 2 if not (fft_plans) else 5.5  # approx!
+            total_memory_needed += (Nx * Na * 4) * factor
+        elif method == "vo":
+            # cupy-based implementation makes many calls to "scipy-like" functions, where the memory usage is not under control
+            # TODO try to estimate this
+            pass
 
     # Reconstruction
     # ---------------
@@ -127,7 +140,14 @@ def estimate_required_memory(process_config, delta_z=None, delta_a=None, max_mem
 
 
 def estimate_max_chunk_size(
-    available_memory_GB, process_config, pipeline_part="all", n_rows=None, step=10, max_mem_allocation_GB=None
+    available_memory_GB,
+    process_config,
+    pipeline_part="all",
+    n_rows=None,
+    step=10,
+    max_mem_allocation_GB=None,
+    fft_plans=True,
+    debug=False,
 ):
     """
     Estimate the maximum size of the data chunk that can be loaded in memory.
@@ -195,7 +215,12 @@ def estimate_max_chunk_size(
     while True:
         try:
             mem = estimate_required_memory(
-                process_config, delta_z=delta_z, delta_a=delta_a, max_mem_allocation_GB=max_mem_allocation_GB
+                process_config,
+                delta_z=delta_z,
+                delta_a=delta_a,
+                max_mem_allocation_GB=max_mem_allocation_GB,
+                fft_plans=fft_plans,
+                debug=debug,
             )
         except ValueError:
             # For very big dataset this function might return "0".
@@ -215,7 +240,17 @@ def estimate_max_chunk_size(
 
     process_config.processing_steps = processing_steps_bak
 
-    res = last_valid_delta_a if pipeline_part == "radios" else last_valid_delta_z
+    if pipeline_part != "radios":
+        if mem / 1e9 < available_memory_GB:
+            res = min(delta_z, process_config.radio_shape()[0])
+        else:
+            res = last_valid_delta_z
+    else:
+        if mem / 1e9 < available_memory_GB:
+            res = min(delta_a, process_config.n_angles())
+        else:
+            res = last_valid_delta_a
+
     # Really not ideal. For very large dataset, "step" should be very small.
     # Otherwise we go from 0 -> OK to 10 -> not OK, and then retain 0...
     if res == 0: