qpms/scattering.py

   1 '''
   2 Object oriented approach for the classical multiple scattering problem.
   3 '''
   4
   5 __TODO__ = '''
   6 - Implement per-scatterer lMax
   7   - This means that Scattering.TMatrices either can not be a single array with a fixed
   8     (N, 2, nelem, 2, nelem) shape but rather list of (2, nelem, 2, nelem) with nelem varying
   9     per particle or some of its elements have to be unused. Anyways, there has to be some kind of
  10     list with the lMaxes.
  11 '''
  12
  13 import numpy as np
  14 nx = np.newaxis
  15 import time
  16 import scipy
  17 import sys
  18 import warnings
  19 import math
  20 from qpms_c import get_mn_y, trans_calculator # TODO be explicit about what is imported
  21 from .qpms_p import cart2sph, nelem2lMax # TODO be explicit about what is imported
  22 from .timetrack import _time_b, _time_e
  23
  24 class Scattering(object):
  25     '''
  26
  27     This is the most general class for a system of scatterers
  28     in a non-lossy homogeneous background
  29     to be solved with the multiple_scattering method. The scatterers,
  30     as long as they comply with the disjoint circumscribed sphere
  31     hypothesis, can each have any position in the 3D space and
  32     any T-matrix.
  33
  34     Note that this object describes the scattering problem only for
  35     a single given frequency, as the T-matrices and wavelenght
  36     otherwise differ and all the computationally demanding
  37     parts have to be done for each frequency. However,
  38     the object can be recycled for many incident field shapes
  39     at the given frequency.
  40
  41     Attributes should be perhaps later redefined to be read-only
  42     (or make descriptors for them).
  43
  44     Args:
  45         positions: (N,3)-shaped real array
  46         TMatrices: (N,2,nelem,2,nelem)-shaped array
  47         k_0 (float): Wave number for the space between scatterers.
  48
  49     Attributes:
  50         positions:
  51         TMatrices:
  52         k_0 (float): Wave number for the space between scatterers.
  53         lMax (int): Absolute maximum l for all scatterers. Depending on implementation,
  54             lMax can be smaller for some individual scatterers in certain subclasses.
  55             FIXME: here it is still implemented as constant lMax for all sites, see #!
  56         prepared (bool): Keeps information whether the interaction matrix has
  57             already been built and factorized.
  58
  59
  60     '''
  61
  62     def __init__(self, positions, TMatrices, k_0, lMax = None, verbose=False, J_scat=3):
  63         self.J_scat = J_scat
  64         self.positions = positions.reshape((-1, positions.shape[-1]))
  65         self.interaction_matrix = None
  66         self.N = self.positions.shape[0]
  67         self.k_0 = k_0
  68         self.lMax = lMax if lMax else nelem2lMax(TMatrices.shape[-1])
  69         self.tc = trans_calculator(self.lMax)
  70         nelem = self.lMax * (self.lMax + 2) #!
  71         self.nelem = nelem #!
  72         self.prepared = False
  73         self.TMatrices = np.broadcast_to(TMatrices, (self.N,2,nelem,2,nelem))
  74         if np.isnan(np.min(TMatrices)):
  75             warnings.warn("Some TMatrices contain NaNs. Expect invalid results")
  76         if np.isnan(np.min(positions)):
  77             warnings.warn("positions contain NaNs. Expect invalid results")
  78         if math.isnan(k_0):
  79             warnings.warn("k_0 is NaN. Expect invalid results")
  80
  81
  82
  83     def prepare(self, keep_interaction_matrix = False, verbose=False):
  84         btime = _time_b(verbose)
  85         if not self.prepared:
  86             if self.interaction_matrix is None:
  87                 self.build_interaction_matrix(verbose=verbose)
  88             self.lupiv = scipy.linalg.lu_factor(self.interaction_matrix,overwrite_a = not keep_interaction_matrix)
  89             if not keep_interaction_matrix:
  90                 self.interaction_matrix = None
  91             self.prepared = True
  92         _time_e(btime, verbose)
  93
  94     def build_interaction_matrix(self,verbose = False):
  95         btime = _time_b(verbose)
  96         N = self.N
  97         my, ny = get_mn_y(self.lMax)
  98         nelem = len(my)
  99         leftmatrix = np.zeros((N,2,nelem,N,2,nelem), dtype=complex)
 100         sbtime = _time_b(verbose, step = 'Calculating interparticle translation coefficients')
 101         """
 102         for i in range(N):
 103             for j in range(N):
 104                 for yi in range(nelem):
 105                     for yj in range(nelem):
 106                         if(i != j):
 107                             d_i2j = cart2sph(self.positions[j]-self.positions[i])
 108                             a = 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)
 109                             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)
 110                             leftmatrix[j,0,yj,i,0,yi] = a
 111                             leftmatrix[j,1,yj,i,1,yi] = a
 112                             leftmatrix[j,0,yj,i,1,yi] = b
 113                             leftmatrix[j,1,yj,i,0,yi] = b
 114         """
 115         kdji = cart2sph(self.positions[:,nx,:] - self.positions[nx,:,:])
 116         kdji[:,:,0] *= self.k_0
 117         # get_AB array structure: [j,yj,i,yi]
 118         a, b = self.tc.get_AB(my[nx,:,nx,nx],ny[nx,:,nx,nx],my[nx,nx,nx,:],ny[nx,nx,nx,:],
 119                 (kdji[:,:,0])[:,nx,:,nx], (kdji[:,:,1])[:,nx,:,nx], (kdji[:,:,2])[:,nx,:,nx],
 120                 False,self.J_scat)
 121         mask = np.broadcast_to(np.eye(N,dtype=bool)[:,nx,:,nx],(N,nelem,N,nelem))
 122         a[mask] = 0 # no self-translations
 123         b[mask] = 0
 124         leftmatrix[:,0,:,:,0,:] = a
 125         leftmatrix[:,1,:,:,1,:] = a
 126         leftmatrix[:,0,:,:,1,:] = b
 127         leftmatrix[:,1,:,:,0,:] = b
 128         _time_e(sbtime, verbose, step = 'Calculating interparticle translation coefficients')
 129         # at this point, leftmatrix is the translation matrix
 130         n2id = np.identity(2*nelem)
 131         n2id.shape = (2,nelem,2,nelem)
 132         for j in range(N):
 133             leftmatrix[j] = - np.tensordot(self.TMatrices[j],leftmatrix[j],axes=([-2,-1],[0,1]))
 134             # at this point, jth row of leftmatrix is that of -MT
 135             leftmatrix[j,:,:,j,:,:] += n2id
 136             # now we are done, 1-MT
 137         leftmatrix.shape=(N*2*nelem,N*2*nelem)
 138         self.interaction_matrix = leftmatrix
 139         _time_e(btime, verbose)
 140
 141     def scatter(self, pq_0, verbose = False):
 142         '''
 143         pq_0 is (N, nelem, 2)-shaped array
 144         '''
 145         btime = _time_b(verbose)
 146         if math.isnan(np.min(pq_0)):
 147             warnings.warn("The incident field expansion contains NaNs. Expect invalid results.")
 148         self.prepare(verbose=verbose)
 149         pq_0 = np.broadcast_to(pq_0, (self.N,2,self.nelem))
 150         MP_0 = np.empty((self.N,2,self.nelem),dtype=np.complex_)
 151         for j in range(self.N):
 152             MP_0[j] = np.tensordot(self.TMatrices[j], pq_0[j],axes=([-2,-1],[-2,-1]))
 153         MP_0.shape = (self.N*2*self.nelem,)
 154         solvebtime = _time_b(verbose,step='Solving the linear equation')
 155         ab = scipy.linalg.lu_solve(self.lupiv, MP_0)
 156         if math.isnan(np.min(ab)):
 157             warnings.warn("Got NaN in the scattering result. Damn.")
 158             raise
 159         _time_e(solvebtime, verbose, step='Solving the linear equation')
 160         ab.shape = (self.N,2,self.nelem)
 161         _time_e(btime, verbose)
 162         return ab
 163
 164     def scatter_constmultipole(self, pq_0_c, verbose = False):
 165         btime = _time_b(verbose)
 166         N = self.N
 167         self.prepare(verbose=verbose)
 168         nelem = self.nelem
 169         if(pq_0_c ==1):
 170             pq_0_c = np.full((2,nelem),1)
 171         ab = np.empty((2,nelem,N*2*nelem), dtype=complex)
 172         for N_or_M in range(2):
 173             for yy in range(nelem):
 174                 pq_0 = np.zeros((2,nelem),dtype=np.complex_)
 175                 pq_0[N_or_M,yy] = pq_0_c[N_or_M,yy]
 176                 pq_0 = np.broadcast_to(pq_0, (N,2,nelem))
 177                 MP_0 = np.empty((N,2,nelem),dtype=np.complex_)
 178                 for j in range(N):
 179                     MP_0[j] = np.tensordot(self.TMatrices[j], pq_0[j],axes=([-2,-1],[-2,-1]))
 180                 MP_0.shape = (N*2*nelem,)
 181                 ab[N_or_M,yy] = scipy.linalg.lu_solve(self.lupiv,MP_0)
 182         ab.shape = (2,nelem,N,2,nelem)
 183         _time_e(btime, verbose)
 184         return ab
 185
 186 class LatticeScattering(Scattering):
 187     def __init__(self, lattice_spec, k_0, zSym = False):
 188         pass
 189
 190
 191 """
 192 class Scattering_2D_lattice_rectcells(Scattering):
 193     def __init__(self, rectcell_dims, rectcell_elem_positions, cellspec, k_0, rectcell_TMatrices = None, TMatrices = None, lMax = None, verbose=False, J_scat=3):
 194         '''
 195         cellspec: dvojice ve tvaru (seznam_zaplněnosti, seznam_pozic)
 196         '''
 197         if (rectcell_TMatrices is None) == (TMatrices is None):
 198             raise ValueError('Either rectcell_TMatrices or TMatrices has to be given')
 199         ###self.positions = ZDE JSEM SKONČIL
 200         self.J_scat = J_scat
 201         self.positions = positions
 202         self.interaction_matrix = None
 203         self.N = positions.shape[0]
 204         self.k_0 = k_0
 205         self.lMax = lMax if lMax else nelem2lMax(TMatrices.shape[-1])
 206         nelem = lMax * (lMax + 2) #!
 207         self.nelem = nelem #!
 208         self.prepared = False
 209         self.TMatrices = np.broadcast_to(TMatrices, (self.N,2,nelem,2,nelem))
 210 """
 211
 212 class Scattering_2D_zsym(Scattering):
 213     def __init__(self, positions, TMatrices, k_0, lMax = None, verbose=False, J_scat=3):
 214         Scattering.__init__(self, positions, TMatrices, k_0, lMax, verbose, J_scat)
 215         #TODO some checks on TMatrices symmetry
 216         self.TE_yz = np.arange(self.nelem)
 217         self.TM_yz = self.TE_yz
 218         self.my, self.ny = get_mn_y(self.lMax)
 219         self.TE_NMz = (self.my + self.ny) % 2
 220         self.TM_NMz = 1 - self.TE_NMz
 221         self.tc = trans_calculator(self.lMax)
 222         # TODO možnost zadávat T-matice rovnou ve zhuštěné podobě
 223         TMatrices_TE = TMatrices[...,self.TE_NMz[:,nx],self.TE_yz[:,nx],self.TE_NMz[nx,:],self.TE_yz[nx,:]]
 224         TMatrices_TM = TMatrices[...,self.TM_NMz[:,nx],self.TM_yz[:,nx],self.TM_NMz[nx,:],self.TM_yz[nx,:]]
 225         self.TMatrices_TE = np.broadcast_to(TMatrices_TE, (self.N, self.nelem, self.nelem))
 226         self.TMatrices_TM = np.broadcast_to(TMatrices_TM, (self.N, self.nelem, self.nelem))
 227         self.prepared_TE = False
 228         self.prepared_TM = False
 229         self.interaction_matrix_TE = None
 230         self.interaction_matrix_TM= None
 231
 232     def prepare_partial(self, TE_or_TM, keep_interaction_matrix = False, verbose=False):
 233         '''
 234         TE is 0, TM is 1.
 235         '''
 236         btime = _time_b(verbose)
 237         if (TE_or_TM == 0): #TE
 238             if not self.prepared_TE:
 239                 if self.interaction_matrix_TE is None:
 240                     self.build_interaction_matrix(0, verbose)
 241                 sbtime = _time_b(verbose, step = 'Calculating LU decomposition of the interaction matrix, TE part')
 242                 self.lupiv_TE = scipy.linalg.lu_factor(self.interaction_matrix_TE, overwrite_a = not keep_interaction_matrix)
 243                 _time_e(sbtime, verbose, step = 'Calculating LU decomposition of the interaction matrix, TE part')
 244                 if(np.isnan(np.min(self.lupiv_TE[0])) or np.isnan(np.min(self.lupiv_TE[1]))):
 245                     warnings.warn("LU factorisation of interaction matrix contains NaNs. Expect invalid results.")
 246                 self.prepared_TE = True
 247         if (TE_or_TM == 1): #TM
 248             if not self.prepared_TM:
 249                 if self.interaction_matrix_TM is None:
 250                     self.build_interaction_matrix(1, verbose)
 251                 sbtime = _time_b(verbose, step = 'Calculating LU decomposition of the interaction matrix, TM part')
 252                 self.lupiv_TM = scipy.linalg.lu_factor(self.interaction_matrix_TM, overwrite_a = not keep_interaction_matrix)
 253                 _time_e(sbtime, verbose, step = 'Calculating LU decomposition of the interaction matrix, TM part')
 254                 if(np.isnan(np.min(self.lupiv_TM[0])) or np.isnan(np.min(self.lupiv_TM[1]))):
 255                     warnings.warn("LU factorisation of interaction matrix contains NaNs. Expect invalid results.")
 256                 self.prepared_TM = True
 257         _time_e(btime, verbose)
 258
 259     def prepare(self, keep_interaction_matrix = False, verbose=False):
 260         btime = _time_b(verbose)
 261         if not self.prepared:
 262             self.prepare_partial(0, keep_interaction_matrix, verbose)
 263             self.prepare_partial(1, keep_interaction_matrix, verbose)
 264             self.prepared = True
 265         _time_e(btime, verbose)
 266
 267     def build_interaction_matrix(self,TE_or_TM = None, verbose = False):
 268         #None means both
 269         btime = _time_b(verbose)
 270         N = self.N
 271         my, ny = get_mn_y(self.lMax)
 272         nelem = len(my)
 273         idm = np.identity(nelem)
 274         if (TE_or_TM == 0):
 275             EoMl = (0,)
 276         elif (TE_or_TM == 1):
 277             EoMl = (1,)
 278         elif (TE_or_TM is None):
 279             EoMl = (0,1)
 280         sbtime = _time_b(verbose, step = 'Calculating interparticle translation coefficients')
 281         kdji = cart2sph(self.positions[:,nx,:] - self.positions[nx,:,:], allow2d=True)
 282         kdji[:,:,0] *= self.k_0
 283         # get_AB array structure: [j,yj,i,yi]
 284         # FIXME I could save some memory by calculating only half of these coefficients
 285         a, b = self.tc.get_AB(my[nx,:,nx,nx],ny[nx,:,nx,nx],my[nx,nx,nx,:],ny[nx,nx,nx,:],
 286                 (kdji[:,:,0])[:,nx,:,nx], (kdji[:,:,1])[:,nx,:,nx], (kdji[:,:,2])[:,nx,:,nx],
 287                 False,self.J_scat)
 288         mask = np.broadcast_to(np.eye(N,dtype=bool)[:,nx,:,nx],(N,nelem,N,nelem))
 289         a[mask] = 0 # no self-translations
 290         b[mask] = 0
 291         if np.isnan(np.min(a)) or np.isnan(np.min(b)):
 292             warnings.warn("Some of the translation coefficients is a NaN. Expect invalid results.")
 293         _time_e(sbtime, verbose, step = 'Calculating interparticle translation coefficients')
 294         for EoM in EoMl:
 295             leftmatrix = np.zeros((N,nelem,N,nelem), dtype=complex)
 296             y = np.arange(nelem)
 297             yi = y[nx,nx,nx,:]
 298             yj = y[nx,:,nx,nx]
 299             mask = np.broadcast_to((((yi - yj) % 2) == 0),(N,nelem,N,nelem))
 300             leftmatrix[mask] = a[mask]
 301             mask = np.broadcast_to((((yi - yj) % 2) != 0),(N,nelem,N,nelem))
 302             leftmatrix[mask] = b[mask]
 303             """ # we use to calculate the AB coefficients here
 304             for i in range(N):
 305                 for j in range(i):
 306                     for yi in range(nelem):
 307                         for yj in range(nelem):
 308                             d_i2j = cart2sph(self.positions[j]-self.positions[i])
 309                             if ((yi - yj) % 2) == 0:
 310                                 tr = 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)
 311                             else:
 312                                 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)
 313                             leftmatrix[j,yj,i,yi] = tr
 314                             leftmatrix[i,yi,j,yj] = tr if (0 == (my[yj]+my[yi]) % 2) else -tr
 315             _time_e(sbtime, verbose, step = 'Calculating interparticle translation coefficients, T%s part' % ('M' if EoM else 'E'))
 316             """
 317             for j in range(N):
 318                 leftmatrix[j] = - np.tensordot(self.TMatrices_TM[j] if EoM else self.TMatrices_TE[j],leftmatrix[j],
 319                         axes = ([-1],[0]))
 320                 leftmatrix[j,:,j,:] += idm
 321             leftmatrix.shape = (self.N*self.nelem, self.N*self.nelem)
 322             if np.isnan(np.min(leftmatrix)):
 323                 warnings.warn("Interaction matrix contains some NaNs. Expect invalid results.")
 324             if EoM == 0:
 325                 self.interaction_matrix_TE = leftmatrix
 326             if EoM == 1:
 327                 self.interaction_matrix_TM = leftmatrix
 328         a = None
 329         b = None
 330         _time_e(btime, verbose)
 331
 332     def scatter_partial(self, TE_or_TM, pq_0, verbose = False):
 333         '''
 334         pq_0 is (N, nelem)-shaped array
 335         '''
 336         btime = _time_b(verbose)
 337         self.prepare_partial(TE_or_TM, verbose = verbose)
 338         pq_0 = np.broadcast_to(pq_0, (self.N, self.nelem))
 339         MP_0 = np.empty((self.N,self.nelem),dtype=np.complex_)
 340         for j in range(self.N):
 341             if TE_or_TM: #TM
 342                 MP_0[j] = np.tensordot(self.TMatrices_TM[j], pq_0[j], axes=([-1],[-1]))
 343             else: #TE
 344                 MP_0[j] = np.tensordot(self.TMatrices_TE[j], pq_0[j], axes=([-1],[-1]))
 345         MP_0.shape = (self.N*self.nelem,)
 346         solvebtime = _time_b(verbose,step='Solving the linear equation')
 347         ab = scipy.linalg.lu_solve(self.lupiv_TM if TE_or_TM else self.lupiv_TE, MP_0)
 348         _time_e(solvebtime, verbose, step='Solving the linear equation')
 349         ab.shape = (self.N, self.nelem)
 350         _time_e(btime,verbose)
 351         return ab
 352
 353     def scatter(self, pq_0, verbose = False):
 354         '''
 355         FI7ME
 356         pq_0 is (N, nelem, 2)-shaped array
 357         '''
 358         btime = _time_b(verbose)
 359         raise Exception('Not yet implemented')
 360         self.prepare(verbose=verbose)
 361         pq_0 = np.broadcast_to(pq_0, (self.N,2,self.nelem))
 362         MP_0 = np.empty((N,2,nelem),dtype=np.complex_)
 363         for j in range(self.N):
 364             MP_0[j] = np.tensordot(self.TMatrices[j], pq_0[j],axes=([-2,-1],[-2,-1]))
 365         MP_0.shape = (N*2*self.nelem,)
 366         solvebtime = _time_b(verbose,step='Solving the linear equation')
 367         ab = scipy.linalg.lu_solve(self.lupiv, MP_0)
 368         _time_e(solvebtime, verbose, step='Solving the linear equation')
 369         ab.shape = (N,2,nelem)
 370         _time_e(btime, verbose)
 371         return ab
 372
 373     def forget_matrices(self):
 374         '''
 375         Free interaction matrices and set the respective flags
 376         (useful when memory is a bottleneck).
 377         '''
 378         self.interaction_matrix_TE = None
 379         self.interaction_matrix_TM = None
 380         self.lupiv_TE = None
 381         self.lupiv_TM = None
 382         self.prepared_TE = False
 383         self.prepared_TM = False
 384         self.prepared = False
 385
 386