nabu 2024.1.0rc4__py3-none-any.whl → 2024.1.1__py3-none-any.whl

build/lib/nabu/pipeline/fullfield/chunked_cuda.py

@@ -0,0 +1,152 @@
+from ...preproc.ccd_cuda import CudaLog, CudaCCDFilter
+from ...preproc.flatfield_cuda import CudaFlatFieldDataUrls
+from ...preproc.shift_cuda import CudaVerticalShift
+from ...preproc.double_flatfield_cuda import CudaDoubleFlatField
+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, 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__
+from ..utils import pipeline_step
+from .chunked import ChunkedPipeline
+
+if __has_pycuda__:
+    import pycuda.gpuarray as garray
+if not (__have_astra__):
+    ConebeamReconstructor = None
+
+
+class CudaChunkedPipeline(ChunkedPipeline):
+    """
+    Cuda backend of ChunkedPipeline
+    """
+
+    backend = "cuda"
+    FlatFieldClass = CudaFlatFieldDataUrls
+    DoubleFlatFieldClass = CudaDoubleFlatField
+    CCDCorrectionClass = CudaCCDFilter
+    PaganinPhaseRetrievalClass = CudaPaganinPhaseRetrieval
+    CTFPhaseRetrievalClass = CudaCTFPhaseRetrieval
+    UnsharpMaskClass = CudaUnsharpMask
+    ImageRotationClass = CudaRotation
+    VerticalShiftClass = CudaVerticalShift
+    MunchDeringerClass = CudaMunchDeringer
+    SinoMeanDeringerClass = CudaSinoMeanDeringer
+    VoDeringerClass = CudaVoDeringer
+    MLogClass = CudaLog
+    SinoBuilderClass = CudaSinoBuilder
+    SinoNormalizationClass = CudaSinoNormalization
+    SinoFilterClass = CudaSinoFilter
+    FBPClass = Backprojector
+    ConebeamClass = ConebeamReconstructor
+    HistogramClass = CudaPartialHistogram
+
+    def __init__(
+        self,
+        process_config,
+        chunk_shape,
+        logger=None,
+        extra_options=None,
+        margin=None,
+        use_grouped_mode=False,
+        cuda_options=None,
+    ):
+        self._init_cuda(cuda_options)
+        super().__init__(
+            process_config,
+            chunk_shape,
+            logger=logger,
+            extra_options=extra_options,
+            use_grouped_mode=use_grouped_mode,
+            margin=margin,
+        )
+        self._allocate_array(self.radios.shape, "f", name="radios")
+        self._determine_when_to_transfer_data_on_gpu()
+
+    def _determine_when_to_transfer_data_on_gpu(self):
+        # Decide when to transfer data to GPU. Normally it's right after reading the data,
+        # But sometimes a part of the processing is done on CPU.
+        self._when_to_transfer_radios_on_gpu = "read_data"
+        if self.flatfield is not None and self.flatfield.distortion_correction is not None:
+            self._when_to_transfer_radios_on_gpu = "flatfield"
+
+    def _init_cuda(self, cuda_options):
+        if not (__has_pycuda__):
+            raise ImportError(__pycuda_error_msg__)
+        cuda_options = cuda_options or {}
+        self.ctx = get_cuda_context(**cuda_options)
+        self._d_radios = None
+        self._d_sinos = None
+        self._d_recs = None
+
+    def _allocate_array(self, shape, dtype, name=None):
+        name = name or "tmp"  # should be mandatory
+        d_name = "_d_" + name
+        d_arr = getattr(self, d_name, None)
+        if d_arr is None:
+            self.logger.debug("Allocating %s: %s" % (name, str(shape)))
+            d_arr = garray.zeros(shape, dtype)
+            setattr(self, d_name, d_arr)
+        return d_arr
+
+    def _transfer_radios_to_gpu(self):
+        self.logger.debug("Transfering radios to GPU")
+        self._d_radios.set(self.radios)
+        self._h_radios = self.radios
+        self.radios = self._d_radios
+
+    def _process_finalize(self):
+        self.radios = self._h_radios
+
+    #
+    # Pipeline execution (class specialization)
+    #
+
+    def _read_data(self):
+        super()._read_data()
+        if self._when_to_transfer_radios_on_gpu == "read_data":
+            self._transfer_radios_to_gpu()
+
+    def _flatfield(self):
+        super()._flatfield()
+        if self._when_to_transfer_radios_on_gpu == "flatfield":
+            self._transfer_radios_to_gpu()
+
+    def _reconstruct(self):
+        super()._reconstruct()
+        if "reconstruction" not in self.processing_steps:
+            return
+        if self.processing_options["reconstruction"]["method"] == "cone":
+            ((U, D), (L, R)) = self.margin
+            U, D = U or None, -D or None
+            # not sure why slicing can't be done before get()
+            self.recs = self.recs.get()[U:D, ...]
+        else:
+            self.recs = self.recs.get()
+
+    def _write_data(self, data=None):
+        super()._write_data(data=data)
+        if "reconstruction" in self.processing_steps:
+            self.recs = self._d_recs
+        self.radios = self._h_radios
+
+    @pipeline_step("histogram", "Computing histogram")
+    def _compute_histogram(self, data=None):
+        if data is None:
+            data = self._d_recs
+        self.recs_histogram = self.histogram.compute_histogram(data)
+
+    def _dump_data_to_file(self, step_name, data=None):
+        if data is None:
+            data = self.radios
+        if step_name not in self.datadump_manager.data_dump:
+            return
+        if isinstance(data, garray.GPUArray):
+            data = data.get()
+        self.datadump_manager.dump_data_to_file(step_name, data=data)

build/lib/nabu/pipeline/fullfield/computations.py

@@ -0,0 +1,260 @@
+from math import ceil
+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, fft_plans=True, debug=False
+):
+    """
+    Estimate the memory (RAM) in Bytes needed for a reconstruction.
+
+    Parameters
+    -----------
+    process_config: `ProcessConfig` object
+        Data structure with the processing configuration
+    delta_z: int, optional
+        How many lines are to be loaded in the each projection image.
+        Default is to load lines [start_z:end_z] (where start_z and end_z come from the user configuration file)
+    delta_a: int, optional
+        How many (partial) projection images to load at the same time. Default is to load all the projection images.
+    max_mem_allocation_GB: float, optional
+        Maximum amount of memory in GB for one single array.
+
+    Returns
+    -------
+    required_memory: float
+        Total required memory (in bytes).
+
+    Raises
+    ------
+    ValueError if one single-array allocation exceeds "max_mem_allocation_GB"
+
+    Notes
+    -----
+    pycuda <= 2022.1 cannot use arrays with more than 2**32 items (i.e 17.18 GB for float32).
+    This was solved in more recent versions.
+    """
+
+    def check_memory_limit(mem_GB, name):
+        if max_mem_allocation_GB is None:
+            return
+        if mem_GB > max_mem_allocation_GB:
+            raise ValueError(
+                "Cannot allocate array '%s' %.3f GB > max_mem_allocation_GB = %.3f GB"
+                % (name, mem_GB, max_mem_allocation_GB)
+            )
+
+    dataset = process_config.dataset_info
+    processing_steps = process_config.processing_steps
+
+    # The "x" dimension (last axis) is always the image width, because
+    #  - Processing is never "cut" along this axis (we either split along frames, or images lines)
+    #  - Even if we want to reconstruct a vertical slice (i.e end_x - start_x == 1),
+    #    the tomography reconstruction will need the full information along this axis.
+    # The only case where reading a x-subregion would be useful is to crop the initial data
+    # (i.e knowing that a part of each image will be completely unused). This is not supported yet.
+    Nx = process_config.radio_shape(binning=True)[-1]
+
+    if delta_z is not None:
+        Nz = delta_z // process_config.binning_z
+    else:
+        Nz = process_config.rec_delta_z  # accounting for binning
+    if delta_a is not None:
+        Na = ceil(delta_a / process_config.subsampling_factor)
+    else:
+        Na = process_config.n_angles(subsampling=True)
+
+    total_memory_needed = 0
+
+    # Read data
+    # ----------
+    data_image_size = Nx * Nz * 4
+    data_volume_size = Na * data_image_size
+    check_memory_limit(data_volume_size / 1e9, "projections")
+    total_memory_needed += data_volume_size
+
+    # CCD processing
+    # ---------------
+    if "flatfield" in processing_steps:
+        # Flat-field is done in-place, but still need to load darks/flats
+        n_darks = len(dataset.darks)
+        n_flats = len(dataset.flats)
+        total_memory_needed += (n_darks + n_flats) * data_image_size
+
+    if "ccd_correction" in processing_steps:
+        # CCD filter is "batched 2D"
+        total_memory_needed += data_image_size
+
+    # Phase retrieval
+    # ---------------
+    if "phase" in processing_steps:
+        # 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 = 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:
+        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
+    # ---------------
+    reconstructed_volume_size = 0
+    if "reconstruction" in processing_steps:
+        rec_config = process_config.processing_options["reconstruction"]
+        Nx_rec = rec_config["end_x"] - rec_config["start_x"] + 1
+        Ny_rec = rec_config["end_y"] - rec_config["start_y"] + 1
+        reconstructed_volume_size = Nz * Nx_rec * Ny_rec * 4
+        check_memory_limit(reconstructed_volume_size / 1e9, "reconstructions")
+        total_memory_needed += reconstructed_volume_size
+        if process_config.rec_params["method"] == "cone":
+            # In cone-beam reconstruction, need both sinograms and reconstruction inside GPU.
+            # That's big!
+            total_memory_needed += 2 * data_volume_size
+
+    if debug:
+        print(
+            "Mem for (delta_z=%s, delta_a=%s)  ==>  (Na=%d, Nz=%d, Nx=%d) : %.3f GB"
+            % (delta_z, delta_a, Na, Nz, Nx, total_memory_needed / 1e9)
+        )
+
+    return total_memory_needed
+
+
+def estimate_max_chunk_size(
+    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.
+
+    Parameters
+    ----------
+    available_memory_GB: float
+        available memory in Giga Bytes (GB - not GiB !).
+    process_config: ProcessConfig
+        ProcessConfig object
+    pipeline_part: str
+        Which pipeline part to consider. Possible options are:
+          - "full": Account for all the pipeline steps (reading data all the way to reconstruction).
+          - "radios": Consider only the processing steps on projection images (ignore sinogram-based steps and reconstruction)
+          - "sinogram": Consider only the processing steps related to sinograms and reconstruction
+    n_rows: int, optional
+        How many lines to load in each projection. Only accounted for pipeline_part="radios".
+    step: int, optional
+        Step size when doing the iterative memory estimation
+    max_mem_allocation_GB: float, optional
+        Maximum size (in GB) for one single array.
+
+    Returns
+    -------
+    n_max: int
+        If pipeline_par is "full" or "sinos": return the maximum number of lines that can be loaded
+        in all the projections while fitting memory, i.e `data[:, 0:n_max, :]`
+        If pipeline_part is "radios", return the maximum number of (partial) images that can be loaded
+        while fitting memory, i.e `data[:, zmin:zmax, 0:n_max]`
+
+
+    Notes
+    -----
+    pycuda <= 2022.1 cannot use arrays with more than 2**32 items (i.e 17.18 GB for float32).
+    This was solved in more recent versions.
+    """
+
+    supported_pipeline_parts = ["all", "radios", "sinos"]
+    check_supported(pipeline_part, supported_pipeline_parts, "pipeline_part")
+
+    processing_steps_bak = process_config.processing_steps.copy()
+    reconstruction_steps = ["sino_rings_correction", "reconstruction"]
+
+    if pipeline_part == "all":
+        # load lines from all the projections
+        delta_a = None
+        delta_z = 0
+    if pipeline_part == "radios":
+        # order should not matter
+        process_config.processing_steps = list(set(process_config.processing_steps) - set(reconstruction_steps))
+        # load lines from only a subset of projections
+        delta_a = 0
+        delta_z = n_rows
+    if pipeline_part == "sinos":
+        process_config.processing_steps = [
+            step for step in process_config.processing_steps if step in reconstruction_steps
+        ]
+        # load lines from all the projections
+        delta_a = None
+        delta_z = 0
+
+    mem = 0
+    last_valid_delta_a = delta_a
+    last_valid_delta_z = delta_z
+    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,
+                fft_plans=fft_plans,
+                debug=debug,
+            )
+        except ValueError:
+            # For very big dataset this function might return "0".
+            # Either start at 1, or use a smaller step...
+            break
+        if mem / 1e9 > available_memory_GB:
+            break
+        if delta_a is not None and delta_a > process_config.n_angles():
+            break
+        if delta_z is not None and delta_z > process_config.radio_shape()[0]:
+            break
+        last_valid_delta_a, last_valid_delta_z = delta_a, delta_z
+        if pipeline_part == "radios":
+            delta_a += step
+        else:
+            delta_z += step
+
+    process_config.processing_steps = processing_steps_bak
+
+    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:
+        res = 1
+    #
+
+    return res

build/lib/nabu/pipeline/fullfield/dataset_validator.py

@@ -0,0 +1,63 @@
+import os
+from ..dataset_validator import DatasetValidatorBase
+
+
+class FullFieldDatasetValidator(DatasetValidatorBase):
+    def _validate(self):
+        self._check_not_empty()
+        self._convert_negative_indices()
+        self._get_output_filename()
+        self._check_can_do_flatfield()
+        self._check_can_do_phase()
+        self._check_can_do_reconstruction()
+        self._check_slice_indices()
+        self._handle_processing_mode()
+        self._handle_binning()
+        self._check_output_file()
+
+    def _check_can_do_flatfield(self):
+        if self.nabu_config["preproc"]["flatfield"]:
+            darks = self.dataset_info.darks
+            assert len(darks) > 0, "Need at least one dark to perform flat-field correction"
+            for dark_id, dark_url in darks.items():
+                assert os.path.isfile(dark_url.file_path()), "Dark file %s not found" % dark_url.file_path()
+            flats = self.dataset_info.flats
+            assert len(flats) > 0, "Need at least one flat to perform flat-field correction"
+            for flat_id, flat_url in flats.items():
+                assert os.path.isfile(flat_url.file_path()), "Flat file %s not found" % flat_url.file_path()
+
+    def _check_slice_indices(self):
+        nx, nz = self.dataset_info.radio_dims
+        rec_params = self.rec_params
+        if rec_params["enable_halftomo"]:
+            ny, nx = self._get_nx_ny()
+        what = (("start_x", "end_x", nx), ("start_y", "end_y", nx), ("start_z", "end_z", nz))
+        for start_name, end_name, numels in what:
+            self._check_start_end_idx(
+                rec_params[start_name], rec_params[end_name], numels, start_name=start_name, end_name=end_name
+            )
+
+    def _check_can_do_phase(self):
+        if self.nabu_config["phase"]["method"] is None:
+            return
+        self.dataset_info.check_defined_attribute("distance")
+        self.dataset_info.check_defined_attribute("pixel_size")
+
+    def _check_can_do_reconstruction(self):
+        rec_options = self.nabu_config["reconstruction"]
+        if rec_options["method"] is None:
+            return
+        self.dataset_info.check_defined_attribute("pixel_size")
+        if rec_options["method"] == "cone":
+            if rec_options["source_sample_dist"] is None:
+                err_msg = "In cone-beam reconstruction, you have to provide 'source_sample_dist' in [reconstruction]"
+                self.logger.fatal(err_msg)
+                raise ValueError(err_msg)
+            if rec_options["sample_detector_dist"] is None:
+                if self.dataset_info.distance is None:
+                    err_msg = "Cone-beam reconstruction: 'sample_detector_dist' was not provided but could not be found in the dataset metadata either. Please provide 'sample_detector_dist'"
+                    self.logger.fatal(err_msg)
+                    raise ValueError(err_msg)
+                self.logger.warning(
+                    "Cone-beam reconstruction: 'sample_detector_dist' not provided, will use the one in dataset metadata"
+                )