69 lines
2.3 KiB
Cython
69 lines
2.3 KiB
Cython
# cythonized functions here
|
|
cimport numpy as np
|
|
import numpy as np
|
|
cimport cython
|
|
|
|
## Auxillary function for retrieving the "meshgrid-like" indices; inc. nmax
|
|
@cython.boundscheck(False)
|
|
def get_mn_y(int nmax):
|
|
"""
|
|
Auxillary function for retreiving the 'meshgrid-like' indices from the flat indexing;
|
|
inc. nmax.
|
|
('y to mn' conversion)
|
|
|
|
Parameters
|
|
----------
|
|
|
|
nmax : int
|
|
The maximum order to which the VSWFs / Legendre functions etc. will be evaluated.
|
|
|
|
Returns
|
|
-------
|
|
|
|
output : (m, n)
|
|
Tuple of two arrays of type np.array(shape=(nmax*nmax + 2*nmax), dtype=np.int),
|
|
where [(m[y],n[y]) for y in range(nmax*nmax + 2*nma)] covers all possible
|
|
integer pairs n >= 1, -n <= m <= n.
|
|
"""
|
|
cdef Py_ssize_t nelems = nmax * nmax + 2 * nmax
|
|
cdef np.ndarray[np.int_t,ndim=1] m_arr = np.empty([nelems], dtype=np.int)
|
|
cdef np.ndarray[np.int_t,ndim=1] n_arr = np.empty([nelems], dtype=np.int)
|
|
cdef Py_ssize_t i = 0
|
|
cdef np.int_t n, m
|
|
for n in range(1,nmax+1):
|
|
for m in range(-n,n+1):
|
|
m_arr[i] = m
|
|
n_arr[i] = n
|
|
i = i + 1
|
|
return (m_arr, n_arr)
|
|
|
|
def get_y_mn_unsigned(int nmax):
|
|
"""
|
|
Auxillary function for mapping 'unsigned m', n indices to the flat y-indexing.
|
|
For use with functions as scipy.special.lpmn, which have to be evaluated separately
|
|
for positive and negative m.
|
|
|
|
Parameters
|
|
----------
|
|
|
|
nmax : int
|
|
The maximum order to which the VSWFs / Legendre functions etc. will be evaluated.
|
|
|
|
output : (ymn_plus, ymn_minus)
|
|
Tuple of two arrays of shape (nmax+1,nmax+1), containing the flat y-indices corresponding
|
|
to the respective (m,n) and (-m,n). The elements for which |m| > n are set to -1.
|
|
(Therefore, the caller must not use those elements equal to -1.)
|
|
"""
|
|
cdef np.ndarray[np.intp_t, ndim=2] ymn_plus = np.full((nmax+1,nmax+1),-1, dtype=np.intp)
|
|
cdef np.ndarray[np.intp_t, ndim=2] ymn_minus = np.full((nmax+1,nmax+1),-1, dtype=np.intp)
|
|
cdef Py_ssize_t i = 0
|
|
cdef np.int_t n, m
|
|
for n in range(1,nmax+1):
|
|
for m in range(-n,0):
|
|
ymn_minus[-m,n] = i
|
|
i = i + 1
|
|
for m in range(0,n+1):
|
|
ymn_plus[m,n] = i
|
|
i = i + 1
|
|
return(ymn_plus, ymn_minus)
|