From 2811cc1fadb094c97a7837ec23179651466c5568 Mon Sep 17 00:00:00 2001 From: Gerardo Suarez Date: Tue, 7 Jan 2025 23:41:46 +0100 Subject: [PATCH] minor formatting, removed code aimed at arbitrary fermionic for now --- qutip/core/environment.py | 209 +++++++++++++++++++++++++++----------- qutip/utilities.py | 75 +++++--------- 2 files changed, 176 insertions(+), 108 deletions(-) diff --git a/qutip/core/environment.py b/qutip/core/environment.py index 33c26d21dc..c7c80fdb6d 100644 --- a/qutip/core/environment.py +++ b/qutip/core/environment.py @@ -776,10 +776,53 @@ def approx_by_aaa( combine: bool = True, tag: Any = None, ) -> tuple[ExponentialBosonicEnvironment, dict[str, 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 to be 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. + """ _, pol, res, _, _ = aaa(self.power_spectrum, wlist, tol=tol, max_iter=N_max*2) + # The *2 is there because half the poles will be filtered out mask = np.imag(pol) < 0 new_pols, new_res = pol[mask], res[mask] @@ -790,7 +833,6 @@ def approx_by_aaa( ckAR = [] vkAR = [] ckAI = [] - vkAI = [] for term in range(len(vk)): a, b, c, d = np.real(ck[term]), -np.real(vk[term]), - \ @@ -812,23 +854,63 @@ def approx_by_mp( combine: bool = True, tag: Any = None, ) -> tuple[ExponentialBosonicEnvironment, dict[str, Any]]: - # amp, phases = matrix_pencil(self.correlation_function(tlist), N) - # ckAR = amp.real - # ckAI = amp.imag - # vk = -((len(tlist)-1)/tlist[-1]) * \ - # (np.log(np.abs(phases))+1j*np.angle(phases)) - amp, phases = matrix_pencil(self.correlation_function(tlist).real, Nr) - amp2, phases2 = matrix_pencil(self.correlation_function(tlist).imag, Ni) - ckAR = amp - ckAI = amp2 - vkAR = -((len(tlist)-1)/tlist[-1]) * \ - (np.log(np.abs(phases))+1j*np.angle(phases)) - vkAI = -((len(tlist)-1)/tlist[-1]) * \ - (np.log(np.abs(phases2))+1j*np.angle(phases2)) + """ + Generates an approximation to this environment by fitting its correlation function + using the Matrix pencil method based on the prony polynomial. + Parameters + ---------- + tlist : array_like + The time range on which to perform the fit. + Nr : optional, int + The number of exponents desired to describe the imaginary part of the correlation function. It defaults to 3 + Nr : optional, int + The number of exponents desired to describe the real part of the correlation function. It defaults to 3 + 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. - cls = ExponentialBosonicEnvironment( - ck_real=ckAR, vk_real=vkAR, ck_imag=ckAI, - vk_imag=vkAI, T=self.T, combine=combine, tag=tag) + Returns + ------- + approx_env : :class:`ExponentialBosonicEnvironment` + The approximated environment with multi-exponential correlation + function. + """ + if Ni != Nr: + amp, phases = matrix_pencil( + self.correlation_function(tlist).real, Nr) + amp2, phases2 = matrix_pencil( + self.correlation_function(tlist).imag, Ni) + ckAR = amp + ckAI = amp2 + vkAR = -((len(tlist)-1)/tlist[-1]) * \ + (np.log(np.abs(phases))+1j*np.angle(phases)) + vkAI = -((len(tlist)-1)/tlist[-1]) * \ + (np.log(np.abs(phases2))+1j*np.angle(phases2)) + cls = ExponentialBosonicEnvironment( + ck_real=ckAR, vk_real=vkAR, ck_imag=ckAI, + vk_imag=vkAI, T=self.T, combine=combine, tag=tag) + else: + vkAR = [] + ckAI = [] + ckAR = [] + + amp, phases = matrix_pencil(self.correlation_function(tlist), Nr) + ck = amp + vk = -((len(tlist)-1)/tlist[-1]) * \ + (np.log(np.abs(phases))+1j*np.angle(phases)) + # for term in range(len(vk)): + # a, b, c, d = np.real(ck[term]), -np.real(vk[term]), - \ + # np.imag(vk[term]), np.imag(ck[term]) + # ckAR.extend([(a + 1j * d) / 2]) + # vkAR.extend([-b - 1j * c]) + # ckAI.extend([-1j * (a + 1j * d) / 2]) + cls = ExponentialBosonicEnvironment( + ck_real=ck.real, vk_real=vk, ck_imag=ck.imag, + vk_imag=vk, T=self.T, combine=combine, tag=tag) return cls, (amp, phases) def approx_by_prony( @@ -839,6 +921,31 @@ def approx_by_prony( combine: bool = True, tag: Any = None, ) -> tuple[ExponentialBosonicEnvironment, dict[str, Any]]: + """ + Generates an approximation to this environment by fitting its correlation function + using the prony method. + Parameters + ---------- + tlist : array_like + The time range on which to perform the fit. + Nr : optional, int + The number of exponents desired to describe the imaginary part of the correlation function. It defaults to 3 + Nr : optional, int + The number of exponents desired to describe the real part of the correlation function. It defaults to 3 + 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. + """ # amp, phases = prony(self.correlation_function(tlist), N) # ckAR = amp.real # ckAI = amp.imag @@ -873,6 +980,31 @@ def approx_by_esprit( combine: bool = True, tag: Any = None, ) -> tuple[ExponentialBosonicEnvironment, dict[str, Any]]: + """ + Generates an approximation to this environment by fitting its correlation function + using the ESPRIT method based on the prony polynomial. + Parameters + ---------- + tlist : array_like + The time range on which to perform the fit. + Nr : optional, int + The number of exponents desired to describe the imaginary part of the correlation function. It defaults to 3 + Nr : optional, int + The number of exponents desired to describe the real part of the correlation function. It defaults to 3 + 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. + """ # amp, phases = esprit(self.correlation_function(tlist), N) # ckAR = amp.real # ckAI = amp.imag @@ -2364,49 +2496,6 @@ def from_spectral_density(cls, **kwargs) -> FermionicEnvironment: "currently not implemented.") -class _FermionicEnvironment_fromSD(BosonicEnvironment): - def __init__(self, J, wlist, wMax, T, tag, args): - super().__init__(T, tag) - self._sd = _real_interpolation(J, wlist, 'spectral density', args) - if wlist is not None: - self.wMax = max(np.abs(wlist[0]), np.abs(wlist[-1])) - else: - self.wMax = wMax - - def correlation_function_plus(self, t, *, eps=1e-10): - if self.T is None: - raise ValueError('The temperature must be specified for this ' - 'operation.') - if self.wMax is None: - raise ValueError('The support of the spectral density (wMax) ' - 'must be specified for this operation.') - return self._cf_from_ps_plus(t, self.wMax, eps=eps) - - def correlation_function_minus(self, t, *, eps=1e-10): - if self.T is None: - raise ValueError('The temperature must be specified for this ' - 'operation.') - if self.wMax is None: - raise ValueError('The support of the spectral density (wMax) ' - 'must be specified for this operation.') - return self._cf_from_ps_minus(t, self.wMax, eps=eps) - - def spectral_density(self, w): - w = np.asarray(w, dtype=float) - - result = np.zeros_like(w) - positive_mask = (w > 0) - result[positive_mask] = self._sd(w[positive_mask]) - - return result.item() if w.ndim == 0 else result - - def power_spectrum_plus(self, w, *, eps=1e-10): - return self._ps_plus_from_sd(w, eps) - - def power_spectrum_minus(self, w, *, eps=1e-10): - return self._ps_minus_from_sd(w, eps) - - class LorentzianEnvironment(FermionicEnvironment): r""" Describes a Lorentzian fermionic environment with the spectral density diff --git a/qutip/utilities.py b/qutip/utilities.py index affef84307..6a542c4a19 100644 --- a/qutip/utilities.py +++ b/qutip/utilities.py @@ -56,45 +56,6 @@ def n_thermal(w, w_th): return result.item() if w.ndim == 0 else result -def fermi_dirac(w, T, mu): - """ - Fermi Dirac distribution - - Parameters - ---------- - - w : float or ndarray - Frequency of the oscillator. - - T : float - The temperature in units of frequency (or the same units as `w`). - mu: float - Chemical potential - - Returns - ------- - - n_avg : float or array - - Return the number of average fermions in thermal equilibrium for a - with the given energy and temperature and chemical potential. - - - """ - - w = np.array(w, dtype=float) - result = np.zeros_like(w) - - if T <= 0: - result[w < 0] = -1 - return result.item() if w.ndim == 0 else result - - non_zero = w != 0 - result[non_zero] = 1 / (np.exp(w[non_zero] / T) + 1) - - return result.item() if w.ndim == 0 else result - - def _factorial_prod(N, arr): arr[:int(N)] += 1 @@ -572,7 +533,7 @@ def aaa(func, z, tol=1e-13, max_iter=100): Computes a rational approximation of the function according to the AAA algorithm as explained in https://doi.org/10.1137/16M1106122 . This implementation is a python adaptation of the matlab version in that paper - + NOTE: I am not sure if this is necessary anymore as scipy 1.15 includes AAA Parameters: ----------- func : callable or np.ndarray @@ -760,7 +721,7 @@ def prz(support_points, values, weights): formula for the residue -https://math.stackexchange.com/questions/2202129/residue-for-quotient-of-functions + https://math.stackexchange.com/questions/2202129/residue-for-quotient-of-functions Parameters: @@ -800,12 +761,10 @@ def prz(support_points, values, weights): return pol, res, zeros -# Matrix Pencil Method (chosen over Prony because of resilience to Noise) - - def matrix_pencil(C: np.ndarray, n: int) -> tuple: """ Estimate amplitudes and frequencies using the Matrix Pencil Method. + Based on the description in https://doi.org/10.1093/imanum/drab108 Args: signal (np.ndarray): The input signal (1D complex array). @@ -830,10 +789,20 @@ def matrix_pencil(C: np.ndarray, n: int) -> tuple: return amplitudes, phases -# Prony +def prony(signal: np.ndarray, n): + """ + Estimate amplitudes and frequencies using the prony Method. + Based on the description in https://doi.org/10.1093/imanum/drab108 + Args: + signal (np.ndarray): The input signal (1D complex array). + n (int): The number of modes to estimate (rank of the signal). -def prony(signal: np.ndarray, n): + Returns: + tuple: A tuple containing: + - amplitudes (np.ndarray): The estimated amplitudes. + - phases (np.ndarray): The estimated complex exponential frequencies. + """ num_freqs = n hankel0 = hankel(c=signal[:num_freqs], r=signal[num_freqs - 1: -1]) hankel1 = hankel(c=signal[1: num_freqs + 1], r=signal[num_freqs:]) @@ -850,10 +819,20 @@ def prony(signal: np.ndarray, n): return np.array(amplitudes), np.array(phases) -# ESPRIT +def esprit(C: np.ndarray, n: int) -> tuple: + """ + Estimate amplitudes and frequencies using the ESPRIT Method. + Based on the description in https://doi.org/10.1093/imanum/drab108 + Args: + signal (np.ndarray): The input signal (1D complex array). + n (int): The number of modes to estimate (rank of the signal). -def esprit(C: np.ndarray, n: int) -> tuple: + Returns: + tuple: A tuple containing: + - amplitudes (np.ndarray): The estimated amplitudes. + - phases (np.ndarray): The estimated complex exponential frequencies. + """ # Step 1: Create the Hankel matrices num_freqs = len(C)-n hankel0 = hankel(c=C[:num_freqs], r=C[num_freqs - 1: -1])