Write vtk files in pure Python + numpy

write_vtk.py

"""
A simple example for writing a valid VTK file in pure Python + numpy (no extra requirements)
"""
import numpy as np

def main(x, y, z, fields, filename, subheader=""):
    nx, ny, nz = shape = (len(_) for _ in (x, y, z))
    ntot = np.product([_-1 for _ in shape])  # number of cells
    with open(filename, 'wb') as fh:
        # two lines of header
        fh.write(b'# vtk DataFile Version 2.0\n')
        fh.write(f"{subheader}\n".encode())

        # declare grid type
        fh.write(f"BINARY\n".encode())
        fh.write("DATASET RECTILINEAR_GRID\n".encode())
        fh.write(f"DIMENSIONS {nx} {ny} {nz}\n".encode())
        fh.write("\n".encode())

        # include grid coordinates
        fh.write(f"X_COORDINATES {nx} float\n".encode())
        x.tofile(fh)
        fh.write("\n".encode())
        fh.write(f"Y_COORDINATES {ny} float\n".encode())
        y.tofile(fh)
        fh.write("\n".encode())
        fh.write(f"Z_COORDINATES {nz} float\n".encode())
        z.tofile(fh)
        fh.write("\n\n".encode())

        # declare on-grid fields      
        fh.write(f"CELL_DATA {ntot}\n".encode())
        fh.write("\n".encode())

        # write fields
        for name, data in fields.items():
            fh.write(f"SCALARS {name} float\n".encode())
            fh.write(f"LOOKUP_TABLE default\n".encode())
            data.tofile(fh)
            fh.write("\n".encode())


if __name__ == "__main__":
    # generate example dataset
    nx, ny, nz = shape = (2, 3, 4)
    ntot = np.product([_-1 for _ in shape])  # number of cells

    x = np.linspace(0, 2, nx, dtype='>f')
    y = np.linspace(0, 2, ny, dtype='>f')
    z = np.linspace(0, 2, nz, dtype='>f')
    data = {"EX": np.ones(ntot, dtype='>f')}
    main(x, y, z, data, "example.vtk")