Skip to content

Commit

Permalink
formatting
Browse files Browse the repository at this point in the history
  • Loading branch information
@gsuarezr
gsuarezr committed Jan 16, 2025
1 parent b96c325 commit 8966d53
Showing 1 changed file with 87 additions and 85 deletions.
172 changes: 87 additions & 85 deletions qutip/core/environment.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -382,22 +382,23 @@ def _cf_from_ps(self, t, wMax, **ps_kwargs):
tMax, tMax=wMax)
result = result_fct(t) / (2 * np.pi)
return result.item() if t.ndim == 0 else result

@overload
def approximate(self,
method: Literal['corr_lsq'],
tlist: ArrayLike,
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,
):
method: Literal['corr_lsq'],
tlist: ArrayLike,
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,
):
...
r"""
Generates an approximation to this environment by fitting its
Expand Down Expand Up @@ -532,17 +533,18 @@ def approximate(self,
A string that summarizes the information about the fit.
"""
...

@overload
def approximate(
self,
method: Literal['spec_lsq'],
wlist: ArrayLike,
Nk: int ,
Nk: int,
target_rmse: float,
Nmax: int ,
guess: list[float] ,
lower: list[float] ,
upper: list[float] ,
Nmax: int,
guess: list[float],
lower: list[float],
upper: list[float],
sigma: float | ArrayLike,
maxfev: int = None,
combine: bool = True,
Expand Down Expand Up @@ -637,20 +639,20 @@ def approximate(

@overload
def approximate(self,
method: Literal['esprit'],
tlist: ArrayLike,
Nr: int ,
Ni: int ,
combine: bool ,
tag: Any,):
method: Literal['esprit'],
tlist: ArrayLike,
Nr: int,
Ni: int,
combine: bool,
tag: Any,):
r"""
Generates an approximation to this environment by fitting its
correlation function using one of the methods based on the Prony
correlation function using one of the methods based on the Prony
polynomial.
- method='prony' For the Prony Method
- method='mp' For the Matrix Pencil Method
- method='esprit' For the Estimation of signal parameters via
- method='prony' For the Prony Method
- method='mp' For the Matrix Pencil Method
- method='esprit' For the Estimation of signal parameters via
rotational invariant techniques method
Expand Down Expand Up @@ -680,58 +682,59 @@ def approximate(self,
"""
@overload
def approximate(self,
method: Literal['aaa'],
wlist: ArrayLike,
tol: float,
N_max: int,
combine: bool,
tag: Any):
"""
Generates an approximation to this environment by fitting its power
spectrum using the AAA algorithm. The function is fit to a rational
polynomial of the form
.. math::
S(\\omega)= 2 \\Re \\left(\\sum_{k} \frac{c_{k}}{\nu_{k}-i \\omega}
\right)
By isolating the poles and residues of a section of the complex plane
the correlation function can be reconstructed as a sum of decaying
exponentials. The main benefit of this method is that it does not
require much knowledge about the function to be fit. On the downside,
if many poles are around the origin, it might require the sample points
to be used for the fit to be a large dense range which makes this
algorithm consume a lot of RAM (it will also be slow if asking for many
exponents)
Parameters
----------
wlist : array_like
The frequency range on which to perform the fit. With this method
typically logarithmic spacing works best.
tol : optional, int
Relative tolerance used to stop the algorithm, if an iteration
contribution is less than the tolerance the fit is stopped.
Nmax : optional, int
The maximum number of exponents desired. Corresponds to the
maximum number of iterations for the AAA algorithm
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
-------
approx_env : :class:`ExponentialBosonicEnvironment`
The approximated environment with multi-exponential correlation
function.
"""
method: Literal['aaa'],
wlist: ArrayLike,
tol: float,
N_max: int,
combine: bool,
tag: Any):
"""
Generates an approximation to this environment by fitting its power
spectrum using the AAA algorithm. The function is fit to a rational
polynomial of the form
.. math::
S(\\omega)= 2 \\Re \\left(\\sum_{k} \frac{c_{k}}{\nu_{k}-i \\omega}
\right)
By isolating the poles and residues of a section of the complex plane
the correlation function can be reconstructed as a sum of decaying
exponentials. The main benefit of this method is that it does not
require much knowledge about the function to be fit. On the downside,
if many poles are around the origin, it might require the sample points
to be used for the fit to be a large dense range which makes this
algorithm consume a lot of RAM (it will also be slow if asking for many
exponents)
Parameters
----------
wlist : array_like
The frequency range on which to perform the fit. With this method
typically logarithmic spacing works best.
tol : optional, int
Relative tolerance used to stop the algorithm, if an iteration
contribution is less than the tolerance the fit is stopped.
Nmax : optional, int
The maximum number of exponents desired. Corresponds to the
maximum number of iterations for the AAA algorithm
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
-------
approx_env : :class:`ExponentialBosonicEnvironment`
The approximated environment with multi-exponential correlation
function.
"""
# --- fitting
def approximate(self,method:str,*args,**kwargs):

def approximate(self, method: str, *args, **kwargs):
dispatch = {
"corr_lsq": self._approx_by_cf_fit,
"spec_lsq": self._approx_by_sd_fit,
Expand All @@ -740,10 +743,10 @@ def approximate(self,method:str,*args,**kwargs):
"esprit": self._approx_by_esprit,
"aaa": self._approx_by_aaa,
}

if method not in dispatch:
raise ValueError(f"Unsupported method: {method}")

return dispatch[method](*args, **kwargs)

def _approx_by_cf_fit(
Expand All @@ -761,7 +764,6 @@ def _approx_by_cf_fit(
combine: bool = True,
tag: Any = None,
) -> tuple[ExponentialBosonicEnvironment, dict[str, Any]]:


# Process arguments
if tag is None and self.tag is not None:
Expand Down Expand Up @@ -862,7 +864,7 @@ def _approx_by_sd_fit(
combine: bool = True,
tag: Any = None,
) -> tuple[ExponentialBosonicEnvironment, dict[str, Any]]:

# Process arguments
if tag is None and self.tag is not None:
tag = (self.tag, "SD Fit")
Expand Down Expand Up @@ -949,7 +951,7 @@ def _approx_by_aaa(
cls = ExponentialBosonicEnvironment(
ck_real=ckAR, vk_real=vkAR, ck_imag=ckAI,
vk_imag=vkAR, T=self.T, combine=combine, tag=tag)
return cls,{}
return cls, {}

def _approx_by_mp(
self,
Expand Down

0 comments on commit 8966d53

Please sign in to comment.