dns-paraview-plugin/paraview-dns-data-reader.py

import numpy as np

from vtkmodules.vtkCommonDataModel import vtkDataSet
from vtkmodules.util.vtkAlgorithm import VTKPythonAlgorithmBase
from vtkmodules.numpy_interface import dataset_adapter as dsa

# new module for ParaView-specific decorators.
from paraview.util.vtkAlgorithm import smproxy, smproperty, smdomain, smhint

# to add a source, instead of a filter, use the `smproxy.source` decorator.
@smproxy.reader(filename_patterns="fort.*", file_description="Turbulent Combustion DNS data", label="POSTECH Turbulent Combustion DNS Reader")
class PythonDnsDataReader(VTKPythonAlgorithmBase):
    """This is dummy VTKPythonAlgorithmBase subclass that
    simply puts out a Superquadric poly data using a vtkSuperquadricSource
    internally"""
    def __init__(self):
        VTKPythonAlgorithmBase.__init__(self,
                nInputPorts=0,
                nOutputPorts=1,
                outputType='vtkRectilinearGrid')
        self._filename = None

    @smproperty.stringvector(name="FileName")
    @smdomain.filelist()
    # @smhint.filechooser(filename_patterns="fort.*", file_description="Numpy CSV files")
    def SetFileName(self, name):
        """Specify filename for the file to read."""
        if self._filename != name:
            self._filename = name
            self._ndata = None
            self._timesteps = None
            self.Modified()


    def RequestData(self, request, inInfo, outInfo):
        import vtk
        from vtkmodules.vtkCommonDataModel import vtkRectilinearGrid
        output = vtkRectilinearGrid.GetData(outInfo, 0)

        if self._filename is None:
            # Note, exceptions are totally fine!
            raise RuntimeError("No filename specified")

        print(self._filename)

        filename = self._filename
        t, nx, ny, nz, U, V, W, Y0, Y1 = self._read_data(filename)

        u = U.ravel()
        v = V.ravel()
        w = W.ravel()
        y1 = Y1.ravel()

        x = np.arange(nx) * 4*np.pi/nx
        y = np.arange(ny) * 2*np.pi/ny
        z = np.arange(nz) * 2*np.pi/nz

        # Create a rectilinear grid by defining three arrays specifying the
        # coordinates in the x-y-z directions.
        xCoords = vtk.vtkFloatArray()
        xCoords.SetNumberOfTuples(len(x))
        for i, xi in enumerate(x):
            xCoords.SetTuple1(i, xi)

        yCoords = vtk.vtkFloatArray()
        yCoords.SetNumberOfTuples(len(y))
        for i, yi in enumerate(y):
            yCoords.SetTuple1(i, yi)

        zCoords = vtk.vtkFloatArray()
        zCoords.SetNumberOfTuples(len(z))
        for i, zi in enumerate(z):
            zCoords.SetTuple1(i, zi)

        # The coordinates are assigned to the rectilinear grid. Make sure that
        # the number of values in each of the XCoordinates, YCoordinates,
        # and ZCoordinates is equal to what is defined in SetDimensions().
        #
        output.SetDimensions(len(x), len(y), len(z))
        output.SetXCoordinates(xCoords)
        output.SetYCoordinates(yCoords)
        output.SetZCoordinates(zCoords)

        numPoints = output.GetNumberOfPoints()

        velocity = vtk.vtkFloatArray();
        velocity.SetNumberOfComponents(3);
        velocity.SetNumberOfTuples(numPoints);
        velocity.SetName("U");

        for i in range(0, numPoints):
            velocity.SetTuple3(i, u[i], v[i], w[i]);

        output.GetPointData().AddArray(velocity)
        del velocity

        c = vtk.vtkFloatArray();
        c.SetNumberOfTuples(numPoints);
        c.SetName("c");

        for i in range(0, numPoints):
            c.SetTuple1(i, 1.0-y1[i]);

        output.GetPointData().AddArray(c)
        del c

        del t, nx, ny, nz, u, v, w, Y0, Y1, x, y, z

        return 1

    # for anything too complex or not yet supported, you can explicitly
    # provide the XML for the method.
    @smproperty.xml("""
        <DoubleVectorProperty name="Center"
            number_of_elements="3"
            default_values="0 0 0"
            command="SetCenter">
            <DoubleRangeDomain name="range" />
            <Documentation>Set center of the superquadric</Documentation>
        </DoubleVectorProperty>""")
    def SetCenter(self, x, y, z):
        self.Modified()

    # In most cases, one can simply use available decorators.
    @smproperty.doublevector(name="Scale", default_values=[1, 1, 1])
    @smdomain.doublerange()
    def SetScale(self, x, y, z):
        self.Modified()

    @smproperty.intvector(name="ThetaResolution", default_values=16)
    def SetThetaResolution(self, x):
        self.Modified()

    @smproperty.intvector(name="PhiResolution", default_values=16)
    @smdomain.intrange(min=0, max=1000)
    def SetPhiResolution(self, x):
        self.Modified()

    @smproperty.doublevector(name="Thickness", default_values=0.3333)
    @smdomain.doublerange(min=1e-24, max=1.0)
    def SetThickness(self, x):
        self.Modified()

    def _read_data (self, fname):
        import struct
        import sys
        import os
        with open(fname, 'rb') as f1 :
            f1.seek(0)

            raw_info = f1.read(4+8*6+4)[4:-4]
            t  = struct.unpack('d', raw_info[ 0: 8])[0]
            nx = struct.unpack('q', raw_info[ 8:16])[0]
            ny = struct.unpack('q', raw_info[16:24])[0]
            nz = struct.unpack('q', raw_info[24:32])[0]
            count = nx*ny*nz
            bSize = count*8 # size in bytes for a variable

            dummy_len = (4+8*3+4) + (4+8*2+4) + (4+8*2+4) + (4+8*2+4) + 4
            dummy = f1.read(dummy_len)
           #dummy = f1.read(4)

           #raw_field = f1.read(4+bSize*5+4)[4:-4]
            V = np.fromfile(f1, dtype=np.double, count=(3*count)).reshape((3,nz,ny,nx))
            s = np.fromfile(f1, dtype=np.double, count=(2*count)).reshape((2,nz,ny,nx))

            u = V[0]
            v = V[1]
            w = V[2]

            Y0 = s[0]
            Y1 = s[1]

        return t, nx, ny, nz, u, v, w, Y0, Y1