import numpy as np

# Parameters
lgrid = 100. # Length of box
ngrid = 32   # Grid size (should be power of 2)
lthres = 0.  # Eigenvalue threshold for environment classification

# Initialisations
nx,ny,nz = ngrid,ngrid,ngrid
lx,ly,lz = lgrid,lgrid,lgrid

# Density field (just set to random values as an example!)
densgrid = np.random.normal(size=(nx,ny,nz))

# Evaluate tidal tensor using Fourier transform methods
densspec = np.fft.rfftn(densgrid)
kx = 2.*np.pi*np.fft.fftfreq(nx,d=lx/nx)
ky = 2.*np.pi*np.fft.fftfreq(ny,d=ly/ny)
kz = 2.*np.pi*np.fft.fftfreq(nz,d=lz/nz)[:nz//2+1]
ksqspec = kx[:,np.newaxis,np.newaxis]**2 + ky[np.newaxis,:,np.newaxis]**2 + kz[np.newaxis,np.newaxis,:]**2
ksqspec[0,0,0] = 1.
kxspec = kx[:,np.newaxis,np.newaxis]*np.ones((nx,ny,nz//2+1))
kyspec = ky[np.newaxis,:,np.newaxis]*np.ones((nx,ny,nz//2+1))
kzspec = kz[np.newaxis,np.newaxis,:]*np.ones((nx,ny,nz//2+1))
txxspec = ((kxspec*kxspec)/ksqspec)*densspec
txyspec = ((kxspec*kyspec)/ksqspec)*densspec
txzspec = ((kxspec*kzspec)/ksqspec)*densspec
tyyspec = ((kyspec*kyspec)/ksqspec)*densspec
tyzspec = ((kyspec*kzspec)/ksqspec)*densspec
tzzspec = ((kzspec*kzspec)/ksqspec)*densspec
txxgrid = np.fft.irfftn(txxspec)
txygrid = np.fft.irfftn(txyspec)
txzgrid = np.fft.irfftn(txzspec)
tyygrid = np.fft.irfftn(tyyspec)
tyzgrid = np.fft.irfftn(tyzspec)
tzzgrid = np.fft.irfftn(tzzspec)

# Loop over the grid, determining the eigenvalues of the tidal tensor
tidaltens = np.empty((3,3))
ienvgrid = np.empty((nx,ny,nz))
for ix in range(nx):
  for iy in range(ny):
    for iz in range(nz):
      tidaltens[0,0] = txxgrid[ix,iy,iz]
      tidaltens[0,1] = txygrid[ix,iy,iz]
      tidaltens[0,2] = txzgrid[ix,iy,iz]
      tidaltens[1,0] = txygrid[ix,iy,iz]
      tidaltens[1,1] = tyygrid[ix,iy,iz]
      tidaltens[1,2] = tyzgrid[ix,iy,iz]
      tidaltens[2,0] = txzgrid[ix,iy,iz]
      tidaltens[2,1] = tyzgrid[ix,iy,iz]
      tidaltens[2,2] = tzzgrid[ix,iy,iz]
      w,v = np.linalg.eig(tidaltens)
      l3,l2,l1 = np.sort(w)
      if (l1 < lthres):
        ienv = 0
      elif ((l1 > lthres) & (l2 < lthres)):
        ienv = 1
      elif ((l2 > lthres) & (l3 < lthres)):
        ienv = 2
      elif (l3 > lthres):
        ienv = 3
      ienvgrid[ix,iy,iz] = ienv