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/nabu/preproc/shift_cuda.py

@@ -0,0 +1,102 @@
+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
+
+
+class CudaVerticalShift(VerticalShift):
+    def __init__(self, radios_shape, shifts, **cuda_options):
+        """
+        Vertical Shifter, Cuda backend.
+        """
+        super().__init__(radios_shape, shifts)
+        self.cuda_processing = CudaProcessing(**(cuda_options or {}))
+        self._init_cuda_arrays()
+
+    def _init_cuda_arrays(self):
+        interp_infos_arr = np.zeros((len(self.interp_infos), 2), "f")
+        self._d_interp_infos = self.cuda_processing.to_device("_d_interp_infos", interp_infos_arr)
+        self._d_radio_new = self.cuda_processing.allocate_array("_d_radio_new", self.radios_shape[1:], "f")
+        self._d_radio = self.cuda_processing.allocate_array("_d_radio", self.radios_shape[1:], "f")
+        self.muladd_kernel = CudaMulAdd(ctx=self.cuda_processing.ctx)
+
+    def apply_vertical_shifts(self, radios, iangles, output=None):
+        """
+        Parameters
+        ----------
+        radios: 3D pycuda.gpuarray.GPUArray
+            The input radios. If the optional parameter is not given, they are modified in-place
+        iangles:  a sequence of integers
+            Must have the same lenght as radios.
+            It contains the index at which the shift is found in `self.shifts`
+            given by `shifts` argument in the initialisation of the object.
+        output: 3D pycuda.gpuarray.GPUArray, optional
+            If given, it will be modified to contain the shifted radios.
+            Must be of the same shape of `radios`.
+        """
+        self._check(radios, iangles)
+        n_a, n_z, n_x = radios.shape
+        assert n_z == self.radios_shape[1]
+        x_slice = slice(0, n_x)  # slice(None, None)
+
+        def nonempty_subregion(region):
+            if region is None:
+                return True
+            z_slice = region[0]
+            return z_slice.stop - z_slice.start > 0
+
+        d_radio_new = self._d_radio_new
+        d_radio = self._d_radio
+
+        for ia in iangles:
+            d_radio_new.fill(0)
+            d_radio[:] = radios[ia, :, :]  # mul-add kernel won't work with pycuda view
+            S0, f = self.interp_infos[ia]
+            f = np.float32(f)
+
+            s0 = S0
+            if s0 > 0:
+                # newradio[:-s0] = radio[s0:] * (1 - f)
+                dst_region = (slice(0, n_z - s0), x_slice)
+                other_region = (slice(s0, n_z), x_slice)
+            elif s0 == 0:
+                # newradio[:] = radio[s0:] * (1 - f)
+                dst_region = None
+                other_region = (slice(s0, n_z), x_slice)
+            else:
+                # newradio[-s0:] = radio[:s0] * (1 - f)
+                dst_region = (slice(-s0, n_z), x_slice)
+                other_region = (slice(0, n_z + s0), x_slice)
+
+            if all([nonempty_subregion(reg) for reg in [dst_region, other_region]]):
+                self.muladd_kernel(
+                    d_radio_new,
+                    d_radio,
+                    1,
+                    1 - f,
+                    dst_region=dst_region,
+                    other_region=other_region,
+                )
+
+            s0 = S0 + 1
+            if s0 > 0:
+                # newradio[:-s0] += radio[s0:] * f
+                dst_region = (slice(0, n_z - s0), x_slice)
+                other_region = (slice(s0, n_z), x_slice)
+            elif s0 == 0:
+                # newradio[:] += radio[s0:] * f
+                dst_region = None
+                other_region = (slice(s0, n_z), x_slice)
+            else:
+                # newradio[-s0:] += radio[:s0] * f
+                dst_region = (slice(-s0, n_z), x_slice)
+                other_region = (slice(0, n_z + s0), x_slice)
+
+            if all([nonempty_subregion(reg) for reg in [dst_region, other_region]]):
+                self.muladd_kernel(d_radio_new, d_radio, 1, f, dst_region=dst_region, other_region=other_region)
+
+            if output is None:
+                radios[ia, :, :] = d_radio_new[:, :]
+            else:
+                output[ia, :, :] = d_radio_new[:, :]

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/preproc/tests/__init__.py

@@ -0,0 +1 @@
+

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/preproc/tests/test_ccd_corr.py

@@ -0,0 +1,59 @@
+import pytest
+import numpy as np
+from nabu.utils import median2 as nabu_median_filter
+from nabu.testutils import get_data
+from nabu.cuda.utils import get_cuda_context, __has_pycuda__
+from nabu.preproc.ccd import CCDFilter
+
+if __has_pycuda__:
+    import pycuda.gpuarray as garray
+    from nabu.preproc.ccd_cuda import CudaCCDFilter
+
+
+@pytest.fixture(scope="class")
+def bootstrap(request):
+    cls = request.cls
+
+    cls.data = get_data("mri_proj_astra.npz")["data"]
+    cls.data /= cls.data.max()
+    cls.put_hotspots_in_data()
+    if __has_pycuda__:
+        cls.ctx = get_cuda_context()
+
+
+@pytest.mark.usefixtures("bootstrap")
+class TestCCDFilter:
+    @classmethod
+    def put_hotspots_in_data(cls):
+        # Put 5 hot spots in the data
+        # (row, column, deviation from median)
+        cls.hotspots = [(50, 51, 0.04), (151, 150, 0.08), (202, 303, 0.12), (322, 203, 0.14)]
+        cls.threshold = 0.1  # parameterize ?
+        data_medfilt = nabu_median_filter(cls.data)
+        for r, c, deviation_from_median in cls.hotspots:
+            cls.data[r, c] = data_medfilt[r, c] + deviation_from_median
+
+    def check_detected_hotspots_locations(self, res):
+        diff = self.data - res
+        rows, cols = np.where(diff > 0)
+        hotspots_arr = np.array(self.hotspots)
+        M = hotspots_arr[:, -1] > self.threshold
+        hotspots_rows = hotspots_arr[M, 0]
+        hotspots_cols = hotspots_arr[M, 1]
+        assert np.allclose(hotspots_rows, rows)
+        assert np.allclose(hotspots_cols, cols)
+
+    def test_median_clip(self):
+        ccd_filter = CCDFilter(self.data.shape, median_clip_thresh=self.threshold)
+        res = np.zeros_like(self.data)
+        res = ccd_filter.median_clip_correction(self.data, output=res)
+        self.check_detected_hotspots_locations(res)
+
+    @pytest.mark.skipif(not (__has_pycuda__), reason="Need cuda/pycuda for this test")
+    def test_cuda_median_clip(self):
+        d_radios = garray.to_gpu(self.data)
+        cuda_ccd_correction = CudaCCDFilter(d_radios.shape, median_clip_thresh=self.threshold)
+        d_out = garray.zeros_like(d_radios)
+        cuda_ccd_correction.median_clip_correction(d_radios, output=d_out)
+        res = d_out.get()
+        self.check_detected_hotspots_locations(res)

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/preproc/tests/test_ctf.py

@@ -0,0 +1,250 @@
+import pytest
+import numpy as np
+import scipy.interpolate
+from nabu.testutils import get_data as nabu_get_data
+from nabu.testutils import __do_long_tests__
+from nabu.preproc.flatfield import FlatFieldArrays
+from nabu.preproc.ccd import CCDFilter
+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__, __has_cufft__, get_cuda_context
+
+if __has_pycuda__ and __has_cufft__:
+    from nabu.preproc.ctf_cuda import CudaCTFPhaseRetrieval
+    import pycuda.gpuarray as garray
+
+
+@pytest.fixture(scope="class")
+def bootstrap_TestCtf(request):
+    cls = request.cls
+    cls.abs_tol = 1.0e-4
+
+    test_data = nabu_get_data("ctf_tests_data_all_pars.npz")
+
+    cls.rand_disp_vh = test_data["rh"]
+    ## the dimension number 1 is over holotomo distances, so far our filter is for one distance only
+    cls.rand_disp_vh.shape = [cls.rand_disp_vh.shape[0], cls.rand_disp_vh.shape[2]]
+    cls.dark = test_data["dark"]
+    cls.flats = [test_data["ref0"], test_data["ref1"]]
+    cls.im = test_data["im"]
+    cls.ipro = int(test_data["ipro"])
+    cls.expected_result = test_data["result"]
+    cls.ref_plain = test_data["ref_plain_float_flat"]
+
+    cls.flats_n = test_data["refns"]
+
+    cls.img_shape_vh = test_data["img_shape_vh"]
+    cls.padded_img_shape_vh = test_data["padded_img_shape_vh"]
+    cls.z1_vh = test_data["z1_vh"]
+    cls.z2 = test_data["z2"]
+    cls.pix_size_det = test_data["pix_size_det"]
+    cls.length_scale = test_data["length_scale"]
+    cls.wavelength = test_data["wave_length"]
+    cls.remove_spikes_threshold = test_data["remove_spikes_threshold"]
+    cls.delta_beta = 27
+
+
+@pytest.mark.usefixtures("bootstrap_TestCtf")
+class TestCtf:
+    def check_result(self, res, ref, error_message):
+        diff = np.abs(res - ref)
+        diff[diff > np.percentile(diff, 99)] = 0
+        assert diff.max() < self.abs_tol * (np.abs(ref).mean()), error_message
+
+    def test_ctf_id16_way(self):
+        """test the ctf phase retrieval.
+        The cft filter, of the CtfFilter class is iniitalised with the geomety informations contained in geo_pars object
+        of the GeoPars class.
+        The geometry encompass the case of astigmatic wavefront with a vertical and horisontal sources which are
+        at distance z1_vh[0] and z1_vh[1] from the object.
+        In the case of parllel geometry put  z1_vh[0]= z1_vh[1] = R
+        where R is a large value ( meters).
+        SI unit system is used. But the same results shudl be obtained with any homogenuous choice of the distance units.
+        The img_shape is the shape of the images which will be processed.
+        padded_img_shape is an intermediate shape which needs to be larger that the img_shape to avoid border effect due
+        to convolutions.
+        Length scale is an internal parameters which should not affect in anyway the result unless there are serious numerical
+        problems involving very small lenghts. You can safely let the default value.
+        """
+        geo_pars = ctf.GeoPars(
+            z1_vh=self.z1_vh,
+            z2=self.z2,
+            pix_size_det=self.pix_size_det,
+            length_scale=self.length_scale,
+            wavelength=self.wavelength,
+        )
+
+        flats = FlatFieldArrays(
+            [1200] + list(self.img_shape_vh), {0: self.flats[0], 1200: self.flats[1]}, {0: self.dark}
+        )
+
+        my_flat = flats.get_flat(self.ipro)
+        my_img = self.im - self.dark
+        my_flat = my_flat - self.dark
+
+        new_coordinates = estimate_flat_distortion(
+            my_flat,
+            my_img,
+            tile_size=100,
+            interpolation_kind="cubic",
+            padding_mode="edge",
+            correction_spike_threshold=3,
+        )
+
+        interpolator = scipy.interpolate.RegularGridInterpolator(
+            (np.arange(my_flat.shape[0]), np.arange(my_flat.shape[1])),
+            my_flat,
+            bounds_error=False,
+            method="linear",
+            fill_value=None,
+        )
+        my_flat = interpolator(new_coordinates)
+
+        my_img = my_img / my_flat
+        my_img = correct_spikes(my_img, self.remove_spikes_threshold)
+
+        my_shift = self.rand_disp_vh[:, self.ipro]
+
+        ctf_filter = ctf.CtfFilter(
+            self.dark.shape,
+            geo_pars,
+            self.delta_beta,
+            padded_shape=self.padded_img_shape_vh,
+            translation_vh=my_shift,
+            normalize_by_mean=True,
+            lim1=1.0e-5,
+            lim2=0.2,
+        )
+        phase = ctf_filter.retrieve_phase(my_img)
+
+        self.check_result(
+            phase, self.expected_result, "retrieved phase and reference result differ beyond the accepted tolerance"
+        )
+
+    @pytest.mark.skipif(not (__do_long_tests__), reason="need environment variable NABU_LONG_TESTS=1")
+    def test_ctf_id16_class(self):
+        geo_pars = ctf.GeoPars(
+            z1_vh=self.z1_vh,
+            z2=self.z2,
+            pix_size_det=self.pix_size_det,
+            length_scale=self.length_scale,
+            wavelength=self.wavelength,
+        )
+        distortion_correction = DistortionCorrection(
+            estimation_method="fft-correlation",
+            estimation_kwargs={
+                "tile_size": 100,
+                "interpolation_kind": "cubic",
+                "padding_mode": "edge",
+                "correction_spike_threshold": 3.0,
+            },
+            correction_method="interpn",
+            correction_kwargs={"fill_value": None},
+        )
+        flats = FlatFieldArrays(
+            [1200] + list(self.img_shape_vh),
+            {0: self.flats[0], 1200: self.flats[1]},
+            {0: self.dark},
+            distortion_correction=distortion_correction,
+        )
+
+        # The "correct_spikes" function is numerically unstable (comparison with a float threshold).
+        # If float32 is used for the image, one spike is detected while it is not in the previous test
+        # (although the max difference between the inputs is about 1e-8).
+        # We use float64 data type for the image to make tests pass.
+        img = self.im.astype(np.float64)
+
+        flats.normalize_single_radio(img, self.ipro)
+        img = correct_spikes(img, self.remove_spikes_threshold)
+
+        shift = self.rand_disp_vh[:, self.ipro]
+
+        ctf_filter = ctf.CtfFilter(
+            img.shape,
+            geo_pars,
+            self.delta_beta,
+            padded_shape=self.padded_img_shape_vh,
+            translation_vh=shift,
+            normalize_by_mean=True,
+            lim1=1.0e-5,
+            lim2=0.2,
+        )
+        phase = ctf_filter.retrieve_phase(img)
+
+        message = "retrieved phase and reference result differ beyond the accepted tolerance"
+        assert np.abs(phase - self.expected_result).max() < self.abs_tol * (
+            np.abs(self.expected_result).mean()
+        ), message
+
+    def test_ctf_plain_way(self):
+        geo_pars = ctf.GeoPars(
+            z1_vh=None,
+            z2=self.z2,
+            pix_size_det=self.pix_size_det,
+            length_scale=self.length_scale,
+            wavelength=self.wavelength,
+        )
+        flatfielder = FlatFieldArrays(
+            [1] + list(self.img_shape_vh),
+            {0: self.flats[0], 1200: self.flats[1]},
+            {0: self.dark},
+            radios_indices=[self.ipro],
+        )
+
+        spikes_corrector = CCDFilter(
+            self.dark.shape, median_clip_thresh=self.remove_spikes_threshold, abs_diff=True, preserve_borders=True
+        )
+
+        img = self.im.astype("f")
+        img = flatfielder.normalize_radios(np.array([img]))[0]
+        img = spikes_corrector.median_clip_correction(img)
+
+        ctf_args = [img.shape, geo_pars, self.delta_beta]
+        ctf_kwargs = {"padded_shape": self.padded_img_shape_vh, "normalize_by_mean": True, "lim1": 1.0e-5, "lim2": 0.2}
+        ctf_filter = ctf.CtfFilter(*ctf_args, **ctf_kwargs)
+        phase = ctf_filter.retrieve_phase(img)
+
+        self.check_result(phase, self.ref_plain, "Something wrong with CtfFilter")
+
+        # Test R2C
+        ctf_numpy = ctf.CtfFilter(*ctf_args, **ctf_kwargs, use_rfft=True)
+        phase_r2c = ctf_numpy.retrieve_phase(img)
+        self.check_result(phase_r2c, self.ref_plain, "Something wrong with CtfFilter-R2C")
+
+        # 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)
+            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")
+    def test_cuda_ctf(self):
+        data = nabu_get_data("brain_phantom.npz")["data"]
+        delta_beta = 50.0
+        energy_kev = 22.0
+        distance_m = 1.0
+        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()
+
+        for normalize in [True, False]:
+            ctf_filter = ctf.CTFPhaseRetrieval(
+                data.shape, geo_pars, delta_beta=delta_beta, normalize_by_mean=normalize, use_rfft=True
+            )
+            cuda_ctf_filter = CudaCTFPhaseRetrieval(
+                data.shape,
+                geo_pars,
+                delta_beta=delta_beta,
+                use_rfft=True,
+                normalize_by_mean=normalize,
+            )
+            ref = ctf_filter.retrieve_phase(data)
+
+            d_data = garray.to_gpu(data)
+            res = cuda_ctf_filter.retrieve_phase(d_data).get()
+            err_max = np.max(np.abs(res - ref))
+
+            assert err_max < 1e-2, "Something wrong with retrieve_phase(normalize_by_mean=%s)" % (str(normalize))

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/preproc/tests/test_double_flatfield.py

@@ -0,0 +1,86 @@
+import os.path as path
+from math import exp
+import tempfile
+import numpy as np
+import pytest
+from silx.io.url import DataUrl
+from tomoscan.esrf.mock import MockNXtomo
+from nabu.io.reader import HDF5Reader
+from nabu.preproc.double_flatfield import DoubleFlatField
+from nabu.cuda.utils import __has_pycuda__, get_cuda_context
+
+if __has_pycuda__:
+    import pycuda.gpuarray as garray
+    from nabu.preproc.double_flatfield_cuda import CudaDoubleFlatField, __has_pycuda__
+
+
+@pytest.fixture(scope="class")
+def bootstrap(request):
+    cls = request.cls
+    cls.tmpdir = tempfile.TemporaryDirectory()
+    dname = cls.tmpdir.name
+    cls.dname = dname
+    radios = MockNXtomo(
+        path.join(dname, "tmp"),
+        10,
+        n_ini_proj=10,
+        dim=100,
+        n_refs=1,
+        scene="increasing value",
+    ).scan.projections
+    reader = HDF5Reader()
+    cls.radios = []
+    Rkeys = list(radios.keys())
+    for k in Rkeys:
+        dataurl = radios[k]
+        data = reader.get_data(dataurl)
+        cls.radios.append(data)
+    cls.radios = np.array(cls.radios)
+    cls.ff_dump_url = DataUrl(
+        file_path=path.join(cls.dname, "dff.h5"), data_path="/entry/double_flatfield/results/data"
+    )
+    cls.ff_cuda_dump_url = DataUrl(
+        file_path=path.join(cls.dname, "dff_cuda.h5"), data_path="/entry/double_flatfield/results/data"
+    )
+    golden = 0
+    for i in range(10):
+        golden += exp(-i)
+    cls.golden = golden / 10
+    cls.tol = 1e-4
+
+    if __has_pycuda__:
+        cls.ctx = get_cuda_context(cleanup_at_exit=False)
+
+    yield
+
+    if __has_pycuda__:
+        cls.ctx.pop()
+
+
+@pytest.mark.usefixtures("bootstrap")
+class TestDoubleFlatField:
+    def test_dff_numpy(self):
+        dff = DoubleFlatField(self.radios.shape, result_url=self.ff_dump_url)
+        mydf = dff.get_double_flatfield(radios=self.radios)
+
+        assert path.isfile(dff.result_url.file_path())
+
+        dff2 = DoubleFlatField(self.radios.shape, result_url=self.ff_dump_url)
+        mydf2 = dff2.get_double_flatfield(radios=self.radios)
+
+        assert np.max(np.abs(mydf2 - mydf)) < self.tol
+        assert np.max(np.abs(mydf - self.golden)) < self.tol
+
+    @pytest.mark.skipif(not (__has_pycuda__), reason="Need pycuda for double flatfield with cuda backend")
+    def test_dff_cuda(self):
+        dff = CudaDoubleFlatField(self.radios.shape, result_url=self.ff_cuda_dump_url, cuda_options={"ctx": self.ctx})
+        d_radios = garray.to_gpu(self.radios)
+        mydf = dff.get_double_flatfield(radios=d_radios).get()
+
+        assert path.isfile(dff.result_url.file_path())
+
+        dff2 = CudaDoubleFlatField(self.radios.shape, result_url=self.ff_cuda_dump_url, cuda_options={"ctx": self.ctx})
+        mydf2 = dff2.get_double_flatfield(radios=d_radios).get()
+
+        assert np.max(np.abs(mydf2 - mydf)) < self.tol
+        assert np.max(np.abs(mydf - self.golden)) < self.tol