Skip to content

Commit

Permalink
Mini refactoring so fitting functions have less arguments. Curve_fit …
Browse files Browse the repository at this point in the history 
…kwargs are now kwargs
  • Loading branch information
@gsuarezr
gsuarezr committed Nov 7, 2024
1 parent e069d1e commit ccfeba4
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 52 deletions.
104 changes: 55 additions & 49 deletions qutip/core/environment.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -360,14 +360,10 @@ def approx_by_cf_fit(
target_rsme: float = 2e-5,
Nr_max: int = 10,
Ni_max: int = 10,
guess: list[float] = None,
lower: list[float] = None,
upper: list[float] = None,
sigma: float | ArrayLike = None,
maxfev: int = None,
full_ansatz: bool = False,
combine: bool = True,
tag: Any = None,
**kwargs
) -> tuple[ExponentialBosonicEnvironment, dict[str, Any]]:
r"""
Generates an approximation to this environment by fitting its
Expand Down Expand Up @@ -416,55 +412,58 @@ def approx_by_cf_fit(
Ni_max : optional, int
The maximum number of modes to use for the fit of the imaginary
part (default 10).
guess : optional, list of float
full_ansatz : optional, bool (default False)
If this is set to False, the parameters :math:`d_k` are all set to
zero. The full ansatz, including :math:`d_k`, usually leads to
significantly slower fits, and some manual tuning of the `guesses`,
`lower` and `upper` is usually needed. On the other hand, the full
ansatz can lead to better fits with fewer exponents, especially
for anomalous spectral densities with
:math:`\operatorname{Im}[C(0)] \neq 0` for which the simplified
ansatz will always give :math:`\operatorname{Im}[C(0)] = 0`.
When using the full ansatz with default values for the guesses and
bounds, if the fit takes too long, we recommend choosing guesses
and bounds manually.
combine : optional, bool (default True)
Whether to combine exponents with the same frequency. See
:meth:`combine <.ExponentialBosonicEnvironment.combine>` for
details.
tag : optional, str, tuple or any other object
An identifier (name) for the approximated environment. If not
provided, a tag will be generated from the tag of this environment.
**kwargs
Optional fitting parameters:
guess : list of float
Initial guesses for the parameters :math:`a_k`, :math:`b_k`, etc.
The same initial guesses are used for all values of k, and for
the real and imaginary parts. If `full_ansatz` is True, `guess` is
a list of size 4, otherwise, it is a list of size 3.
If none of `guess`, `lower` and `upper` are provided, these
parameters will be chosen automatically.
lower : optional, list of float
lower : list of float
Lower bounds for the parameters :math:`a_k`, :math:`b_k`, etc.
The same lower bounds are used for all values of k, and for
the real and imaginary parts. If `full_ansatz` is True, `lower` is
a list of size 4, otherwise, it is a list of size 3.
If none of `guess`, `lower` and `upper` are provided, these
parameters will be chosen automatically.
upper : optional, list of float
upper : list of float
Upper bounds for the parameters :math:`a_k`, :math:`b_k`, etc.
The same upper bounds are used for all values of k, and for
the real and imaginary parts. If `full_ansatz` is True, `upper` is
a list of size 4, otherwise, it is a list of size 3.
If none of `guess`, `lower` and `upper` are provided, these
parameters will be chosen automatically.
sigma : optional, float or list of float
sigma : float or list of float
Adds an uncertainty to the correlation function of the environment,
i.e., adds a leeway to the fit. This parameter is useful to adjust
if the correlation function is very small in parts of the time
range. For more details, see the documentation of
``scipy.optimize.curve_fit``.
maxfev : optional, int
maxfev : int
Number of times the parameters of the fit are allowed to vary
during the optimization (per fit).
full_ansatz : optional, bool (default False)
If this is set to False, the parameters :math:`d_k` are all set to
zero. The full ansatz, including :math:`d_k`, usually leads to
significantly slower fits, and some manual tuning of the `guesses`,
`lower` and `upper` is usually needed. On the other hand, the full
ansatz can lead to better fits with fewer exponents, especially
for anomalous spectral densities with
:math:`\operatorname{Im}[C(0)] \neq 0` for which the simplified
ansatz will always give :math:`\operatorname{Im}[C(0)] = 0`.
When using the full ansatz with default values for the guesses and
bounds, if the fit takes too long, we recommend choosing guesses
and bounds manually.
combine : optional, bool (default True)
Whether to combine exponents with the same frequency. See
:meth:`combine <.ExponentialBosonicEnvironment.combine>` for
details.
tag : optional, str, tuple or any other object
An identifier (name) for the approximated environment. If not
provided, a tag will be generated from the tag of this environment.
Returns
-------
Expand Down Expand Up @@ -516,7 +515,11 @@ def approx_by_cf_fit(
Nr_min, Ni_min = Nr_max, Ni_max
else:
Nr_min, Ni_min = 1, 1

lower = kwargs.get("lower", None)
guess = kwargs.get("guess", None)
upper = kwargs.get("upper", None)
sigma = kwargs.get("sigma", None)
maxfev = kwargs.get("maxfev", None)
clist = self.correlation_function(tlist)
if guess is None and lower is None and upper is None:
guess_re, lower_re, upper_re = _default_guess_cfreal(
Expand All @@ -531,8 +534,8 @@ def approx_by_cf_fit(
start_real = time()
rmse_real, params_real = iterated_fit(
_cf_real_fit_model, num_params, tlist, np.real(clist), target_rsme,
guess_re, Nr_min, Nr_max, lower_re, upper_re, sigma=sigma,
maxfev=maxfev
Nr_min, Nr_max, guess=guess_re, lower=lower_re, upper=upper_re,
sigma=sigma, maxfev=maxfev
)
end_real = time()
fit_time_real = end_real - start_real
Expand All @@ -541,7 +544,7 @@ def approx_by_cf_fit(
start_imag = time()
rmse_imag, params_imag = iterated_fit(
_cf_imag_fit_model, num_params, tlist, np.imag(clist), target_rsme,
guess_im, Ni_min, Ni_max, lower_im, upper_im,
Ni_min, Ni_max, guess=guess_im, lower=lower_im, upper=upper_im,
sigma=sigma, maxfev=maxfev
)
end_imag = time()
Expand Down Expand Up @@ -593,13 +596,9 @@ def approx_by_sd_fit(
Nk: int = 1,
target_rmse: float = 5e-6,
Nmax: int = 10,
guess: list[float] = None,
lower: list[float] = None,
upper: list[float] = None,
sigma: float | ArrayLike = None,
maxfev: int = None,
combine: bool = True,
tag: Any = None,
**kwargs
) -> tuple[ExponentialBosonicEnvironment, dict[str, Any]]:
r"""
Generates an approximation to this environment by fitting its spectral
Expand Down Expand Up @@ -630,6 +629,16 @@ def approx_by_sd_fit(
modes (`Nmax`).
Nmax : optional, int
The maximum number of modes to use for the fit (default 10).
combine : optional, bool (default True)
Whether to combine exponents with the same frequency. See
:meth:`combine <.ExponentialBosonicEnvironment.combine>` for
details.
tag : optional, str, tuple or any other object
An identifier (name) for the approximated environment. If not
provided, a tag will be generated from the tag of this environment.
**kwargs
Optional fitting parameters:
guess : optional, list of float
Initial guesses for the parameters :math:`a_k`, :math:`b_k` and
:math:`c_k`. The same initial guesses are used for all values of
Expand Down Expand Up @@ -657,13 +666,6 @@ def approx_by_sd_fit(
maxfev : optional, int
Number of times the parameters of the fit are allowed to vary
during the optimization (per fit).
combine : optional, bool (default True)
Whether to combine exponents with the same frequency. See
:meth:`combine <.ExponentialBosonicEnvironment.combine>` for
details.
tag : optional, str, tuple or any other object
An identifier (name) for the approximated environment. If not
provided, a tag will be generated from the tag of this environment.
Returns
-------
Expand All @@ -686,7 +688,11 @@ def approx_by_sd_fit(
"summary"
A string that summarizes the information about the fit.
"""

lower = kwargs.get("lower", None)
guess = kwargs.get("guess", None)
upper = kwargs.get("upper", None)
sigma = kwargs.get("sigma", None)
maxfev = kwargs.get("maxfev", None)
# Process arguments
if tag is None and self.tag is not None:
tag = (self.tag, "SD Fit")
Expand All @@ -704,8 +710,8 @@ def approx_by_sd_fit(
# Fit
start = time()
rmse, params = iterated_fit(
_sd_fit_model, 3, wlist, jlist, target_rmse, guess,
Nmin, Nmax, lower, upper, sigma=sigma, maxfev=maxfev
_sd_fit_model, 3, wlist, jlist, target_rmse, Nmin, Nmax,
guess=guess, lower=lower, upper=upper, sigma=sigma, maxfev=maxfev
)
end = time()
fit_time = end - start
Expand Down Expand Up @@ -849,7 +855,7 @@ class DrudeLorentzEnvironment(BosonicEnvironment):
An identifier (name) for this environment.
Nk : optional, int, defaults to 10
The number of Pade exponents to be used for the calculation of the
correlation function.
correlation function.
"""

def __init__(
Expand Down
13 changes: 10 additions & 3 deletions qutip/utilities.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -354,10 +354,8 @@ def iterated_fit(
fun: Callable[..., complex], num_params: int,
xdata: ArrayLike, ydata: ArrayLike,
target_rmse: float = 1e-5,
guess: ArrayLike | Callable[[int], ArrayLike] = None,
Nmin: int = 1, Nmax: int = 10,
lower: ArrayLike = None, upper: ArrayLike = None,
sigma: float | ArrayLike = None, maxfev: int = None
**kwargs
) -> tuple[float, ArrayLike]:
r"""
Iteratively tries to fit the given data with a model of the form
Expand Down Expand Up @@ -396,6 +394,9 @@ def iterated_fit(
The maximum number of terms to be used for the fit (default 10).
If the number `Nmax` of terms is reached, the function returns even if
the target rmse has not been reached yet.
**kwargs
Scipy's curve_fit keyword arguments:
lower : optional, list of length `num_params`
Lower bounds on the parameters for the fit.
upper : optional, list of length `num_params`
Expand All @@ -414,6 +415,12 @@ def iterated_fit(
The model parameters in the form
`[[p11, ..., p1n], [p21, ..., p2n], ..., [pN1, ..., pNn]]`.
"""
# Extract kwargs and set to None if not provided
lower = kwargs.get("lower", None)
guess = kwargs.get("guess", None)
upper = kwargs.get("upper", None)
sigma = kwargs.get("sigma", None)
maxfev = kwargs.get("maxfev", None)
if len(xdata) != len(ydata):
raise ValueError(
"The shape of the provided fit data is not consistent")
Expand Down

0 comments on commit ccfeba4

Please sign in to comment.