Models/shifts: added utility function to combine vertical and horizon…

…tal shifts

Signed-off-by: Nicola VIGANO <[email protected]>

corrct/models.py

@@ -519,6 +519,37 @@ def get_default_from_volume(volume: NDArray) -> "VolumeGeometry":
         return get_vol_geom_from_volume(volume=volume)
 
 
+def combine_shifts_vu(shifts_v: NDArray, shifts_u: NDArray, cor: Union[float, NDArray, None] = None) -> NDArray:
+    """Combine vertical and horizontal shifts.
+
+    Parameters
+    ----------
+    shifts_v : NDArray
+        The vertical shifts
+    shifts_u : NDArray
+        The horizontal shifts
+    cor : Union[float, NDArray, None], optional
+        The center-of-rotation, if applicable, by default None
+
+    Returns
+    -------
+    NDArray
+        The combined shifts
+    """
+    if np.sum(np.array(shifts_v.shape) > 1) > 1:
+        raise ValueError(f"Expected 1-dimensional array for vertical shifts, but an array {shifts_v.shape = } was passed")
+    if np.sum(np.array(shifts_u.shape) > 1) > 1:
+        raise ValueError(f"Expected 1-dimensional array for horizontal shifts, but an array {shifts_u.shape = } was passed")
+    if shifts_v.size != shifts_u.size:
+        raise ValueError(f"Number of vertical shifts ({shifts_v.size}) and horizontal shifts ({shifts_u.size}) should match")
+    if cor is not None:
+        cor = np.array(cor)
+        if cor.ndim > 0 and np.any(np.array(shifts_u.shape) > 1):
+            raise ValueError(f"Expected 0-dimensional array for center-of-rotation, but an array {cor.shape = } was passed")
+
+    return np.stack([shifts_v, shifts_u + cor if cor is not None else 0], axis=-1)
+
+
 def get_rot_axis_dir(rot_axis_dir: Union[str, ArrayLike, NDArray] = "clockwise") -> NDArray:
     """Process the requested rotation axis direction and return a meaningful value.