Alignment: allow fitting of multiple 1D peaks in one call

Signed-off-by: Nicola VIGANÒ <nicola.vigano@esrf.fr>

nabu/preproc/alignment.py

@@ -110,41 +110,53 @@ class AlignmentBase(object):
         float
             Estimated function max, according to the coordinates in fx.
         """
-        if not len(f_vals.shape) == 1:
+        if not len(f_vals.shape) in (1, 2):
             raise ValueError(
-                "The fitted values should form a 1-dimensional array. Array of shape: [%s] was given."
+                "The fitted values should be either one or a collection of 1-dimensional arrays. Array of shape: [%s] was given."
                 % (" ".join(("%d" % s for s in f_vals.shape)))
             )
+        num_vals = f_vals.shape[0]
+
         if fx is None:
-            fx_half_size = (f_vals.size - 1) / 2
-            fx = np.linspace(-fx_half_size, fx_half_size, f_vals.size)
-        elif not (len(fx.shape) == 1 and np.all(fx.size == f_vals.size)):
+            fx_half_size = (num_vals - 1) / 2
+            fx = np.linspace(-fx_half_size, fx_half_size, num_vals)
+        elif not (len(fx.shape) == 1 and np.all(fx.size == num_vals)):
             raise ValueError(
                 "Base coordinates should have the same length as values array. Sizes of fx: %d, f_vals: %d"
-                % (fx.size, f_vals.size)
+                % (fx.size, num_vals)
             )
 
-        # using Polynomial.fit, because supposed to be more numerically stable
-        # than previous solutions (according to numpy).
-        poly = Polynomial.fit(fx, f_vals, deg=2)
-        coeffs = poly.convert().coef
+        if len(f_vals.shape) == 1:
+            # using Polynomial.fit, because supposed to be more numerically
+            # stable than previous solutions (according to numpy).
+            poly = Polynomial.fit(fx, f_vals, deg=2)
+            coeffs = poly.convert().coef
+        else:
+            coords = np.array([np.ones(num_vals), fx, fx ** 2])
+            coeffs = np.linalg.lstsq(coords.T, f_vals, rcond=None)[0]
 
         # For a 1D parabola `f(x) = c + bx + ax^2`, the vertex position is:
         # x_v = -b / 2a.
-        vertex_x = -coeffs[1] / (2 * coeffs[2])
+        vertex_x = -coeffs[1, :] / (2 * coeffs[2, :])
 
         vertex_min_x = np.min(fx)
         vertex_max_x = np.max(fx)
-        if not (vertex_min_x < vertex_x < vertex_max_x):
-            raise ValueError(
-                "Fitted x: {} position is outide the margins of input: x: [{}, {}]".format(
+        lower_bound_ok = vertex_min_x < vertex_x
+        upper_bound_ok = vertex_x < vertex_max_x
+        if not np.all(lower_bound_ok * upper_bound_ok):
+            if len(f_vals.shape) == 1:
+                message = "Fitted position {} is outide the input margins [{}, {}]".format(
                     vertex_x, vertex_min_x, vertex_max_x
                 )
-            )
+            else:
+                message = "Fitted positions outide the input margins [{}, {}]: %d below and %d above".format(
+                    vertex_min_x, vertex_max_x, np.sum(1 - lower_bound_ok), np.sum(1 - upper_bound_ok)
+                )
+            raise ValueError(message)
         return vertex_x
 
     @staticmethod
-    def extract_peak_region(cc, peak_radius=1, cc_vs=None, cc_hs=None):
+    def extract_peak_region_2d(cc, peak_radius=1, cc_vs=None, cc_hs=None):
         """
         Extracts a region around the maximum value.
 
@@ -440,7 +452,7 @@ class CenterOfRotation(AlignmentBase):
         cc_vs = np.fft.fftfreq(img_shape[-2], 1 / img_shape[-2])
         cc_hs = np.fft.fftfreq(img_shape[-1], 1 / img_shape[-1])
 
-        (f_vals, fv, fh) = self.extract_peak_region(cc, peak_radius=peak_fit_radius, cc_vs=cc_vs, cc_hs=cc_hs)
+        (f_vals, fv, fh) = self.extract_peak_region_2d(cc, peak_radius=peak_fit_radius, cc_vs=cc_vs, cc_hs=cc_hs)
         fitted_shifts_vh = self.refine_max_position_2d(f_vals, fv, fh)
 
         return fitted_shifts_vh[-1] / 2.0
@@ -595,7 +607,7 @@ class DetectorTranslationAlongBeam(AlignmentBase):
 
         shifts_vh = np.empty((num_imgs - 1, 2))
         for ii, cc in enumerate(ccs):
-            (f_vals, fv, fh) = self.extract_peak_region(cc, peak_radius=peak_fit_radius, cc_vs=cc_vs, cc_hs=cc_hs)
+            (f_vals, fv, fh) = self.extract_peak_region_2d(cc, peak_radius=peak_fit_radius, cc_vs=cc_vs, cc_hs=cc_hs)
             shifts_vh[ii, :] = self.refine_max_position_2d(f_vals, fv, fh)
 
         if equispaced_increments: