[broken] Jit všude možně (rozbil jsem to)
Former-commit-id: c88f554cee3aed543c657a27982139b3ae1ad0ce
This commit is contained in:
Marek Nečada
parent
26cabb1e87
commit
2b64a50244
|
|
@ -152,7 +152,7 @@ def sph_loccart_basis(sph, sphaxis=-1, cartaxis=None):
|
|||
out = np.concatenate((x,y,z),axis=cartaxis)
|
||||
return out
|
||||
|
||||
@jit
|
||||
@jit(u=False)
|
||||
def lpy(nmax, z):
|
||||
"""
|
||||
Associated legendre function and its derivatative at z in the 'y-indexing'.
|
||||
|
|
@ -249,6 +249,7 @@ _sph_zn = [_sph_zn_1,_sph_zn_2,_sph_zn_3,_sph_zn_4]
|
|||
|
||||
# computes bessel/hankel functions for orders from 0 up to n; drops
|
||||
# the derivatives which are also included in scipy.special.sph_jn/yn
|
||||
@jit
|
||||
def zJn(n, z, J=1):
|
||||
return _sph_zn[J-1](n=n,z=z)
|
||||
|
||||
|
|
@ -258,6 +259,7 @@ def zJn(n, z, J=1):
|
|||
|
||||
# FIXME: this can be expressed simply as:
|
||||
# $$ -\frac{1}{2}\sqrt{\frac{2n+1}{4\pi}n\left(n+1\right)}(\delta_{m,1}+\delta_{m,-1}) $$
|
||||
@jit(u=True)
|
||||
def π̃_zerolim(nmax): # seems OK
|
||||
"""
|
||||
lim_{θ→ 0-} π̃(cos θ)
|
||||
|
|
@ -275,6 +277,7 @@ def π̃_zerolim(nmax): # seems OK
|
|||
π̃_y = prenorm * π̃_y
|
||||
return π̃_y
|
||||
|
||||
@jit(u=True)
|
||||
def π̃_pilim(nmax): # Taky OK, jen to možná není kompatibilní se vzorečky z mathematiky
|
||||
"""
|
||||
lim_{θ→ π+} π̃(cos θ)
|
||||
|
|
@ -294,6 +297,7 @@ def π̃_pilim(nmax): # Taky OK, jen to možná není kompatibilní se vzorečky
|
|||
|
||||
# FIXME: this can be expressed simply as
|
||||
# $$ -\frac{1}{2}\sqrt{\frac{2n+1}{4\pi}n\left(n+1\right)}(\delta_{m,1}-\delta_{m,-1}) $$
|
||||
@jit(u=True)
|
||||
def τ̃_zerolim(nmax):
|
||||
"""
|
||||
lim_{θ→ 0-} τ̃(cos θ)
|
||||
|
|
@ -304,6 +308,7 @@ def τ̃_zerolim(nmax):
|
|||
p0[minus1mmask] = -p0[minus1mmask]
|
||||
return p0
|
||||
|
||||
@jit(u=True)
|
||||
def τ̃_pilim(nmax):
|
||||
"""
|
||||
lim_{θ→ π+} τ̃(cos θ)
|
||||
|
|
@ -314,6 +319,7 @@ def τ̃_pilim(nmax):
|
|||
t[plus1mmask] = -t[plus1mmask]
|
||||
return t
|
||||
|
||||
@jit(u=True)
|
||||
def get_π̃τ̃_y1(θ,nmax):
|
||||
# TODO replace with the limit functions (below) when θ approaches
|
||||
# the extreme values at about 1e-6 distance
|
||||
|
|
@ -390,7 +396,7 @@ def vswf_yr1(pos_sph,nmax,J=1):
|
|||
# return 1j**ny * np.sqrt((2*ny+1)*factorial(ny-my) /
|
||||
# (ny*(ny+1)*factorial(ny+my))
|
||||
# )
|
||||
|
||||
@jit(u=True)
|
||||
def zplane_pq_y(nmax, betap = 0):
|
||||
"""
|
||||
The z-propagating plane wave expansion coefficients as in [1, (1.12)]
|
||||
|
|
@ -409,6 +415,7 @@ def zplane_pq_y(nmax, betap = 0):
|
|||
|
||||
|
||||
#import warnings
|
||||
@jit(u=True)
|
||||
def plane_pq_y(nmax, kdir_cart, E_cart):
|
||||
"""
|
||||
The plane wave expansion coefficients for any direction kdir_cart
|
||||
|
|
@ -465,11 +472,13 @@ def plane_pq_y(nmax, kdir_cart, E_cart):
|
|||
# Functions copied from scattering_xu, additionaly normalized
|
||||
from py_gmm.gmm import vec_trans as vc
|
||||
|
||||
@jit(u=True)
|
||||
def q_max(m,n,μ,ν):
|
||||
return min(n,ν,(n+ν-abs(m+μ))/2)
|
||||
|
||||
# returns array with indices corresponding to q
|
||||
# argument q does nothing for now
|
||||
@jit(u=True)
|
||||
def a_q(m,n,μ,ν,q = None):
|
||||
qm=q_max(m,n,μ,ν)
|
||||
res, err= vc.gaunt_xu(m,n,μ,ν,qm)
|
||||
|
|
@ -480,7 +489,7 @@ def a_q(m,n,μ,ν,q = None):
|
|||
|
||||
# All arguments are single numbers (for now)
|
||||
# ZDE VYCHÁZEJÍ DIVNÁ ZNAMÉNKA
|
||||
#@jit
|
||||
@jit(u=True)
|
||||
def Ã(m,n,μ,ν,kdlj,θlj,φlj,r_ge_d,J):
|
||||
"""
|
||||
The à translation coefficient for spherical vector waves.
|
||||
|
|
@ -604,6 +613,7 @@ def B̃(m,n,μ,ν,kdlj,θlj,φlj,r_ge_d,J):
|
|||
# In[7]:
|
||||
|
||||
# Material parameters
|
||||
@jit(u=True)
|
||||
def ε_drude(ε_inf, ω_p, γ_p, ω): # RELATIVE permittivity, of course
|
||||
return ε_inf - ω_p*ω_p/(ω*(ω+1j*γ_p))
|
||||
|
||||
|
|
@ -611,6 +621,7 @@ def ε_drude(ε_inf, ω_p, γ_p, ω): # RELATIVE permittivity, of course
|
|||
# In[8]:
|
||||
|
||||
# Mie scattering
|
||||
@jit(u=True)
|
||||
def mie_coefficients(a, nmax, #ω, ε_i, ε_e=1, J_ext=1, J_scat=3
|
||||
k_i, k_e, μ_i=1, μ_e=1, J_ext=1, J_scat=3):
|
||||
"""
|
||||
|
|
@ -782,6 +793,7 @@ def G_Mie_scat_precalc_cart(source_cart, dest_cart, RH, RV, a, nmax, k_i, k_e,
|
|||
G_source_dest = sph_loccart2cart(G_source_dest, sph=orig2dest_sph, axis=-1)
|
||||
return G_source_dest
|
||||
|
||||
@jit(u=True)
|
||||
def G_Mie_scat_cart(source_cart, dest_cart, a, nmax, k_i, k_e, μ_i=1, μ_e=1, J_ext=1, J_scat=3):
|
||||
"""
|
||||
TODO
|
||||
|
|
@ -801,6 +813,7 @@ def cross_section_Mie(a, nmax, k_i, k_e, μ_i, μ_e,):
|
|||
# In[9]:
|
||||
|
||||
# From PRL 112, 253601 (1)
|
||||
@jit(u=True)
|
||||
def Grr_Delga(nmax, a, r, k, ε_m, ε_b):
|
||||
om = k * c
|
||||
z = (r-a)/a
|
||||
|
|
@ -822,6 +835,7 @@ def Grr_Delga(nmax, a, r, k, ε_m, ε_b):
|
|||
# Test if the decomposition of plane wave works also for absorbing environment (complex k).
|
||||
|
||||
# From PRL 112, 253601 (1)
|
||||
@jit(u=True)
|
||||
def Grr_Delga(nmax, a, r, k, ε_m, ε_b):
|
||||
om = k * c
|
||||
z = (r-a)/a
|
||||
|
|
@ -830,7 +844,7 @@ def Grr_Delga(nmax, a, r, k, ε_m, ε_b):
|
|||
s = np.sum( (n+1)**2 * (ε_m-ε_b) / ((1+z)**(2*n+4) * (ε_m + ((n+1)/n)*ε_b)))
|
||||
return (g0 + s * c**2/(4*π*om**2*ε_b*a**3))
|
||||
|
||||
|
||||
@jit(u=True)
|
||||
def G0_dip_1(r_cart,k):
|
||||
"""
|
||||
Free-space dyadic Green's function in terms of the spherical vector waves.
|
||||
|
|
@ -847,12 +861,15 @@ def G0_dip_1(r_cart,k):
|
|||
|
||||
# Free-space dyadic Green's functions from RMP 70, 2, 447 =: [1]
|
||||
# (The numerical value is correct only at the regular part, i.e. r != 0)
|
||||
@jit(u=True)
|
||||
def _P(z):
|
||||
return (1-1/z+1/(z*z))
|
||||
@jit(u=True)
|
||||
def _Q(z):
|
||||
return (-1+3/z-3/(z*z))
|
||||
|
||||
# [1, (9)] FIXME The sign here is most likely wrong!!!
|
||||
@jit(u=True)
|
||||
def G0_analytical(r #cartesian!
|
||||
, k):
|
||||
I=np.identity(3)
|
||||
|
|
@ -866,6 +883,7 @@ def G0_analytical(r #cartesian!
|
|||
))
|
||||
|
||||
# [1, (11)]
|
||||
@jit(u=True)
|
||||
def G0L_analytical(r, k):
|
||||
I=np.identity(3)
|
||||
rn = sph_loccart2cart(np.array([1.,0.,0.]), cart2sph(r), axis=-1)
|
||||
|
|
@ -874,10 +892,11 @@ def G0L_analytical(r, k):
|
|||
return (I-3*rnxrn)/(4*π*k*k*r**3)[...,ň,ň]
|
||||
|
||||
# [1,(10)]
|
||||
@jit
|
||||
def G0T_analytical(r, k):
|
||||
return G0_analytical(r,k) - G0L_analytical(r,k)
|
||||
|
||||
|
||||
@jit(u=True)
|
||||
def G0_sum_1_slow(source_cart, dest_cart, k, nmax):
|
||||
my, ny = get_mn_y(nmax)
|
||||
nelem = len(my)
|
||||
|
|
@ -888,6 +907,7 @@ def G0_sum_1_slow(source_cart, dest_cart, k, nmax):
|
|||
|
||||
|
||||
# Transformations of spherical bases
|
||||
@jit
|
||||
def WignerD_mm(l, quat):
|
||||
"""
|
||||
Calculates Wigner D matrix (as an numpy (2*l+1,2*l+1)-shaped array)
|
||||
|
|
@ -901,6 +921,7 @@ def WignerD_mm(l, quat):
|
|||
Delems = sf.Wigner_D_element(quat, indices).reshape(2*l+1,2*l+1)
|
||||
return Delems
|
||||
|
||||
@jit
|
||||
def WignerD_mm_fromvector(l, vect):
|
||||
"""
|
||||
TODO doc
|
||||
|
|
@ -908,6 +929,7 @@ def WignerD_mm_fromvector(l, vect):
|
|||
return WignerD_mm(l, quaternion.from_rotation_vector(vect))
|
||||
|
||||
|
||||
@jit
|
||||
def WignerD_yy(lmax, quat):
|
||||
"""
|
||||
TODO doc
|
||||
|
|
@ -922,6 +944,7 @@ def WignerD_yy(lmax, quat):
|
|||
b_in = e_in
|
||||
return Delems
|
||||
|
||||
@jit
|
||||
def WignerD_yy_fromvector(lmax, vect):
|
||||
"""
|
||||
TODO doc
|
||||
|
|
@ -929,7 +952,7 @@ def WignerD_yy_fromvector(lmax, vect):
|
|||
return WignerD_yy(lmax, quaternion.from_rotation_vector(vect))
|
||||
|
||||
|
||||
|
||||
@jit
|
||||
def xflip_yy(lmax):
|
||||
"""
|
||||
TODO doc
|
||||
|
|
@ -946,10 +969,12 @@ def xflip_yy(lmax):
|
|||
b_in = e_in
|
||||
return elems
|
||||
|
||||
@jit
|
||||
def xflip_tyy(lmax):
|
||||
fl_yy = xflip_yy(lmax)
|
||||
return np.array([fl_yy,-fl_yy])
|
||||
|
||||
@jit
|
||||
def xflip_tyty(lmax):
|
||||
fl_yy = xflip_yy(lmax)
|
||||
nelem = fl_yy.shape[0]
|
||||
|
|
@ -958,6 +983,7 @@ def xflip_tyty(lmax):
|
|||
fl_tyty[1,:,1,:] = -fl_yy
|
||||
return fl_tyty
|
||||
|
||||
@jit
|
||||
def yflip_yy(lmax):
|
||||
"""
|
||||
TODO doc
|
||||
|
|
@ -969,10 +995,12 @@ def yflip_yy(lmax):
|
|||
elems[(my % 2)==1] = elems[(my % 2)==1] * -1 # Obvious sign of tiredness (this is correct but ugly; FIXME)
|
||||
return elems
|
||||
|
||||
@jit
|
||||
def yflip_tyy(lmax):
|
||||
fl_yy = yflip_yy(lmax)
|
||||
return np.array([fl_yy,-fl_yy])
|
||||
|
||||
@jit
|
||||
def yflip_tyty(lmax):
|
||||
fl_yy = yflip_yy(lmax)
|
||||
nelem = fl_yy.shape[0]
|
||||
|
|
@ -981,6 +1009,7 @@ def yflip_tyty(lmax):
|
|||
fl_tyty[1,:,1,:] = -fl_yy
|
||||
return fl_tyty
|
||||
|
||||
@jit
|
||||
def zflip_yy(lmax):
|
||||
"""
|
||||
TODO doc
|
||||
|
|
@ -996,10 +1025,12 @@ def zflip_yy(lmax):
|
|||
b_in = e_in
|
||||
return elems
|
||||
|
||||
@jit
|
||||
def zflip_tyy(lmax):
|
||||
fl_yy = zflip_yy(lmax)
|
||||
return np.array([fl_yy,-fl_yy])
|
||||
|
||||
@jit
|
||||
def zflip_tyty(lmax):
|
||||
fl_yy = zflip_yy(lmax)
|
||||
nelem = fl_yy.shape[0]
|
||||
|
|
@ -1008,6 +1039,7 @@ def zflip_tyty(lmax):
|
|||
fl_tyty[1,:,1,:] = -fl_yy
|
||||
return fl_tyty
|
||||
|
||||
@jit
|
||||
def parity_yy(lmax):
|
||||
"""
|
||||
Parity operator (flip in x,y,z)
|
||||
|
|
@ -1025,6 +1057,7 @@ def parity_yy(lmax):
|
|||
#----------------------------------------------------#
|
||||
|
||||
# We don't really need this particular function anymore, but...
|
||||
@jit
|
||||
def _scuffTMatrixConvert_EM_01(EM):
|
||||
#print(EM)
|
||||
if (EM == b'E'):
|
||||
|
|
@ -1034,6 +1067,7 @@ def _scuffTMatrixConvert_EM_01(EM):
|
|||
else:
|
||||
return None
|
||||
|
||||
@jit(u=True)
|
||||
def loadScuffTMatrices(fileName):
|
||||
"""
|
||||
TODO doc
|
||||
|
|
@ -1069,7 +1103,7 @@ def loadScuffTMatrices(fileName):
|
|||
|
||||
|
||||
# misc tensor maniputalion
|
||||
|
||||
@jit
|
||||
def apply_matrix_left(matrix, tensor, axis):
|
||||
"""
|
||||
TODO doc
|
||||
|
|
@ -1080,6 +1114,7 @@ def apply_matrix_left(matrix, tensor, axis):
|
|||
tmp = np.tensordot(matrix, tensor, axes=(-1,axis))
|
||||
return np.moveaxis(tmp, 0, axis)
|
||||
|
||||
@jit
|
||||
def apply_ndmatrix_left(matrix,tensor,axes):
|
||||
"""
|
||||
Generalized apply_matrix_left, the matrix can have more (2N) abstract dimensions,
|
||||
|
|
@ -1095,7 +1130,7 @@ def apply_ndmatrix_left(matrix,tensor,axes):
|
|||
# Array simulations
|
||||
####################
|
||||
|
||||
|
||||
@jit
|
||||
def nelem2lMax(nelem):
|
||||
"""
|
||||
Auxiliary inverse function to nelem(lMax) = (lMax + 2) * lMax. Returns 0 if
|
||||
|
|
@ -1151,6 +1186,7 @@ def scatter_plane_wave(omega, epsilon_b, positions, Tmatrices, k_dirs, E_0s, #sa
|
|||
pass
|
||||
|
||||
import warnings
|
||||
@jit(u=True)
|
||||
def scatter_plane_wave_rectarray(omega, epsilon_b, xN, yN, xd, yd, TMatrices, k_dirs, E_0s,
|
||||
return_pq_0 = False, return_pq= False, return_xy = False, watch_time = False):
|
||||
"""
|
||||
|
|
@ -1382,6 +1418,7 @@ def scatter_plane_wave_rectarray(omega, epsilon_b, xN, yN, xd, yd, TMatrices, k_
|
|||
|
||||
|
||||
import warnings
|
||||
@jit(u=True)
|
||||
def scatter_constmultipole_rectarray(omega, epsilon_b, xN, yN, xd, yd, TMatrices, pq_0_c = 1,
|
||||
return_pq= False, return_xy = False, watch_time = False):
|
||||
"""
|
||||
|
|
|
|||
Loading…
Reference in New Issue