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

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/pipeline/helical/tests/test_pipeline_elements_full.py

@@ -0,0 +1,355 @@
+from nabu.pipeline.helical import gridded_accumulator, span_strategy
+from nabu.testutils import get_data, __do_long_tests__
+import os
+import numpy as np
+import pytest
+from nabu.preproc.ccd import Log, CCDFilter
+from nabu.preproc.phase import PaganinPhaseRetrieval
+from nabu.cuda.utils import get_cuda_context, __has_pycuda__
+from nabu.pipeline.helical import gridded_accumulator, span_strategy
+from nabu.pipeline.helical.weight_balancer import WeightBalancer
+from nabu.pipeline.helical.helical_utils import find_mirror_indexes
+
+from nabu.cuda.utils import get_cuda_context, __has_pycuda__, __pycuda_error_msg__, replace_array_memory
+
+if __has_pycuda__:
+    import pycuda.gpuarray as garray
+    from nabu.pipeline.helical.fbp import BackprojectorHelical as FBPClass
+
+
+@pytest.fixture(scope="class")
+def bootstrap(request):
+    cls = request.cls
+
+    # This is a helical dataset derived
+    # from "crayon" dataset, using 5 slices and covering 2.5 x 360 angular span
+    # in halftomo, with vertical translations.
+
+    #  >>> d=load("small_sparse_helical_dataset.npz")
+    #  dd=dict( d.items() )
+    #  >>> dd["median_clip_threshold"] = 0.04
+    #  >>> savez("retouched_test.npz",**dd)
+
+    helical_dataset = get_data("small_sparse_helical_dataset.npz")
+
+    helical_dataset = dict(list(helical_dataset.items()))
+    # the radios, in the dataset file, are stored by swapping angular and x dimension
+    # so that the fast running dimension runs over the projections.
+    # Due to the sparsity of the dataset, where only an handful of slices
+    # has signal, this gives a much better compression even when the axis is translating
+    # vertically
+    helical_dataset["radios"] = np.array(np.swapaxes(helical_dataset["radios_transposed"], 0, 2))
+    del helical_dataset["radios_transposed"]
+
+    # adding members: radios,  dark, flats, z_pix_per_proj, x_pix_per_proj, projection_angles_deg,
+    # pixel_size_mm, phase_margin_pix,
+    #  weigth_field=weigth_field, double_flat
+    dataset_keys = [
+        "dark",
+        "flats",
+        "z_pix_per_proj",
+        "x_pix_per_proj",
+        "projection_angles_deg",
+        "pixel_size_mm",
+        "phase_margin_pix",
+        "weights_field",
+        "double_flat",
+        "rotation_axis_position",
+        "detector_shape_vh",
+        "result_inset",
+        "radios",
+        # further  parameters, added on top of the test data which was originally made for gridded_accumulator
+        "median_clip_threshold",
+        "distance_m",
+        "energy_kev",
+        "delta_beta",
+        "pixel_size_m",
+        "padding_type",
+        "phase_margin_for_pag",
+        "rec_reference",
+        "ref_tol",
+        "ref_start",
+        "ref_end",
+    ]
+    # the test dataset is the original one from the accumulator test
+    # plus some patched metadeta information for phase retrieval.
+    # The original dataset had phase_margin_pix and had 3 usefule slices.
+    # Here, to test phase retrieval, we redefine the phase margin to the maximum
+    # that we can do with such a small dataset
+    helical_dataset["phase_margin_pix"] = helical_dataset["phase_margin_for_pag"]
+
+    for key in dataset_keys:
+        setattr(cls, key, helical_dataset[key])
+
+    cls.padding_type = str(cls.padding_type)
+    cls.rotation_axis_position = float(cls.rotation_axis_position)
+    cls.rtol_regridded = 1.0e-6
+
+    cls.projection_angles_rad = np.rad2deg(cls.projection_angles_deg)
+
+
+@pytest.mark.skipif(not (__do_long_tests__), reason="need environment variable NABU_LONG_TESTS=1")
+@pytest.mark.skipif(not (__has_pycuda__), reason="Needs pycuda for this test")
+@pytest.mark.usefixtures("bootstrap")
+class TestGriddedAccumulator:
+    """
+    Test the GriddedAccumulator.
+    Rebuilds the sinogram for some selected slices of the crayon dataset
+    """
+
+    def test_regridding(self):
+        span_info = span_strategy.SpanStrategy(
+            z_pix_per_proj=self.z_pix_per_proj,
+            x_pix_per_proj=self.x_pix_per_proj,
+            detector_shape_vh=self.detector_shape_vh,
+            phase_margin_pix=self.phase_margin_pix,
+            projection_angles_deg=self.projection_angles_deg,
+            require_redundancy=True,
+            pixel_size_mm=self.pixel_size_mm,
+            logger=None,
+        )
+
+        # I would like to reconstruct from feaseable height 15 to feaseable  height 18
+        # relatively to the first doable slice in the vertical translation direction
+        # I get the heights in the detector frame of the first and of the last
+
+        self.reconstruction_space = gridded_accumulator.get_reconstruction_space(
+            span_info=span_info, min_scanwise_z=15, end_scanwise_z=18, phase_margin_pix=self.phase_margin_pix
+        )
+
+        chunk_info = span_info.get_chunk_info((self.reconstruction_space.my_z_min, self.reconstruction_space.my_z_end))
+
+        sub_region = (
+            self.reconstruction_space.my_z_min - self.phase_margin_pix,
+            self.reconstruction_space.my_z_end + self.phase_margin_pix,
+        )
+
+        # useful projections
+        proj_num_start, proj_num_end = chunk_info.angle_index_span
+
+        # the first of the chunk angular range
+        my_first_pnum = proj_num_start
+
+        self.accumulator = gridded_accumulator.GriddedAccumulator(
+            gridded_radios=self.reconstruction_space.gridded_radios,
+            gridded_weights=self.reconstruction_space.gridded_cumulated_weights,
+            diagnostic_radios=self.reconstruction_space.diagnostic_radios,
+            diagnostic_weights=self.reconstruction_space.diagnostic_weights,
+            diagnostic_angles=self.reconstruction_space.diagnostic_proj_angle,
+            dark=self.dark,
+            flat_indexes=[0, 7501],
+            flats=self.flats,
+            weights=self.weights_field,
+            double_flat=self.double_flat,
+        )
+
+        # splitting in sub ranges of 100 projections
+        n_granularity = 100
+        pnum_start_list = list(np.arange(proj_num_start, proj_num_end, n_granularity))
+        pnum_end_list = pnum_start_list[1:] + [proj_num_end]
+
+        for pnum_start, pnum_end in zip(pnum_start_list, pnum_end_list):
+            start_in_chunk = pnum_start - my_first_pnum
+            end_in_chunk = pnum_end - my_first_pnum
+
+            self._read_data_and_apply_flats(
+                slice(pnum_start, pnum_end), slice(start_in_chunk, end_in_chunk), chunk_info, sub_region, span_info
+            )
+
+        res_flatfielded = self.reconstruction_space.gridded_radios / self.reconstruction_space.gridded_cumulated_weights
+
+        # but in real pipeline the radio_shape is obtained from the pipeline get_shape utility method
+        self._init_ccd_corrections(res_flatfielded.shape[1:])
+
+        # but in the actual pipeline the argument is not given, and the processed stack is the one internally
+        # kept by the pipeline object ( self.gridded_radios in the pipeline )
+        self._ccd_corrections(res_flatfielded)
+
+        self._init_phase(res_flatfielded.shape[1:])
+        processed_radios = self._retrieve_phase(res_flatfielded)
+
+        self._init_mlog(processed_radios.shape)
+        self._take_log(processed_radios)
+
+        top_margin = -self.phase_margin_pix if self.phase_margin_pix else None
+        processed_weights = self.reconstruction_space.gridded_cumulated_weights[
+            :, self.phase_margin_pix : top_margin, :
+        ]
+
+        self._init_weight_balancer()
+        self._balance_weights(processed_weights)
+
+        self._init_reconstructor(processed_radios.shape)
+
+        i_slice = 0
+
+        self.d_radios_slim.set(processed_radios[:, i_slice, :])
+        self._filter()
+
+        self._apply_weights(i_slice, processed_weights)
+
+        res = self._reconstruct()
+
+        test_slicing = slice(self.ref_start, self.ref_end)
+        tested_inset = res[test_slicing, test_slicing]
+        assert np.max(np.abs(tested_inset - self.rec_reference)) < self.ref_tol
+
+        # uncomment the four following lines to get the slice image
+        # import fabio
+        # edf = fabio.edfimage.edfimage()
+        # edf.data = res
+        # edf.write("reconstructed_slice.edf")
+
+        # put the test here
+
+    def _reconstruct(self):
+        axis_corrections = np.zeros_like(self.reconstruction_space.gridded_angles_rad)
+        self.reconstruction.set_custom_angles_and_axis_corrections(
+            self.reconstruction_space.gridded_angles_rad, axis_corrections
+        )
+
+        self.reconstruction.backprojection(self.d_radios_slim, output=self.d_rec_res)
+
+        self.d_rec_res.get(self.rec_res)
+
+        return self.rec_res
+
+    def _apply_weights(self, i_slice, weights):
+        """d_radios_slim is on gpu"""
+        n_provided_angles = self.d_radios_slim.shape[0]
+
+        for first_angle_index in range(0, n_provided_angles, self.num_weight_radios_per_app):
+            end_angle_index = min(n_provided_angles, first_angle_index + self.num_weight_radios_per_app)
+            self._d_radios_weights[: end_angle_index - first_angle_index].set(
+                weights[first_angle_index:end_angle_index, i_slice]
+            )
+            self.d_radios_slim[first_angle_index:end_angle_index] *= self._d_radios_weights[
+                : end_angle_index - first_angle_index
+            ]
+
+    def _filter(self):
+        self.mirror_angle_relative_indexes = find_mirror_indexes(self.reconstruction_space.gridded_angles_deg)
+
+        self.reconstruction.sino_filter.filter_sino(
+            self.d_radios_slim,
+            mirror_indexes=self.mirror_angle_relative_indexes,
+            rot_center=self.rotation_axis_position,
+            output=self.d_radios_slim,
+        )
+
+    def _init_reconstructor(self, processed_radios_shape):
+        one_slice_data_shape = processed_radios_shape[:1] + processed_radios_shape[2:]
+
+        self.d_radios_slim = garray.zeros(one_slice_data_shape, np.float32)
+
+        # let's make room for loading chunck of weights without necessarily doubling the memory footprint.
+        # The weights will be used to multiplied the d_radios_slim.
+        # We will proceed by bunches
+        self.num_weight_radios_per_app = 200
+        self._d_radios_weights = garray.zeros((self.num_weight_radios_per_app,) + one_slice_data_shape[1:], np.float32)
+
+        pixel_size_cm = self.pixel_size_m * 100
+
+        radio_size_h = processed_radios_shape[-1]
+
+        assert (
+            2 * self.rotation_axis_position > radio_size_h
+        ), """The code of this test is adapted for HA on the right. This seems to be a case
+        of HA on the left because  self.rotation_axis_position={self.rotation_axis_position} and radio_size_h = {radio_size_h}
+        """
+
+        rec_dim = int(round(2 * self.rotation_axis_position))
+
+        self.reconstruction = FBPClass(
+            one_slice_data_shape,
+            angles=np.zeros(processed_radios_shape[0], "f"),
+            rot_center=self.rotation_axis_position,
+            filter_name=None,
+            slice_roi=(0, rec_dim, 0, rec_dim),
+            extra_options={
+                "scale_factor": 1.0 / pixel_size_cm,
+                "axis_correction": np.zeros(processed_radios_shape[0], "f"),
+                "padding_mode": "edge",
+            },
+        )
+        self.reconstruction.fbp = self.reconstruction.backproj
+
+        self.d_rec_res = garray.zeros((rec_dim, rec_dim), np.float32)
+        self.rec_res = np.zeros((rec_dim, rec_dim), np.float32)
+
+    def _init_weight_balancer(self):
+        self.weight_balancer = WeightBalancer(self.rotation_axis_position, self.reconstruction_space.gridded_angles_rad)
+
+    def _balance_weights(self, weights):
+        self.weight_balancer.balance_weights(weights)
+
+    def _retrieve_phase(self, radios):
+        processed_radios = np.zeros(
+            (radios.shape[0],) + (radios.shape[1] - 2 * self.phase_margin_pix,) + (radios.shape[2],), radios.dtype
+        )
+        for i in range(radios.shape[0]):
+            processed_radios[i] = self.phase_retrieval.apply_filter(radios[i])
+        return processed_radios
+
+    def _read_data_and_apply_flats(self, sub_total_prange_slice, subchunk_slice, chunk_info, sub_region, span_info):
+        my_integer_shifts_v = chunk_info.integer_shift_v[subchunk_slice]
+
+        subr_start_z, subr_end_z = sub_region
+
+        subr_start_z_list = subr_start_z - my_integer_shifts_v
+        subr_end_z_list = subr_end_z - my_integer_shifts_v + 1
+
+        dtasrc_start_z = max(0, subr_start_z_list.min())
+        dtasrc_end_z = min(span_info.detector_shape_vh[0], subr_end_z_list.max())
+
+        data_raw = self.radios[sub_total_prange_slice, slice(dtasrc_start_z, dtasrc_end_z), :]
+
+        subsampling_file_slice = sub_total_prange_slice
+
+        # my_subsampled_indexes = self.chunk_reader._sorted_files_indices[subsampling_file_slice]
+        my_subsampled_indexes = (np.arange(10000))[subsampling_file_slice]
+
+        self.accumulator.extract_preprocess_with_flats(
+            subchunk_slice,
+            my_subsampled_indexes,
+            chunk_info,
+            np.array((subr_start_z, subr_end_z), "i"),
+            np.array((dtasrc_start_z, dtasrc_end_z), "i"),
+            data_raw,
+        )
+
+    def _init_ccd_corrections(self, radio_shape):
+        # but in real pipeline the radio_shape is obtained from the pipeline get_shape utility method
+        self.ccd_correction = CCDFilter(radio_shape, median_clip_thresh=self.median_clip_threshold)
+
+    def _ccd_corrections(self, radios):
+        _tmp_radio = np.empty_like(radios[0])
+        for i in range(radios.shape[0]):
+            self.ccd_correction.median_clip_correction(radios[i], output=_tmp_radio)
+            radios[i][:] = _tmp_radio[:]
+
+    def _take_log(self, radios):
+        self.mlog.take_logarithm(radios)
+
+    def _init_mlog(self, radios_shape):
+        log_shape = radios_shape
+        clip_min = 1.0e-6
+        clip_max = 1.1
+        self.mlog = Log(log_shape, clip_min=clip_min, clip_max=clip_max)
+
+    def _init_phase(self, raw_shape):
+        self.phase_retrieval = PaganinPhaseRetrieval(
+            raw_shape,
+            distance=self.distance_m,
+            energy=self.energy_kev,
+            delta_beta=self.delta_beta,
+            pixel_size=self.pixel_size_m,
+            padding=self.padding_type,
+            margin=((self.phase_margin_pix,) * 2, (0, 0)),
+            use_R2C=True,
+            fftw_num_threads=True,  # TODO tune in advanced params of nabu config file
+        )
+        if self.phase_retrieval.use_fftw:
+            self.logger.debug(
+                "PaganinPhaseRetrieval using FFTW with %d threads" % self.phase_retrieval.fftw.num_threads
+            )

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/pipeline/helical/tests/test_strategy.py

@@ -0,0 +1,61 @@
+import pytest
+import numpy as np
+from nabu.testutils import get_data as nabu_get_data
+from nabu.pipeline.helical.span_strategy import SpanStrategy
+
+
+@pytest.fixture(scope="class")
+def bootstrap_TestStrategy(request):
+    cls = request.cls
+    cls.abs_tol = 1.0e-6
+
+    # from the Paul telephone dataset
+    test_data = nabu_get_data("data_test_strategy.npz")
+
+    cls.z_pix_per_proj = test_data["z_pix_per_proj"]
+    cls.x_pix_per_proj = test_data["x_pix_per_proj"]
+    cls.detector_shape_vh = test_data["detector_shape_vh"]
+    cls.phase_margin_pix = test_data["phase_margin_pix"]
+    cls.projection_angles_deg = test_data["projection_angles_deg"]
+    cls.require_redundancy = test_data["require_redundancy"]
+    cls.pixel_size_mm = test_data["pixel_size_mm"]
+
+    cls.result_angle_index_span = test_data["result_angle_index_span"]
+    cls.result_angles_rad = test_data["result_angles_rad"]
+    cls.result_fract_complement_to_integer_shift_v = test_data["result_fract_complement_to_integer_shift_v"]
+    cls.result_integer_shift_v = test_data["result_integer_shift_v"]
+    cls.result_span_v = test_data["result_span_v"]
+    cls.result_x_pix_per_proj = test_data["result_x_pix_per_proj"]
+    cls.result_z_pix_per_proj = test_data["result_z_pix_per_proj"]
+
+    cls.test_data = test_data
+
+
+@pytest.mark.usefixtures("bootstrap_TestStrategy")
+class TestStrategy:
+    def test_strategy(self):
+        # the python implementation is slow. so we take only a p[art of the scan
+        limit = 4000
+        span_info = SpanStrategy(
+            z_pix_per_proj=self.z_pix_per_proj[:limit],
+            x_pix_per_proj=self.x_pix_per_proj[:limit],
+            detector_shape_vh=self.detector_shape_vh,
+            phase_margin_pix=self.phase_margin_pix,
+            projection_angles_deg=self.projection_angles_deg[:limit],
+            pixel_size_mm=self.pixel_size_mm,
+            require_redundancy=self.require_redundancy,
+        )
+
+        print(span_info.get_informative_string())
+        chunk_info = span_info.get_chunk_info(self.result_span_v)
+
+        for key, val in chunk_info.__dict__.items():
+            reference = getattr(self, "result_" + key)
+            ref_array = np.array(reference)
+            res_array = np.array(val)
+            if res_array.dtype in [bool, np.int32, np.int64]:
+                message = f" different result for {key} attribute in the chunk_info returned value  "
+                assert np.array_equal(res_array, ref_array), message
+            elif res_array.dtype in [np.float32, np.float64]:
+                message = f" different result for {key} attribute in the chunk_info returned value  "
+                assert np.all(np.isclose(res_array, ref_array, atol=self.abs_tol)), message

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/pipeline/helical/utils.py

@@ -0,0 +1,51 @@
+from ...utils import *
+from ...io.writer import Writers, NXProcessWriter
+from ...io.tiffwriter_zmm import TIFFWriter
+from ...resources.logger import LoggerOrPrint
+from ...resources.utils import is_hdf5_extension
+from os import path, mkdir
+from ...utils import check_supported
+from ..fallback_utils import WriterConfigurator
+from ..params import files_formats
+
+Writers["tif"] = TIFFWriter
+Writers["tiff"] = TIFFWriter
+
+
+class WriterConfiguratorHelical(WriterConfigurator):
+    def _get_initial_writer_kwarg(self):
+        if self.file_format in ["tif", "tiff"]:
+            return {"heights_above_stage_mm": self.heights_above_stage_mm}
+        else:
+            return {}
+
+    def __init__(
+        self,
+        output_dir,
+        file_prefix,
+        file_format="hdf5",
+        overwrite=False,
+        start_index=None,
+        logger=None,
+        nx_info=None,
+        write_histogram=False,
+        histogram_entry="entry",
+        writer_options=None,
+        extra_options=None,
+        heights_above_stage_mm=None,
+    ):
+        self.heights_above_stage_mm = heights_above_stage_mm
+        self.file_format = file_format
+        super().__init__(
+            output_dir,
+            file_prefix,
+            file_format=file_format,
+            overwrite=overwrite,
+            start_index=start_index,
+            logger=logger,
+            nx_info=nx_info,
+            write_histogram=write_histogram,
+            histogram_entry=histogram_entry,
+            writer_options=writer_options,
+            extra_options=extra_options,
+        )

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/pipeline/helical/weight_balancer.py

@@ -0,0 +1,101 @@
+import numpy as np
+import math
+
+
+class WeightBalancer:
+    def __init__(self, rot_center, angles_rad):
+        """This  class contains the method for rebalancing the weight  prior to
+        backprojection. The weights of halfomo redundacy data ( the central part) are rebalanced.
+        In a pipeline, the weights rebalanced by  the method balance_weight, have to be applied to the ramp-filtered
+        data prior to backprojection.
+        As a matter of fact the method balance_weights  could be called as a function, but in order to be  conformant
+        to Nabu, we create this class and follow the scheme initialisation + application.
+
+        Parameters
+        ==========
+
+        rot_center : float
+                    the center of rotation in pixel units
+        angles_rad :
+                   the angles corresponding to the to be rebalanced projections
+
+        """
+        self.rot_center = rot_center
+        self.my_angles_rad = angles_rad
+
+    def balance_weights(self, radios_weights):
+        """
+        The parameter radios_weights is a stack having having the same weight as the stack of projections. It is modified in place, correcting
+        the value of overlapping data, so that the sum is always one
+        """
+        radios_weights[:] = self._rebalance(radios_weights)
+
+    def _rebalance(self, radios_weights):
+        """rebalance the  weights,  within groups of equivalent (up to multiple of 180), data pixels"""
+
+        balanced = np.zeros_like(radios_weights)
+        n_span = int(math.ceil(self.my_angles_rad[-1] - self.my_angles_rad[0]) / np.pi)
+        center = (radios_weights.shape[-1] - 1) / 2
+        nloop = balanced.shape[0]
+
+        for i in range(nloop):
+            w_res = balanced[i]
+            angle = self.my_angles_rad[i]
+
+            for i_half_turn in range(-n_span - 1, n_span + 2):
+                if i_half_turn == 0:
+                    w_res[:] += radios_weights[i]
+                    continue
+
+                shifted_angle = angle + i_half_turn * np.pi
+
+                insertion_index = np.searchsorted(self.my_angles_rad, shifted_angle)
+
+                if insertion_index in [0, self.my_angles_rad.shape[0]]:
+                    if insertion_index == 0:
+                        if abs(self.my_angles_rad[0] - shifted_angle) > np.pi / 100:
+                            continue
+                        myimage = radios_weights[0]
+                    else:
+                        if abs(self.my_angles_rad[-1] - shifted_angle) > np.pi / 100:
+                            continue
+                        myimage = radios_weights[-1]
+                else:
+                    partition = shifted_angle - self.my_angles_rad[insertion_index - 1]
+                    myimage = (1.0 - partition) * radios_weights[insertion_index - 1] + partition * radios_weights[
+                        insertion_index
+                    ]
+
+                if i_half_turn % 2 == 0:
+                    w_res[:] += myimage
+                else:
+                    myimage = np.fliplr(myimage)
+                    w_res[:] += shift(myimage, (2 * self.rot_center - 2 * center))
+
+        mask = np.equal(0, radios_weights)
+        balanced[:] = radios_weights / balanced
+        balanced[mask] = 0
+        return balanced
+
+
+def shift(arr, shift, fill_value=0.0):
+    """trivial horizontal shift.
+    Contrarily to scipy.ndimage.interpolation.shift, this shift does not cut the tails abruptly, but by interpolation
+    """
+    result = np.zeros_like(arr)
+
+    num1 = int(math.floor(shift))
+    num2 = num1 + 1
+    partition = shift - num1
+
+    for num, factor in zip([num1, num2], [(1 - partition), partition]):
+        if num > 0:
+            result[:, :num] += fill_value * factor
+            result[:, num:] += arr[:, :-num] * factor
+        elif num < 0:
+            result[:, num:] += fill_value * factor
+            result[:, :num] += arr[:, -num:] * factor
+        else:
+            result[:] += arr * factor
+
+    return result