Refactor & add tests

Author: Tom-Willemsen

src/ibex_bluesky_core/fitting.py

@@ -514,6 +514,19 @@ def _center_of_mass_and_mass(
     return com, total_area
 
 
+def _guess_cen_and_width(
+    x: npt.NDArray[np.float64], y: npt.NDArray[np.float64]
+) -> tuple[float, float]:
+    """Guess the center and width of a positive peak."""
+    com, total_area = _center_of_mass_and_mass(x, y)
+    y_range = np.max(y) - np.min(y)
+    if y_range == 0.0:
+        width = (np.max(x) - np.min(x)) / 2
+    else:
+        width = total_area / y_range
+    return com, width
+
+
 class TopHat(Fit):
     """Top Hat Fitting."""
 
@@ -541,27 +554,31 @@ def guess(
         def guess(
             x: npt.NDArray[np.float64], y: npt.NDArray[np.float64]
         ) -> dict[str, lmfit.Parameter]:
-            com, total_area = _center_of_mass_and_mass(x, y)
-            width = (np.max(y) - np.min(y)) / total_area
+            cen, width = _guess_cen_and_width(x, y)
 
             init_guess = {
-                "cen": lmfit.Parameter("cen", com, min=np.min(x), max=np.max(x)),
-                "width": lmfit.Parameter("width", width, min=0, max=np.max(x) - np.min(x)),
+                "cen": lmfit.Parameter("cen", cen),
+                "width": lmfit.Parameter("width", width, min=0),
                 "height": lmfit.Parameter(
                     "height",
-                    (max(y) - min(y)),
-                    min=-(np.max(y) - np.min(y)),
-                    max=np.max(y) - np.min(y),
+                    np.max(y) - np.min(y),
                 ),
-                "background": lmfit.Parameter("background", min(y), min=np.min(y), max=np.max(y)),
-                "method": "basinhopping",
+                "background": lmfit.Parameter("background", np.min(y)),
             }
 
             return init_guess
 
         return guess
 
 
+def _guess_trapezoid_gradient(x: npt.NDArray[np.float64], y: npt.NDArray[np.float64]) -> float:
+    gradients = np.zeros_like(x[1:], dtype=np.float64)
+    x_diffs = x[:-1] - x[1:]
+    y_diffs = y[:-1] - y[1:]
+    np.divide(y_diffs, x_diffs, out=gradients, where=x_diffs != 0)
+    return np.max(np.abs(gradients))
+
+
 class Trapezoid(Fit):
     """Trapezoid Fitting."""
 
@@ -597,26 +614,15 @@ def guess(
         def guess(
             x: npt.NDArray[np.float64], y: npt.NDArray[np.float64]
         ) -> dict[str, lmfit.Parameter]:
-            com, total_mass = _center_of_mass_and_mass(x, y)
-            y_range = np.max(y) - np.min(y)
-            if y_range != 0:
-                width_guess = total_mass / y_range
-            else:
-                width_guess = (np.max(x) - np.min(x)) / 2
-
-            # Guess that the gradient is the maximum gradient between any two points
-            gradients = np.zeros_like(x[1:], dtype=np.float64)
-            x_diffs = x[:-1] - x[1:]
-            y_diffs = y[:-1] - y[1:]
-            np.divide(y_diffs, x_diffs, out=gradients, where=x_diffs != 0)
-            gradient_guess = np.max(np.abs(gradients))
+            cen, width = _guess_cen_and_width(x, y)
+            gradient_guess = _guess_trapezoid_gradient(x, y)
 
             height = np.max(y) - np.min(y)
             background = np.min(y)
-            y_offset = gradient_guess * width_guess / 2.0
+            y_offset = gradient_guess * width / 2.0
 
             init_guess = {
-                "cen": lmfit.Parameter("cen", com, min=np.min(x), max=np.max(x)),
+                "cen": lmfit.Parameter("cen", cen, min=np.min(x), max=np.max(x)),
                 "gradient": lmfit.Parameter("gradient", gradient_guess, min=0),
                 "height": lmfit.Parameter("height", height, min=0),
                 "background": lmfit.Parameter("background", background),
@@ -663,22 +669,15 @@ def guess(
         def guess(
             x: npt.NDArray[np.float64], y: npt.NDArray[np.float64]
         ) -> dict[str, lmfit.Parameter]:
-            com, total_mass = _center_of_mass_and_mass(x, -y)
-            width_guess = total_mass / (np.max(y) - np.min(y))
-
-            # Guess that the gradient is the maximum gradient between any two points
-            gradients = np.zeros_like(x[1:], dtype=np.float64)
-            x_diffs = x[:-1] - x[1:]
-            y_diffs = y[:-1] - y[1:]
-            np.divide(y_diffs, x_diffs, out=gradients, where=x_diffs != 0)
-            gradient_guess = np.max(np.abs(gradients))
+            cen, width = _guess_cen_and_width(x, -y)
+            gradient_guess = _guess_trapezoid_gradient(x, y)
 
             height = np.max(y) - np.min(y)
             background = np.max(y)
-            y_offset = -gradient_guess * width_guess / 2.0
+            y_offset = -gradient_guess * width / 2.0
 
             init_guess = {
-                "cen": lmfit.Parameter("cen", com, min=np.min(x), max=np.max(x)),
+                "cen": lmfit.Parameter("cen", cen, min=np.min(x), max=np.max(x)),
                 "gradient": lmfit.Parameter("gradient", gradient_guess, min=0),
                 "height": lmfit.Parameter("height", height, min=0),
                 "background": lmfit.Parameter("background", background),

tests/callbacks/fitting/test_fitting_methods.py

@@ -17,6 +17,7 @@
     Gaussian,
     Linear,
     Lorentzian,
+    NegativeTrapezoid,
     Polynomial,
     SlitScan,
     TopHat,
@@ -599,7 +600,7 @@ def test_width_guess(self):
 
             outp = TopHat.guess()(x, y)
 
-            assert outp["width"] == pytest.approx(1.0, rel=1e-2)
+            assert outp["width"] == pytest.approx(2.0, rel=1e-2)
 
         def test_guess_given_flat_data(self):
             x = np.arange(-5.0, 5.0, 1.0, dtype=np.float64)
@@ -688,15 +689,15 @@ def test_gradient_guess(self):
 
             outp = Trapezoid.guess()(x, y)
 
-            assert outp["gradient"] == pytest.approx(10.0, rel=1e-2)
+            assert outp["gradient"] == pytest.approx(8.0, rel=1e-2)
 
         def test_y_offset_guess(self):
-            x = np.arange(-5.0, 5.0, 1.0, dtype=np.float64)
+            x = np.linspace(-5.0, 5.0, num=11, dtype=np.float64)
             y = np.array([1.0, 1.0, 1.0, 2.0, 3.0, 3.0, 3.0, 2.0, 1.0, 1.0, 1.0], dtype=np.float64)
 
             outp = Trapezoid.guess()(x, y)
 
-            x1 = np.arange(-6.0, 6.0, 1.0, dtype=np.float64)
+            x1 = np.linspace(-6.0, 6.0, num=13, dtype=np.float64)
             y1 = np.array(
                 [1.0, 1.0, 1.0, 2.0, 3.0, 3.0, 3.0, 3.0, 3.0, 2.0, 1.0, 1.0, 1.0], dtype=np.float64
             )
@@ -713,3 +714,112 @@ def test_guess_given_flat_data(self):
             outp = Trapezoid.guess()(x, y)
             # check that with flat data gradient guess is 0
             assert outp["gradient"] == pytest.approx(0.0, rel=1e-2)
+
+
+class TestNegativeTrapezoid:
+    class TestNegativeTrapezoidModel:
+        def test_negative_trapezoid_model(self):
+            x = np.arange(-5.0, 6.0, 1.0, dtype=np.float64)
+            cen = 0
+            y_offset = -1
+            height = 1
+            background = 1
+            gradient = 1
+
+            outp = NegativeTrapezoid.model().func(
+                x,
+                cen=cen,
+                y_offset=y_offset,
+                height=height,
+                background=background,
+                gradient=gradient,
+            )
+
+            assert background == pytest.approx(np.max(outp), rel=1e-2)
+            assert background - height == pytest.approx(np.min(outp), rel=1e-2)
+
+            outp1 = NegativeTrapezoid.model().func(
+                x,
+                cen=cen + 3,
+                y_offset=y_offset,
+                height=height,
+                background=background,
+                gradient=gradient - 0.5,
+            )
+
+            # check centre moves when data is shifted
+            assert np.mean(x[np.where(outp < background)]) < np.mean(
+                x[np.where(outp1 < background)]
+            )
+
+            # check gradient: a greater gradient means smaller average y values as wider
+            assert np.mean(outp) > np.mean(outp1)
+
+            outp2 = NegativeTrapezoid.model().func(
+                x,
+                cen=cen,
+                y_offset=y_offset - 5,
+                height=height,
+                background=background,
+                gradient=gradient,
+            )
+
+            # check y_offset: a smaller y_offset means smaller average y values as wider
+            assert np.mean(outp) > np.mean(outp2)
+
+    class TestNegativeTrapezoidGuess:
+        def test_background_guess(self):
+            x = np.array([-1.0, 0.0, 1.0, 2.0, 3.0], dtype=np.float64)
+            y = np.array([-1.0, -2.0, -2.0, -2.0, -1.0], dtype=np.float64)
+
+            outp = NegativeTrapezoid.guess()(x, y)
+
+            assert outp["background"] == pytest.approx(-1.0, rel=1e-2)
+
+        def test_cen_height_guess(self):
+            x = np.array([-1.0, 0.0, 1.0, 2.0, 3.0], dtype=np.float64)
+            y = np.array([-1.0, -1.0, -2.0, -1.0, -1.0], dtype=np.float64)
+
+            outp = NegativeTrapezoid.guess()(x, y)
+
+            assert outp["cen"] == pytest.approx(1.0, rel=1e-2)
+            assert outp["height"] == pytest.approx(1.0, rel=1e-2)
+
+        def test_gradient_guess(self):
+            x = np.array([-4.0, -3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0, 4.0], dtype=np.float64)
+            y = np.array([1.0, 2.0, 4.0, 8.0, 16.0, 8.0, 4.0, 2.0, 1.0], dtype=np.float64)
+
+            # Should choose x = -1.0 as x1 and x = -2.5 as x0
+            # height = 16 - 1 = 15
+            # gradient = 15 / (-1 - -2.5) = 10
+
+            outp = NegativeTrapezoid.guess()(x, y)
+
+            assert outp["gradient"] == pytest.approx(8.0, rel=1e-2)
+
+        def test_y_offset_guess(self):
+            x = np.linspace(-5.0, 5.0, num=11, dtype=np.float64)
+            y = np.array(
+                [-1.0, -1.0, -1.0, -2.0, -3.0, -3.0, -3.0, -2.0, -1.0, -1.0, -1.0], dtype=np.float64
+            )
+
+            outp = NegativeTrapezoid.guess()(x, y)
+
+            x1 = np.linspace(-6.0, 6.0, num=13, dtype=np.float64)
+            y1 = np.array(
+                [-1.0, -1.0, -1.0, -2.0, -3.0, -3.0, -3.0, -3.0, -3.0, -2.0, -1.0, -1.0, -1.0],
+                dtype=np.float64,
+            )
+
+            outp1 = NegativeTrapezoid.guess()(x1, y1)
+
+            # Assert that with a greater top width, y_offset decreases
+            assert outp["y_offset"] > outp1["y_offset"]
+
+        def test_guess_given_flat_data(self):
+            x = np.arange(-5.0, 5.0, 1.0, dtype=np.float64)
+            y = np.zeros_like(x, dtype=np.float64) + 1
+
+            outp = NegativeTrapezoid.guess()(x, y)
+            # check that with flat data gradient guess is 0
+            assert outp["gradient"] == pytest.approx(0.0, rel=1e-2)