nabu 2025.1.0.dev14__py3-none-any.whl → 2025.1.0rc1__py3-none-any.whl

nabu/estimation/tests/test_motion_estimation.py

@@ -0,0 +1,471 @@
+from math import cos, sin
+import numpy as np
+from nabu.utils import search_sorted
+from nabu.estimation.motion import MotionEstimation
+from nabu.estimation.translation import estimate_shifts
+from nabu.testutils import get_data, __do_long_tests__
+from scipy.ndimage import shift
+import pytest
+
+try:
+    import astra
+except (ImportError, RuntimeError):
+    astra = None
+
+
+def test_estimate_shifts():
+    image = get_data("chelsea.npz")["data"]
+    n_tests = 10
+    shift_min = -50
+    shift_max = 50
+    max_tol = 0.2  # in pixel unit
+    shifts = np.random.rand(n_tests, 2) * (shift_max - shift_min) - shift_min
+    for s in shifts:
+        image_shifted = shift(image, s)
+        estimated_shifts = estimate_shifts(image, image_shifted)
+        abs_diff = np.abs(np.array(estimated_shifts) - np.array(s))
+        assert np.all(abs_diff < max_tol), "Wrong estimation for shifts=%s : got %s" % (s, estimated_shifts)
+
+
+def _bootstrap_test_motion_combined(cls):
+    cls._fdesc = get_data("motion/test_correct_motion_combined.npz")
+    cls.projs = cls._fdesc["projs"]
+    cls.angles = cls._fdesc["angles_rad"]
+    cls.shifts = cls._fdesc["shifts"]
+    cls.cor = -2.81  # center of rotation in pixel units, (relative to the middle of the detector)
+    if __do_long_tests__:
+        cls._fdesc = get_data("motion/test_correct_motion_horizontal.npz")
+        cls.projs_horiz = cls._fdesc["projs"]
+        cls.angles_horiz = cls._fdesc["angles_rad"]
+        cls.shifts_horiz = cls._fdesc["shifts"]
+        cls._fdesc = get_data("motion/test_correct_motion_vertical.npz")
+        cls.projs_vertic = cls._fdesc["projs"]
+        cls.angles_vertic = cls._fdesc["angles_rad"]
+        cls.shifts_vertic = cls._fdesc["shifts"]
+
+
+@pytest.fixture(scope="class")
+def bootstrap_test_motion_combined(request):
+    cls = request.cls
+    _bootstrap_test_motion_combined(cls)
+
+
+@pytest.mark.usefixtures("bootstrap_test_motion_combined")
+class TestMotionEstimation:
+
+    def _test_estimate_motion_360(
+        self,
+        projs,
+        angles,
+        reference_shifts,
+        deg=2,
+        deg_z=2,
+        tol_x=1e-5,
+        tol_y=1e-5,
+        tol_z=1e-5,
+        estimate_cor=False,
+        verbose=True,
+    ):
+        """
+        Test MotionEstimation using pairs of projs
+        """
+        n_a = projs.shape[0]
+        projs1 = projs[: n_a // 2]
+        projs2 = projs[n_a // 2 :]
+        angles1 = angles[: n_a // 2]
+        angles2 = angles[n_a // 2 :]
+        estimator = MotionEstimation(projs1, projs2, angles1, angles2, shifts_estimator="DetectorTranslationAlongBeam")
+        estimated_shifts_v = estimator.estimate_vertical_motion(degree=deg_z)
+        cor = None if estimate_cor else self.cor
+        estimated_shifts_h, cor_est = estimator.estimate_horizontal_motion(degree=deg, cor=cor)
+
+        shifts = reference_shifts
+        shifts_xy = -shifts[:, :2]
+        shifts_z = -shifts[:, 2]
+
+        # mind the signs, nabu vs RTK geometry
+        abs_diff_x = np.max(np.abs(estimated_shifts_h[:, 0] - shifts_xy[:360, 0]))
+        abs_diff_y = np.max(np.abs(estimated_shifts_h[:, 1] - shifts_xy[:360, 1]))
+        abs_diff_z = np.max(np.abs(estimated_shifts_v - shifts_z[:360]))
+
+        if verbose:
+            estimator.plot_detector_shifts(cor=cor)
+            estimator.plot_movements(cor=cor, angles_rad=angles, gt_xy=shifts_xy, gt_z=shifts_z)
+
+        assert abs_diff_x < tol_x, "Wrong x-movement estimation (estimate_cor=%s)" % estimate_cor
+        assert abs_diff_y < tol_y, "Wrong y-movement estimation (estimate_cor=%s)" % estimate_cor
+        assert abs_diff_z < tol_z, "Wrong z-movement estimation (estimate_cor=%s)" % estimate_cor
+
+    def test_estimate_motion_360(self, verbose=False):
+        """
+        Test with a synthetic dataset that underwent horizontal AND vertical translations
+        """
+        params_dict = {
+            "estimate_cor": [False, True],
+            "tol_x": [0.05, 0.1],
+            "tol_y": [0.05, 0.2],
+            "tol_z": [0.005, 0.01],
+        }
+        params_names = params_dict.keys()
+        params_values = np.array(list(params_dict.values()), dtype="O")
+        for i in range(params_values.shape[1]):
+            params = {k: v for k, v in zip(params_names, params_values[:, i].tolist())}
+            self._test_estimate_motion_360(
+                self.projs[:360],
+                self.angles[:360],
+                self.shifts[:360],
+                verbose=verbose,
+                **params,
+            )
+
+    def test_estimate_motion_360_return(self, verbose=False):
+        return_angles = np.deg2rad([360, 270, 180, 90, 0])
+        outward_angles_indices = np.array([search_sorted(self.angles[:360], ra) for ra in return_angles])
+        outward_angles = self.angles[outward_angles_indices]
+        projs1 = self.projs[outward_angles_indices]
+        projs2 = self.projs[-5:]
+
+        estimator = MotionEstimation(
+            projs1,
+            projs2,
+            outward_angles,
+            return_angles,
+            indices1=outward_angles_indices,
+            indices2=np.arange(self.projs.shape[0])[-5:],
+            n_projs_tot=self.projs.shape[0],
+        )
+        estimated_shifts_z = estimator.estimate_vertical_motion(degree=2)
+        estimated_shifts_xy, c = estimator.estimate_horizontal_motion(degree=2, cor=self.cor)
+
+        shifts = self.shifts
+        gt_xy = -shifts[:, :2][:-5]
+        gt_z = -shifts[:, 2][:-5]
+
+        if verbose:
+            estimator.plot_detector_shifts(cor=self.cor)
+            estimator.plot_movements(cor=self.cor, angles_rad=self.angles[:-5], gt_xy=gt_xy, gt_z=gt_z)
+
+        estimated_shifts_xy = estimator.apply_fit_horiz(angles=self.angles[:-5])
+        estimated_shifts_z = estimator.apply_fit_vertic(angles=self.angles[:-5])
+
+        abs_diff_x = np.max(np.abs(estimated_shifts_xy[:360, 0] - gt_xy[:360, 0]))
+        abs_diff_y = np.max(np.abs(estimated_shifts_xy[:360, 1] - gt_xy[:360, 1]))
+        abs_diff_z = np.max(np.abs(estimated_shifts_z[:360] - gt_z[:360]))
+
+        max_fit_err_vu = estimator.get_max_fit_error(cor=self.cor)
+        assert (
+            max_fit_err_vu[0] < 0.7
+        ), "Max difference between detector_v_shifts and fit(detector_v_shifts) is too high"  # can't do better for z estimation ?!
+        assert (
+            max_fit_err_vu[1] < 0.2
+        ), "Max difference between detector_u_shifts and fit(detector_u_shifts) is too high"
+
+        assert np.max(abs_diff_x) < 0.3, "Wrong x-movement estimation"
+        assert np.max(abs_diff_y) < 0.5, "Wrong y-movement estimation"
+        assert np.max(abs_diff_z) < 0.5, "Wrong z-movement estimation"
+
+
+@pytest.fixture(scope="class")
+def bootstrap_test_motion_estimation2(request):
+    cls = request.cls
+    cls.volume = get_data("motion/mri_volume_subsampled.npy")
+
+
+def _create_translations_vector(a, alpha, beta):
+    return alpha * a**2 + beta * a
+
+
+def project_volume(vol, angles, tx, ty, tz, cor=0, orig_det_dist=0):
+    """
+    Forward-project a volume with translations (tx, ty, tz) of the sample
+    """
+    n_y, n_x, n_z = vol.shape
+    vol_geom = astra.create_vol_geom(n_y, n_x, n_z)
+
+    det_row_count = n_z
+    det_col_count = max(n_y, n_x)
+
+    vectors = np.zeros((len(angles), 12))
+    for i in range(len(angles)):
+        theta = angles[i]
+        # ray direction
+        vectors[i, 0] = sin(theta)
+        vectors[i, 1] = -cos(theta)
+        vectors[i, 2] = 0
+
+        # center of detector
+        vectors[i, 3] = (
+            (cos(theta) ** 2) * tx[i] + cos(theta) * sin(theta) * ty[i] - orig_det_dist * sin(theta) + cos(theta) * cor
+        )
+        vectors[i, 4] = (
+            sin(theta) * cos(theta) * tx[i] + (sin(theta) ** 2) * ty[i] + orig_det_dist * cos(theta) + sin(theta) * cor
+        )
+        vectors[i, 5] = tz[i]
+
+        # vector from detector pixel (0,0) to (0,1)
+        vectors[i, 6] = cos(theta)  # uX
+        vectors[i, 7] = sin(theta)  # uY
+        vectors[i, 8] = 0  # uZ
+
+        # vector from detector pixel (0,0) to (1,0)
+        vectors[i, 9] = 0
+        vectors[i, 10] = 0
+        vectors[i, 11] = 1
+
+    proj_geom = astra.create_proj_geom("parallel3d_vec", det_row_count, det_col_count, vectors)
+
+    sinogram_id, sinogram = astra.create_sino3d_gpu(vol, proj_geom, vol_geom)
+    return sinogram
+
+
+def check_motion_estimation(
+    motion_estimator,
+    angles,
+    cor,
+    gt_xy,
+    gt_z,
+    fit_error_shifts_tol_vu=(1e-5, 1e-5),
+    fit_error_det_tol_vu=(1e-5, 1e-5),
+    fit_error_tol_xyz=(1e-5, 1e-5, 1e-5),
+    fit_error_det_all_angles_tol_vu=(1e-5, 1e-5),
+):
+    """
+    Assess the fit quality of a MotionEstimation object, given known ground truth sample movements.
+
+    Parameters
+    ----------
+    motion_estimator: MotionEstimation
+        MotionEstimation object
+    angles: numpy.ndarray
+        Array containing rotation angles in radians. For 180° scan it must not contain the "return projections".
+    cor: float
+        Center of rotation
+    gt_xy: numpy.ndarray
+        Ground-truth sample movements in (x, y). Shape (n_angles, 2)
+    gt_z: numpy.ndarray
+        Ground-truth sample movements in z. Shape (n_angles,)
+    fit_error_shifts_tol_vu: tuple of float
+        Maximum error tolerance when assessing the fit of radios movements,
+        i.e see how well the cross-correlation shifts between pairs of radios  are fitted
+    fit_error_det_tol_vu: tuple of float
+        Maximum error tolerance when assessing the sample movements fit on used angles, projected on the detector
+    fit_error_tol_xyz: tuple of float
+        Maximum error tolerance when assessing the fit of sample movements, in the sample domain
+    fit_error_det_all_angles_tol_vu: tuple of float
+        Maximum error tolerance when assessing the sample movements fit on ALL angles, projected on the detector
+    """
+
+    is_return = motion_estimator.is_return
+    n_angles = angles.size
+    outward_angles = motion_estimator.angles1
+    return_angles = motion_estimator.angles1
+    ma = lambda x: np.max(np.abs(x))
+    outward_angles_indices = motion_estimator.indices1
+
+    # (1) Check the fit in the detector domain, wrt measured shifts between pair of projs
+    # ------------------------------------------------------------------------------------
+    for fit_err, fit_max_error_tol, coord_name in zip(
+        motion_estimator.get_max_fit_error(cor), fit_error_shifts_tol_vu, ["v", "u"]
+    ):
+        assert (
+            fit_err < fit_max_error_tol
+        ), f"Max error between (measured detector {coord_name}-shifts) and (estimated fit) is too high: {fit_err} > {fit_max_error_tol}"
+
+    # (2) Check the fit wrt the ground truth motion projected onto detector
+    # ----------------------------------------------------------------------
+    shifts_vu = motion_estimator.apply_fit_to_get_detector_displacements(cor)
+    if is_return:
+        gt_shifts_vu_1 = motion_estimator.convert_sample_motion_to_detector_shifts(
+            gt_xy[outward_angles_indices], gt_z[outward_angles_indices], outward_angles, cor=cor
+        )
+        gt_shifts_vu_2 = motion_estimator.convert_sample_motion_to_detector_shifts(
+            gt_xy[n_angles:], gt_z[n_angles:], return_angles, cor=cor
+        )
+    else:
+        gt_shifts_vu_1 = motion_estimator.convert_sample_motion_to_detector_shifts(
+            gt_xy[: n_angles // 2], gt_z[: n_angles // 2], outward_angles, cor=cor
+        )
+        gt_shifts_vu_2 = motion_estimator.convert_sample_motion_to_detector_shifts(
+            gt_xy[n_angles // 2 :], gt_z[n_angles // 2 :], outward_angles, cor=cor
+        )
+    gt_shifts_vu = gt_shifts_vu_2 - gt_shifts_vu_1
+    if not (is_return):
+        gt_shifts_vu[:, 1] += (
+            2 * cor
+        )  # when comparing pairs of opposits projs, delta_u = u(theta+) - u(theta) + 2*cor !
+
+    err_max_vu = np.max(np.abs(gt_shifts_vu - shifts_vu), axis=0)
+    assert (
+        err_max_vu[0] < fit_error_det_tol_vu[0]
+    ), f"Max difference of fit(used_angles) on 'v' coordinate is too high: {err_max_vu[0]} > {fit_error_det_tol_vu[0]}"
+    assert (
+        err_max_vu[1] < fit_error_det_tol_vu[1]
+    ), f"Max difference of fit(used_angles) on 'u'' coordinate is too high: {err_max_vu[1]} > {fit_error_det_tol_vu[1]}"
+
+    # (3) Check the fit wrt ground truth motion of the sample
+    # This is less precise when using only return projections,
+    # because the movements were estimated from the detector shifts between a few pairs of projections
+    # Here apply the shift on all angles (above was only on the angles used to calculate the fit)
+    # --------------------------------------------------------------------------------------------
+    txy_est_all_angles = motion_estimator.apply_fit_horiz(angles=angles)
+    tz_est_all_angles = motion_estimator.apply_fit_vertic(angles=angles)
+    err_max_x = ma(txy_est_all_angles[:, 0] - gt_xy[:n_angles, 0])
+    err_max_y = ma(txy_est_all_angles[:, 1] - gt_xy[:n_angles, 1])
+    err_max_z = ma(tz_est_all_angles - gt_z[:n_angles])
+    assert (
+        err_max_x < fit_error_tol_xyz[0]
+    ), f"Max error for x coordinate is too high: {err_max_x} > {fit_error_tol_xyz[0]}"
+    assert (
+        err_max_y < fit_error_tol_xyz[1]
+    ), f"Max error for y coordinate is too high: {err_max_y} > {fit_error_tol_xyz[1]}"
+    assert (
+        err_max_z < fit_error_tol_xyz[2]
+    ), f"Max error for z coordinate is too high: {err_max_z} > {fit_error_tol_xyz[2]}"
+
+    # (4) Check the fit wrt ground truth motion, projected onto the detector, for all angles
+    # ---------------------------------------------------------------------------------------
+    estimated_shifts_vu_all_angles = motion_estimator.convert_sample_motion_to_detector_shifts(
+        txy_est_all_angles, tz_est_all_angles, angles, cor=cor
+    )
+    gt_shifts_vu_all_angles = motion_estimator.convert_sample_motion_to_detector_shifts(
+        gt_xy[:n_angles], gt_z[:n_angles], angles, cor=cor
+    )
+    err_max_vu_all_angles = np.max(np.abs(estimated_shifts_vu_all_angles - gt_shifts_vu_all_angles), axis=0)
+    assert (
+        err_max_vu_all_angles[0] < fit_error_det_all_angles_tol_vu[0]
+    ), f"Max error of fit(all_angles) on 'v' coordinate is too high: {err_max_vu_all_angles[0]} > {fit_error_det_all_angles_tol_vu[0]}"
+    assert (
+        err_max_vu_all_angles[1] < fit_error_det_all_angles_tol_vu[1]
+    ), f"Max error  of fit(all_angles) on 'u' coordinate is too high: {err_max_vu_all_angles[1]} > {fit_error_det_all_angles_tol_vu[1]}"
+
+
+@pytest.mark.skipif(not (__do_long_tests__), reason="Need NABU_LONG_TESTS=1 for this test")
+@pytest.mark.skipif(astra is None, reason="Need astra for this test")
+@pytest.mark.usefixtures("bootstrap_test_motion_estimation2")
+class TestMotionEstimationFromProjectedVolume:
+
+    def test_180_with_return_projections(self, verbose=False):
+        n_angles = 250
+        cor = -10
+
+        alpha_x = 4
+        beta_x = 3
+        alpha_y = -5
+        beta_y = 10
+        beta_z = 0
+        orig_det_dist = 0
+
+        angles0 = np.linspace(0, np.pi, n_angles, False)
+        return_angles = np.deg2rad([180.0, 135.0, 90.0, 45.0, 0.0])
+        angles = np.hstack([angles0, return_angles]).ravel()
+        a = np.arange(angles0.size + return_angles.size) / angles0.size
+
+        tx = _create_translations_vector(a, alpha_x, beta_x)
+        ty = _create_translations_vector(a, alpha_y, beta_y)
+        tz = _create_translations_vector(a, 0, beta_z)
+
+        if not (hasattr(self, "volume")):
+            self.volume = get_data("motion/mri_volume_subsampled.npy")
+        sinos = project_volume(self.volume, angles, -tx, -ty, -tz, cor=-cor, orig_det_dist=orig_det_dist)
+        radios = np.moveaxis(sinos, 1, 0)
+
+        n_return_angles = return_angles.size
+        projs_stack2 = radios[-n_return_angles:]
+        outward_angles_indices = np.array([search_sorted(angles0, ra) for ra in return_angles])
+        outward_angles = angles0[outward_angles_indices]
+        projs_stack1 = radios[:-n_return_angles][outward_angles_indices]
+
+        motion_estimator = MotionEstimation(
+            projs_stack1,
+            projs_stack2,
+            outward_angles,
+            return_angles,
+            indices1=outward_angles_indices,
+            indices2=np.arange(n_angles, n_angles + n_return_angles),
+            shifts_estimator="DetectorTranslationAlongBeam",
+        )
+        motion_estimator.estimate_horizontal_motion(degree=2, cor=cor)
+        motion_estimator.estimate_vertical_motion()
+
+        gt_h = np.stack([-tx, ty], axis=1)
+        gt_v = -tz
+        if verbose:
+            motion_estimator.plot_detector_shifts(cor=cor)
+            motion_estimator.plot_movements(cor, angles_rad=angles0, gt_xy=gt_h[:n_angles], gt_z=gt_v[:n_angles])
+
+        check_motion_estimation(
+            motion_estimator,
+            angles0,
+            cor,
+            gt_h,
+            gt_v,
+            fit_error_shifts_tol_vu=(0.01, 0.05),
+            fit_error_det_tol_vu=(0.02, 0.5),
+            fit_error_tol_xyz=(0.5, 1, 0.05),
+            fit_error_det_all_angles_tol_vu=(0.02, 1),
+        )
+
+    def test_360(self, verbose=False):
+        n_angles = 250
+        cor = -5.5
+
+        alpha_x = -2
+        beta_x = 7.0
+        alpha_y = -2
+        beta_y = 3
+        beta_z = 100
+        orig_det_dist = 0
+
+        angles = np.linspace(0, 2 * np.pi, n_angles, False)
+        a = np.linspace(0, 1, angles.size, endpoint=False)  # theta/theta_max
+
+        tx = _create_translations_vector(a, alpha_x, beta_x)
+        ty = _create_translations_vector(a, alpha_y, beta_y)
+        tz = _create_translations_vector(a, 0, beta_z)
+
+        if not (hasattr(self, "volume")):
+            self.volume = get_data("motion/mri_volume_subsampled.npy")
+
+        sinos = project_volume(self.volume, angles, -tx, -ty, -tz, cor=-cor, orig_det_dist=orig_det_dist)
+        radios = np.moveaxis(sinos, 1, 0)
+
+        projs_stack1 = radios[: n_angles // 2]
+        projs_stack2 = radios[n_angles // 2 :]
+        angles1 = angles[: n_angles // 2]
+        angles2 = angles[n_angles // 2 :]
+
+        motion_estimator = MotionEstimation(
+            projs_stack1, projs_stack2, angles1, angles2, shifts_estimator="phase_cross_correlation"
+        )
+        motion_estimator.estimate_horizontal_motion(degree=2, cor=cor)
+        motion_estimator.estimate_vertical_motion()
+
+        gt_xy = np.stack([-tx, ty], axis=1)
+        gt_z = -tz
+
+        if verbose:
+            motion_estimator.plot_detector_shifts(cor=cor)
+            motion_estimator.plot_movements(cor=cor, angles_rad=angles, gt_xy=gt_xy, gt_z=gt_z)
+
+        check_motion_estimation(
+            motion_estimator,
+            angles,
+            cor,
+            gt_xy,
+            gt_z,
+            fit_error_shifts_tol_vu=(1e-5, 0.2),
+            fit_error_det_tol_vu=(1e-5, 0.05),
+            fit_error_tol_xyz=(0.1, 0.1, 1e-5),
+            fit_error_det_all_angles_tol_vu=(1e-5, 0.1),
+        )
+
+
+if __name__ == "__main__":
+    test1 = TestMotionEstimationFromProjectedVolume()
+    test1.test_180_with_return_projections(verbose=True)
+    test1.test_360(verbose=True)
+
+    test2 = TestMotionEstimation()
+    _bootstrap_test_motion_combined(test2)
+    test2.test_estimate_motion_360(verbose=True)
+    test2.test_estimate_motion_360_return(verbose=True)

nabu/estimation/tilt.py

@@ -131,7 +131,7 @@ class CameraTilt(CenterOfRotation):
             img_1, img_2, padding_mode, axes=(-1,), high_pass=high_pass, low_pass=low_pass
         )
 
-        img_shape = img_2.shape
+        img_shape = cc.shape  # cc.shape can differ from img_2.shape, e.g. in case of odd nb of cols
         cc_h_coords = np.fft.fftfreq(img_shape[-1], 1 / img_shape[-1])
 
         (f_vals, fh) = self.extract_peak_regions_1d(cc, peak_radius=peak_fit_radius, cc_coords=cc_h_coords)

nabu/estimation/translation.py

@@ -1,5 +1,8 @@
 import numpy as np
 from numpy.polynomial.polynomial import Polynomial, polyval
+from scipy import fft
+
+from ..utils import calc_padding_lengths
 from .alignment import AlignmentBase, plt
 
 
@@ -151,7 +154,7 @@ class DetectorTranslationAlongBeam(AlignmentBase):
             for ii in range(1, num_imgs)
         ]
 
-        img_shape = img_stack.shape[-2:]
+        img_shape = ccs[0].shape  # cc.shape can differ from img.shape, e.g. in case of odd number of cols.
         cc_vs = np.fft.fftfreq(img_shape[-2], 1 / img_shape[-2])
         cc_hs = np.fft.fftfreq(img_shape[-1], 1 / img_shape[-1])
 
@@ -191,3 +194,46 @@ class DetectorTranslationAlongBeam(AlignmentBase):
             return coeffs_v[1], coeffs_h[1], shifts_vh
         else:
             return coeffs_v[1], coeffs_h[1]
+
+
+def _fft_pad(i, axes=None, padding_mode="constant"):
+    pad_len = calc_padding_lengths(i.shape, np.array(i.shape) * 2)
+    i_p = np.pad(i, pad_len, mode=padding_mode)
+    return fft.fftn(i_p)
+
+
+def estimate_shifts(im1, im2):
+    """
+    Simple implementation of shift estimation between two images, based on phase cross correlation.
+    """
+    pr = _fft_pad(im1) * _fft_pad(im2).conjugate()
+    pr_n = pr / np.maximum(1e-7, np.abs(pr))
+    corr = np.fft.fftshift(fft.ifftn(pr_n).real)
+    argmax = np.array(np.unravel_index(np.argmax(corr), pr.shape))
+    shp = np.array(pr.shape)
+    argmax_refined = refine_parabola_2D(corr, argmax)
+    argmax = argmax + argmax_refined
+    return shp // 2 - np.array(argmax)
+
+
+def refine_parabola_2D(im_vals, argmax):
+    argmax = tuple(argmax)
+    maxval = im_vals[argmax]
+    ny, nx = im_vals.shape
+
+    iy, ix = np.array(argmax, dtype=np.intp)
+    ixm, ixp = (ix - 1) % nx, (ix + 1) % nx
+    iym, iyp = (iy - 1) % ny, (iy + 1) % ny
+
+    F = maxval
+    A = (im_vals[iy, ixp] + im_vals[iy, ixm]) / 2 - F
+    D = (im_vals[iy, ixp] - im_vals[iy, ixm]) / 2
+    B = (im_vals[iyp, ix] + im_vals[iym, ix]) / 2 - F
+    E = (im_vals[iyp, ix] - im_vals[iym, ix]) / 2
+    C = (im_vals[iyp, ixp] - im_vals[iym, ixp] - im_vals[iyp, ixm] + im_vals[iym, ixm]) / 4
+    det = C**2 - 4 * A * B
+    dx = (2 * B * D - C * E) / det
+    dy = (2 * A * E - C * D) / det
+    dx = np.clip(dx, -0.5, 0.5)
+    dy = np.clip(dy, -0.5, 0.5)
+    return (dy, dx)

nabu/io/cast_volume.py

@@ -1,10 +1,12 @@
 import os
-
-from tomoscan.esrf.volume.singleframebase import VolumeSingleFrameBase
-from nabu.misc.utils import rescale_data
-from nabu.pipeline.params import files_formats
+import logging
+from glob import glob
+from shutil import rmtree
+import numpy
+from silx.io.utils import get_data
+from silx.io.url import DataUrl
 from tomoscan.volumebase import VolumeBase
-from tomoscan.scanbase import TomoScanBase
+from tomoscan.esrf.volume.singleframebase import VolumeSingleFrameBase
 from tomoscan.esrf.volume import (
     EDFVolume,
     HDF5Volume,
@@ -13,11 +15,10 @@ from tomoscan.esrf.volume import (
     TIFFVolume,
 )
 from tomoscan.io import HDF5File
-from silx.io.utils import get_data
-import numpy
-from silx.io.url import DataUrl
-from typing import Optional
-import logging
+from ..utils import first_generator_item
+from ..misc.utils import rescale_data
+from ..pipeline.params import files_formats
+from .reader import get_hdf5_file_all_virtual_sources, list_hdf5_entries
 
 _logger = logging.getLogger(__name__)
 
@@ -133,7 +134,7 @@ def cast_volume(
     output_data_type: numpy.dtype,
     data_min=None,
     data_max=None,
-    scan: Optional[TomoScanBase] = None,
+    scan=None,
     rescale_min_percentile=RESCALE_MIN_PERCENTILE,
     rescale_max_percentile=RESCALE_MAX_PERCENTILE,
     save=True,
@@ -283,7 +284,7 @@ def clamp_and_rescale_data(
     return rescaled_data
 
 
-def find_histogram(volume: VolumeBase, scan: Optional[TomoScanBase] = None) -> Optional[DataUrl]:
+def find_histogram(volume: VolumeBase, scan=None):
     """
     Look for histogram of the provided url. If found one return the DataUrl of the nabu histogram
     """
@@ -330,7 +331,7 @@ def find_histogram(volume: VolumeBase, scan: Optional[TomoScanBase] = None) -> O
             data_path = getattr(scan, "entry/histogram/results/data", "entry/histogram/results/data")
         else:
 
-            def get_file_entries(file_path: str) -> Optional[tuple]:
+            def get_file_entries(file_path: str):
                 if os.path.exists(file_path):
                     with HDF5File(file_path, mode="r") as h5s:
                         return tuple(h5s.keys())
@@ -408,3 +409,83 @@ def _min_max_from_histo(url: DataUrl, rescale_min_percentile: int, rescale_max_p
         return _get_hst_saturations(
             hist, bins, numpy.float32(rescale_min_percentile), numpy.float32(rescale_max_percentile)
         )
+
+
+def _remove_volume_singleframe(volume, check=True):
+    volume_directory = volume.data_url.file_path()
+    if check:
+        volume_files = set(volume.browse_data_files())
+        files_names_pattern = os.path.join(volume_directory, "*." + volume.data_extension)
+        files_on_disk = set(glob(files_names_pattern))
+        # Don't check strict equality here, as some files on disk might be already removed.
+        # i.e, there should be no more files on disk than expected files in the volume
+        if not (files_on_disk.issubset(volume_files)):
+            raise RuntimeError(f"Unexpected files present in {volume_directory}: {files_on_disk - volume_files}")
+        # TODO also check for metadata file(s) ?
+    rmtree(volume_directory)
+
+
+def _remove_volume_multiframe(volume, check=True):
+    file_path = volume.data_url.file_path()
+    if check:
+        if not (os.path.isfile(file_path)):
+            raise RuntimeError(f"Expected a file: {file_path}")
+    os.remove(file_path)
+
+
+def _remove_volume_hdf5(volume, check=True):
+    file_path = volume.data_url.file_path()
+    entry = volume.data_url.data_path().lstrip("/").split("/")[0]
+
+    # Nabu HDF5 reconstructions have a folder alongside the HDF5 file, with the same prefix
+    # For example the HDF5 file "/path/to/rec.hdf5" has an associated directory "/path/to/rec"
+    associated_dir, _ = os.path.splitext(os.path.basename(file_path))
+    associated_dir_abs = os.path.join(os.path.dirname(file_path), associated_dir)
+
+    with HDF5File(file_path, "r") as f:
+        fdesc = f[entry]
+        virtual_sources = get_hdf5_file_all_virtual_sources(fdesc, return_only_filenames=True)
+
+    # TODO check if this is legitimate. Nabu reconstruction will only do one VS (for entry/reconstruction/results/data).
+    # Bliss/Lima do have multiple VS (flats/darks/projs), but we generally don't want to remove raw data ?
+    if len(virtual_sources) > 1:
+        raise ValueError("Found more than one virtual source - this looks weird. Interrupting.")
+    #
+    if len(virtual_sources) > 0:
+        h5path, virtual_source_files_paths = first_generator_item(virtual_sources[0].items())
+        if len(virtual_source_files_paths) == 1:
+            target_dir = os.path.dirname(virtual_source_files_paths[0])
+        else:
+            target_dir = os.path.commonpath(virtual_source_files_paths)
+        target_dir_abs = os.path.join(os.path.dirname(file_path), target_dir)
+        if check and (target_dir_abs != associated_dir_abs):
+            raise ValueError(
+                f"The virtual sources in {file_path}:{h5path} reference the directory {target_dir}, but expected was {associated_dir}"
+            )
+        if os.path.isdir(target_dir_abs):
+            rmtree(associated_dir_abs)
+    os.remove(file_path)
+
+
+def remove_volume(volume, check=True):
+    """
+    Remove files belonging to a volume, claim disk space.
+
+    Parameters
+    ----------
+    volume: tomoscan.esrf.volume
+        Volume object
+    check: bool, optional
+        Whether to check if the files that would be removed do not have extra other files ; interrupt the operation if so.
+
+    """
+    if isinstance(volume, (EDFVolume, JP2KVolume, TIFFVolume)):
+        _remove_volume_singleframe(volume, check=check)
+    elif isinstance(volume, MultiTIFFVolume):
+        _remove_volume_multiframe(volume, check=check)
+    elif isinstance(volume, HDF5Volume):
+        if len(list_hdf5_entries(volume.file_path)) > 1:
+            raise NotImplementedError("Removing a HDF5 volume with more than one entry is not supported")
+        _remove_volume_hdf5(volume, check=check)
+    else:
+        raise TypeError("Unknown type of volume")