nmreval/nmreval/distributions/loggaussian.py

124 lines
3.7 KiB
Python

from multiprocessing import Pool, cpu_count
from itertools import product
import numpy as np
from ..lib.decorator import adjust_dims
try:
from scipy.integrate import simpson
except ImportError:
from scipy.integrate import simps as simpson
from scipy.integrate import quad
from .base import Distribution
__all__ = ['LogGaussian']
class LogGaussian(Distribution):
name = 'Log-Gaussian'
parameter = [r'\sigma']
bounds = [(0, 10)]
@staticmethod
def distribution(tau, tau0, sigma: float):
return np.exp(-0.5*(np.log(tau/tau0)/sigma)**2)/np.sqrt(2*np.pi)/sigma
@staticmethod
def correlation(t, tau0, sigma: float):
_t = np.atleast_1d(t)
_tau = np.atleast_1d(tau0)
pool = Pool(processes=min(cpu_count(), 4))
integration_ranges = [(omega_i, tau_j, sigma) for (omega_i, tau_j) in product(_t, _tau)]
with np.errstate(divide='ignore'):
res = np.array(pool.map(_integrate_process_time, integration_ranges))
ret_val = res.reshape((_t.shape[0], _tau.shape[0]))
return ret_val.squeeze()
@staticmethod
def susceptibility(omega, tau0, sigma: float):
_omega = np.atleast_1d(omega)
_tau = np.atleast_1d(tau0)
pool = Pool(processes=min(cpu_count(), 4))
integration_ranges = [(omega_i, tau_j, sigma) for (omega_i, tau_j) in product(_omega, _tau)]
with np.errstate(divide='ignore'):
res_real = np.array(pool.map(_integrate_process_imag, integration_ranges))
res_imag = np.array(pool.map(_integrate_process_real, integration_ranges))
ret_val = (res_real+1j*res_imag).reshape((_omega.shape[0], _tau.shape[0]))
return ret_val.squeeze()
@staticmethod
def specdens(omega, tau0, sigma):
_omega = np.atleast_1d(omega)
_tau = np.atleast_1d(tau0)
pool = Pool(processes=min(cpu_count(), 4))
integration_ranges = [(omega_i, tau_j, sigma) for (omega_i, tau_j) in product(_omega, _tau)]
with np.errstate(divide='ignore'):
res = np.array(pool.map(_integrate_process_imag, integration_ranges))
ret_val = res.reshape((_omega.shape[0], _tau.shape[0]))
ret_val /= _omega[:, None]
return ret_val.squeeze()
def mean(*args):
return args[0]*np.exp(args[1]**2 / 2)
def _integrate_process_imag(args):
omega_i, tau_j, sigma = args
area = quad(_integrand_freq_imag_high, 0, 50, args=(omega_i, tau_j, sigma))[0]
area += quad(_integrand_freq_imag_low, -50, 0, args=(omega_i, tau_j, sigma))[0]
return area
def _integrate_process_real(args):
omega_i, tau_j, sigma = args
area = quad(_integrand_freq_real_high, 0, 50, args=(omega_i, tau_j, sigma))[0]
area += quad(_integrand_freq_real_low, -50, 0, args=(omega_i, tau_j, sigma))[0]
return area
def _integrate_process_time(args):
omega_i, tau_j, sigma = args
return quad(_integrand_time, -50, 50, args=(omega_i, tau_j, sigma))[0]
def _integrand_time(u, t, tau, sigma):
uu = np.exp(u)
return LogGaussian.distribution(uu, tau, sigma) * np.exp(-t/uu)
# integrands
def _integrand_freq_imag_low(u, omega, tau, sigma):
# integrand
uu = np.exp(u)
return LogGaussian.distribution(uu, tau, sigma) * omega * uu / (1 + (omega*uu)**2)
def _integrand_freq_imag_high(u, omega, tau, sigma):
uu = np.exp(-u)
return LogGaussian.distribution(1/uu, tau, sigma) * omega * uu / (uu**2 + omega**2)
def _integrand_freq_real_low(u, omega, tau, sigma):
uu = np.exp(u)
return LogGaussian.distribution(uu, tau, sigma) / (1 + (omega*uu)**2)
def _integrand_freq_real_high(u, omega, tau, sigma):
uu = np.exp(-2*u)
return LogGaussian.distribution(np.exp(u), tau, sigma) * uu / (uu + omega**2)