Skip to content

Commit

Permalink
minor formatting, removed code aimed at arbitrary fermionic for now
Browse files Browse the repository at this point in the history
  • Loading branch information
gsuarezr committed Jan 7, 2025
1 parent 4c82b42 commit 2811cc1
Show file tree
Hide file tree
Showing 2 changed files with 176 additions and 108 deletions.
209 changes: 149 additions & 60 deletions qutip/core/environment.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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]), - \
Expand All @@ -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(
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
75 changes: 27 additions & 48 deletions qutip/utilities.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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).
Expand All @@ -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:])
Expand All @@ -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])
Expand Down

0 comments on commit 2811cc1

Please sign in to comment.