diff --git a/oqupy/bath_correlations.py b/oqupy/bath_correlations.py index 3bc40b4..55b3f8a 100644 --- a/oqupy/bath_correlations.py +++ b/oqupy/bath_correlations.py @@ -13,7 +13,7 @@ Module for environment correlations. """ -from typing import Callable, Optional, Text +from typing import Callable, Optional, Union, Text from typing import Any as ArrayLike from functools import lru_cache @@ -309,6 +309,30 @@ def _gaussian_cutoff(omega: ArrayLike, omega_c: float) -> ArrayLike: # --- the spectral density classes -------------------------------------------- +def _weighted_integral( + re_integrand: Callable[[float], float], + im_integrand: Callable[[float], float], + w: float, + a: Optional[float] = 0.0, + b: Optional[float] = 1.0, + epsrel: Optional[float] = INTEGRATE_EPSREL, + limit: Optional[int] = SUBDIV_LIMIT) -> complex: + re_int = integrate.quad(re_integrand, + a=a, + b=b, + epsrel=epsrel, + limit=limit, + weight='cos', + wvar=w)[0] + im_int = integrate.quad(im_integrand, + a=a, + b=b, + epsrel=epsrel, + limit=limit, + weight='sin', + wvar=w)[0] + return re_int + 1j * im_int + def _complex_integral( integrand: Callable[[float], complex], a: Optional[float] = 0.0, @@ -362,6 +386,25 @@ class CustomSD(BaseCorrelations): ``'gaussian'``} temperature: float The environment's temperature. + alt_integrator: bool + Whether or not to use a sine/cosine weighted version of the + integrals that could be better suited if you're encountering + numerical difficulties, especially for long memory times. + The integral computed in `eta_function` is cut in two: one + handling the divergent part near 0 using the default integrator + and the other taking care of the fast oscillations at angular + frequency :math:`\tau` using an appropriate integration scheme. + The cutoff point can be modified with `num_oscillations`. + num_oscillations: int, float, callable + When using `alt_integrator`, specifies how many oscillations to keep + for the non-weighted part of the integral of `eta_function`. This should + be either an integer or specified as a function of :math:`\tau`. To help + choose: the higher this number is, the more the default integrator will + struggle to integrate the non-weighted part. The lower it is, the more + the weighted integrators will struggle to integrate the near divergent + part close to 0. A constant value is a good choice for most use cases. + Otherwise, it is possible to input a function to target possible issues + at specific memory times :math:`\tau`. name: str An optional name for the correlations. description: str @@ -374,6 +417,9 @@ def __init__( cutoff: float, cutoff_type: Optional[Text] = 'exponential', temperature: Optional[float] = 0.0, + alt_integrator: Optional[bool] = False, + num_oscillations: Optional[ + Union[int, float, Callable[[float], float]]] = None, name: Optional[Text] = None, description: Optional[Text] = None) -> None: """Create a CustomFunctionSD (spectral density) object. """ @@ -409,6 +455,28 @@ def __init__( tmp_temperature)) self.temperature = tmp_temperature + # input check for alt_integrator + assert isinstance(alt_integrator, bool) + self._alt_integrator = alt_integrator + + # create the first_cutoff function and input check for num_oscillations + if num_oscillations is None: + num_oscillations = 3 + if isinstance(num_oscillations, (int, float)) and num_oscillations > 0: + tmp_n = num_oscillations + # This form is a step version of the first cutoff which prevents + # useless computations of _truncated_eta + self.first_cutoff = lambda t: 2*np.pi*tmp_n/(1+1/tmp_n)**np.floor( + np.log(t)/np.log(1+1/tmp_n)) + else: + try: + float(num_oscillations(1.0)) + except Exception as e: + raise AssertionError("num_oscillations must be a positive" \ + "float, int or a function returning floats or ints.") from e + self.first_cutoff = lambda t: max(min(2*np.pi*num_oscillations(t)/t, + self.cutoff), 0.0) + self._cutoff_function = \ lambda omega: CUTOFF_DICT[self.cutoff_type](omega, self.cutoff) self._spectral_density = \ @@ -515,23 +583,75 @@ def integrand(w): integral = integral.real return integral + @lru_cache(maxsize=2 ** 10, typed=False) + def _truncated_eta(self, + a: float, + epsrel: float, + limit: int) -> complex: + """Computes the parts of `eta_function` that are tau independant when + `alt_integrator` is set to True. + """ + if self.temperature == 0.0: + def re_integrand(w): + return self._spectral_density(w) / w ** 2 + else: + def re_integrand(w): + # this is to stop overflow + if np.exp(-w / self.temperature) > np.finfo(float).eps: + inte = self._spectral_density(w) / w ** 2 \ + * 1/np.tanh(w / (2*self.temperature)) + else: + inte = self._spectral_density(w) / w ** 2 + return inte + + def im_integrand(w): + return self._spectral_density(w)/w + + im_constant = integrate.quad(im_integrand, + a=a, + b=self.cutoff, + epsrel=epsrel, + limit=limit)[0] + + re_constant = integrate.quad(re_integrand, + a=a, + b=self.cutoff, + epsrel=epsrel, + limit=limit)[0] + + if self.cutoff_type != "hard": + im_constant += integrate.quad(im_integrand, + a=self.cutoff, + b=np.inf, + epsrel=epsrel, + limit=limit)[0] + + re_constant += integrate.quad(re_integrand, + a=self.cutoff, + b=np.inf, + epsrel=epsrel, + limit=limit)[0] + return re_constant, im_constant + @lru_cache(maxsize=2 ** 10, typed=False) def eta_function( self, tau: ArrayLike, epsrel: Optional[float] = INTEGRATE_EPSREL, subdiv_limit: Optional[int] = SUBDIV_LIMIT, + alt_integrator: Optional[Union[bool, None]] = None, matsubara: Optional[bool] = False) -> ArrayLike: r""" - Auto-correlation function associated to the spectral density at the - given temperature :math:`T` + :math:`\eta` function associated to the spectral density at the + given temperature :math:`T` as given in [Strathearn2017] .. math:: - C(\tau) = \int_0^{\infty} J(\omega) \ - \left[ \cos(\omega \tau) \ + \eta(\tau) = \int_0^{\infty} \frac{J(\omega)}{\omega^2} \ + \left[ ( \cos(\omega \tau) - 1 ) \ \coth\left( \frac{\omega}{2 T}\right) \ - - i \sin(\omega \tau) \right] \mathrm{d}\omega . + - i ( \sin(\omega \tau) - \omega \tau ) \right] \ + \mathrm{d}\omega . with time difference `tau` :math:`\tau`. @@ -543,12 +663,23 @@ def eta_function( Relative error tolerance. subdiv_limit: int Maximal number of interval subdivisions for numerical integration. - + alt_integrator: Union[bool, None] + Whether or not to use a sine/cosine weighted version of the + integrals that could be better suited if you're encountering + numerical difficulties, especially for long memory times. If the + value is `None`, the value of `alt_integrator` set during + initialization of the `CustomSD` object is used instead. Returns ------- correlation : ndarray - The auto-correlation function :math:`C(\tau)` at time :math:`\tau`. + The function :math:`\eta(\tau)` at time :math:`\tau`. """ + if alt_integrator is None: + alt_integrator = self._alt_integrator + check_true(not (matsubara is True and alt_integrator is True), + "Can't use weighted integrator with matsubara") + check_true(isinstance(alt_integrator, bool), + "alt_integrator has to be either True, False or None") # real and imaginary part of the integrand if matsubara: tau = -1j * tau @@ -574,6 +705,50 @@ def integrand(w): * (np.exp(-1j * w * tau) - 1 + 1j * w * tau) return inte + # Alternative computation with weighted integrators + if alt_integrator and tau*self.cutoff > 1.: + im_integrand = lambda w: - self._spectral_density(w) / w ** 2 + if self.temperature == 0.0: + re_integrand = lambda w: self._spectral_density(w) / w ** 2 + else: + def re_integrand(w): + # this is to stop overflow + if np.exp(-w / self.temperature) > np.finfo(float).eps: + inte = self._spectral_density(w) / w ** 2 \ + * 1/np.tanh(w / (2*self.temperature)) + else: + inte = self._spectral_density(w) / w ** 2 + return inte + + omega_tilde = self.first_cutoff(tau) + + integral = _complex_integral(integrand, + a=0.0, + b=omega_tilde, + epsrel=epsrel, + limit=subdiv_limit) + + integral += _weighted_integral(re_integrand, im_integrand, + w=tau, + a=omega_tilde, + b=self.cutoff, + epsrel=epsrel, + limit=subdiv_limit) + + if self.cutoff_type != "hard": + integral += _weighted_integral(re_integrand, im_integrand, + w=tau, + a=self.cutoff, + b=np.inf, + epsrel=epsrel, + limit=subdiv_limit) + + re_constant, im_constant = self._truncated_eta(a=omega_tilde, + epsrel=epsrel, + limit=subdiv_limit) + return - (integral + 1.j*tau*im_constant - re_constant) + + # Default computation of eta_function when alt_integrator is False integral = _complex_integral(integrand, a=0.0, b=self.cutoff, @@ -598,6 +773,7 @@ def correlation_2d_integral( shape: Optional[Text] = 'square', epsrel: Optional[float] = INTEGRATE_EPSREL, subdiv_limit: Optional[int] = SUBDIV_LIMIT, + alt_integrator: Optional[Union[bool, None]] = None, matsubara: Optional[bool] = False) -> complex: r""" 2D integrals of the correlation function @@ -632,6 +808,10 @@ def correlation_2d_integral( Relative error tolerance. subdiv_limit: int Maximal number of interval subdivisions for numerical integration. + alt_integrator: Union[bool, None] + Whether or not to use a sine/cosine weighted version of the + integrals that could be better suited if you're encountering + numerical difficulties, especially for long times. Returns ------- @@ -642,6 +822,7 @@ def correlation_2d_integral( kwargs = { 'epsrel': epsrel, 'subdiv_limit': subdiv_limit, + 'alt_integrator': alt_integrator, 'matsubara': matsubara} if shape == 'upper-triangle': @@ -701,6 +882,25 @@ class PowerLawSD(CustomSD): ``'gaussian'``} temperature: float The environment's temperature. + alt_integrator: bool + Whether or not to use a sine/cosine weighted version of the + integrals that could be better suited if you're encountering + numerical difficulties, especially for long memory times. + The integral computed in `eta_function` is cut in two: one + handling the divergent part near 0 using the default integrator + and the other taking care of the fast oscillations at angular + frequency :math:`\tau` using an appropriate integration scheme. + The cutoff point can be modified with `num_oscillations`. + num_oscillations: int, float, callable + When using `alt_integrator`, specifies how many oscillations to keep + for the non-weighted part of the integral of `eta_function`. This should + be either an integer or specified as a function of :math:`\tau`. To help + choose: the higher this number is, the more the default integrator will + struggle to integrate the non-weighted part. The lower it is, the more + the weighted integrators will struggle to integrate the near divergent + part close to 0. A constant value is a good choice for most use cases. + Otherwise, it is possible to input a function to target possible issues + at specific memory times :math:`\tau`. name: str An optional name for the correlations. description: str @@ -714,6 +914,9 @@ def __init__( cutoff: float, cutoff_type: Text = 'exponential', temperature: Optional[float] = 0.0, + alt_integrator: Optional[bool] = False, + num_oscillations: Optional[ + Union[int, float, Callable[[float], float]]] = None, name: Optional[Text] = None, description: Optional[Text] = None) -> None: """Create a StandardSD (spectral density) object. """ @@ -747,6 +950,7 @@ def __init__( cutoff=cutoff, cutoff_type=cutoff_type, temperature=temperature, + alt_integrator=alt_integrator, name=name, description=description) diff --git a/tests/coverage/correlations_test.py b/tests/coverage/correlations_test.py index 6a67a53..01fa1e4 100644 --- a/tests/coverage/correlations_test.py +++ b/tests/coverage/correlations_test.py @@ -63,22 +63,25 @@ def test_custom_correlations_bad_input(): def test_custom_s_d(): for cutoff_type in ["hard", "exponential", "gaussian"]: for temperature in [0.0, 2.0]: - sd = CustomSD(square_function, - cutoff=2.0, - temperature=temperature, - cutoff_type=cutoff_type) - str(sd) - w = np.linspace(0, 8.0 * sd.cutoff, 10) - y = sd.spectral_density(w) - t = np.linspace(0, 4.0 / sd.cutoff, 10) - [sd.correlation(tt) for tt in t] - for shape in ["square", "upper-triangle"]: - sd.correlation_2d_integral(time_1=0.25, - delta=0.05, - shape=shape) - sd.correlation_2d_integral(time_1=0.25, - delta=0.05, - shape=shape) + for alt_integrator in [True, False, None]: + sd = CustomSD(square_function, + cutoff=2.0, + temperature=temperature, + cutoff_type=cutoff_type) + str(sd) + w = np.linspace(0, 8.0 * sd.cutoff, 10) + y = sd.spectral_density(w) + t = np.linspace(0, 4.0 / sd.cutoff, 10) + [sd.correlation(tt) for tt in t] + for shape in ["square", "upper-triangle"]: + sd.correlation_2d_integral(time_1=0.25, + delta=0.05, + shape=shape, + alt_integrator=alt_integrator) + sd.correlation_2d_integral(time_1=0.5, + delta=0.05, + shape=shape, + alt_integrator=alt_integrator) def test_matsubara_custom_s_d(): @@ -155,7 +158,9 @@ def test_matsubara_custom_s_d_bad_input(): delta=0.1, shape=shape, matsubara=True) - + correlations = CustomSD(square_function, cutoff=2.0, cutoff_type="gaussian", temperature=1.0, alt_integrator=True) + with pytest.raises(ValueError): + correlations.correlation_2d_integral(time_1=2, delta=0.1, shape='upper-triangle', matsubara=True, alt_integrator=True) @@ -171,6 +176,15 @@ def test_custom_s_d_bad_input(): with pytest.raises(AssertionError): CustomSD(square_function, cutoff=2.0, cutoff_type="gaussian", \ temperature="bla") + with pytest.raises(AssertionError): + CustomSD(square_function, cutoff=2.0, cutoff_type="gaussian", \ + temperature=0.0, alt_integrator="bla") + with pytest.raises(AssertionError): + CustomSD(square_function, cutoff=2.0, cutoff_type="gaussian", \ + temperature=0.0, num_oscillations=-1.0) + with pytest.raises(AssertionError): + CustomSD(square_function, cutoff=2.0, cutoff_type="gaussian", \ + temperature=0.0, num_oscillations=lambda t: "a") with pytest.raises(ValueError): CustomSD(square_function, cutoff=2.0, cutoff_type="hard", \ temperature=-2.0) diff --git a/tests/physics/bath_correlations_test.py b/tests/physics/bath_correlations_test.py new file mode 100644 index 0000000..12f07df --- /dev/null +++ b/tests/physics/bath_correlations_test.py @@ -0,0 +1,90 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Tests for the oqupy.bath_correlations module. +""" + +from itertools import product +import pytest +import warnings +import numpy as np +from scipy.integrate import IntegrationWarning +from scipy.special import loggamma + +from oqupy.bath_correlations import BaseCorrelations +from oqupy.bath_correlations import CustomCorrelations +from oqupy.bath_correlations import CustomSD +from oqupy.bath_correlations import PowerLawSD + +def exact_2d_correlations(alpha, omega_cutoff, temperature, k, dt, shape): + # Exact 2d correlations for an Ohmic spectral density and exponential cutoff + if temperature == 0.0: + eta_func = lambda k: np.log(1/omega_cutoff+1.j*k*dt)-1.j*omega_cutoff*dt*k + else: + eta_func = lambda k: -(1/2*np.log(1+(k*dt*omega_cutoff)**2)+ + loggamma(temperature*(1/omega_cutoff+1.j*k*dt))+ + loggamma(temperature*(1/omega_cutoff-1.j*k*dt))- + 2*loggamma(temperature/omega_cutoff)+ + 1j*(omega_cutoff*k*dt-np.arctan(omega_cutoff*k*dt))) + if shape == "square": + correl = eta_func(k-1)-2*eta_func(k)+eta_func(k+1) + elif shape == "upper-triangle": + correl = eta_func(k+1)-eta_func(k) + return 2*alpha*correl + +def test(): + kranges = [np.arange(1, 5), np.arange(395, 400)] + + omega_cutoff = 3.0 + temperature = 0.0 + alpha = 5e-2 + dt = 0.4 + + for shape, temperature, krange in product(["square", "upper-triangle"], + [0.0, 2.0], + kranges): + exact_correls_2d = exact_2d_correlations(alpha, + omega_cutoff, + temperature, + krange, + dt, + shape) + for alt_integrator in [False, True]: + correls = PowerLawSD(alpha=alpha, + zeta=1.0, + cutoff=omega_cutoff, + cutoff_type="exponential", + temperature=temperature, + alt_integrator=alt_integrator) + with warnings.catch_warnings(): + warnings.simplefilter(action="error", + category=IntegrationWarning) + try: + correls_2d = [correls.correlation_2d_integral(dt, k*dt, shape=shape) + for k in krange] + except IntegrationWarning: + pass + else: + try: + np.testing.assert_allclose(exact_correls_2d, + correls_2d, + rtol=1e-3, atol=1e-6) + except AssertionError: + pytest.fail("correlation_2d_integral should "\ + "either give precise results or at least "\ + "warn the user that there could be some "\ + f"numerical error, failed: shape={shape}, "\ + f"temperature={temperature}, ", + f"ks={krange}, "\ + f"alt_integrator={alt_integrator}") + +