eval_spectra.py 930 B

123456789101112131415161718192021222324252627282930313233
  1. from ast import Constant
  2. from numba import njit, prange, complex128, float64
  3. import numpy as np
  4. from scattnlay import mesomie
  5. params_count = 4
  6. @njit(complex128(float64, float64[:]),
  7. fastmath=True, cache=True)
  8. def lorentzian(omega, xvec):
  9. pc = params_count
  10. res = 0
  11. poles = len(xvec)//params_count
  12. factor = 1
  13. for i in range(poles):
  14. gamma = xvec[pc*i+0] / (factor**i)
  15. omega_n = xvec[pc*i+1] / (factor**i)
  16. f = (xvec[pc*i+2] + 1j*xvec[pc*i+3]) / (factor**i)
  17. if (np.abs(f) == 0):
  18. return res
  19. res = res + f / (omega * (omega + 1j*gamma) - omega_n**2)
  20. return res
  21. @njit(complex128[:](float64[:], float64[:]),
  22. fastmath=True, cache=True)
  23. def multi_lorentzian(omega, xvec):
  24. poles = len(xvec)//params_count
  25. val = np.zeros(omega.size, dtype=np.cdouble)
  26. for i in range(omega.size):
  27. val[i] = lorentzian(omega[i], xvec)
  28. return val