From 27d0795bd362e38dddf4e36f9b6a266afbff4935 Mon Sep 17 00:00:00 2001 From: Ryuichi Arafune Date: Sun, 21 Apr 2024 18:18:42 +0900 Subject: [PATCH] =?UTF-8?q?=F0=9F=92=AC=20=20Update=20=5F=5Fdoc=5F=5F=20fo?= =?UTF-8?q?r=20fit=20models?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/arpes/fits/fit_models/decay.py | 9 +++-- src/arpes/fits/fit_models/dirac.py | 4 ++- src/arpes/fits/fit_models/fermi_edge.py | 38 +++++++++++++------- src/arpes/fits/fit_models/misc.py | 6 ++-- src/arpes/fits/fit_models/peaks.py | 4 +-- src/arpes/fits/fit_models/two_dimensional.py | 4 +-- src/arpes/fits/fit_models/x_model_mixin.py | 16 +++++---- 7 files changed, 52 insertions(+), 29 deletions(-) diff --git a/src/arpes/fits/fit_models/decay.py b/src/arpes/fits/fit_models/decay.py index ec27bac0..a733ff66 100644 --- a/src/arpes/fits/fit_models/decay.py +++ b/src/arpes/fits/fit_models/decay.py @@ -79,7 +79,10 @@ def guess( return update_param_vals(pars, self.prefix, **kwargs) - __init__.doc = lf.models.COMMON_INIT_DOC + __init__.__doc__ = ( + "A model for fitting an exponential decay with a constant background." + + lf.models.COMMON_INIT_DOC + ) guess.__doc__ = lf.models.COMMON_GUESS_DOC @@ -136,5 +139,7 @@ def guess( return update_param_vals(pars, self.prefix, **kwargs) - __init__.doc = lf.models.COMMON_INIT_DOC + __init__.__doc__ = ( + "Like `exponential_decay_c`, except with two timescales." + lf.models.COMMON_INIT_DOC + ) guess.__doc__ = lf.models.COMMON_GUESS_DOC diff --git a/src/arpes/fits/fit_models/dirac.py b/src/arpes/fits/fit_models/dirac.py index fd279899..d71c5267 100644 --- a/src/arpes/fits/fit_models/dirac.py +++ b/src/arpes/fits/fit_models/dirac.py @@ -70,5 +70,7 @@ def guess(self, **kwargs: float) -> lf.Parameters: pars = self.make_params() return update_param_vals(pars, self.prefix, **kwargs) - __init__.doc = lf.models.COMMON_INIT_DOC + __init__.__doc__ = ( + "Model for dirac_dispersion symmetric about the dirac point." + lf.models.COMMON_INIT_DOC + ) guess.__doc__ = lf.models.COMMON_GUESS_DOC diff --git a/src/arpes/fits/fit_models/fermi_edge.py b/src/arpes/fits/fit_models/fermi_edge.py index 4c399490..64b9833f 100644 --- a/src/arpes/fits/fit_models/fermi_edge.py +++ b/src/arpes/fits/fit_models/fermi_edge.py @@ -109,7 +109,9 @@ def guess(self, data: XrTypes, **kwargs: float) -> lf.Parameters: return update_param_vals(pars, self.prefix, **kwargs) - __init__.doc = lf.models.COMMON_INIT_DOC + __init__.__doc__ = ( + "affine density of states broadened by Fermi-Dirac " + lf.models.COMMON_INIT_DOC + ) guess.__doc__ = lf.models.COMMON_GUESS_DOC @@ -175,7 +177,9 @@ def guess( return update_param_vals(pars, self.prefix, **kwargs) - __init__.doc = lf.models.COMMON_INIT_DOC + __init__.__doc__ = ( + "Lorentzian multiplied by a gstepb background model" + lf.models.COMMON_INIT_DOC + ) guess.__doc__ = lf.models.COMMON_GUESS_DOC @@ -201,7 +205,7 @@ def guess(self, data: DataType, **kwargs: Incomplete) -> lf.Parameters: return update_param_vals(pars, self.prefix, **kwargs) - __init__.doc = lf.models.COMMON_INIT_DOC + __init__.__doc__ = "Fermi-Dirc distribution model" + lf.models.COMMON_INIT_DOC guess.__doc__ = lf.models.COMMON_GUESS_DOC @@ -247,7 +251,9 @@ def guess( return update_param_vals(pars, self.prefix, **kwargs) - __init__.doc = lf.models.COMMON_INIT_DOC + __init__.__doc__ = ( + """Fermi functions with a linear background model""" + lf.models.COMMON_INIT_DOC + ) guess.__doc__ = lf.models.COMMON_GUESS_DOC @@ -389,10 +395,6 @@ def __init__(self, **kwargs: Unpack[ModelArgs]) -> None: self.set_param_hint("offset", min=-10) self.set_param_hint("center", vary=False) - """ - - """ - def guess( self, data: XrTypes, @@ -493,7 +495,10 @@ def guess( return update_param_vals(pars, self.prefix, **kwargs) - __init__.doc = lf.models.COMMON_INIT_DOC + __init__.__doc__ = ( + "Fermi Dirac function with affine background multiplied, then all convolved with Gaussian" + + lf.models.COMMON_INIT_DOC + ) guess.__doc__ = lf.models.COMMON_GUESS_DOC @@ -553,7 +558,10 @@ def guess( return update_param_vals(pars, self.prefix, **kwargs) - __init__.doc = lf.models.COMMON_INIT_DOC + __init__.__doc__ = ( + "Fermi-Dirac distribution function with a linear background model" + + lf.models.COMMON_INIT_DOC + ) guess.__doc__ = lf.models.COMMON_GUESS_DOC @@ -611,7 +619,10 @@ def guess( return update_param_vals(pars, self.prefix, **kwargs) - __init__.doc = lf.models.COMMON_INIT_DOC + __init__.__doc__ = ( + """A model for fitting Fermi functions with a linear background.""" + + lf.models.COMMON_INIT_DOC + ) guess.__doc__ = lf.models.COMMON_GUESS_DOC @@ -682,5 +693,8 @@ def guess( return update_param_vals(pars, self.prefix, **kwargs) - __init__.doc = lf.models.COMMON_INIT_DOC + __init__.__doc__ = ( + "A model for (two lorentzians with an affine background) multiplied by a gstepb." + + lf.models.COMMON_INIT_DOC + ) guess.__doc__ = lf.models.COMMON_GUESS_DOC diff --git a/src/arpes/fits/fit_models/misc.py b/src/arpes/fits/fit_models/misc.py index 69bf318d..ff94c64f 100644 --- a/src/arpes/fits/fit_models/misc.py +++ b/src/arpes/fits/fit_models/misc.py @@ -57,7 +57,7 @@ def guess( return update_param_vals(pars, self.prefix, **kwargs) - __init__.doc = lf.models.COMMON_INIT_DOC + __init__.__doc__ = "Quadratic model" + lf.models.COMMON_INIT_DOC guess.__doc__ = lf.models.COMMON_GUESS_DOC @@ -100,7 +100,7 @@ def guess(self, **kwargs: float) -> lf.Parameters: return update_param_vals(pars, self.prefix, **kwargs) - __init__.doc = lf.models.COMMON_INIT_DOC + __init__.__doc__ = "Fermi velocity renormalization model" + lf.models.COMMON_INIT_DOC guess.__doc__ = lf.models.COMMON_GUESS_DOC @@ -155,5 +155,5 @@ def guess(self, **kwargs: float) -> lf.Parameters: return update_param_vals(pars, self.prefix, **kwargs) - __init__.doc = lf.models.COMMON_INIT_DOC + __init__.__doc__ = "Log renormalization model" + lf.models.COMMON_INIT_DOC guess.__doc__ = lf.models.COMMON_GUESS_DOC diff --git a/src/arpes/fits/fit_models/peaks.py b/src/arpes/fits/fit_models/peaks.py index d4caefe5..6b51b7bb 100644 --- a/src/arpes/fits/fit_models/peaks.py +++ b/src/arpes/fits/fit_models/peaks.py @@ -78,7 +78,7 @@ def guess( return update_param_vals(pars, self.prefix, **kwargs) - __init__.doc = lf.models.COMMON_INIT_DOC + __init__.__doc__ = "Two gaussian model" + lf.models.COMMON_INIT_DOC guess.__doc__ = lf.models.COMMON_GUESS_DOC @@ -121,5 +121,5 @@ def guess( return update_param_vals(pars, self.prefix, **kwargs) - __init__.doc = lf.models.COMMON_INIT_DOC + __init__.__doc__ = "Two lorenzian model" + lf.models.COMMON_INIT_DOC guess.__doc__ = lf.models.COMMON_GUESS_DOC diff --git a/src/arpes/fits/fit_models/two_dimensional.py b/src/arpes/fits/fit_models/two_dimensional.py index ec81d6d7..ba27b521 100644 --- a/src/arpes/fits/fit_models/two_dimensional.py +++ b/src/arpes/fits/fit_models/two_dimensional.py @@ -28,7 +28,7 @@ class Gaussian2DModel(XModelMixin): """2D Gaussian fitting.""" n_dims = 2 - dimension_order: ClassVar = [any_dim_sentinel, any_dim_sentinel] + dimension_order: ClassVar[list[str | None]] = [any_dim_sentinel, any_dim_sentinel] @staticmethod def gaussian_2d_bkg( # noqa: PLR0913 @@ -169,5 +169,5 @@ def guess( return update_param_vals(pars, self.prefix, **kwargs) - __init__.doc = lf.models.COMMON_INIT_DOC + __init__.__doc__ = lf.models.COMMON_INIT_DOC guess.__doc__ = lf.models.COMMON_GUESS_DOC diff --git a/src/arpes/fits/fit_models/x_model_mixin.py b/src/arpes/fits/fit_models/x_model_mixin.py index 1de3fc6b..e786578a 100644 --- a/src/arpes/fits/fit_models/x_model_mixin.py +++ b/src/arpes/fits/fit_models/x_model_mixin.py @@ -5,7 +5,7 @@ import operator import warnings from logging import DEBUG, INFO, WARNING, Formatter, StreamHandler, getLogger -from typing import TYPE_CHECKING, Required, TypedDict +from typing import TYPE_CHECKING, ClassVar, Required, TypedDict import lmfit as lf import numpy as np @@ -109,7 +109,7 @@ class XModelMixin(lf.Model): """ n_dims = 1 - dimension_order = None + dimension_order: ClassVar[list[str | None]] = [None] def guess_fit( # noqa: PLR0913 self, @@ -282,7 +282,7 @@ def _real_data_etc_from_xarray( real_data, flat_data = data.values, data.values assert len(real_data.shape) == self.n_dims coord_values = {} - new_dim_order = None + new_dim_order: list[str] if self.n_dims == 1: coord_values["x"] = data.coords[next(iter(data.indexes))].values else: @@ -297,14 +297,16 @@ def find_appropriate_dimension(dim_or_dim_list: str | list[str]) -> str: return next(iter(intersect)) # resolve multidimensional parameters - if self.dimension_order is None or all(d is None for d in self.dimension_order): - new_dim_order = data.dims + if all(d is None for d in self.dimension_order): + new_dim_order = [str(dim) for dim in data.dims] else: new_dim_order = [ - find_appropriate_dimension(dim_options) for dim_options in self.dimension_order + find_appropriate_dimension(dim_options) + for dim_options in self.dimension_order + if dim_options is not None ] - if list(new_dim_order) != list(data.dims): + if new_dim_order != list(data.dims): warnings.warn("Transposing data for multidimensional fit.", stacklevel=2) data = data.transpose(*new_dim_order)