feat(analysis.fit): Add background option to MultiPeakModel and `Mu…

…ltiPeakFunction`

docs/source/user-guide/curve-fitting.ipynb

@@ -31,8 +31,7 @@
                 "fit multidimensional xarray objects. Finally, we will show how to use `iminuit\n",
                 "<https://github.com/scikit-hep/iminuit>`_ with lmfit models.\n",
                 "\n",
-                "Basic fitting with ``lmfit``\n",
-                "----------------------------"
+                "If you are familiar with lmfit, you can skip the first section."
             ]
         },
         {
@@ -325,11 +324,17 @@
                 "Fitting multiple peaks\n",
                 "~~~~~~~~~~~~~~~~~~~~~~\n",
                 "\n",
-                "One example is :class:`MultiPeakModel <erlab.analysis.fit.models.MultiPeakModel>`, which is\n",
-                "a composite model of multiple Gaussian or Lorentzian peaks on a linear background. By\n",
-                "supplying keyword arguments, you can specify the number of peaks, their shapes, whether\n",
-                "to multiply with a Fermi-Dirac distribution, and whether to convolve the result with\n",
-                "experimental resolution."
+                "One example is :class:`MultiPeakModel <erlab.analysis.fit.models.MultiPeakModel>`, which\n",
+                "works like a composite model of multiple Gaussian or Lorentzian peaks.\n",
+                "\n",
+                "By supplying keyword arguments, you can specify the number of peaks, their shapes,\n",
+                "whether to multiply with a Fermi-Dirac distribution, the shape of the background, and\n",
+                "whether to convolve the result with experimental resolution.\n",
+                "\n",
+                "For a detailed explanation of all the arguments, see its `documentation\n",
+                "<erlab.analysis.fit.models.MultiPeakModel>`.\n",
+                "\n",
+                "The model can be constructed as follows:"
             ]
         },
         {
@@ -469,6 +474,8 @@
                 }
             },
             "source": [
+                "\n",
+                "\n",
                 "Fitting ``xarray`` objects\n",
                 "--------------------------\n",
                 "\n",

environment.yml

@@ -14,6 +14,7 @@ dependencies:
   - pip>=19.3
   - pyperclip>=1.8.2
   - pyqtgraph>=0.13.1
+  - pyside6>=6.7.0
   - python>=3.11
   - qtawesome>=1.3.1
   - qtpy>=2.4.1
@@ -22,4 +23,3 @@ dependencies:
   - uncertainties>=3.1.4
   - varname>=0.13.0
   - xarray>=2024.05.0
-  - pyside6>=6.7.0

src/erlab/analysis/fit/functions/dynamic.py

@@ -14,7 +14,7 @@
 import functools
 import inspect
 from collections.abc import Callable, Sequence
-from typing import Any, ClassVar, TypedDict, no_type_check
+from typing import Any, ClassVar, Literal, TypedDict, no_type_check
 
 import numpy as np
 import numpy.typing as npt
@@ -154,6 +154,16 @@ class MultiPeakFunction(DynamicFunction):
         distribution. This adds three parameters to the model: `efermi`, `temp`, and
         `offset`, each corresponding to the Fermi level, temperature in K, and constant
         background.
+    background
+        The type of background to include in the model. The options are: ``'constant'``,
+        ``'linear'``, ``'polynomial'``, or ``'none'``. If ``'constant'``, adds a
+        ``const_bkg`` parameter. If ``'linear'``, adds a ``lin_bkg`` parameter and a
+        ``const_bkg`` parameter. If ``'polynomial'``, adds  ``c0``, ``c1``, ...
+        corresponding to the polynomial coefficients. The polynomial degree can be
+        specified with `degree`. If ``'none'``, no background is added.
+    degree
+        The degree of the polynomial background. Only used if `background` is
+        ``'polynomial'``. Default is 2.
     convolve
         Flag indicating whether the model should be convolved with a gaussian kernel. If
         `True`, adds a `resolution` parameter to the model, corresponding to the FWHM of
@@ -173,20 +183,28 @@ def __init__(
         npeaks: int,
         peak_shapes: list[str] | str | None = None,
         fd: bool = True,
+        background: Literal["constant", "linear", "polynomial", "none"] = "linear",
+        degree: int = 2,
         convolve: bool = True,
     ):
         super().__init__()
         self.npeaks = npeaks
         self.fd = fd
         self.convolve = convolve
+        self.background = background
+
+        if self.background == "polynomial":
+            self.bkg_degree = degree
 
         if peak_shapes is None:
             peak_shapes = [self.DEFAULT_PEAK] * self.npeaks
+
         if isinstance(peak_shapes, str):
             peak_shapes = peak_shapes.split(" ")
 
         if len(peak_shapes) == 1:
             peak_shapes = peak_shapes * self.npeaks
+
         elif len(peak_shapes) != self.npeaks:
             raise ValueError("Number of peaks does not match given peak shapes")
 
@@ -230,16 +248,22 @@ def argnames(self) -> list[str]:
 
     @property
     def kwargs(self):
-        kws = [
-            ("lin_bkg", 0.0),
-            ("const_bkg", 0.0),
-        ]
+        kws: list[tuple[str, float]] = []
+
+        if self.background == "constant" or self.background == "linear":
+            kws.append(("const_bkg", 0.0))
+        if self.background == "linear":
+            kws.append(("lin_bkg", 0.0))
+        elif self.background == "polynomial":
+            kws += [(f"c{i}", 0.0) for i in range(self.bkg_degree + 1)]
+
         if self.fd:
             kws += [
                 ("efermi", 0.0),  # fermi level
                 ("temp", 30.0),  # temperature
                 ("offset", 0.0),
             ]
+
         if self.convolve:
             kws += [("resolution", 0.02)]
 
@@ -275,7 +299,17 @@ def eval_peak(self, index: int, x, **params):
         )
 
     def eval_bkg(self, x, **params):
-        return params["lin_bkg"] * x + params["const_bkg"]
+        match self.background:
+            case "constant":
+                return 0.0 * x + params["const_bkg"]
+            case "linear":
+                return params["lin_bkg"] * x + params["const_bkg"]
+            case "polynomial":
+                return PolynomialFunction(self.bkg_degree)(
+                    x, **{f"c{i}": params[f"c{i}"] for i in range(self.bkg_degree + 1)}
+                )
+            case "none":
+                return 0.0 * x
 
     def pre_call(self, x, **params):
         x = np.asarray(x).copy()
@@ -286,7 +320,7 @@ def pre_call(self, x, **params):
                 x, **{arg: params[f"p{i}_{arg}"] for arg in self.peak_argnames[func]}
             )
 
-        y += params["lin_bkg"] * x + params["const_bkg"]
+        y += self.eval_bkg(x, **params)
 
         if self.fd:
             y *= fermi_dirac(x, center=params["efermi"], temp=params["temp"])
@@ -309,7 +343,7 @@ def __call__(self, x, **params):
 
 
 class FermiEdge2dFunction(DynamicFunction):
-    def __init__(self, degree=1) -> None:
+    def __init__(self, degree: int = 1):
         super().__init__()
         self.poly = PolynomialFunction(degree)
 
@@ -318,7 +352,7 @@ def argnames(self) -> list[str]:
         return ["eV", "alpha"] + self.poly.argnames[1:]
 
     @property
-    def kwargs(self):
+    def kwargs(self) -> list[tuple[str, float]]:
         return [
             ("temp", 30.0),
             ("lin_bkg", 0.0),

src/erlab/analysis/fit/models.py

@@ -10,6 +10,8 @@
     "StepEdgeModel",
 ]
 
+from typing import Literal
+
 import lmfit
 import numba
 import numpy as np
@@ -111,6 +113,34 @@ def fit_edges_linear(x, data, len_fit) -> tuple[float, float, float, float]:
     return n0, m0, n1, m1
 
 
+def _get_edges(x, y, fraction=0.2):
+    """
+    Get the edges of the input arrays based on a given fraction.
+
+    Primarily used to make initial guesses for backgrounds.
+
+    Parameters
+    ----------
+    x
+        The input x-values.
+    y
+        The input y-values.
+    fraction
+        The fraction of the array length to consider as edges. Defaults to 0.2.
+
+    """
+    if len(x) != len(y):
+        raise ValueError("x and y must have the same length")
+    if fraction < 0 or fraction > 0.5:
+        raise ValueError("Fraction must be between 0 and 0.5")
+    if len(x) < 10:
+        raise ValueError("Array must have at least 10 elements")
+
+    num = max(5, round(len(x) * fraction))
+
+    return np.r_[x[:num], x[-num:]], np.r_[y[:num], y[-num:]]
+
+
 class FermiEdgeModel(lmfit.Model):
     """Model for fitting a Fermi edge with a linear background.
 
@@ -238,13 +268,20 @@ def __init__(
         npeaks: int = 1,
         peak_shapes: list[str] | str | None = None,
         fd: bool = True,
+        background: Literal["constant", "linear", "polynomial", "none"] = "linear",
+        degree: int = 2,
         convolve: bool = True,
         **kwargs,
     ):
         kwargs.setdefault("name", f"{npeaks}Peak")
         super().__init__(
             MultiPeakFunction(
-                npeaks, peak_shapes=peak_shapes, fd=fd, convolve=convolve
+                npeaks,
+                peak_shapes=peak_shapes,
+                fd=fd,
+                background=background,
+                degree=degree,
+                convolve=convolve,
             ),
             **kwargs,
         )
@@ -271,12 +308,21 @@ def guess(self, data, x=None, **kwargs):
         if self.func.fd:
             pars[f"{self.prefix}offset"].set(value=float(data[x >= 0].mean()))
 
-        poly1 = PolynomialModel(1).guess(data, x)
-        pars[f"{self.prefix}lin_bkg"].set(poly1["c1"].value)
-        pars[f"{self.prefix}const_bkg"].set(poly1["c0"].value)
-
-        # for i, func in enumerate(self.func.peak_funcs):
-        # self.func.peak_argnames
+        # Sample edges for background estimation
+        xc, yc = _get_edges(x, data, fraction=0.2)
+        if self.func.background == "constant":
+            poly = PolynomialModel(0).guess(yc, xc)
+            pars[f"{self.prefix}const_bkg"].set(poly["c0"].value)
+
+        elif self.func.background == "linear":
+            poly = PolynomialModel(1).guess(yc, xc)
+            pars[f"{self.prefix}const_bkg"].set(poly["c0"].value)
+            pars[f"{self.prefix}lin_bkg"].set(poly["c1"].value)
+
+        elif self.func.background == "polynomial":
+            poly = PolynomialModel(self.func.bkg_degree).guess(yc, xc)
+            for i in range(self.func.bkg_degree + 1):
+                pars[f"{self.prefix}c{i}"].set(poly[f"c{i}"].value)
 
         xrange = float(x.max() - x.min())
 

tests/analysis/fit/test_functions_dynamic.py

@@ -62,18 +62,20 @@ def test_multi_peak_function_call():
         + params["const_bkg"]
     )
     assert np.allclose(
-        MultiPeakFunction(npeaks, peak_shapes, fd, convolve)(x, **params),
+        MultiPeakFunction(npeaks, peak_shapes, fd=fd, convolve=convolve)(x, **params),
         expected_result,
     )
     # Test case 2: With different peak shape signature
     assert np.allclose(
-        MultiPeakFunction(npeaks, "l g", fd, convolve)(x, **params),
+        MultiPeakFunction(npeaks, "l g", fd=fd, convolve=convolve)(x, **params),
         expected_result,
     )
 
     # Test case 3: Evaluate multi-peak function with xarray input
     x = xr.DataArray(np.linspace(-5, 5, 20, dtype=np.float64), dims="x")
-    result_xr = MultiPeakFunction(npeaks, peak_shapes, fd, convolve)(x, **params)
+    result_xr = MultiPeakFunction(npeaks, peak_shapes, fd=fd, convolve=convolve)(
+        x, **params
+    )
     assert isinstance(result_xr, xr.DataArray)
     assert np.allclose(result_xr, expected_result)
 
@@ -101,13 +103,65 @@ def test_multi_peak_function_call():
         "resolution": 0.1,
     }
 
-    expected_result = MultiPeakFunction(npeaks, "lorentzian", fd, convolve)(x, **params)
+    expected_result = MultiPeakFunction(npeaks, "lorentzian", fd=fd, convolve=convolve)(
+        x, **params
+    )
 
     for peak_shapes in ("l l l", ["lorentzian"] * 3):
         assert np.allclose(
-            MultiPeakFunction(npeaks, peak_shapes, fd, convolve)(x, **params),
+            MultiPeakFunction(npeaks, peak_shapes, fd=fd, convolve=convolve)(
+                x, **params
+            ),
             expected_result,
         )
+    # Test case 5: Evaluate multi-peak function with different background shapes
+    x = np.linspace(-5, 5, 20, dtype=np.float64)
+    npeaks = 1
+    peak_shapes = ["lorentzian"]
+    fd = False
+    convolve = False
+    degree = 2
+
+    params = {
+        "p0_center": 0.0,
+        "p0_height": 1.0,
+        "p0_width": 0.5,
+    }
+
+    for background in ("constant", "linear", "polynomial", "none"):
+        expected_result = lorentzian_wh(
+            x, params["p0_center"], params["p0_width"], params["p0_height"]
+        )
+
+        if background == "constant":
+            params["const_bkg"] = 0.1
+            expected_bkg = params["const_bkg"]
+
+        elif background == "linear":
+            params["const_bkg"] = 0.1
+            params["lin_bkg"] = 0.2
+            expected_bkg = params["lin_bkg"] * x + params["const_bkg"]
+
+        elif background == "polynomial":
+            for d in range(degree + 1):
+                params[f"c{d}"] = 0.1 * degree
+
+            coeffs = tuple(params[f"c{d}"] for d in range(degree + 1))
+            expected_bkg = np.polynomial.polynomial.polyval(x, coeffs)
+        else:
+            expected_bkg = 0.0
+
+        assert np.allclose(
+            MultiPeakFunction(
+                npeaks,
+                peak_shapes,
+                background=background,
+                degree=degree,
+                fd=fd,
+                convolve=convolve,
+            )(x, **params),
+            expected_result + expected_bkg,
+        )
 
 
 @pytest.mark.filterwarnings("ignore:overflow encountered in exp.*:RuntimeWarning")

tests/analysis/fit/test_minuit.py

@@ -24,11 +24,11 @@ def test_minuit_from_lmfit():
 
     # lmfit model
     model = MultiPeakModel(
-        npeaks=2, peak_shapes=["lorentzian"], fd=False, convolve=True
+        npeaks=2, peak_shapes=["lorentzian"], background="none", fd=False, convolve=True
     )
 
-    m = Minuit.from_lmfit(model, yval, xval, yerr)
-
+    m = Minuit.from_lmfit(model, yval, xval, yerr, p0_center=-0.5, p1_center=0.5)
+    m.scipy()
     m.migrad()
     m.minos()
     m.hesse()

tests/analysis/fit/test_models.py

@@ -118,6 +118,12 @@ def test_multi_peak_model():
         components["2Peak_bkg"], model.func.eval_bkg(x, **result.params.valuesdict())
     )
 
+    # Make sure guesses work for different backgrounds
+    for background in ["constant", "linear", "polynomial", "none"]:
+        model = models.MultiPeakModel(npeaks=2, background=background)
+        params = model.guess(y, x=x)
+        model.fit(y, params, x=x)
+
 
 def test_polynomial_model():
     # Create test data