From 0aa4bd4b48940f10bbfbd92f10a8920e5e98a05c Mon Sep 17 00:00:00 2001 From: Bing Li Date: Mon, 18 Nov 2024 17:44:31 -0500 Subject: [PATCH] updated Fit1D class --- src/tavi/data/fit.py | 226 +++++++++++++++--------------------- src/tavi/data/fit_backup.py | 207 +++++++++++++++++++++++++++++++++ src/tavi/plotter.py | 11 +- tests/test_fit.py | 52 ++++++++- 4 files changed, 359 insertions(+), 137 deletions(-) create mode 100644 src/tavi/data/fit_backup.py diff --git a/src/tavi/data/fit.py b/src/tavi/data/fit.py index 6979100..afad0aa 100644 --- a/src/tavi/data/fit.py +++ b/src/tavi/data/fit.py @@ -4,7 +4,19 @@ from lmfit import Parameters, models from tavi.data.scan_data import ScanData1D -from tavi.plotter import Plot1D + + +class FitData1D(object): + + def __init__( + self, + x: np.ndarray, + y: np.ndarray, + ) -> None: + + self.x = x + self.y = y + self.fmt: dict = {} class Fit1D(object): @@ -41,7 +53,6 @@ def __init__( ): """initialize a fit model, mask based on fit_range if given""" - self.NUM_PTS: int = 100 self.x: np.ndarray = data.x self.y: np.ndarray = data.y self.err: Optional[np.ndarray] = data.err @@ -51,7 +62,8 @@ def __init__( self.pars = Parameters() self._num_backgrounds = 0 self._num_signals = 0 - self.fit_result = None + self.result = None + self.fit_data: Optional[FitData1D] = None self.PLOT_SEPARATELY = False @@ -67,141 +79,91 @@ def set_range(self, fit_range: tuple[float, float]): if self.err is not None: self.err = self.err[mask] - @property - def x_plot(self): - return np.linspace(self.x.min(), self.x.max(), num=self.NUM_PTS) - - def add_background( - self, - model: Literal["Constant", "Linear", "Quadratic", "Polynomial", "Exponential", "PowerLaw"] = "Constant", - values=None, - vary=None, - mins=None, - maxs=None, - exprs=None, - ): - """Set the model for background - - Args: - model (str): Constant, Linear, Quadratic, Polynomial, - Exponential, PowerLaw - p0 (tuple | None): inital parameters - min (tuple | None): minimum - max (tuple | None): maximum - fixed (tuple | None): tuple of flags - expr (tuple| None ): constraint expressions - """ - self._num_backgrounds += 1 - - # add prefix if more than one background - if self._num_backgrounds > 1: - prefix = f"b{self._num_backgrounds}_" - else: - prefix = "" - - model = Fit1D.models[model](prefix=prefix, nan_policy="propagate") - param_names = model.param_names - # guess initials - pars = model.guess(self.y, x=self.x) - - # overwrite with user input - if values is not None: - for idx, v in enumerate(values): - if v is not None: - pars[param_names[idx]].set(value=v) - - if vary is not None: - for idx, v in enumerate(vary): - if v is not None: - pars[param_names[idx]].set(vary=v) - - if mins is not None: - for idx, v in enumerate(mins): - if v is not None: - pars[param_names[idx]].set(min=v) - if maxs is not None: - for idx, v in enumerate(maxs): - if v is not None: - pars[param_names[idx]].set(max=v) - - if exprs is not None: - for idx, v in enumerate(exprs): - if v is not None: - pars[param_names[idx]].set(expr=v) - - for param_name in param_names: - self.pars.add(pars[param_name]) - - self.background_models.append(model) + @staticmethod + def _add_model(model, prefix): + model = Fit1D.models[model] + return model(prefix=prefix, nan_policy="propagate") def add_signal( self, - model="Gaussian", - values=None, - vary=None, - mins=None, - maxs=None, - exprs=None, + model_name: Literal[ + "Gaussian", "Lorentzian", "Voigt", "PseudoVoigt", "DampedOscillator", "DampedHarmonicOscillator" + ], ): - """Set the model for signal - - Args: - model (str): Constant, Linear, Quadratic, Polynomial, - Exponential, PowerLaw - p0 (tuple | None): inital parameters - min (tuple | None): minimum - max (tuple | None): maximum - expr (str| None ): constraint expression - """ self._num_signals += 1 prefix = f"s{self._num_signals}_" - model = Fit1D.models[model](prefix=prefix, nan_policy="propagate") - param_names = model.param_names - # guess initials - pars = model.guess(self.y, x=self.x) - - # overwrite with user input - if values is not None: - for idx, v in enumerate(values): - if v is not None: - pars[param_names[idx]].set(value=v) - - if vary is not None: - for idx, v in enumerate(vary): - if v is not None: - pars[param_names[idx]].set(vary=v) - - if mins is not None: - for idx, v in enumerate(mins): - if v is not None: - pars[param_names[idx]].set(min=v) - if maxs is not None: - for idx, v in enumerate(maxs): - if v is not None: - pars[param_names[idx]].set(max=v) - - if exprs is not None: - for idx, v in enumerate(exprs): - if v is not None: - pars[param_names[idx]].set(expr=v) - - for param_name in param_names: - self.pars.add(pars[param_name]) - self.signal_models.append(model) - - def perform_fit(self) -> None: - model = np.sum(self.signal_models) - - if self._num_backgrounds > 0: - model += np.sum(self.background_models) - - if self.err is None: - out = model.fit(self.y, self.pars, x=self.x) + self.signal_models.append(Fit1D._add_model(model_name, prefix)) + + def add_background( + self, model_name: Literal["Constant", "Linear", "Quadratic", "Polynomial", "Exponential", "PowerLaw"] + ): + self._num_backgrounds += 1 + prefix = f"b{self._num_backgrounds}_" + self.background_models.append(Fit1D._add_model(model_name, prefix)) + + @staticmethod + def _get_model_params(models): + params = [] + for model in models: + params.append(model.param_names) + return params + + def get_signal_params(self): + return Fit1D._get_model_params(self.signal_models) + + def get_background_params(self): + return Fit1D._get_model_params(self.background_models) + + def guess(self): + pars = Parameters() + for signal in self.signal_models: + pars += signal.guess(self.y, x=self.x) + for bkg in self.background_models: + pars += bkg.guess(self.y, x=self.x) + self.pars = pars + return pars + + @property + def x_to_plot(self): + return + + def eval(self, pars: Parameters, num_of_pts: Optional[int] = 100) -> FitData1D: + mod = self.signal_models[0] + if (sz := len(self.signal_models)) > 1: + for i in range(1, sz): + mod += self.signal_models[i] + + for bkg in self.background_models: + mod += bkg + + if num_of_pts is None: + x_to_plot = self.x + elif isinstance(num_of_pts, int): + x_to_plot = np.linspace(self.x.min(), self.x.max(), num=num_of_pts) + else: + raise ValueError(f"num_of_points={num_of_pts} needs to be an integer.") + y_to_plot = mod.eval(pars, x=x_to_plot) + return FitData1D(x_to_plot, y_to_plot) + + def fit(self, pars: Parameters, num_of_pts: Optional[int] = 100) -> FitData1D: + mod = self.signal_models[0] + if (sz := len(self.signal_models)) > 1: + for i in range(1, sz): + mod += self.signal_models[i] + + for bkg in self.background_models: + mod += bkg + + result = mod.fit(self.y, pars, x=self.x, weights=self.err) + self.result = result + + if num_of_pts is None: + x_to_plot = self.x + elif isinstance(num_of_pts, int): + x_to_plot = np.linspace(self.x.min(), self.x.max(), num=num_of_pts) else: - out = model.fit(self.y, self.pars, x=self.x, weights=self.err) + raise ValueError(f"num_of_points={num_of_pts} needs to be an integer.") - self.result = out - self.y_plot = model.eval(out.params, x=self.x_plot) + y_to_plot = mod.eval(result.params, x=x_to_plot) - self.fit_report = out.fit_report(min_correl=0.25) - self.fit_plot = Plot1D(x=self.x_plot, y=self.y_plot) + return FitData1D(x_to_plot, y_to_plot) diff --git a/src/tavi/data/fit_backup.py b/src/tavi/data/fit_backup.py new file mode 100644 index 0000000..6979100 --- /dev/null +++ b/src/tavi/data/fit_backup.py @@ -0,0 +1,207 @@ +from typing import Literal, Optional + +import numpy as np +from lmfit import Parameters, models + +from tavi.data.scan_data import ScanData1D +from tavi.plotter import Plot1D + + +class Fit1D(object): + """Fit a 1d curve + + Attributes: + + """ + + models = { + # ---------- peak models --------------- + "Gaussian": models.GaussianModel, + "Lorentzian": models.LorentzianModel, + "Voigt": models.VoigtModel, + "PseudoVoigt": models.PseudoVoigtModel, + "DampedOscillator": models.DampedOscillatorModel, + "DampedHarmonicOscillator": models.DampedHarmonicOscillatorModel, + # ---------- background models ------------ + "Constant": models.ConstantModel, + "Linear": models.LinearModel, + "Quadratic": models.QuadraticModel, + "Polynomial": models.PolynomialModel, + "Exponential": models.ExponentialModel, + "PowerLaw": models.PowerLawModel, + # --------- expression --------------- + "Expression": models.ExpressionModel, + "Spline": models.SplineModel, + } + + def __init__( + self, + data: ScanData1D, + fit_range: Optional[tuple[float, float]] = None, + ): + """initialize a fit model, mask based on fit_range if given""" + + self.NUM_PTS: int = 100 + self.x: np.ndarray = data.x + self.y: np.ndarray = data.y + self.err: Optional[np.ndarray] = data.err + + self.background_models: models = [] + self.signal_models: models = [] + self.pars = Parameters() + self._num_backgrounds = 0 + self._num_signals = 0 + self.fit_result = None + + self.PLOT_SEPARATELY = False + + if fit_range is not None: + self.set_range(fit_range) + + def set_range(self, fit_range: tuple[float, float]): + """set the range used for fitting""" + fit_min, fit_max = fit_range + mask = np.bitwise_and(self.x >= fit_min, self.x <= fit_max) + self.x = self.x[mask] + self.y = self.y[mask] + if self.err is not None: + self.err = self.err[mask] + + @property + def x_plot(self): + return np.linspace(self.x.min(), self.x.max(), num=self.NUM_PTS) + + def add_background( + self, + model: Literal["Constant", "Linear", "Quadratic", "Polynomial", "Exponential", "PowerLaw"] = "Constant", + values=None, + vary=None, + mins=None, + maxs=None, + exprs=None, + ): + """Set the model for background + + Args: + model (str): Constant, Linear, Quadratic, Polynomial, + Exponential, PowerLaw + p0 (tuple | None): inital parameters + min (tuple | None): minimum + max (tuple | None): maximum + fixed (tuple | None): tuple of flags + expr (tuple| None ): constraint expressions + """ + self._num_backgrounds += 1 + + # add prefix if more than one background + if self._num_backgrounds > 1: + prefix = f"b{self._num_backgrounds}_" + else: + prefix = "" + + model = Fit1D.models[model](prefix=prefix, nan_policy="propagate") + param_names = model.param_names + # guess initials + pars = model.guess(self.y, x=self.x) + + # overwrite with user input + if values is not None: + for idx, v in enumerate(values): + if v is not None: + pars[param_names[idx]].set(value=v) + + if vary is not None: + for idx, v in enumerate(vary): + if v is not None: + pars[param_names[idx]].set(vary=v) + + if mins is not None: + for idx, v in enumerate(mins): + if v is not None: + pars[param_names[idx]].set(min=v) + if maxs is not None: + for idx, v in enumerate(maxs): + if v is not None: + pars[param_names[idx]].set(max=v) + + if exprs is not None: + for idx, v in enumerate(exprs): + if v is not None: + pars[param_names[idx]].set(expr=v) + + for param_name in param_names: + self.pars.add(pars[param_name]) + + self.background_models.append(model) + + def add_signal( + self, + model="Gaussian", + values=None, + vary=None, + mins=None, + maxs=None, + exprs=None, + ): + """Set the model for signal + + Args: + model (str): Constant, Linear, Quadratic, Polynomial, + Exponential, PowerLaw + p0 (tuple | None): inital parameters + min (tuple | None): minimum + max (tuple | None): maximum + expr (str| None ): constraint expression + """ + self._num_signals += 1 + prefix = f"s{self._num_signals}_" + model = Fit1D.models[model](prefix=prefix, nan_policy="propagate") + param_names = model.param_names + # guess initials + pars = model.guess(self.y, x=self.x) + + # overwrite with user input + if values is not None: + for idx, v in enumerate(values): + if v is not None: + pars[param_names[idx]].set(value=v) + + if vary is not None: + for idx, v in enumerate(vary): + if v is not None: + pars[param_names[idx]].set(vary=v) + + if mins is not None: + for idx, v in enumerate(mins): + if v is not None: + pars[param_names[idx]].set(min=v) + if maxs is not None: + for idx, v in enumerate(maxs): + if v is not None: + pars[param_names[idx]].set(max=v) + + if exprs is not None: + for idx, v in enumerate(exprs): + if v is not None: + pars[param_names[idx]].set(expr=v) + + for param_name in param_names: + self.pars.add(pars[param_name]) + self.signal_models.append(model) + + def perform_fit(self) -> None: + model = np.sum(self.signal_models) + + if self._num_backgrounds > 0: + model += np.sum(self.background_models) + + if self.err is None: + out = model.fit(self.y, self.pars, x=self.x) + else: + out = model.fit(self.y, self.pars, x=self.x, weights=self.err) + + self.result = out + self.y_plot = model.eval(out.params, x=self.x_plot) + + self.fit_report = out.fit_report(min_correl=0.25) + self.fit_plot = Plot1D(x=self.x_plot, y=self.y_plot) diff --git a/src/tavi/plotter.py b/src/tavi/plotter.py index 69c22da..0c8e2ac 100644 --- a/src/tavi/plotter.py +++ b/src/tavi/plotter.py @@ -6,6 +6,7 @@ from mpl_toolkits.axisartist.grid_finder import MaxNLocator from mpl_toolkits.axisartist.grid_helper_curvelinear import GridHelperCurveLinear +from tavi.data.fit import FitData1D from tavi.data.scan_data import ScanData1D, ScanData2D from tavi.instrument.resolution.ellipse import ResoEllipse @@ -37,6 +38,7 @@ class Plot1D(object): def __init__(self) -> None: # self.ax = None self.scan_data: list[ScanData1D] = [] + self.fit_data: list[FitData1D] = [] self.title = "" self.xlabel = None self.ylabel = None @@ -52,12 +54,19 @@ def add_scan(self, scan_data: ScanData1D, **kwargs): for key, val in kwargs.items(): scan_data.fmt.update({key: val}) + def add_fit(self, fit_data: FitData1D, **kwargs): + self.fit_data.append(fit_data) + for key, val in kwargs.items(): + fit_data.fmt.update({key: val}) + def plot(self, ax): for data in self.scan_data: if data.err is None: ax.plot(data.x, data.y, **data.fmt) else: ax.errorbar(x=data.x, y=data.y, yerr=data.err, **data.fmt) + for fit in self.fit_data: + ax.plot(fit.x, fit.y, **fit.fmt) if self.xlim is not None: ax.set_xlim(left=self.xlim[0], right=self.xlim[1]) @@ -80,7 +89,7 @@ def plot(self, ax): ax.set_ylabel(",".join(set(ylabels))) ax.grid(alpha=0.6) - for data in self.scan_data: + for data in self.scan_data + self.fit_data: if "label" in data.fmt.keys(): ax.legend() break diff --git a/tests/test_fit.py b/tests/test_fit.py index df7a732..25031bb 100644 --- a/tests/test_fit.py +++ b/tests/test_fit.py @@ -13,19 +13,21 @@ def test_fit_single_peak_external_model(): path_to_spice_folder = "./test_data/exp424" scan42 = Scan.from_spice(path_to_spice_folder=path_to_spice_folder, scan_num=42) - s1_scan = scan42.get_data(norm_channel="mcu", norm_val=30) + s1_scan = scan42.get_data(norm_to=(30, "mcu")) f1 = Fit1D(s1_scan, fit_range=(0.5, 4.0)) bkg = ConstantModel(prefix="bkg_", nan_policy="propagate") peak = GaussianModel(prefix="peak_", nan_policy="propagate") model = peak + bkg + pars = peak.guess(f1.y, x=f1.x) pars += bkg.make_params(c=0) + out = model.fit(f1.y, pars, x=f1.x, weight=f1.err) assert np.allclose(out.values["peak_center"], 3.54, atol=0.01) assert np.allclose(out.values["peak_fwhm"], 0.39, atol=0.01) - assert np.allclose(out.redchi, 10.012, atol=0.01) + assert np.allclose(out.redchi, 2.50, atol=0.01) # p1 = Plot1D() # p1.add_scan(s1_scan, fmt="o") @@ -44,14 +46,56 @@ def test_fit_single_peak_external_model(): # plt.show() +def test_get_fitting_variables(): + path_to_spice_folder = "./test_data/exp424" + scan42 = Scan.from_spice(path_to_spice_folder=path_to_spice_folder, scan_num=42) + + s1_scan = scan42.get_data(norm_to=(30, "mcu")) + f1 = Fit1D(s1_scan, fit_range=(0.5, 4.0)) + + f1.add_signal(model_name="Gaussian") + f1.add_background(model_name="Constant") + + assert f1.get_signal_params() == [["s1_amplitude", "s1_center", "s1_sigma"]] + assert f1.get_background_params() == [["b1_c"]] + + +def test_guess_initial(): + path_to_spice_folder = "./test_data/exp424" + scan42 = Scan.from_spice(path_to_spice_folder=path_to_spice_folder, scan_num=42) + + s1_scan = scan42.get_data(norm_to=(30, "mcu")) + f1 = Fit1D(s1_scan, fit_range=(0.5, 4.0)) + + f1.add_signal(model_name="Gaussian") + f1.add_background(model_name="Constant") + pars = f1.guess() + inital = f1.eval(pars) + fit = f1.fit(pars) + + p1 = Plot1D() + p1.add_scan(s1_scan, fmt="o", label="data") + p1.add_fit(inital, label="guess") + p1.add_fit(fit, label="fit") + + fig, ax = plt.subplots() + p1.plot(ax) + plt.show() + + def test_fit_single_peak_internal_model(): path_to_spice_folder = "./test_data/exp424" scan42 = Scan.from_spice(path_to_spice_folder=path_to_spice_folder, scan_num=42) - s1_scan = scan42.get_data(norm_channel="mcu", norm_val=30) + s1_scan = scan42.get_data(norm_to=(30, "mcu")) f1 = Fit1D(s1_scan, fit_range=(0.5, 4.0)) + f1.add_background(model="Constant") + f1.add_signal(model_name="Gaussian") + f1.eval() + f1.fit() + bkg = ConstantModel(prefix="bkg_", nan_policy="propagate") peak = GaussianModel(prefix="peak_", nan_policy="propagate") model = peak + bkg @@ -91,7 +135,7 @@ def test_fit_two_peak(): f1.add_background(values=(0.7,)) f1.add_signal(values=(None, 3.5, 0.29), vary=(True, True, True)) f1.add_signal( - model="Gaussian", + model_name="Gaussian", values=(None, 0, None), vary=(True, False, True), mins=(0, 0, 0.1),