qpms/hexpoints.py

   1 import math
   2 import numpy as np
   3 nx = None
   4
   5 _s3 = math.sqrt(3)
   6
   7
   8 def generate_trianglepoints(maxlayer, include_origin = False, v3d = True, circular = True, sixthindices = False, mirrorindices = False):
   9     _e6 = np.array([[math.cos(2*math.pi*i/6),math.sin(2*math.pi*i/6),0] if v3d else [math.cos(2*math.pi*i/6),math.sin(2*math.pi*i/6)] for i in range(6)])
  10     points = np.empty((3*maxlayer*(maxlayer+1)+(1 if include_origin else 0), 3 if v3d else 2))
  11     point_i = 0
  12     if (include_origin):
  13         points[0] = np.array((0,0,0) if v3d else (0,0))
  14         point_i = 1
  15     if sixthindices:
  16         si = np.empty((6,(maxlayer*(maxlayer+1))//2), dtype=int)
  17         sii = [0,0,0,0,0,0]
  18     if mirrorindices:
  19         if(maxlayer < 3):
  20             mi = np.empty((2,0,), dtype=int)
  21         else:
  22             #layer indices start from one!
  23             ilayer = np.arange(1, maxlayer+1) # We need first to "copy" layer indices to correspond to the Muster count
  24             mustercount = (ilayer - 1)//2
  25             mustercum = np.cumsum(mustercount)
  26             layerstart = np.zeros((mustercum[maxlayer - 1]), dtype=int)
  27             layerstart[mustercum[:(maxlayer-1)]] = 1
  28             layer = np.cumsum(layerstart) + 2 # That's it
  29             lb = 3*layer*(layer-1) # layer base (lowest) index
  30             li = np.arange(len(layer)) - mustercum[layer-2] # muster indices for each layers
  31             mi = np.empty((2, len(layer)), dtype=int)
  32             mi[0] = lb + 1 + li + include_origin
  33             mi[1] = lb + layer - (1 + li) + include_origin
  34             # there are two non-musters in each even layer, one non-muster in each odd
  35         layer = (2*np.arange(((3*maxlayer)//2))+1)//3 + 1
  36         nmi = 3*layer*(layer-1)
  37         nmi[2::3] += layer[2::3] // 2 # second non-musters in even layers
  38         nmi += include_origin
  39
  40     for layer in range(1,maxlayer+1):
  41         for i in range(6):
  42             base = _e6[i]*layer
  43             shift = _e6[(i+2)%6]
  44             ar = np.arange(layer)
  45             points[point_i:(point_i+layer)] = base[nx,:] + ar[:,nx] * shift[nx,:]
  46             if sixthindices:
  47                 si[i, sii[i]:sii[i]+layer] = point_i + ar
  48                 sii[i] += layer
  49             point_i += layer
  50     if (circular):
  51         mask = (np.sum(points * points, axis = -1) <= maxlayer * maxlayer * 3/ 4 + 0.1)  # UGLY FIX OF ASYMMETRY BECAUSE OF ROUNDING ERROR
  52         points = points[mask]
  53         if sixthindices:
  54             cum = np.cumsum(mask) - 1
  55             mask0 = mask[si[0]]
  56             si_ = si[:,mask0]
  57             si = cum[si_]
  58         if mirrorindices:
  59             cum = np.cumsum(mask) - 1
  60             mask0 = mask[mi[0]]
  61             mi_ = mi[:,mask0]
  62             mi = cum[mi_]
  63             mask0 = mask[nmi]
  64             nmi_ = nmi[mask0]
  65             nmi = cum[nmi_]
  66     if not (mirrorindices or sixthindices):
  67         return points
  68     else:
  69         return {'points': points,
  70                 'si' : si if sixthindices else None,
  71                 'mi' : mi if mirrorindices else None,
  72                 'nmi' : nmi if mirrorindices else None}
  73
  74 def generate_trianglepoints_hexcomplement(maxlayer, v3d = True, circular = True, thirdindices = False, mirrorindices=False):
  75     _e6 = np.array([[math.cos(2*math.pi*i/6),math.sin(2*math.pi*i/6),0] if v3d else [math.cos(2*math.pi*i/6),math.sin(2*math.pi*i/6)] for i in range(6)])
  76     _f6 = np.array([[-math.sin(2*math.pi*i/6),math.cos(2*math.pi*i/6),0] if v3d else [math.sin(2*math.pi*i/6),-math.cos(2*math.pi*i/6)] for i in range(6)])
  77     points = np.empty((3*maxlayer*maxlayer, 3 if v3d else 2))
  78     point_i = 0
  79     # 3 * layer ** 2 is the basis index for a layer, a layer contains 3 * (2*layer + 1) points
  80     if thirdindices:
  81         ii = np.arange(maxlayer**2)
  82         layer = np.empty((maxlayer**2), dtype=int)
  83         layer = np.sqrt(ii, out=layer, casting='unsafe')
  84         #ti0 = 2*layer**2 + ii
  85         ti = np.arange(3)[:, nx] * (2*layer + 1)[nx, :] + (2*layer**2 + ii)[nx,:]
  86     if mirrorindices:
  87         ii = np.arange(((maxlayer-1)*maxlayer)//2)
  88         layer = np.empty((((maxlayer-1)*maxlayer)//2), dtype=int)
  89         layer = (np.sqrt(1+8*ii, out=layer, casting='unsafe')+1)//2
  90         li = ii - ((layer ) * (layer-1))//2# numbers indices in a each layer
  91         lb = 3*layer **2  # base index of a layer
  92         mi = np.empty((2,len(layer)), dtype=int)
  93         mi[0] = lb + li + layer % 2
  94         mi[1] = lb + 2*layer - li
  95         # indices of non-mirrored/self-mirrored
  96         layer = np.arange(maxlayer)
  97         lb = 3 * layer**2
  98         nmi = lb + ((layer + 1) % 2) * layer
  99     for layer in range(0,maxlayer):
 100         if (layer % 2): # odd layer
 101             for i in range(3):
 102                 base = _f6[(2*i-1)%6] * ((0.5 + 1.5 * layer) / _s3)
 103                 shift = _e6[(2*i+2)%6]
 104                 count = (layer + 1) // 2
 105                 ar = np.arange(count)
 106                 points[point_i:point_i+count,:] = base + ar[:,nx]*shift[nx,:]
 107                 point_i += count
 108
 109                 base = _e6[(2*i+1)%6]*layer + _f6[(2*i)%6] / _s3
 110                 shift = _e6[(2*i+3)%6]
 111                 count = layer
 112                 ar = np.arange(count)
 113                 points[point_i:point_i+count,:] = base + ar[:,nx]*shift[nx,:]
 114                 point_i += count
 115
 116                 base = _e6[(2*i+2)%6]*layer + _f6[(2*i)%6] / _s3
 117                 shift = _e6[(2*i+4)%6]
 118                 count = (layer + 1) // 2
 119                 ar = np.arange(count)
 120                 points[point_i:point_i+count,:] = base + ar[:,nx]*shift[nx,:]
 121                 point_i += count
 122         else: # even layer:
 123             for i in range(3):
 124                 shift = _e6[(2*i+2)%6]
 125                 base = _f6[(2*i-1)%6] * ((0.5 + 1.5 * layer) / _s3) + shift / 2
 126                 count = layer // 2
 127                 ar = np.arange(count)
 128                 points[point_i:point_i+count,:] = base + ar[:,nx]*shift[nx,:]
 129                 point_i += count
 130
 131                 base = _e6[(2*i+1)%6]*layer + _f6[(2*i)%6] / _s3
 132                 shift = _e6[(2*i+3)%6]
 133                 count = layer
 134                 ar = np.arange(count)
 135                 points[point_i:point_i+count,:] = base + ar[:,nx]*shift[nx,:]
 136                 point_i += count
 137
 138                 base = _e6[(2*i+2)%6]*layer + _f6[(2*i)%6] / _s3
 139                 shift = _e6[(2*i+4)%6]
 140                 count = (layer + 2) // 2
 141                 ar = np.arange(count)
 142                 points[point_i:point_i+count,:] = base + ar[:,nx]*shift[nx,:]
 143                 point_i += count
 144     #if (mirrorindices):
 145
 146     if (circular):
 147         mask = (np.sum(points * points, axis = -1) <= maxlayer * maxlayer * 3/ 4 + 0.01)  # UGLY FIX OF ASYMMETRY BECAUSE OF ROUNDING ERROR
 148         points = points[mask]
 149         if thirdindices:
 150             cum = np.cumsum(mask) - 1
 151             mask0 = mask[ti[0]]
 152             ti_ = ti[:,mask0]
 153             ti = cum[ti_]
 154         if mirrorindices:
 155             cum = np.cumsum(mask) - 1
 156             mask0 = mask[mi[0]]
 157             mi_ = mi[:,mask0]
 158             mi = cum[mi_]
 159             mask0 = mask[nmi]
 160             nmi_ = nmi[mask0]
 161             nmi = cum[nmi_]
 162     if not (mirrorindices or thirdindices):
 163         return points
 164     else:
 165         return {'points': points,
 166                 'ti' : ti if thirdindices else None,
 167                 'mi' : mi if mirrorindices else None,
 168                 'nmi' : nmi if mirrorindices else None
 169                }
 170
 171
 172 from .cycommon import get_mn_y
 173 from .cytranslations import trans_calculator
 174 from .qpms_p import cart2sph
 175
 176 def hexlattice_precalc_AB_save(file, lMax, k_hexside, maxlayer, circular=True, savepointinfo = False, J_scat=3):
 177     params = {
 178         'lMax' : lMax,
 179         'k_hexside' : k_hexside,
 180         'maxlayer' : maxlayer,
 181         'circular' : circular,
 182         'savepointinfo' : savepointinfo,
 183         'J_scat' : J_scat
 184     }
 185     tpdict = generate_trianglepoints(maxlayer, v3d=True, circular=circular, sixthindices=True, mirrorindices=True)
 186     tphcdict = generate_trianglepoints_hexcomplement(maxlayer, v3d=True, circular=circular, thirdindices=True, mirrorindices=True)
 187     my, ny = get_mn_y(lMax)
 188     nelem = len(my)
 189     a_self_nm = np.empty((tpdict['nmi'].shape[0],nelem,nelem), dtype=complex)
 190     b_self_nm = np.empty((tpdict['nmi'].shape[0],nelem,nelem), dtype=complex)
 191     a_self_m0 = np.empty((tpdict['mi'].shape[1],nelem,nelem), dtype=complex)
 192     b_self_m0 = np.empty((tpdict['mi'].shape[1],nelem,nelem), dtype=complex)
 193     a_d2u_nm = np.empty((tphcdict['nmi'].shape[0],nelem,nelem), dtype=complex)
 194     b_d2u_nm = np.empty((tphcdict['nmi'].shape[0],nelem,nelem), dtype=complex)
 195     a_d2u_m0 = np.empty((tphcdict['mi'].shape[1],nelem,nelem), dtype=complex)
 196     b_d2u_m0 = np.empty((tphcdict['mi'].shape[1],nelem,nelem), dtype=complex)
 197
 198     k_0 = k_hexside*_s3 # not really a wave vector here because of the normalisation!
 199     tc = trans_calculator(lMax)
 200
 201     y = np.arange(nelem)
 202
 203     points = tpdict['points'][tpdict['nmi']]
 204     d_i2j = cart2sph(points)
 205     a_self_nm, b_self_nm = tc.get_AB_arrays(k_0*d_i2j[:,0],d_i2j[:,1],d_i2j[:,2],np.array([False]),J_scat)
 206
 207     points = tpdict['points'][tpdict['mi'][0]]
 208     d_i2j = cart2sph(points)
 209     a_self_m0, b_self_m0 = tc.get_AB_arrays(k_0*d_i2j[:,0],d_i2j[:,1],d_i2j[:,2],np.array([False]),J_scat)
 210
 211     points = tphcdict['points'][tphcdict['nmi']]
 212     d_i2j = cart2sph(points)
 213     a_d2u_nm, b_d2u_nm = tc.get_AB_arrays(k_0*d_i2j[:,0],d_i2j[:,1],d_i2j[:,2],np.array([False]),J_scat)
 214
 215     points = tphcdict['points'][tphcdict['mi'][0]]
 216     d_i2j = cart2sph(points)
 217     a_d2u_m0, b_d2u_m0 = tc.get_AB_arrays(k_0*d_i2j[:,0],d_i2j[:,1],d_i2j[:,2],np.array([False]),J_scat)
 218
 219     tosave = {
 220         'a_self_nm' : a_self_nm,
 221         'a_self_m0' : a_self_m0,
 222         'b_self_nm' : b_self_nm,
 223         'b_self_m0' : b_self_m0,
 224         'a_d2u_nm' : a_d2u_nm,
 225         'a_d2u_m0' : a_d2u_m0,
 226         'b_d2u_nm' : b_d2u_nm,
 227         'b_d2u_m0' : b_d2u_m0,
 228         'precalc_params' : params
 229     }
 230     if savepointinfo:
 231         tosave['tp_points'] = tpdict['points'],
 232         tosave['tp_si'] = tpdict['si'],
 233         tosave['tp_mi'] = tpdict['mi'],
 234         tosave['tp_nmi'] = tpdict['nmi']
 235         tosave['tphc_points'] = tphcdict['points'],
 236         tosave['tphc_ti'] = tphcdict['ti'],
 237         tosave['tphc_mi'] = tphcdict['mi'],
 238         tosave['tphc_nmi'] = tphcdict['nmi']
 239     np.savez(file, **tosave)
 240
 241 def hexlattice_precalc_AB_loadunwrap(file, tpdict = None, tphcdict = None, return_points = False):
 242     npz = np.load(file)
 243     precalc_params = npz['precalc_params'][()]
 244     my, ny = get_mn_y(precalc_params['lMax'])
 245     nelem = len(my)
 246     # this I should have made more universal...
 247     if precalc_params['savepointinfo']:
 248         if not tpdict:
 249             tpdict = {
 250                 'points' : npz['tp_points'],
 251                 'si' : npz['tp_si'],
 252                 'mi' : npz['tp_mi'],
 253                 'nmi' : npz['tp_nmi'],
 254             }
 255         if not tphcdict:
 256             tphcdict = {
 257                 'points' : npz['tphc_points'],
 258                 'ti' : npz['tphc_ti'],
 259                 'mi' : npz['tphc_mi'],
 260                 'nmi' : npz['tphc_nmi']
 261         }
 262     else:
 263         if not tpdict:
 264             tpdict = generate_trianglepoints(maxlayer = precalc_params['maxlayer'], v3d=True,
 265                                              circular=precalc_params['circular'], sixthindices=True, mirrorindices=True)
 266         if not tphcdict:
 267             tphcdict = generate_trianglepoints_hexcomplement(maxlayer=precalc_params['maxlayer'], v3d=True,
 268                                                              circular=precalc_params['circular'], thirdindices=True, mirrorindices=True)
 269
 270     # For some obscure reason, I keep getting trailing single-dimension in the beginning for these arrays
 271     for a in (tpdict['points'], tphcdict['points'], tpdict['si'], tpdict['mi'],
 272              tphcdict['ti'], tphcdict['mi']):
 273         if len(a.shape) > 2:
 274             a.shape = a.shape[1::]
 275
 276     self_tr = tpdict['points']
 277     d2u_tr = tphcdict['points']
 278     if len(self_tr.shape)>2:
 279         self_tr = np.reshape(self_tr, self_tr.shape[1::])
 280     if len(d2u_tr.shape)>2:
 281         d2u_tr = np.reshape(d2u_tr, d2u_tr.shape[1::])
 282     u2d_tr = -d2u_tr
 283     a_self = np.empty((self_tr.shape[0],nelem,nelem), dtype=complex)
 284     b_self = np.empty((self_tr.shape[0],nelem,nelem), dtype=complex)
 285     a_d2u  = np.empty(( d2u_tr.shape[0],nelem,nelem), dtype=complex)
 286     b_d2u  = np.empty(( d2u_tr.shape[0],nelem,nelem), dtype=complex)
 287     a_self[tpdict['nmi']]=npz['a_self_nm']
 288     a_self[tpdict['mi'][0]]=npz['a_self_m0']
 289     b_self[tpdict['nmi']]=npz['b_self_nm']
 290     b_self[tpdict['mi'][0]]=npz['b_self_m0']
 291     mirrorangles = cart2sph(self_tr[tpdict['mi'][1]])[:,2] - cart2sph(self_tr[tpdict['mi'][0]])[:,2]
 292     a_self[tpdict['mi'][1],:,:] = a_self[tpdict['mi'][0],:,:] * np.exp(1j*mirrorangles[:,nx,nx]*(my[nx,nx,:]-my[nx,:,nx]))
 293     b_self[tpdict['mi'][1],:,:] = b_self[tpdict['mi'][0],:,:] * np.exp(1j*mirrorangles[:,nx,nx]*(my[nx,nx,:]-my[nx,:,nx]))
 294     for i in range(1,6):
 295         a_self[tpdict['si'][i],:,:] = a_self[tpdict['si'][0],:,:] * np.exp(1j*math.pi/3*i*(my[nx,:]-my[:,nx]))
 296         b_self[tpdict['si'][i],:,:] = b_self[tpdict['si'][0],:,:] * np.exp(1j*math.pi/3*i*(my[nx,:]-my[:,nx]))
 297     a_d2u[tphcdict['nmi']]=npz['a_d2u_nm']
 298     a_d2u[tphcdict['mi'][0]]=npz['a_d2u_m0']
 299     b_d2u[tphcdict['nmi']]=npz['b_d2u_nm']
 300     b_d2u[tphcdict['mi'][0]]=npz['b_d2u_m0']
 301     mirrorangles = cart2sph(self_tr[tphcdict['mi'][1]])[:,2] - cart2sph(self_tr[tphcdict['mi'][0]])[:,2]
 302     a_d2u[tphcdict['mi'][1],:,:] = a_d2u[tphcdict['mi'][0],:,:] * np.exp(1j*mirrorangles[:,nx,nx]*(my[nx,nx,:]-my[nx,:,nx]))
 303     b_d2u[tphcdict['mi'][1],:,:] = b_d2u[tphcdict['mi'][0],:,:] * np.exp(1j*mirrorangles[:,nx,nx]*(my[nx,nx,:]-my[nx,:,nx]))
 304     for i in (1,-1):
 305         a_d2u[tphcdict['ti'][i],:,:] = a_d2u[tphcdict['ti'][0],:,:] * np.exp(i*2j*math.pi/3*(my[nx,:]-my[:,nx]))
 306         b_d2u[tphcdict['ti'][i],:,:] = b_d2u[tphcdict['ti'][0],:,:] * np.exp(i*2j*math.pi/3*(my[nx,:]-my[:,nx]))
 307     a_u2d = a_d2u * (-1)**(my[nx,:]-my[:,nx])
 308     b_u2d = b_d2u * (-1)**(my[nx,:]-my[:,nx])
 309     d = {
 310         'a_self' : a_self,
 311         'b_self' : b_self,
 312         'a_d2u'  : a_d2u,
 313         'b_d2u'  : b_d2u,
 314         'a_u2d'  : a_u2d,
 315         'b_u2d'  : b_u2d,
 316     }
 317     for k in precalc_params.keys():
 318         d[k] = precalc_params[k]
 319     if return_points:
 320         d['d2u_tr'] = tphcdict['points']
 321         d['u2d_tr'] = -tphcdict['points']
 322         d['self_tr'] = tpdict['points']
 323     return d
 324
 325 def hexlattice_get_AB(lMax, k_hexside, maxlayer, circular=True, return_points = True, J_scat=3):
 326     params = {
 327         'lMax' : lMax,
 328         'k_hexside' : k_hexside,
 329         'maxlayer' : maxlayer,
 330         'circular' : circular,
 331         'savepointinfo' : return_points, # should I delete this key?
 332         'J_scat' : J_scat
 333     }
 334     tpdict = generate_trianglepoints(maxlayer, v3d=True, circular=circular, sixthindices=True, mirrorindices=True)
 335     tphcdict = generate_trianglepoints_hexcomplement(maxlayer, v3d=True, circular=circular, thirdindices=True, mirrorindices=True)
 336     my, ny = get_mn_y(lMax)
 337     nelem = len(my)
 338     a_self_nm = np.empty((tpdict['nmi'].shape[0],nelem,nelem), dtype=complex)
 339     b_self_nm = np.empty((tpdict['nmi'].shape[0],nelem,nelem), dtype=complex)
 340     a_self_m0 = np.empty((tpdict['mi'].shape[1],nelem,nelem), dtype=complex)
 341     b_self_m0 = np.empty((tpdict['mi'].shape[1],nelem,nelem), dtype=complex)
 342     a_d2u_nm = np.empty((tphcdict['nmi'].shape[0],nelem,nelem), dtype=complex)
 343     b_d2u_nm = np.empty((tphcdict['nmi'].shape[0],nelem,nelem), dtype=complex)
 344     a_d2u_m0 = np.empty((tphcdict['mi'].shape[1],nelem,nelem), dtype=complex)
 345     b_d2u_m0 = np.empty((tphcdict['mi'].shape[1],nelem,nelem), dtype=complex)
 346
 347     k_0 = k_hexside*_s3 # not really a wave vector here because of the normalisation!
 348     tc = trans_calculator(lMax)
 349
 350     y = np.arange(nelem)
 351
 352     points = tpdict['points'][tpdict['nmi']]
 353     d_i2j = cart2sph(points)
 354     a_self_nm, b_self_nm = tc.get_AB_arrays(k_0*d_i2j[:,0],d_i2j[:,1],d_i2j[:,2],np.array([False]),J_scat)
 355
 356     points = tpdict['points'][tpdict['mi'][0]]
 357     d_i2j = cart2sph(points)
 358     a_self_m0, b_self_m0 = tc.get_AB_arrays(k_0*d_i2j[:,0],d_i2j[:,1],d_i2j[:,2],np.array([False]),J_scat)
 359
 360     points = tphcdict['points'][tphcdict['nmi']]
 361     d_i2j = cart2sph(points)
 362     a_d2u_nm, b_d2u_nm = tc.get_AB_arrays(k_0*d_i2j[:,0],d_i2j[:,1],d_i2j[:,2],np.array([False]),J_scat)
 363
 364     points = tphcdict['points'][tphcdict['mi'][0]]
 365     d_i2j = cart2sph(points)
 366     a_d2u_m0, b_d2u_m0 = tc.get_AB_arrays(k_0*d_i2j[:,0],d_i2j[:,1],d_i2j[:,2],np.array([False]),J_scat)
 367     '''
 368     tosave = {
 369         'a_self_nm' : a_self_nm,
 370         'a_self_m0' : a_self_m0,
 371         'b_self_nm' : b_self_nm,
 372         'b_self_m0' : b_self_m0,
 373         'a_d2u_nm' : a_d2u_nm,
 374         'a_d2u_m0' : a_d2u_m0,
 375         'b_d2u_nm' : b_d2u_nm,
 376         'b_d2u_m0' : b_d2u_m0,
 377         'precalc_params' : params
 378     }
 379     if savepointinfo:
 380         tosave['tp_points'] = tpdict['points'],
 381         tosave['tp_si'] = tpdict['si'],
 382         tosave['tp_mi'] = tpdict['mi'],
 383         tosave['tp_nmi'] = tpdict['nmi']
 384         tosave['tphc_points'] = tphcdict['points'],
 385         tosave['tphc_ti'] = tphcdict['ti'],
 386         tosave['tphc_mi'] = tphcdict['mi'],
 387         tosave['tphc_nmi'] = tphcdict['nmi']
 388     np.savez(file, **tosave)
 389     '''
 390     self_tr = tpdict['points']
 391     d2u_tr = tphcdict['points']
 392     if len(self_tr.shape)>2:
 393         self_tr = np.reshape(self_tr, self_tr.shape[1::])
 394     if len(d2u_tr.shape)>2:
 395         d2u_tr = np.reshape(d2u_tr, d2u_tr.shape[1::])
 396     u2d_tr = -d2u_tr
 397     a_self = np.empty((self_tr.shape[0],nelem,nelem), dtype=complex)
 398     b_self = np.empty((self_tr.shape[0],nelem,nelem), dtype=complex)
 399     a_d2u  = np.empty(( d2u_tr.shape[0],nelem,nelem), dtype=complex)
 400     b_d2u  = np.empty(( d2u_tr.shape[0],nelem,nelem), dtype=complex)
 401     a_self[tpdict['nmi']]=a_self_nm
 402     a_self[tpdict['mi'][0]]=a_self_m0
 403     b_self[tpdict['nmi']]=b_self_nm
 404     b_self[tpdict['mi'][0]]=b_self_m0
 405     mirrorangles = cart2sph(self_tr[tpdict['mi'][1]])[:,2] - cart2sph(self_tr[tpdict['mi'][0]])[:,2]
 406     a_self[tpdict['mi'][1],:,:] = a_self[tpdict['mi'][0],:,:] * np.exp(1j*mirrorangles[:,nx,nx]*(my[nx,nx,:]-my[nx,:,nx]))
 407     b_self[tpdict['mi'][1],:,:] = b_self[tpdict['mi'][0],:,:] * np.exp(1j*mirrorangles[:,nx,nx]*(my[nx,nx,:]-my[nx,:,nx]))
 408     for i in range(1,6):
 409         a_self[tpdict['si'][i],:,:] = a_self[tpdict['si'][0],:,:] * np.exp(1j*math.pi/3*i*(my[nx,:]-my[:,nx]))
 410         b_self[tpdict['si'][i],:,:] = b_self[tpdict['si'][0],:,:] * np.exp(1j*math.pi/3*i*(my[nx,:]-my[:,nx]))
 411     a_d2u[tphcdict['nmi']]=a_d2u_nm
 412     a_d2u[tphcdict['mi'][0]]=a_d2u_m0
 413     b_d2u[tphcdict['nmi']]=b_d2u_nm
 414     b_d2u[tphcdict['mi'][0]]=b_d2u_m0
 415     mirrorangles = cart2sph(self_tr[tphcdict['mi'][1]])[:,2] - cart2sph(self_tr[tphcdict['mi'][0]])[:,2]
 416     a_d2u[tphcdict['mi'][1],:,:] = a_d2u[tphcdict['mi'][0],:,:] * np.exp(1j*mirrorangles[:,nx,nx]*(my[nx,nx,:]-my[nx,:,nx]))
 417     b_d2u[tphcdict['mi'][1],:,:] = b_d2u[tphcdict['mi'][0],:,:] * np.exp(1j*mirrorangles[:,nx,nx]*(my[nx,nx,:]-my[nx,:,nx]))
 418     for i in (1,-1):
 419         a_d2u[tphcdict['ti'][i],:,:] = a_d2u[tphcdict['ti'][0],:,:] * np.exp(i*2j*math.pi/3*(my[nx,:]-my[:,nx]))
 420         b_d2u[tphcdict['ti'][i],:,:] = b_d2u[tphcdict['ti'][0],:,:] * np.exp(i*2j*math.pi/3*(my[nx,:]-my[:,nx]))
 421     a_u2d = a_d2u * (-1)**(my[nx,:]-my[:,nx])
 422     b_u2d = b_d2u * (-1)**(my[nx,:]-my[:,nx])
 423     d = {
 424         'a_self' : a_self,
 425         'b_self' : b_self,
 426         'a_d2u'  : a_d2u,
 427         'b_d2u'  : b_d2u,
 428         'a_u2d'  : a_u2d,
 429         'b_u2d'  : b_u2d,
 430     }
 431     for k in params.keys():
 432         d[k] = params[k]
 433     if return_points:
 434         d['d2u_tr'] = tphcdict['points']
 435         d['u2d_tr'] = -tphcdict['points']
 436         d['self_tr'] = tpdict['points']
 437     return d
 438
 439 from scipy.constants import c
 440 from .timetrack import _time_b, _time_e
 441 from .tmatrices import symz_indexarrays
 442
 443 def hexlattice_zsym_getSVD(lMax, TMatrices_om, epsilon_b, hexside, maxlayer, omega, klist, gaussianSigma=False, onlyNmin=0, verbose=False):
 444     btime = _time_b(verbose)
 445     nelem = lMax * (lMax + 2)
 446     n2id = np.identity(2*nelem)
 447     n2id.shape = (2,nelem,2,nelem)
 448     nan = float('nan')
 449     k_0 = omega * math.sqrt(epsilon_b) / c
 450     tdic = hexlattice_get_AB(lMax,k_0*hexside,maxlayer)
 451     a_self = tdic['a_self'][:,:nelem,:nelem]
 452     b_self = tdic['b_self'][:,:nelem,:nelem]
 453     a_u2d = tdic['a_u2d'][:,:nelem,:nelem]
 454     b_u2d = tdic['b_u2d'][:,:nelem,:nelem]
 455     a_d2u = tdic['a_d2u'][:,:nelem,:nelem]
 456     b_d2u = tdic['b_d2u'][:,:nelem,:nelem]
 457     unitcell_translations = tdic['self_tr']*hexside*_s3
 458     u2d_translations = tdic['u2d_tr']*hexside*_s3
 459     d2u_translations = tdic['d2u_tr']*hexside*_s3
 460
 461     if gaussianSigma:
 462         sbtime = _time_b(verbose, step='Calculating gaussian envelope')
 463         unitcell_envelope = np.exp(-np.sum(tdic['self_tr']**2,axis=-1)/(2*gaussianSigma**2))
 464         u2d_envelope = np.exp(-np.sum(tdic['u2d_tr']**2,axis=-1)/(2*gaussianSigma**2))
 465         d2u_envelope = np.exp(-np.sum(tdic['d2u_tr']**2,axis=-1)/(2*gaussianSigma**2))
 466         _time_e(sbtime, verbose, step='Calculating gaussian envelope')
 467
 468     #TMatrices_om = TMatrices_interp(omega)
 469     if(not onlyNmin):
 470         svUfullTElist = np.full((klist.shape[0], 2*nelem, 2*nelem), np.nan, dtype=complex)
 471         svVfullTElist = np.full((klist.shape[0], 2*nelem, 2*nelem), np.nan, dtype=complex)
 472         svSfullTElist = np.full((klist.shape[0], 2*nelem), np.nan, dtype=complex)
 473         svUfullTMlist = np.full((klist.shape[0], 2*nelem, 2*nelem), np.nan, dtype=complex)
 474         svVfullTMlist = np.full((klist.shape[0], 2*nelem, 2*nelem), np.nan, dtype=complex)
 475         svSfullTMlist = np.full((klist.shape[0], 2*nelem), np.nan, dtype=complex)
 476     else:
 477         minsvTElist = np.full((klist.shape[0], onlyNmin),np.nan)
 478         minsvTMlist = np.full((klist.shape[0], onlyNmin),np.nan)
 479
 480     leftmatrixlist = np.full((klist.shape[0],2,2,nelem,2,2,nelem),np.nan,dtype=complex)
 481     #isNaNlist = np.zeros((klist.shape[0]), dtype=bool)
 482     isNaNlist = (k_0*k_0 - klist[:,0]**2 - klist[:,1]**2 < 0)
 483     nnlist = np.logical_not(isNaNlist)
 484
 485     sbtime = _time_b(verbose, step='Initialization of matrices for SVD for a given list of k\'s')
 486
 487
 488     #ki = np.arange(klist.shape[0])[k_0*k_0 - klist[:,0]**2 - klist[:,1]**2 >= 0]
 489     k = klist[nnlist]
 490
 491     phases_self = np.exp(1j*np.tensordot(k,unitcell_translations,axes=(-1,-1)))
 492     phases_u2d = np.exp(1j*np.tensordot(k,u2d_translations,axes=(-1,-1)))
 493     phases_d2u = np.exp(1j*np.tensordot(k,d2u_translations,axes=(-1,-1)))
 494     if gaussianSigma:
 495         phases_self *= unitcell_envelope
 496         phases_u2d *= u2d_envelope
 497         phases_d2u *= d2u_envelope
 498     leftmatrix = np.zeros((k.shape[0],2,2,nelem, 2,2,nelem), dtype=complex)
 499     #       0:[u,E<--u,E  ]
 500     #       1:[d,M<--d,M  ]
 501     leftmatrix[:,0,0,:,0,0,:] = np.tensordot(phases_self,a_self, axes=(-1,0)) # u2u, E2E
 502     leftmatrix[:,1,0,:,1,0,:] = leftmatrix[:,0,0,:,0,0,:] # d2d, E2E
 503     leftmatrix[:,0,1,:,0,1,:] = leftmatrix[:,0,0,:,0,0,:] # u2u, M2M
 504     leftmatrix[:,1,1,:,1,1,:] = leftmatrix[:,0,0,:,0,0,:] # d2d, M2M
 505     leftmatrix[:,0,0,:,0,1,:] = np.tensordot(phases_self,b_self, axes=(-1,0)) # u2u, M2E
 506     leftmatrix[:,0,1,:,0,0,:] = leftmatrix[:,0,0,:,0,1,:] # u2u, E2M
 507     leftmatrix[:,1,1,:,1,0,:] = leftmatrix[:,0,0,:,0,1,:] # d2d, E2M
 508     leftmatrix[:,1,0,:,1,1,:] = leftmatrix[:,0,0,:,0,1,:] # d2d, M2E
 509     leftmatrix[:,0,0,:,1,0,:] = np.tensordot(phases_d2u, a_d2u,axes=(-1,0)) #d2u,E2E
 510     leftmatrix[:,0,1,:,1,1,:] = leftmatrix[:,0,0,:,1,0,:] #d2u, M2M
 511     leftmatrix[:,1,0,:,0,0,:] = np.tensordot(phases_u2d, a_u2d,axes=(-1,0)) #u2d,E2E
 512     leftmatrix[:,1,1,:,0,1,:] = leftmatrix[:,1,0,:,0,0,:] #u2d, M2M
 513     leftmatrix[:,0,0,:,1,1,:] = np.tensordot(phases_d2u, b_d2u,axes=(-1,0)) #d2u,M2E
 514     leftmatrix[:,0,1,:,1,0,:] = leftmatrix[:,0,0,:,1,1,:] #d2u, E2M
 515     leftmatrix[:,1,0,:,0,1,:] = np.tensordot(phases_u2d, b_u2d,axes=(-1,0)) #u2d,M2E
 516     leftmatrix[:,1,1,:,0,0,:] = leftmatrix[:,1,0,:,0,1,:] #u2d, E2M
 517     #leftmatrix is now the translation matrix T
 518     for j in range(2):
 519         leftmatrix[:,j] = np.rollaxis(-np.tensordot(TMatrices_om[j], leftmatrix[:,j], axes=([-2,-1],[1,2])),2)
 520         # at this point, jth row of leftmatrix is that of -MT
 521         leftmatrix[:,j,:,:,j,:,:] += n2id
 522     #now we are done, 1-MT
 523
 524     leftmatrixlist[nnlist] = leftmatrix
 525
 526     '''
 527     # sem nějaká rozumná smyčka
 528     for ki in range(klist.shape[0]):
 529         k = klist[ki]
 530         if (k_0*k_0 - k[0]*k[0] - k[1]*k[1] < 0):
 531             isNaNlist[ki] = True
 532             continue
 533
 534         phases_self = np.exp(1j*np.tensordot(k,unitcell_translations,axes=(0,-1)))
 535         phases_u2d = np.exp(1j*np.tensordot(k,u2d_translations,axes=(0,-1)))
 536         phases_d2u = np.exp(1j*np.tensordot(k,d2u_translations,axes=(0,-1)))
 537         if gaussianSigma:
 538             phases_self *= unitcell_envelope
 539             phases_u2d *= u2d_envelope
 540             phases_d2u *= d2u_envelope
 541         leftmatrix = np.zeros((2,2,nelem, 2,2,nelem), dtype=complex)
 542         #       0:[u,E<--u,E  ]
 543         #       1:[d,M<--d,M  ]
 544         leftmatrix[0,0,:,0,0,:] = np.tensordot(a_self,phases_self, axes=(0,-1)) # u2u, E2E
 545         leftmatrix[1,0,:,1,0,:] = leftmatrix[0,0,:,0,0,:] # d2d, E2E
 546         leftmatrix[0,1,:,0,1,:] = leftmatrix[0,0,:,0,0,:] # u2u, M2M
 547         leftmatrix[1,1,:,1,1,:] = leftmatrix[0,0,:,0,0,:] # d2d, M2M
 548         leftmatrix[0,0,:,0,1,:] = np.tensordot(b_self,phases_self, axes=(0,-1)) # u2u, M2E
 549         leftmatrix[0,1,:,0,0,:] = leftmatrix[0,0,:,0,1,:] # u2u, E2M
 550         leftmatrix[1,1,:,1,0,:] = leftmatrix[0,0,:,0,1,:] # d2d, E2M
 551         leftmatrix[1,0,:,1,1,:] = leftmatrix[0,0,:,0,1,:] # d2d, M2E
 552         leftmatrix[0,0,:,1,0,:] = np.tensordot(a_d2u, phases_d2u,axes=(0,-1)) #d2u,E2E
 553         leftmatrix[0,1,:,1,1,:] = leftmatrix[0,0,:,1,0,:] #d2u, M2M
 554         leftmatrix[1,0,:,0,0,:] = np.tensordot(a_u2d, phases_u2d,axes=(0,-1)) #u2d,E2E
 555         leftmatrix[1,1,:,0,1,:] = leftmatrix[1,0,:,0,0,:] #u2d, M2M
 556         leftmatrix[0,0,:,1,1,:] = np.tensordot(b_d2u, phases_d2u,axes=(0,-1)) #d2u,M2E
 557         leftmatrix[0,1,:,1,0,:] = leftmatrix[0,0,:,1,1,:] #d2u, E2M
 558         leftmatrix[1,0,:,0,1,:] = np.tensordot(b_u2d, phases_u2d,axes=(0,-1)) #u2d,M2E
 559         leftmatrix[1,1,:,0,0,:] = leftmatrix[1,0,:,0,1,:] #u2d, E2M
 560         #leftmatrix is now the translation matrix T
 561         for j in range(2):
 562             leftmatrix[j] = -np.tensordot(TMatrices_om[j], leftmatrix[j], axes=([-2,-1],[0,1]))
 563             # at this point, jth row of leftmatrix is that of -MT
 564             leftmatrix[j,:,:,j,:,:] += n2id
 565         #now we are done, 1-MT
 566
 567         leftmatrixlist[ki] = leftmatrix
 568     '''
 569
 570     leftmatrixlist_s = np.reshape(leftmatrixlist,(klist.shape[0], 2*2*nelem,2*2*nelem))[nnlist]
 571     TEč, TMč = symz_indexarrays(lMax, 2)
 572     leftmatrixlist_TE = leftmatrixlist_s[np.ix_(np.arange(leftmatrixlist_s.shape[0]),TEč,TEč)]
 573     leftmatrixlist_TM = leftmatrixlist_s[np.ix_(np.arange(leftmatrixlist_s.shape[0]),TMč,TMč)]
 574     _time_e(sbtime, verbose, step='Initializing matrices for SVD for a given list of k\'s')
 575
 576     sbtime = _time_b(verbose, step='Calculating SVDs for a given list of k\'s.')
 577     if(not onlyNmin):
 578         svUfullTElist[nnlist], svSfullTElist[nnlist], svVfullTElist[nnlist] = np.linalg.svd(leftmatrixlist_TE, compute_uv=True)
 579         svUfullTMlist[nnlist], svSfullTMlist[nnlist], svVfullTMlist[nnlist] = np.linalg.svd(leftmatrixlist_TM, compute_uv=True)
 580         _time_e(sbtime, verbose, step='Calculating SVDs for a given list of k\'s.')
 581         _time_e(btime, verbose)
 582         return ((svUfullTElist, svSfullTElist, svVfullTElist), (svUfullTMlist, svSfullTMlist, svVfullTMlist))
 583     else:
 584         minsvTElist[nnlist] = np.linalg.svd(leftmatrixlist_TE, compute_uv=False)[...,-onlyNmin:]
 585         minsvTMlist[nnlist] = np.linalg.svd(leftmatrixlist_TM, compute_uv=False)[...,-onlyNmin:]
 586         _time_e(sbtime, verbose, step='Calculating SVDs for a given list of k\'s.')
 587         _time_e(btime, verbose)
 588         return (minsvTElist, minsvTMlist)