nabu 2024.1.10__py3-none-any.whl → 2024.2.0__py3-none-any.whl

nabu/preproc/tests/test_flatfield.py

@@ -1,16 +1,11 @@
-from tempfile import mkdtemp
 import os
 import numpy as np
 import pytest
-from silx.io.url import DataUrl
-from silx.io import get_data
-from silx.io.dictdump import dicttoh5
 from nabu.cuda.utils import get_cuda_context, __has_pycuda__
-from nabu.preproc.flatfield import FlatField, FlatFieldDataUrls
+from nabu.preproc.flatfield import FlatField
 
 if __has_pycuda__:
-    import pycuda.gpuarray as garray
-    from nabu.preproc.flatfield_cuda import CudaFlatFieldDataUrls, CudaFlatField
+    from nabu.preproc.flatfield_cuda import CudaFlatField
 
 
 # Flats values should be O(k) so that linear interpolation between flats gives exact results
@@ -84,7 +79,6 @@ def generate_test_flatfield_generalized(
     flats_values,
     darks_indices,
     darks_values,
-    h5_fname,
     dtype=np.uint16,
 ):
     """
@@ -112,14 +106,11 @@ def generate_test_flatfield_generalized(
     -------
     radios: numpy.ndarray
         3D array with raw radios
-    darks: dict of DataUrls
-        Dictionary where each key is the dark indice, and value is a DataUrl
-    flats: dict of DataUrls
-        Dictionary where each key is the flat indice, and value is a DataUrl
+    darks: dict of arrays
+        Dictionary where each key is the dark indice, and value is an array
+    flats: dict of arrays
+        Dictionary where each key is the flat indice, and value is an array
     """
-    tempdir = mkdtemp(prefix="nabu_")
-    testffname = os.path.join(tempdir, h5_fname)
-
     # Radios
     radios = np.zeros((len(radios_values),) + image_shape, dtype="f")
     n_radios = radios.shape[0]
@@ -129,26 +120,15 @@ def generate_test_flatfield_generalized(
 
     # Flats
     flats = {}
-    flats_urls = {}
     for i, flat_idx in enumerate(flats_indices):
-        flats["flats_%06d" % flat_idx] = np.zeros(img_shape, dtype=dtype) + flats_values[i]
-        flats_urls[flat_idx] = DataUrl(
-            file_path=testffname, data_path=str("/flats/flats_%06d" % flat_idx), scheme="silx"
-        )
+        flats[flat_idx] = np.zeros(img_shape, dtype=dtype) + flats_values[i]
 
     # Darks
     darks = {}
-    darks_urls = {}
     for i, dark_idx in enumerate(darks_indices):
-        darks["darks_%06d" % dark_idx] = np.zeros(img_shape, dtype=dtype) + darks_values[i]
-        darks_urls[dark_idx] = DataUrl(
-            file_path=testffname, data_path=str("/darks/darks_%06d" % dark_idx), scheme="silx"
-        )
+        darks[dark_idx] = np.zeros(img_shape, dtype=dtype) + darks_values[i]
 
-    dicttoh5(flats, testffname, h5path="/flats", mode="w")
-    dicttoh5(darks, testffname, h5path="/darks", mode="a")
-
-    return radios, flats_urls, darks_urls
+    return radios, flats, darks
 
 
 @pytest.fixture(scope="class")
@@ -173,7 +153,7 @@ def bootstrap(request):
 class TestFlatField:
     def get_test_elements(self, case_name):
         config = flatfield_tests_cases[case_name]
-        radios_stack, flats_url, darks_url = generate_test_flatfield_generalized(
+        radios_stack, flats, darks = generate_test_flatfield_generalized(
             config["image_shape"],
             config["radios_indices"],
             config["radios_values"],
@@ -181,12 +161,11 @@ class TestFlatField:
             config["flats_values"],
             config["darks_indices"],
             config["darks_values"],
-            "test_ff.h5",
         )
-        fname = flats_url[list(flats_url.keys())[0]].file_path()
-        self.tmp_files.append(fname)
-        self.tmp_dirs.append(os.path.dirname(fname))
-        return radios_stack, flats_url, darks_url, config
+        # fname = flats_url[list(flats_url.keys())[0]].file_path()
+        # self.tmp_files.append(fname)
+        # self.tmp_dirs.append(os.path.dirname(fname))
+        return radios_stack, flats, darks, config
 
     @staticmethod
     def check_normalized_radios(radios_corr, expected_values):
@@ -203,9 +182,9 @@ class TestFlatField:
         (I - D)/(F - D)   where I = (1, 2, ...), D = 1, F = 0.5
         = (0, -2, -4, -6, ...)
         """
-        radios_stack, flats_url, darks_url, config = self.get_test_elements("simple_nearest_interp")
+        radios_stack, flats, darks, config = self.get_test_elements("simple_nearest_interp")
 
-        flatfield = FlatFieldDataUrls(radios_stack.shape, flats_url, darks_url)
+        flatfield = FlatField(radios_stack.shape, flats, darks)
         radios_corr = flatfield.normalize_radios(np.copy(radios_stack))
         self.check_normalized_radios(radios_corr, config["expected_result"])
 
@@ -213,16 +192,17 @@ class TestFlatField:
         """
         Same as test_flatfield_simple, but in a vertical subregion of the radios.
         """
-        radios_stack, flats_url, darks_url, config = self.get_test_elements("simple_nearest_interp")
+        radios_stack, flats, darks, config = self.get_test_elements("simple_nearest_interp")
         end_z = 51
+        flats = {k: arr[:end_z, :] for k, arr in flats.items()}
+        darks = {k: arr[:end_z, :] for k, arr in darks.items()}
         radios_chunk = np.copy(radios_stack[:, :end_z, :])
         # we only have a chunk in memory. Instantiate the class with the
         # corresponding subregion to only load the relevant part of dark/flat
-        flatfield = FlatFieldDataUrls(
+        flatfield = FlatField(
             radios_chunk.shape,
-            flats_url,
-            darks_url,
-            sub_region=(None, None, None, end_z),  # start_x, end_x, start_z, end_z
+            flats,
+            darks,
         )
         radios_corr = flatfield.normalize_radios(radios_chunk)
         self.check_normalized_radios(radios_corr, config["expected_result"])
@@ -239,8 +219,8 @@ class TestFlatField:
             = (I-D)/(F-D)
             = (I-1)/I
         """
-        radios_stack, flats_url, darks_url, config = self.get_test_elements("two_flats_no_radios_indices")
-        flatfield = FlatFieldDataUrls(radios_stack.shape, flats_url, darks_url)
+        radios_stack, flats, darks, config = self.get_test_elements("two_flats_no_radios_indices")
+        flatfield = FlatField(radios_stack.shape, flats, darks)
         radios_corr = flatfield.normalize_radios(np.copy(radios_stack))
         self.check_normalized_radios(radios_corr, config["expected_result"])
 
@@ -249,10 +229,10 @@ class TestFlatField:
         # F   = 2                       11
         # F_i = 2   3.8  5.6 7.4  9.2   11   11   11   11    11
         # R   = 0   .357  .435 .469 .488 .5  .6   .7  .8     .9
-        flats_url = flats_url.copy()
-        flats_url[5] = flats_url[9]
-        flats_url.pop(9)
-        flatfield = FlatFieldDataUrls(radios_stack.shape, flats_url, darks_url)
+        flats = {k: v.copy() for k, v in flats.items()}
+        flats[5] = flats[9]
+        flats.pop(9)
+        flatfield = FlatField(radios_stack.shape, flats, darks)
         radios_corr = flatfield.normalize_radios(np.copy(radios_stack))
         self.check_normalized_radios(
             radios_corr, [0.0, 0.35714286, 0.43478261, 0.46875, 0.48780488, 0.5, 0.6, 0.7, 0.8, 0.9]
@@ -263,13 +243,13 @@ class TestFlatField:
         """
         Test the flat-field with cuda back-end.
         """
-        radios_stack, flats_url, darks_url, config = self.get_test_elements("two_flats_no_radios_indices")
-        d_radios = garray.to_gpu(radios_stack.astype("f"))
-        cuda_flatfield = CudaFlatFieldDataUrls(
-            d_radios.shape,
-            flats_url,
-            darks_url,
+        radios_stack, flats, darks, config = self.get_test_elements("two_flats_no_radios_indices")
+        cuda_flatfield = CudaFlatField(
+            radios_stack.shape,
+            flats,
+            darks,
         )
+        d_radios = cuda_flatfield.cuda_processing.to_device("d_radios", radios_stack.astype("f"))
         cuda_flatfield.normalize_radios(d_radios)
         radios_corr = d_radios.get()
         self.check_normalized_radios(radios_corr, config["expected_result"])
@@ -277,27 +257,27 @@ class TestFlatField:
     # Linear interpolation, two flats, one dark
     def test_twoflats_simple(self):
         radios, flats, darks, config = self.get_test_elements("two_flats_with_radios_indices")
-        FF = FlatFieldDataUrls(radios.shape, flats, darks, radios_indices=config["radios_indices"])
+        FF = FlatField(radios.shape, flats, darks, radios_indices=config["radios_indices"])
         FF.normalize_radios(radios)
         self.check_normalized_radios(radios, config["expected_result"])
 
     def _setup_numerical_issue(self):
         radios, flats, darks, config = self.get_test_elements("two_flats_with_radios_indices")
-
-        # Retrieve the actual data for radios/darks/flats to use FlatField instead of FlatFieldDataUrl.
-        # Create a setting yielding "0/0": one pixel such that flat==dark and radio==dark
-        for flat_idx, flat_url in flats.items():
-            flats[flat_idx] = get_data(flat_url)
-            flats[flat_idx][0, 0] = 99
-        for dark_idx, dark_url in darks.items():
-            darks[dark_idx] = get_data(dark_url)
-            darks[dark_idx][0, 0] = 99
+        flats_copy = {}
+        darks_copy = {}
+
+        for flat_idx, flat in flats.items():
+            flats_copy[flat_idx] = flat.copy()
+            flats_copy[flat_idx][0, 0] = 99
+        for dark_idx, dark in darks.items():
+            darks_copy[dark_idx] = dark.copy()
+            darks_copy[dark_idx][0, 0] = 99
         radios[:, 0, 0] = 99
-        return radios, flats, darks, config
+        return radios, flats_copy, darks_copy, config
 
     def _check_numerical_issue(self, radios, expected_result, nan_value=None):
         if nan_value is None:
-            assert np.alltrue(np.logical_not(np.isfinite(radios[:, 0, 0]))), "First pixel should be nan or inf"
+            assert np.all(np.logical_not(np.isfinite(radios[:, 0, 0]))), "First pixel should be nan or inf"
             radios[:, 0, 0] = radios[:, 1, 1]
             self.check_normalized_radios(radios, expected_result)
         else:
@@ -341,10 +321,10 @@ class TestFlatField:
         """
         radios, flats, darks, config = self._setup_numerical_issue()
         radios0 = radios.copy()
-        d_radios = garray.to_gpu(radios)
         FF_no_nan_handling = CudaFlatField(
             radios.shape, flats, darks, radios_indices=config["radios_indices"], nan_value=None
         )
+        d_radios = FF_no_nan_handling.cuda_processing.to_device("radios", radios)
         # In a cuda kernel, no one can hear you scream
         FF_no_nan_handling.normalize_radios(d_radios)
         radios = d_radios.get()
@@ -363,7 +343,7 @@ class TestFlatField:
     def test_srcurrent(self):
         radios, flats, darks, config = self.get_test_elements("three_flats_srcurrent")
 
-        FF = FlatFieldDataUrls(
+        FF = FlatField(
             radios.shape,
             flats,
             darks,
@@ -377,9 +357,8 @@ class TestFlatField:
     @pytest.mark.skipif(not (__has_pycuda__), reason="Need cuda/pycuda for this test")
     def test_srcurrent_cuda(self):
         radios, flats, darks, config = self.get_test_elements("three_flats_srcurrent")
-        d_radios = garray.to_gpu(radios)
 
-        FF = CudaFlatFieldDataUrls(
+        FF = CudaFlatField(
             radios.shape,
             flats,
             darks,
@@ -387,6 +366,7 @@ class TestFlatField:
             radios_srcurrent=config["radios_srcurrent"],
             flats_srcurrent=config["flats_srcurrent"],
         )
+        d_radios = FF.cuda_processing.to_device("radios", radios)
         FF.normalize_radios(d_radios)
         radios_corr = d_radios.get()
         self.check_normalized_radios(radios_corr, config["expected_result"])
@@ -429,8 +409,6 @@ class FlatFieldTestDataset:
 
     def __init__(self):
         self._generate_projections()
-        self._dump_to_h5()
-        self._generate_dataurls()
 
     def get_flat_idx(self, proj_idx):
         flats_idx = sorted(list(self.flats.keys()))
@@ -461,26 +439,6 @@ class FlatFieldTestDataset:
             self.projs[str(proj_idx)] = np.zeros(self.shp, "f") + proj_val
             self.projs_data[i] = self.projs[str(proj_idx)]
 
-    def _dump_to_h5(self):
-        self.tempdir = mkdtemp(prefix="nabu_")
-        self.fname = os.path.join(self.tempdir, "projs_flats.h5")
-        dicttoh5(
-            {
-                "projs": self.projs,
-                "flats": {str(k): v for k, v in self.flats.items()},
-                "darks": {str(k): v for k, v in self.darks.items()},
-            },
-            h5file=self.fname,
-        )
-
-    def _generate_dataurls(self):
-        self.flats_urls = {}
-        for idx in self.flats.keys():
-            self.flats_urls[int(idx)] = DataUrl(file_path=self.fname, data_path="/flats/%d" % idx)
-        self.darks_urls = {}
-        for idx in self.darks.keys():
-            self.darks_urls[int(idx)] = DataUrl(file_path=self.fname, data_path="/darks/0")
-
 
 @pytest.fixture(scope="class")
 def bootstraph5(request):
@@ -496,9 +454,6 @@ def bootstraph5(request):
     cls.tol_std = 1e-3
 
     yield
-    # tear-down
-    os.remove(cls.dataset.fname)
-    os.rmdir(cls.dataset.tempdir)
 
 
 @pytest.mark.usefixtures("bootstraph5")
@@ -512,10 +467,10 @@ class TestFlatFieldH5:
         assert np.max(np.abs(errs)) < self.tol, "Something wrong with flat-field normalization"
 
     def test_flatfield(self):
-        flatfield = FlatFieldDataUrls(
+        flatfield = FlatField(
             self.dataset.projs_data.shape,
-            self.dataset.flats_urls,
-            self.dataset.darks_urls,
+            self.dataset.flats,
+            self.dataset.darks,
             radios_indices=self.dataset.projs_idx,
             interpolation="linear",
         )
@@ -525,13 +480,13 @@ class TestFlatFieldH5:
 
     @pytest.mark.skipif(not (__has_pycuda__), reason="Need cuda/pycuda for this test")
     def test_cuda_flatfield(self):
-        d_projs = garray.to_gpu(self.dataset.projs_data)
-        cuda_flatfield = CudaFlatFieldDataUrls(
+        cuda_flatfield = CudaFlatField(
             self.dataset.projs_data.shape,
-            self.dataset.flats_urls,
-            self.dataset.darks_urls,
+            self.dataset.flats,
+            self.dataset.darks,
             radios_indices=self.dataset.projs_idx,
         )
+        d_projs = cuda_flatfield.cuda_processing.to_device("d_projs", self.dataset.projs_data)
         cuda_flatfield.normalize_radios(d_projs)
         projs = d_projs.get()
         self.check_normalization(projs)
@@ -551,24 +506,17 @@ class TestFlatFieldH5:
 def generate_test_flatfield(n_radios, radio_shape, flat_interval, h5_fname):
     radios = np.zeros((n_radios,) + radio_shape, "f")
     dark_data = np.ones(radios.shape[1:], "f")
-    tempdir = mkdtemp(prefix="nabu_")
-    testffname = os.path.join(tempdir, h5_fname)
     flats = {}
-    flats_urls = {}
     # F_i = i + 2
     # R_i = i*(F_i - 1) + 1
     # N_i = (R_i - D)/(F_i - D) = i*(F_i - 1)/( F_i - 1) = i
     for i in range(n_radios):
         f_i = i + 2
         if (i % flat_interval) == 0:
-            flats["flats_%06d" % i] = np.zeros(radio_shape, "f") + f_i
-            flats_urls[i] = DataUrl(file_path=testffname, data_path=str("/flats/flats_%06d" % i), scheme="silx")
+            flats[i] = np.zeros(radio_shape, "f") + f_i
         radios[i] = i * (f_i - 1) + 1
-    dark = {"dark_0000": dark_data}
-    dicttoh5(flats, testffname, h5path="/flats", mode="w")
-    dicttoh5(dark, testffname, h5path="/dark", mode="a")
-    dark_url = {0: DataUrl(file_path=testffname, data_path="/dark/dark_0000", scheme="silx")}
-    return radios, flats_urls, dark_url
+    darks = {0: dark_data}
+    return radios, flats, darks
 
 
 @pytest.fixture(scope="class")
@@ -582,17 +530,14 @@ def bootstrap_multiflats(request):
 
     radios, flats, dark = generate_test_flatfield(n_radios, radio_shape, cls.flat_interval, h5_fname)
     cls.radios = radios
-    cls.flats_urls = flats
-    cls.darks_urls = dark
+    cls.flats = flats
+    cls.darks = dark
     cls.expected_results = np.arange(n_radios)
 
     cls.tol = 5e-4
     cls.tol_std = 1e-4
 
     yield
-    # tear down
-    os.remove(dark[0].file_path())
-    os.rmdir(os.path.dirname(dark[0].file_path()))
 
 
 @pytest.mark.usefixtures("bootstrap_multiflats")
@@ -607,7 +552,7 @@ class TestFlatFieldMultiFlat:
         assert np.max(np.abs(errs)) < self.tol, "Something wrong with flat-field normalization"
 
     def test_flatfield(self):
-        flatfield = FlatFieldDataUrls(self.radios.shape, self.flats_urls, self.darks_urls, interpolation="linear")
+        flatfield = FlatField(self.radios.shape, self.flats, self.darks, interpolation="linear")
         projs = np.copy(self.radios)
         flatfield.normalize_radios(projs)
         print(projs[:, 0, 0])
@@ -615,12 +560,12 @@ class TestFlatFieldMultiFlat:
 
     @pytest.mark.skipif(not (__has_pycuda__), reason="Need cuda/pycuda for this test")
     def test_cuda_flatfield(self):
-        d_projs = garray.to_gpu(self.radios)
-        cuda_flatfield = CudaFlatFieldDataUrls(
+        cuda_flatfield = CudaFlatField(
             self.radios.shape,
-            self.flats_urls,
-            self.darks_urls,
+            self.flats,
+            self.darks,
         )
+        d_projs = cuda_flatfield.cuda_processing.to_device("radios", self.radios)
         cuda_flatfield.normalize_radios(d_projs)
         projs = d_projs.get()
         self.check_normalization(projs)

nabu/preproc/tests/test_paganin.py

@@ -1,32 +1,35 @@
 import pytest
 import numpy as np
 from nabu.preproc.phase import PaganinPhaseRetrieval
-from nabu.testutils import get_data
+from nabu.processing.fft_cuda import get_available_fft_implems
+from nabu.testutils import generate_tests_scenarios, get_data
 from nabu.thirdparty.tomopy_phase import retrieve_phase
-from nabu.cuda.utils import __has_pycuda__, __has_cufft__
+from nabu.cuda.utils import __has_pycuda__
 
+__has_cufft__ = False
 if __has_pycuda__:
     from nabu.preproc.phase_cuda import CudaPaganinPhaseRetrieval
 
-scenarios = [
-    {
-        "distance": 1,
-        "energy": 35,
-        "delta_beta": 1e1,
-        "margin": ((50, 50), (0, 0)),
-    }
-]
+    avail_fft = get_available_fft_implems()
+    __has_cufft__ = len(avail_fft) > 0
 
+scenarios = {
+    "distance": [1],
+    "energy": [35],
+    "delta_beta": [1e1],
+    "margin": [((50, 50), (0, 0)), None],
+}
 
-@pytest.fixture(scope="class", params=scenarios)
+scenarios = generate_tests_scenarios(scenarios)
+
+
+@pytest.fixture(scope="class")
 def bootstrap(request):
     cls = request.cls
-    cls.paganin_config = request.param
 
     cls.data = get_data("mri_proj_astra.npz")["data"]
     cls.rtol = 1.1e-6
     cls.rtol_pag = 5e-3
-    cls.paganin = PaganinPhaseRetrieval(cls.data.shape, **cls.paganin_config)
 
 
 @pytest.mark.usefixtures("bootstrap")
@@ -36,32 +39,53 @@ class TestPaganin:
     The reference implementation is tomopy.
     """
 
-    def crop_to_margin(self, data):
-        s0, s1 = self.paganin.shape_inner
-        ((U, _), (L, _)) = self.paganin.margin
-        return data[U : U + s0, L : L + s1]
+    @staticmethod
+    def get_paganin_instance_and_data(cfg, data):
+        pag_kwargs = cfg.copy()
+        margin = pag_kwargs.pop("margin")
+        if margin is not None:
+            data = np.pad(data, margin, mode="edge")
+        paganin = PaganinPhaseRetrieval(data.shape, **pag_kwargs)
+        return paganin, data, pag_kwargs
 
-    def test_paganin(self):
-        data_tomopy = np.atleast_3d(np.copy(self.data)).T
+    @staticmethod
+    def crop_to_margin(data, margin):
+        if margin is None:
+            return data
+        ((U, D), (L, R)) = margin
+        D = None if D == 0 else -D
+        R = None if R == 0 else -R
+        return data[U:D, L:R]
 
+    @pytest.mark.parametrize("config", scenarios)
+    def test_paganin(self, config):
+        paganin, data, _ = self.get_paganin_instance_and_data(config, self.data)
+        res = paganin.apply_filter(data)
+
+        data_tomopy = np.atleast_3d(np.copy(data)).T
         res_tomopy = retrieve_phase(
             data_tomopy,
-            pixel_size=self.paganin.pixel_size_micron * 1e-4,
-            dist=self.paganin.distance_cm,
-            energy=self.paganin.energy_kev,
-            alpha=1.0 / (4 * 3.141592**2 * self.paganin.delta_beta),
+            pixel_size=paganin.pixel_size_xy_micron[0] * 1e-4,
+            dist=paganin.distance_cm,
+            energy=paganin.energy_kev,
+            alpha=1.0 / (4 * 3.141592**2 * paganin.delta_beta),
         )
-        res_tomopy = self.crop_to_margin(res_tomopy[0].T)
 
-        res = self.paganin.apply_filter(self.data)
+        res_tomopy = self.crop_to_margin(res_tomopy[0].T, config["margin"])
+        res = self.crop_to_margin(res, config["margin"])
 
         errmax = np.max(np.abs(res - res_tomopy) / np.max(res_tomopy))
         assert errmax < self.rtol_pag, "Max error is too high"
 
-    @pytest.mark.skipif(not (__has_pycuda__ and __has_cufft__), reason="Need pycuda and scikit-cuda for this test")
-    def test_gpu_paganin(self):
-        gpu_paganin = CudaPaganinPhaseRetrieval(self.data.shape, **self.paganin_config)
-        ref = self.paganin.apply_filter(self.data)
-        res = gpu_paganin.apply_filter(self.data)
+    @pytest.mark.skipif(
+        not (__has_pycuda__ and __has_cufft__), reason="Need pycuda and (scikit-cuda or vkfft) for this test"
+    )
+    @pytest.mark.parametrize("config", scenarios)
+    def test_gpu_paganin(self, config):
+        paganin, data, pag_kwargs = self.get_paganin_instance_and_data(config, self.data)
+
+        gpu_paganin = CudaPaganinPhaseRetrieval(data.shape, **pag_kwargs)
+        ref = paganin.apply_filter(data)
+        res = gpu_paganin.apply_filter(data)
         errmax = np.max(np.abs((res - ref) / np.max(ref)))
         assert errmax < self.rtol, "Max error is too high"