386 lines
18 KiB
Python
386 lines
18 KiB
Python
'''
|
|
Object oriented approach for the classical multiple scattering problem.
|
|
'''
|
|
|
|
__TODO__ = '''
|
|
- Implement per-scatterer lMax
|
|
- This means that Scattering.TMatrices either can not be a single array with a fixed
|
|
(N, 2, nelem, 2, nelem) shape but rather list of (2, nelem, 2, nelem) with nelem varying
|
|
per particle or some of its elements have to be unused. Anyways, there has to be some kind of
|
|
list with the lMaxes.
|
|
'''
|
|
|
|
import numpy as np
|
|
nx = np.newaxis
|
|
import time
|
|
import scipy
|
|
import sys
|
|
import warnings
|
|
import math
|
|
from qpms_c import get_mn_y, trans_calculator # TODO be explicit about what is imported
|
|
from .qpms_p import cart2sph, nelem2lMax # TODO be explicit about what is imported
|
|
from .timetrack import _time_b, _time_e
|
|
|
|
class Scattering(object):
|
|
'''
|
|
|
|
This is the most general class for a system of scatterers
|
|
in a non-lossy homogeneous background
|
|
to be solved with the multiple_scattering method. The scatterers,
|
|
as long as they comply with the disjoint circumscribed sphere
|
|
hypothesis, can each have any position in the 3D space and
|
|
any T-matrix.
|
|
|
|
Note that this object describes the scattering problem only for
|
|
a single given frequency, as the T-matrices and wavelenght
|
|
otherwise differ and all the computationally demanding
|
|
parts have to be done for each frequency. However,
|
|
the object can be recycled for many incident field shapes
|
|
at the given frequency.
|
|
|
|
Attributes should be perhaps later redefined to be read-only
|
|
(or make descriptors for them).
|
|
|
|
Args:
|
|
positions: (N,3)-shaped real array
|
|
TMatrices: (N,2,nelem,2,nelem)-shaped array
|
|
k_0 (float): Wave number for the space between scatterers.
|
|
|
|
Attributes:
|
|
positions:
|
|
TMatrices:
|
|
k_0 (float): Wave number for the space between scatterers.
|
|
lMax (int): Absolute maximum l for all scatterers. Depending on implementation,
|
|
lMax can be smaller for some individual scatterers in certain subclasses.
|
|
FIXME: here it is still implemented as constant lMax for all sites, see #!
|
|
prepared (bool): Keeps information whether the interaction matrix has
|
|
already been built and factorized.
|
|
|
|
|
|
'''
|
|
|
|
def __init__(self, positions, TMatrices, k_0, lMax = None, verbose=False, J_scat=3):
|
|
self.J_scat = J_scat
|
|
self.positions = positions.reshape((-1, positions.shape[-1]))
|
|
self.interaction_matrix = None
|
|
self.N = self.positions.shape[0]
|
|
self.k_0 = k_0
|
|
self.lMax = lMax if lMax else nelem2lMax(TMatrices.shape[-1])
|
|
self.tc = trans_calculator(self.lMax)
|
|
nelem = self.lMax * (self.lMax + 2) #!
|
|
self.nelem = nelem #!
|
|
self.prepared = False
|
|
self.TMatrices = np.broadcast_to(TMatrices, (self.N,2,nelem,2,nelem))
|
|
if np.isnan(np.min(TMatrices)):
|
|
warnings.warn("Some TMatrices contain NaNs. Expect invalid results")
|
|
if np.isnan(np.min(positions)):
|
|
warnings.warn("positions contain NaNs. Expect invalid results")
|
|
if math.isnan(k_0):
|
|
warnings.warn("k_0 is NaN. Expect invalid results")
|
|
|
|
|
|
|
|
def prepare(self, keep_interaction_matrix = False, verbose=False):
|
|
btime = _time_b(verbose)
|
|
if not self.prepared:
|
|
if self.interaction_matrix is None:
|
|
self.build_interaction_matrix(verbose=verbose)
|
|
self.lupiv = scipy.linalg.lu_factor(self.interaction_matrix,overwrite_a = not keep_interaction_matrix)
|
|
if not keep_interaction_matrix:
|
|
self.interaction_matrix = None
|
|
self.prepared = True
|
|
_time_e(btime, verbose)
|
|
|
|
def build_interaction_matrix(self,verbose = False):
|
|
btime = _time_b(verbose)
|
|
N = self.N
|
|
my, ny = get_mn_y(self.lMax)
|
|
nelem = len(my)
|
|
leftmatrix = np.zeros((N,2,nelem,N,2,nelem), dtype=complex)
|
|
sbtime = _time_b(verbose, step = 'Calculating interparticle translation coefficients')
|
|
"""
|
|
for i in range(N):
|
|
for j in range(N):
|
|
for yi in range(nelem):
|
|
for yj in range(nelem):
|
|
if(i != j):
|
|
d_i2j = cart2sph(self.positions[j]-self.positions[i])
|
|
a = Ã(my[yj],ny[yj],my[yi],ny[yi],kdlj=d_i2j[0]*self.k_0,θlj=d_i2j[1],φlj=d_i2j[2],r_ge_d=False,J=self.J_scat)
|
|
b = B̃(my[yj],ny[yj],my[yi],ny[yi],kdlj=d_i2j[0]*self.k_0,θlj=d_i2j[1],φlj=d_i2j[2],r_ge_d=False,J=self.J_scat)
|
|
leftmatrix[j,0,yj,i,0,yi] = a
|
|
leftmatrix[j,1,yj,i,1,yi] = a
|
|
leftmatrix[j,0,yj,i,1,yi] = b
|
|
leftmatrix[j,1,yj,i,0,yi] = b
|
|
"""
|
|
kdji = cart2sph(self.positions[:,nx,:] - self.positions[nx,:,:])
|
|
kdji[:,:,0] *= self.k_0
|
|
# get_AB array structure: [j,yj,i,yi]
|
|
a, b = self.tc.get_AB(my[nx,:,nx,nx],ny[nx,:,nx,nx],my[nx,nx,nx,:],ny[nx,nx,nx,:],
|
|
(kdji[:,:,0])[:,nx,:,nx], (kdji[:,:,1])[:,nx,:,nx], (kdji[:,:,2])[:,nx,:,nx],
|
|
False,self.J_scat)
|
|
mask = np.broadcast_to(np.eye(N,dtype=bool)[:,nx,:,nx],(N,nelem,N,nelem))
|
|
a[mask] = 0 # no self-translations
|
|
b[mask] = 0
|
|
leftmatrix[:,0,:,:,0,:] = a
|
|
leftmatrix[:,1,:,:,1,:] = a
|
|
leftmatrix[:,0,:,:,1,:] = b
|
|
leftmatrix[:,1,:,:,0,:] = b
|
|
_time_e(sbtime, verbose, step = 'Calculating interparticle translation coefficients')
|
|
# at this point, leftmatrix is the translation matrix
|
|
n2id = np.identity(2*nelem)
|
|
n2id.shape = (2,nelem,2,nelem)
|
|
for j in range(N):
|
|
leftmatrix[j] = - np.tensordot(self.TMatrices[j],leftmatrix[j],axes=([-2,-1],[0,1]))
|
|
# at this point, jth row of leftmatrix is that of -MT
|
|
leftmatrix[j,:,:,j,:,:] += n2id
|
|
# now we are done, 1-MT
|
|
leftmatrix.shape=(N*2*nelem,N*2*nelem)
|
|
self.interaction_matrix = leftmatrix
|
|
_time_e(btime, verbose)
|
|
|
|
def scatter(self, pq_0, verbose = False):
|
|
'''
|
|
pq_0 is (N, nelem, 2)-shaped array
|
|
'''
|
|
btime = _time_b(verbose)
|
|
if math.isnan(np.min(pq_0)):
|
|
warnings.warn("The incident field expansion contains NaNs. Expect invalid results.")
|
|
self.prepare(verbose=verbose)
|
|
pq_0 = np.broadcast_to(pq_0, (self.N,2,self.nelem))
|
|
MP_0 = np.empty((self.N,2,self.nelem),dtype=np.complex_)
|
|
for j in range(self.N):
|
|
MP_0[j] = np.tensordot(self.TMatrices[j], pq_0[j],axes=([-2,-1],[-2,-1]))
|
|
MP_0.shape = (self.N*2*self.nelem,)
|
|
solvebtime = _time_b(verbose,step='Solving the linear equation')
|
|
ab = scipy.linalg.lu_solve(self.lupiv, MP_0)
|
|
if math.isnan(np.min(ab)):
|
|
warnings.warn("Got NaN in the scattering result. Damn.")
|
|
raise
|
|
_time_e(solvebtime, verbose, step='Solving the linear equation')
|
|
ab.shape = (self.N,2,self.nelem)
|
|
_time_e(btime, verbose)
|
|
return ab
|
|
|
|
def scatter_constmultipole(self, pq_0_c, verbose = False):
|
|
btime = _time_b(verbose)
|
|
N = self.N
|
|
self.prepare(verbose=verbose)
|
|
nelem = self.nelem
|
|
if(pq_0_c ==1):
|
|
pq_0_c = np.full((2,nelem),1)
|
|
ab = np.empty((2,nelem,N*2*nelem), dtype=complex)
|
|
for N_or_M in range(2):
|
|
for yy in range(nelem):
|
|
pq_0 = np.zeros((2,nelem),dtype=np.complex_)
|
|
pq_0[N_or_M,yy] = pq_0_c[N_or_M,yy]
|
|
pq_0 = np.broadcast_to(pq_0, (N,2,nelem))
|
|
MP_0 = np.empty((N,2,nelem),dtype=np.complex_)
|
|
for j in range(N):
|
|
MP_0[j] = np.tensordot(self.TMatrices[j], pq_0[j],axes=([-2,-1],[-2,-1]))
|
|
MP_0.shape = (N*2*nelem,)
|
|
ab[N_or_M,yy] = scipy.linalg.lu_solve(self.lupiv,MP_0)
|
|
ab.shape = (2,nelem,N,2,nelem)
|
|
_time_e(btime, verbose)
|
|
return ab
|
|
|
|
class LatticeScattering(Scattering):
|
|
def __init__(self, lattice_spec, k_0, zSym = False):
|
|
|
|
|
|
"""
|
|
class Scattering_2D_lattice_rectcells(Scattering):
|
|
def __init__(self, rectcell_dims, rectcell_elem_positions, cellspec, k_0, rectcell_TMatrices = None, TMatrices = None, lMax = None, verbose=False, J_scat=3):
|
|
'''
|
|
cellspec: dvojice ve tvaru (seznam_zaplněnosti, seznam_pozic)
|
|
'''
|
|
if (rectcell_TMatrices is None) == (TMatrices is None):
|
|
raise ValueError('Either rectcell_TMatrices or TMatrices has to be given')
|
|
###self.positions = ZDE JSEM SKONČIL
|
|
self.J_scat = J_scat
|
|
self.positions = positions
|
|
self.interaction_matrix = None
|
|
self.N = positions.shape[0]
|
|
self.k_0 = k_0
|
|
self.lMax = lMax if lMax else nelem2lMax(TMatrices.shape[-1])
|
|
nelem = lMax * (lMax + 2) #!
|
|
self.nelem = nelem #!
|
|
self.prepared = False
|
|
self.TMatrices = np.broadcast_to(TMatrices, (self.N,2,nelem,2,nelem))
|
|
"""
|
|
|
|
class Scattering_2D_zsym(Scattering):
|
|
def __init__(self, positions, TMatrices, k_0, lMax = None, verbose=False, J_scat=3):
|
|
Scattering.__init__(self, positions, TMatrices, k_0, lMax, verbose, J_scat)
|
|
#TODO some checks on TMatrices symmetry
|
|
self.TE_yz = np.arange(self.nelem)
|
|
self.TM_yz = self.TE_yz
|
|
self.my, self.ny = get_mn_y(self.lMax)
|
|
self.TE_NMz = (self.my + self.ny) % 2
|
|
self.TM_NMz = 1 - self.TE_NMz
|
|
self.tc = trans_calculator(self.lMax)
|
|
# TODO možnost zadávat T-matice rovnou ve zhuštěné podobě
|
|
TMatrices_TE = TMatrices[...,self.TE_NMz[:,nx],self.TE_yz[:,nx],self.TE_NMz[nx,:],self.TE_yz[nx,:]]
|
|
TMatrices_TM = TMatrices[...,self.TM_NMz[:,nx],self.TM_yz[:,nx],self.TM_NMz[nx,:],self.TM_yz[nx,:]]
|
|
self.TMatrices_TE = np.broadcast_to(TMatrices_TE, (self.N, self.nelem, self.nelem))
|
|
self.TMatrices_TM = np.broadcast_to(TMatrices_TM, (self.N, self.nelem, self.nelem))
|
|
self.prepared_TE = False
|
|
self.prepared_TM = False
|
|
self.interaction_matrix_TE = None
|
|
self.interaction_matrix_TM= None
|
|
|
|
def prepare_partial(self, TE_or_TM, keep_interaction_matrix = False, verbose=False):
|
|
'''
|
|
TE is 0, TM is 1.
|
|
'''
|
|
btime = _time_b(verbose)
|
|
if (TE_or_TM == 0): #TE
|
|
if not self.prepared_TE:
|
|
if self.interaction_matrix_TE is None:
|
|
self.build_interaction_matrix(0, verbose)
|
|
sbtime = _time_b(verbose, step = 'Calculating LU decomposition of the interaction matrix, TE part')
|
|
self.lupiv_TE = scipy.linalg.lu_factor(self.interaction_matrix_TE, overwrite_a = not keep_interaction_matrix)
|
|
_time_e(sbtime, verbose, step = 'Calculating LU decomposition of the interaction matrix, TE part')
|
|
if(np.isnan(np.min(self.lupiv_TE[0])) or np.isnan(np.min(self.lupiv_TE[1]))):
|
|
warnings.warn("LU factorisation of interaction matrix contains NaNs. Expect invalid results.")
|
|
self.prepared_TE = True
|
|
if (TE_or_TM == 1): #TM
|
|
if not self.prepared_TM:
|
|
if self.interaction_matrix_TM is None:
|
|
self.build_interaction_matrix(1, verbose)
|
|
sbtime = _time_b(verbose, step = 'Calculating LU decomposition of the interaction matrix, TM part')
|
|
self.lupiv_TM = scipy.linalg.lu_factor(self.interaction_matrix_TM, overwrite_a = not keep_interaction_matrix)
|
|
_time_e(sbtime, verbose, step = 'Calculating LU decomposition of the interaction matrix, TM part')
|
|
if(np.isnan(np.min(self.lupiv_TM[0])) or np.isnan(np.min(self.lupiv_TM[1]))):
|
|
warnings.warn("LU factorisation of interaction matrix contains NaNs. Expect invalid results.")
|
|
self.prepared_TM = True
|
|
_time_e(btime, verbose)
|
|
|
|
def prepare(self, keep_interaction_matrix = False, verbose=False):
|
|
btime = _time_b(verbose)
|
|
if not self.prepared:
|
|
self.prepare_partial(0, keep_interaction_matrix, verbose)
|
|
self.prepare_partial(1, keep_interaction_matrix, verbose)
|
|
self.prepared = True
|
|
_time_e(btime, verbose)
|
|
|
|
def build_interaction_matrix(self,TE_or_TM = None, verbose = False):
|
|
#None means both
|
|
btime = _time_b(verbose)
|
|
N = self.N
|
|
my, ny = get_mn_y(self.lMax)
|
|
nelem = len(my)
|
|
idm = np.identity(nelem)
|
|
if (TE_or_TM == 0):
|
|
EoMl = (0,)
|
|
elif (TE_or_TM == 1):
|
|
EoMl = (1,)
|
|
elif (TE_or_TM is None):
|
|
EoMl = (0,1)
|
|
sbtime = _time_b(verbose, step = 'Calculating interparticle translation coefficients')
|
|
kdji = cart2sph(self.positions[:,nx,:] - self.positions[nx,:,:], allow2d=True)
|
|
kdji[:,:,0] *= self.k_0
|
|
# get_AB array structure: [j,yj,i,yi]
|
|
# FIXME I could save some memory by calculating only half of these coefficients
|
|
a, b = self.tc.get_AB(my[nx,:,nx,nx],ny[nx,:,nx,nx],my[nx,nx,nx,:],ny[nx,nx,nx,:],
|
|
(kdji[:,:,0])[:,nx,:,nx], (kdji[:,:,1])[:,nx,:,nx], (kdji[:,:,2])[:,nx,:,nx],
|
|
False,self.J_scat)
|
|
mask = np.broadcast_to(np.eye(N,dtype=bool)[:,nx,:,nx],(N,nelem,N,nelem))
|
|
a[mask] = 0 # no self-translations
|
|
b[mask] = 0
|
|
if np.isnan(np.min(a)) or np.isnan(np.min(b)):
|
|
warnings.warn("Some of the translation coefficients is a NaN. Expect invalid results.")
|
|
_time_e(sbtime, verbose, step = 'Calculating interparticle translation coefficients')
|
|
for EoM in EoMl:
|
|
leftmatrix = np.zeros((N,nelem,N,nelem), dtype=complex)
|
|
y = np.arange(nelem)
|
|
yi = y[nx,nx,nx,:]
|
|
yj = y[nx,:,nx,nx]
|
|
mask = np.broadcast_to((((yi - yj) % 2) == 0),(N,nelem,N,nelem))
|
|
leftmatrix[mask] = a[mask]
|
|
mask = np.broadcast_to((((yi - yj) % 2) != 0),(N,nelem,N,nelem))
|
|
leftmatrix[mask] = b[mask]
|
|
""" # we use to calculate the AB coefficients here
|
|
for i in range(N):
|
|
for j in range(i):
|
|
for yi in range(nelem):
|
|
for yj in range(nelem):
|
|
d_i2j = cart2sph(self.positions[j]-self.positions[i])
|
|
if ((yi - yj) % 2) == 0:
|
|
tr = Ã(my[yj],ny[yj],my[yi],ny[yi],kdlj=d_i2j[0]*self.k_0,θlj=d_i2j[1],φlj=d_i2j[2],r_ge_d=False,J=self.J_scat)
|
|
else:
|
|
tr = B̃(my[yj],ny[yj],my[yi],ny[yi],kdlj=d_i2j[0]*self.k_0,θlj=d_i2j[1],φlj=d_i2j[2],r_ge_d=False,J=self.J_scat)
|
|
leftmatrix[j,yj,i,yi] = tr
|
|
leftmatrix[i,yi,j,yj] = tr if (0 == (my[yj]+my[yi]) % 2) else -tr
|
|
_time_e(sbtime, verbose, step = 'Calculating interparticle translation coefficients, T%s part' % ('M' if EoM else 'E'))
|
|
"""
|
|
for j in range(N):
|
|
leftmatrix[j] = - np.tensordot(self.TMatrices_TM[j] if EoM else self.TMatrices_TE[j],leftmatrix[j],
|
|
axes = ([-1],[0]))
|
|
leftmatrix[j,:,j,:] += idm
|
|
leftmatrix.shape = (self.N*self.nelem, self.N*self.nelem)
|
|
if np.isnan(np.min(leftmatrix)):
|
|
warnings.warn("Interaction matrix contains some NaNs. Expect invalid results.")
|
|
if EoM == 0:
|
|
self.interaction_matrix_TE = leftmatrix
|
|
if EoM == 1:
|
|
self.interaction_matrix_TM = leftmatrix
|
|
a = None
|
|
b = None
|
|
_time_e(btime, verbose)
|
|
|
|
def scatter_partial(self, TE_or_TM, pq_0, verbose = False):
|
|
'''
|
|
pq_0 is (N, nelem)-shaped array
|
|
'''
|
|
btime = _time_b(verbose)
|
|
self.prepare_partial(TE_or_TM, verbose = verbose)
|
|
pq_0 = np.broadcast_to(pq_0, (self.N, self.nelem))
|
|
MP_0 = np.empty((self.N,self.nelem),dtype=np.complex_)
|
|
for j in range(self.N):
|
|
if TE_or_TM: #TM
|
|
MP_0[j] = np.tensordot(self.TMatrices_TM[j], pq_0[j], axes=([-1],[-1]))
|
|
else: #TE
|
|
MP_0[j] = np.tensordot(self.TMatrices_TE[j], pq_0[j], axes=([-1],[-1]))
|
|
MP_0.shape = (self.N*self.nelem,)
|
|
solvebtime = _time_b(verbose,step='Solving the linear equation')
|
|
ab = scipy.linalg.lu_solve(self.lupiv_TM if TE_or_TM else self.lupiv_TE, MP_0)
|
|
_time_e(solvebtime, verbose, step='Solving the linear equation')
|
|
ab.shape = (self.N, self.nelem)
|
|
_time_e(btime,verbose)
|
|
return ab
|
|
|
|
def scatter(self, pq_0, verbose = False):
|
|
'''
|
|
FI7ME
|
|
pq_0 is (N, nelem, 2)-shaped array
|
|
'''
|
|
btime = _time_b(verbose)
|
|
raise Exception('Not yet implemented')
|
|
self.prepare(verbose=verbose)
|
|
pq_0 = np.broadcast_to(pq_0, (self.N,2,self.nelem))
|
|
MP_0 = np.empty((N,2,nelem),dtype=np.complex_)
|
|
for j in range(self.N):
|
|
MP_0[j] = np.tensordot(self.TMatrices[j], pq_0[j],axes=([-2,-1],[-2,-1]))
|
|
MP_0.shape = (N*2*self.nelem,)
|
|
solvebtime = _time_b(verbose,step='Solving the linear equation')
|
|
ab = scipy.linalg.lu_solve(self.lupiv, MP_0)
|
|
_time_e(solvebtime, verbose, step='Solving the linear equation')
|
|
ab.shape = (N,2,nelem)
|
|
_time_e(btime, verbose)
|
|
return ab
|
|
|
|
def forget_matrices(self):
|
|
'''
|
|
Free interaction matrices and set the respective flags
|
|
(useful when memory is a bottleneck).
|
|
'''
|
|
self.interaction_matrix_TE = None
|
|
self.interaction_matrix_TM = None
|
|
self.lupiv_TE = None
|
|
self.lupiv_TM = None
|
|
self.prepared_TE = False
|
|
self.prepared_TM = False
|
|
self.prepared = False
|
|
|
|
|