diff --git a/docs/source/user-guide/curve-fitting.ipynb b/docs/source/user-guide/curve-fitting.ipynb index 7e2419c0..8d19bf72 100644 --- a/docs/source/user-guide/curve-fitting.ipynb +++ b/docs/source/user-guide/curve-fitting.ipynb @@ -31,8 +31,7 @@ "fit multidimensional xarray objects. Finally, we will show how to use `iminuit\n", "`_ 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 `, 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 `, 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", + "`.\n", + "\n", + "The model can be constructed as follows:" ] }, { @@ -469,6 +474,8 @@ } }, "source": [ + "\n", + "\n", "Fitting ``xarray`` objects\n", "--------------------------\n", "\n", diff --git a/environment.yml b/environment.yml index f3eda705..70a88623 100644 --- a/environment.yml +++ b/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 diff --git a/src/erlab/analysis/fit/functions/dynamic.py b/src/erlab/analysis/fit/functions/dynamic.py index f50f12ad..54914ab4 100644 --- a/src/erlab/analysis/fit/functions/dynamic.py +++ b/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), diff --git a/src/erlab/analysis/fit/models.py b/src/erlab/analysis/fit/models.py index 96eab70f..0b40bd6a 100644 --- a/src/erlab/analysis/fit/models.py +++ b/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()) diff --git a/tests/analysis/fit/test_functions_dynamic.py b/tests/analysis/fit/test_functions_dynamic.py index ba53abcd..ce55d835 100644 --- a/tests/analysis/fit/test_functions_dynamic.py +++ b/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") diff --git a/tests/analysis/fit/test_minuit.py b/tests/analysis/fit/test_minuit.py index 15ef0249..407852dd 100644 --- a/tests/analysis/fit/test_minuit.py +++ b/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() diff --git a/tests/analysis/fit/test_models.py b/tests/analysis/fit/test_models.py index bc2b1d34..d9c2b437 100644 --- a/tests/analysis/fit/test_models.py +++ b/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