170 lines
7.2 KiB
Cython
170 lines
7.2 KiB
Cython
from .qpms_cdefs cimport *
|
|
from .qpms_c cimport ScatteringSystem, _ScatteringSystemAtOmega, _ScatteringSystemAtOmegaK
|
|
from .cycommon import EwaldPart, BesselType
|
|
from libc.stdlib cimport malloc, free, calloc
|
|
import numpy as np
|
|
from cython.parallel cimport prange, parallel
|
|
from cython cimport boundscheck, wraparound
|
|
|
|
cdef extern from "ewald.h":
|
|
void ewald3_2_sigma_long_Delta(cdouble *target, double *err, int maxn, cdouble x, int xbranch, cdouble z)
|
|
void ewald3_2_sigma_long_Delta_series(cdouble *target, double *err, int maxn, cdouble x, int xbranch, cdouble z)
|
|
void ewald3_2_sigma_long_Delta_recurrent(cdouble *target, double *err, int maxn, cdouble x, int xbranch, cdouble z, bint bigimz)
|
|
int complex_gamma_inc_e(double a, cdouble x, int xbranch, qpms_csf_result *result)
|
|
extern int ewald_factor_ipow_l
|
|
extern int ewald_factor_ipow_m
|
|
|
|
cdef extern from "translations_dbg.h":
|
|
int qpms_trans_calculator_test_sswf(const qpms_trans_calculator *c,
|
|
cdouble *dest, const csph_t kdlj, const qpms_bessel_t J) nogil
|
|
int qpms_trans_calculator_test_sswf_lc3p(const qpms_trans_calculator *c,
|
|
cdouble *dest, const cdouble k, const cart3_t r, const qpms_bessel_t J) nogil
|
|
int qpms_trans_calculator_test_sswf_e32(const qpms_trans_calculator *c,
|
|
cdouble * const sigmas_total, double * serr_total,
|
|
const double eta, const cdouble k,
|
|
const cart2_t b1, const cart2_t b2,
|
|
const cart2_t beta,
|
|
const cart3_t particle_shift,
|
|
double maxR, double maxK,
|
|
const qpms_ewald_part parts) nogil
|
|
|
|
cdef extern from "scatsystem_dbg.h":
|
|
int qpms_scatsyswk_test_sswf_basis_e(cdouble *target, const qpms_scatsys_at_omega_k_t *sswk,
|
|
qpms_bessel_t typ, cart3_t evalpoint, qpms_ewald_part parts) nogil
|
|
|
|
def e32_Delta(int maxn, cdouble x, cdouble z, int xbranch = 0, get_err=True, method='auto'):
|
|
cdef np.ndarray[double, ndim=1] err_np
|
|
cdef double[::1] err_view
|
|
cdef np.ndarray[np.complex_t, ndim=1] target_np = np.empty((maxn+1,), dtype=complex, order='C')
|
|
cdef cdouble[::1] target_view = target_np
|
|
if get_err:
|
|
err_np = np.empty((maxn+1,), order='C')
|
|
err_view = err_np
|
|
if method == 'recurrent':
|
|
ewald3_2_sigma_long_Delta_recurrent(&target_view[0], &err_view[0] if get_err else NULL, maxn, x, xbranch, z, False)
|
|
if method == 'recurrent_bigimz':
|
|
ewald3_2_sigma_long_Delta_recurrent(&target_view[0], &err_view[0] if get_err else NULL, maxn, x, xbranch, z, True)
|
|
elif method == 'series':
|
|
ewald3_2_sigma_long_Delta_series(&target_view[0], &err_view[0] if get_err else NULL, maxn, x, xbranch, z)
|
|
else:
|
|
ewald3_2_sigma_long_Delta(&target_view[0], &err_view[0] if get_err else NULL, maxn, x, xbranch, z)
|
|
if get_err:
|
|
return target_np, err_np
|
|
else:
|
|
return target_np
|
|
|
|
def scatsyswk_sswf_basis(_ScatteringSystemAtOmegaK sswk, evalpos, btyp=QPMS_HANKEL_PLUS, ewaldparts=EwaldPart.FULL):
|
|
"""Evaluate Ewald-summed scalar SWFs that are used to compose the tranlation operators.
|
|
|
|
Parameters
|
|
----------
|
|
evalpos: array_like of floats and shape (..., 3)
|
|
Evaluation points in cartesian coordinates.
|
|
|
|
Returns
|
|
-------
|
|
ndarray of complex, with the shape `evalpos.shape[:-1] + (particle_count, pnelem)`
|
|
"""
|
|
if(btyp != QPMS_HANKEL_PLUS):
|
|
raise NotImplementedError("Only first kind Bessel function-based fields are supported")
|
|
cdef qpms_bessel_t btyp_c = BesselType(btyp)
|
|
cdef qpms_ewald_part ewaldparts_c = EwaldPart(ewaldparts)
|
|
evalpos = np.array(evalpos, dtype=float, copy=False)
|
|
if evalpos.shape[-1] != 3:
|
|
raise ValueError("Last dimension of evalpos has to be 3")
|
|
|
|
cdef ScatteringSystem ss = sswk.ssw_pyref.ss_pyref
|
|
cdef qpms_scatsys_t *ss_c = ss.s
|
|
cdef qpms_scatsys_at_omega_k_t *sswk_c = sswk.rawpointer()
|
|
cdef const qpms_trans_calculator *c = ss_c[0].c
|
|
cdef qpms_l_t sswf_lMax = 2*c[0].lMax+1
|
|
cdef qpms_y_t pnelem = qpms_lMax2nelem_sc(sswf_lMax)
|
|
cdef qpms_ss_pi_t particle_count = sswk.ssw_pyref.ss_pyref.particle_count
|
|
|
|
cdef np.ndarray[double,ndim=2] evalpos_a = evalpos.reshape(-1,3)
|
|
cdef np.ndarray[complex, ndim=3] results = np.empty((evalpos_a.shape[0], particle_count, pnelem), dtype=complex)
|
|
cdef cdouble *res
|
|
cdef cart3_t pos
|
|
cdef Py_ssize_t i, pi, j
|
|
with nogil, wraparound(False), parallel():
|
|
res = <cdouble *> malloc(particle_count * pnelem * sizeof(cdouble))
|
|
for i in prange(evalpos_a.shape[0]):
|
|
pos.x = evalpos_a[i,0]
|
|
pos.y = evalpos_a[i,1]
|
|
pos.z = evalpos_a[i,2]
|
|
qpms_scatsyswk_test_sswf_basis_e(res, sswk_c, btyp_c, pos, ewaldparts_c)
|
|
for pi in range(particle_count):
|
|
for j in range(pnelem):
|
|
results[i,pi,j] = res[pi * pnelem + j]
|
|
free(res)
|
|
return results.reshape(evalpos.shape[:-1] + (particle_count, pnelem))
|
|
|
|
def scatsysw_eval_sswf(_ScatteringSystemAtOmega ssw, evalpos, btyp=QPMS_HANKEL_PLUS):
|
|
"""Evaluate sswfs by internal qpms_trans_calculator at given locations.
|
|
|
|
Parameters
|
|
----------
|
|
ssw: _ScatteringSystemAtOmega
|
|
We take the wavenumber and qpms_trans_calculator from this, nothing else.
|
|
(In a way, this is an overkill, but there is no up-to-date standalone cython
|
|
wrapper for qpms_trans_calculator.)
|
|
evalpos: array_like of floats and shape (..., 3)
|
|
Evaluation points in cartesian coordinates.
|
|
btyp: BesselType
|
|
Kind of Bessel functions to be used
|
|
|
|
Returns
|
|
-------
|
|
ndarray of complex, with the shape `evalpos.shape[:-1] + (pnelem,)`
|
|
where `pnelem` is TODO
|
|
"""
|
|
cdef qpms_bessel_t btyp_c = BesselType(btyp)
|
|
evalpos = np.array(evalpos, dtype=float, copy=False)
|
|
if evalpos.shape[-1] != 3:
|
|
raise ValueError("Last dimension of evalpos has to be 3")
|
|
|
|
cdef ScatteringSystem ss = ssw.ss_pyref
|
|
cdef qpms_scatsys_t *ss_c = ss.s
|
|
cdef const qpms_trans_calculator *c = ss_c[0].c
|
|
cdef qpms_l_t sswf_lMax = 2*c[0].lMax+1
|
|
cdef qpms_y_t pnelem = qpms_lMax2nelem_sc(sswf_lMax)
|
|
|
|
cdef cdouble k = ssw.wavenumber
|
|
|
|
cdef np.ndarray[double,ndim=2] evalpos_a = evalpos.reshape(-1,3)
|
|
cdef np.ndarray[complex,ndim=2] results = np.empty((evalpos_a.shape[0], pnelem), dtype=complex)
|
|
cdef cdouble *res
|
|
cdef cart3_t pos
|
|
cdef csph_t kdlj
|
|
cdef Py_ssize_t i, j
|
|
with nogil, wraparound(False), parallel():
|
|
res = <cdouble *> malloc(pnelem * sizeof(cdouble))
|
|
for i in prange(evalpos_a.shape[0]):
|
|
pos.x = evalpos_a[i,0]
|
|
pos.y = evalpos_a[i,1]
|
|
pos.z = evalpos_a[i,2]
|
|
kdlj = cart2csph(pos)
|
|
kdlj.r *= k
|
|
qpms_trans_calculator_test_sswf(c, res, kdlj, btyp_c)
|
|
for j in range(pnelem):
|
|
results[i,j] = res[j]
|
|
free(res)
|
|
return results.reshape(evalpos.shape[:-1] + (pnelem,))
|
|
|
|
def gamma_inc(double a, cdouble x, int xbranch=0):
|
|
cdef qpms_csf_result res
|
|
complex_gamma_inc_e(a, x, xbranch, &res)
|
|
return res.val
|
|
|
|
def ipow_l_factor(n = None):
|
|
global ewald_factor_ipow_l
|
|
if n is not None:
|
|
ewald_factor_ipow_l = n
|
|
return ewald_factor_ipow_l
|
|
|
|
def ipow_m_factor(n = None):
|
|
global ewald_factor_ipow_m
|
|
if n is not None:
|
|
ewald_factor_ipow_m = n
|
|
return ewald_factor_ipow_m
|