AnalysisFunctions.py 文件源码

def kdiv(arx,ary,arz,kf=None,lx=2*pi,ly=2*pi,lz=2*pi):
   """
   Function to compute divergence in Fourier space
   """

   #  COMPUTE THE ARRAY SIZE
   kx, ky, kz, km = create_kgrid(*np.shape(arx), lx=lx, ly=ly, lz=lz)

   #  FOURIER TRANSFORM THE ARRAY
   far = [nf.fftshift(nf.fftn(a)) for a in (arx, ary, arz)]

   # COMPUTE div=i*(kx*ax+ky*ay+kz*az)
   mg = np.meshgrid(kx,ky,kz)

   divarf = 1.j*reduce(operator.add, [a*b for a,b in zip(mg, far)])

   #  SET VALUES ABOVE kf AS 0+0i if kf non zero
   if kf is not None:
      divarf = (np.sign(km - kf) - 1.)/(-2.)*divarf

   divar = np.real(nf.ifftn(nf.ifftshift(divarf)))

   return divar