Reads vtk binary data for 3D MHD runs from SNOOPY

readvtk_snoopy.py

# Reads output vtk files from the spectral code, SNOOPY
# Currently only for 3D MHD runs
# Output data is by default float 32 binary

# Big endian vs Little endian: depends on the machine
# <f4 means read as np.float32 type little endian byter order
# >f4 means read as np.float32 type big    endian byter order

import numpy as np

# Usage:
#   time, V, B, N, D = readvtk_snoopy(fname='v1000.vtk')
# time = Simulation time
# N = Grid points vector (Nx, Ny, Nz)
# D = Difference vector (Dx, Dy, Dz)
# V = Velocity field vector (3, Nx, Ny, Nz)
# B = Magnetic field vector (3, Nx, Ny, Nz)

def readvtk_sno(fname='v0001.vtk'):
  
  with open(fname,'rb') as ff:
    ff.seek(0, 0) # start of file

    ff.readline() # l1 Out[70]: b'# vtk DataFile Version 2.0\n'
    
    # Read time
    l2 = ff.readline() # l2 Out[71]: b't= 5.025051671818386e-01 Snoopy Code v5.0\n'
    tt = np.float32(l2.split()[1])

    ff.readline() # l3 Out[72]: b'BINARY\n'
    ff.readline() # l4 Out[73]: b'DATASET STRUCTURED_POINTS\n'
    
    # Read dimensions
    l5 = ff.readline() # l5 Out[74]: b'DIMENSIONS 32 64 32\n'
    dims = np.array(l5.split()[1:], dtype=np.int32)

    ff.readline() # l6 Out[75]: b'ORIGIN -0.5 -1 -0.5\n'

    # Read spacing
    l7 = ff.readline() # l7 Out[76]: b'SPACING 0.03125 0.03125 0.03125\n'
    dgrid = np.array(l7.split()[1:], dtype=np.float32)

    ff.readline() # l8 Out[77]: b'POINT_DATA 65536\n'
    ff.readline() # l9 Out[78]: b'SCALARS vx float\n'
    ff.readline() # l10 Out[79]: b'LOOKUP_TABLE default\n'

    ###
    # IMPORTANT to read fields as FLOAT32 since its binary; 
    # FLOAT64 reads more than it should!!
    ###
    vx = np.fromfile(ff, dtype='>f4', count=np.prod(dims)).reshape(dims, order='f')
    ff.readline() # l11 Out[81]: b'FIELD FieldData 5\n'
    ff.readline() # l12 Out[82]: b'vy 1 65536 float\n'

    vy = np.fromfile(ff, dtype='>f4', count=np.prod(dims)).reshape(dims,order='f')
    ff.readline() # l13 Out[84]: b'vz 1 65536 float\n'

    vz = np.fromfile(ff, dtype='>f4', count=np.prod(dims)).reshape(dims,order='f')
    ff.readline() # Out[86]: b'bx 1 65536 float\n'

    bx = np.fromfile(ff, dtype='>f4', count=np.prod(dims)).reshape(dims,order='f')
    ff.readline() #Out[88]: b'by 1 65536 float\n'

    by = np.fromfile(ff, dtype='>f4', count=np.prod(dims)).reshape(dims,order='f')
    ff.readline() #Out[90]: b'bz 1 65536 float\n'

    bz = np.fromfile(ff, dtype='>f4', count=np.prod(dims)).reshape(dims,order='f')

  V = np.stack((vx, vy, vz),axis=0)
  B = np.stack((bx, by, bz),axis=0)

  return tt, V, B, dims, dgrid