Merge pull request #224 from ISISComputingGroup/MomentumFit

Momentum fit

Author: rerpha

doc/_static/images_fits/muons_momentum.png

@@ -183,3 +183,22 @@ g(x) = \max(f(x), \text{background} - \text{height})
 ```{math}
 y = \min(g(x), \text{background})
 ```
+
+## Muon Momentum
+
+API Reference: [`MuonMomentum`](ibex_bluesky_core.fitting.MuonMomentum)
+
+Fits data from a momentum scan, it is designed for the specific use case of scanning over magnet current on muon instruments.
+
+- `x0` - The center (x) of the model
+- `w` - The horizontal stretch factor of the model
+- `R` - The amplitude of the model
+- `b` - The minimum value (y) of the model
+- `p` - Changes the gradient of the tail ends of the model
+
+``` {math}
+y = \left (\text{erfc} \mathopen{} \left(\frac{x-x_0}{w} \right) \mathclose{} \cdot \frac{R}{2} + b \right) \cdot \left (\frac{x}{x_0} \right)^p
+```
+
+![MuonMomentumModel](../_static/images_fits/muons_momentum.png)
+ 

doc/fitting/standard_fits.md

@@ -20,6 +20,7 @@
     "Gaussian",
     "Linear",
     "Lorentzian",
+    "MuonMomentum",
     "NegativeTrapezoid",
     "Polynomial",
     "SlitScan",
@@ -687,3 +688,65 @@ def guess(
             return init_guess
 
         return guess
+
+
+class MuonMomentum(Fit):
+    """Muon momentum fitting."""
+
+    equation = """
+        y=(erfc((x-x0/w))*(r/2)+b)*((x/x0)**p)"""
+
+    @classmethod
+    def model(cls, *args: int) -> lmfit.Model:
+        """Momentum scan model."""
+
+        def model(
+            x: npt.NDArray[np.float64], x0: float, r: float, w: float, p: float, b: float
+        ) -> npt.NDArray[np.float64]:
+            return (scipy.special.erfc((x - x0) / w) * (r / 2) + b) * ((x / x0) ** p)
+
+        return lmfit.Model(model, name=f"{cls.__name__}  [{cls.equation}]")
+
+    @classmethod
+    def guess(
+        cls, *args: int
+    ) -> Callable[[npt.NDArray[np.float64], npt.NDArray[np.float64]], dict[str, lmfit.Parameter]]:
+        """Momentum Scan Fit Guessing."""
+
+        def guess(
+            x: npt.NDArray[np.float64], y: npt.NDArray[np.float64]
+        ) -> dict[str, lmfit.Parameter]:
+            index_array = np.argsort(x)
+            x = x[index_array]
+            y = y[index_array]
+
+            index_min_y = np.argmin(y)
+            index_max_y = np.argmax(y)
+
+            b = np.min(y)
+            r = np.max(y) - b
+
+            x_slope = x[
+                index_max_y:index_min_y
+            ]  # Gets all x values between the maximum and minimum y
+
+            if len(x_slope) != 0:
+                x0 = np.mean(x_slope)
+            else:
+                x0 = x[-1]  # Picked as it can't be 0
+
+            p = 1  # Expected value, not likely to change
+
+            const = 3  # largest difference in erfc function is between -1.5 and 1.5
+            w = np.abs(x[index_max_y] - x[index_min_y]) / const
+
+            init_guess = {
+                "b": lmfit.Parameter("b", b),
+                "r": lmfit.Parameter("r", r, min=0),
+                "x0": lmfit.Parameter("x0", x0, min=0),
+                "p": lmfit.Parameter("p", p, min=0),
+                "w": lmfit.Parameter("w", w, min=0),
+            }
+            return init_guess
+
+        return guess

src/ibex_bluesky_core/fitting.py

@@ -1,3 +1,4 @@
+import typing
 import warnings
 from collections.abc import Callable
 from unittest import mock
@@ -17,6 +18,7 @@
     Gaussian,
     Linear,
     Lorentzian,
+    MuonMomentum,
     NegativeTrapezoid,
     Polynomial,
     SlitScan,
@@ -823,3 +825,132 @@ def test_guess_given_flat_data(self):
             outp = NegativeTrapezoid.guess()(x, y)
             # check that with flat data gradient guess is 0
             assert outp["gradient"] == pytest.approx(0.0, rel=1e-2)
+
+
+class TestMuonMomentum:
+    values: typing.ClassVar[dict[str, float]] = {
+        "x0": 21.89,
+        "w": 0.301,
+        "r": 110,
+        "b": 3.37,
+        "p": 1.82,
+    }
+    x: typing.ClassVar[npt.NDArray[np.float64]] = np.linspace(20.5, 23, 15, dtype=np.float64)
+
+    out: typing.ClassVar[npt.NDArray[np.float64]] = MuonMomentum.model().func(
+        x=x,
+        x0=values["x0"],
+        w=values["w"],
+        r=values["r"],
+        b=values["b"],
+        p=values["p"],
+    )
+
+    class TestMuonMomentumModel:
+        def test_muon_momentum_model_w(self):
+            out2 = MuonMomentum.model().func(
+                x=TestMuonMomentum.x,
+                x0=TestMuonMomentum.values["x0"],
+                w=TestMuonMomentum.values["w"] + 1,
+                r=TestMuonMomentum.values["r"],
+                b=TestMuonMomentum.values["b"],
+                p=TestMuonMomentum.values["p"],
+            )
+
+            assert np.mean(TestMuonMomentum.out) < np.mean(out2)
+
+        def test_muon_momentum_model_b(self):
+            out2 = MuonMomentum.model().func(
+                x=TestMuonMomentum.x,
+                x0=TestMuonMomentum.values["x0"],
+                w=TestMuonMomentum.values["w"],
+                r=TestMuonMomentum.values["r"],
+                b=TestMuonMomentum.values["b"] + 1,
+                p=TestMuonMomentum.values["p"],
+            )
+
+            assert np.min(TestMuonMomentum.out) < np.min(out2)
+
+        def test_muon_momentum_model_r(self):
+            out2 = MuonMomentum.model().func(
+                x=TestMuonMomentum.x,
+                x0=TestMuonMomentum.values["x0"],
+                w=TestMuonMomentum.values["w"],
+                r=TestMuonMomentum.values["r"] + 10,
+                b=TestMuonMomentum.values["b"],
+                p=TestMuonMomentum.values["p"],
+            )
+
+            assert np.max(TestMuonMomentum.out) - np.min(TestMuonMomentum.out) < np.max(
+                out2
+            ) - np.min(out2)
+
+        def test_muon_momentum_model_x0(self):
+            out2 = MuonMomentum.model().func(
+                x=TestMuonMomentum.x,
+                x0=TestMuonMomentum.values["x0"] + 1,
+                w=TestMuonMomentum.values["w"],
+                r=TestMuonMomentum.values["r"],
+                b=TestMuonMomentum.values["b"],
+                p=TestMuonMomentum.values["p"],
+            )
+
+            assert (
+                TestMuonMomentum.x[np.argmax(TestMuonMomentum.out)]
+                < TestMuonMomentum.x[np.argmax(out2)]
+            )
+            assert (
+                TestMuonMomentum.x[np.argmin(TestMuonMomentum.out)]
+                < TestMuonMomentum.x[np.argmin(out2)]
+            )
+
+    class TestMuonMomentumGuess:
+        def test_muon_momentum_guess_b(self):
+            x = np.array([0, 1, 2], dtype=np.float64)
+            y = np.array([20, 20, 10], dtype=np.float64)
+
+            out = MuonMomentum.guess()(x, y)
+
+            assert out["b"].value is not None
+            assert np.min(y) == out["b"].value
+
+        def test_muon_momentum_guess_r(self):
+            x = np.array([0, 1, 2], dtype=np.float64)
+            y = np.array([20, 20, 10], dtype=np.float64)
+
+            out = MuonMomentum.guess()(x, y)
+
+            assert out["r"].value is not None
+            assert np.max(y) - np.min(y) == out["r"].value
+
+        def test_muon_momentum_guess_w(self):
+            x = np.array([0, 1, 2], dtype=np.float64)
+            y = np.array([20, 20, 10], dtype=np.float64)
+            x1 = np.array([0, 1, 1.5], dtype=np.float64)
+
+            out = MuonMomentum.guess()(x, y)
+            out1 = MuonMomentum.guess()(x1, y)
+
+            assert out["w"].value is not None
+            assert out1["w"].value is not None
+            assert out["w"].value > out1["w"].value
+
+        def test_muon_momentum_guess_x0(self):
+            x = np.array([0, 1, 2], dtype=np.float64)
+            y = np.array([20, 20, 10], dtype=np.float64)
+            x1 = np.array([3, 4, 5], dtype=np.float64)
+
+            out = MuonMomentum.guess()(x, y)
+            out1 = MuonMomentum.guess()(x1, y)
+
+            assert out["x0"].value is not None
+            assert out1["x0"].value is not None
+            assert out1["x0"].value > out["x0"].value
+
+        def test_muon_momentum_guess_x0_noslope(self):
+            x = np.array([0, 1, 2], dtype=np.float64)
+            y = np.array([20, 20, 20], dtype=np.float64)
+
+            out = MuonMomentum.guess()(x, y)
+
+            assert out["x0"].value == x[-1]