AnalysisFunctions.py 文件源码

def vecpot(arx,ary,arz, kf, lx=2*np.pi, ly=2*np.pi, lz=2*np.pi):
   """
   Function to compute vector potential of a 3D array
   """
   nx,ny,nz=arx.shape

   #  COMPUTE THE ARRAY SIZE
   kx, ky, kz, km = create_kgrid(*arx.shape, lx=lx, ly=ly, lz=lz)
  #kx, ky, kz, k2 = kx[:,nna,nna],ky[nna,:,nna],kz[nna,nna,:], km**2
   k2=km**2
   k2[nx/2,ny/2,nz/2]=1.

   #  FOURIER TRANSFORM THE ARRAY
   farx = nf.fftshift(nf.fftn(arx))
   fary = nf.fftshift(nf.fftn(ary))
   farz = nf.fftshift(nf.fftn(arz))

   #  SET VALUES ABOVE kf AS 0+0i
   farx = (np.sign(km - kf) - 1.)/(-2.)*farx
   fary = (np.sign(km - kf) - 1.)/(-2.)*fary
   farz = (np.sign(km - kf) - 1.)/(-2.)*farz

   #  FIND THE CORRESPONDING VECTOR POTENTIAL A = -ik x B /k^2
   axf = -eye*(ky*farz-kz*fary)/k2
   ayf = -eye*(kz*farx-kx*farz)/k2
   azf = -eye*(kx*fary-ky*farx)/k2

   #  BACK TRANSFORM TO REAL SPACE
   ax  = np.real(nf.ifftn(nf.ifftshift(axf)))
   ay  = np.real(nf.ifftn(nf.ifftshift(ayf)))
   az  = np.real(nf.ifftn(nf.ifftshift(azf)))
   return ax,ay,az