Source code for Muscat.IO.VtkReader

# -*- coding: utf-8 -*-
#
# This file is subject to the terms and conditions defined in
# file 'LICENSE.txt', which is part of this source code package.
#


"""Vtu and Vtk file reader
"""
import os
import numpy as np

import Muscat.Containers.ElementsDescription as ED
from Muscat.Containers.Mesh import Mesh
from Muscat.IO.ReaderBase import ReaderBase
from Muscat.Bridges.vtkBridge import VtkToMesh
from Muscat.IO.IOFactory import RegisterReaderClass


class VtkReader(ReaderBase):
    """Vtk Reader class"""

    def __init__(self, fileName: str = "") -> None:
        super().__init__(fileName=fileName)

    def Read(self, fileName: str = None) -> Mesh:
        """Function that performs the reading mesh using vtk,
        the reader is selected based on the file_extension.
        Current files supported are vtk, stl

        Parameters
        ----------
        fileName : str, optional
            name of the file to be read, by default None

        Returns
        -------
        Mesh
            output unstructured mesh object containing reading result
        """
        from vtkmodules.vtkCommonDataModel import vtkUnstructuredGrid, vtkPolyData
        from vtkmodules.vtkIOLegacy import vtkGenericDataObjectReader
        from vtkmodules.vtkIOGeometry import vtkSTLReader
        from vtkmodules.vtkIOXML import vtkXMLUnstructuredGridReader, vtkXMLImageDataReader, vtkXMLStructuredGridReader, vtkXMLRectilinearGridReader, vtkXMLPolyDataReader

        VTK_ReaderByExtention = {
            ".stl": vtkSTLReader,
            ".vtu": vtkXMLUnstructuredGridReader,
            ".vti": vtkXMLImageDataReader,
            ".vts": vtkXMLStructuredGridReader,
            ".vtr": vtkXMLRectilinearGridReader,
            ".vtp": vtkXMLPolyDataReader,
        }

        if fileName is not None:
            self.SetFileName(fileName)

        filename, file_extension = os.path.splitext(self.fileName)
        reader = VTK_ReaderByExtention.get(file_extension, vtkGenericDataObjectReader)()
        reader.SetFileName(self.fileName)
        reader.Update()
        res = reader.GetOutput()

        output = VtkToMesh(res)
        output.ConvertDataForNativeTreatment()
        return output


try:
    import vtkmodules.vtkIOLegacy

    RegisterReaderClass(".vtk", VtkReader)
    RegisterReaderClass(".vti", VtkReader)
    RegisterReaderClass(".vtp", VtkReader)
    RegisterReaderClass(".vtr", VtkReader)
    RegisterReaderClass(".vts", VtkReader)
    RegisterReaderClass(".vtu", VtkReader)
except:  # pragma: no cover
    pass


def CheckIntegrity():

    try:
        import vtkmodules.vtkIOLegacy
    except:
        return "skip"

    from Muscat.TestData import GetTestDataPath

    filesToTest = [
        ("Cube222.vtk", 27, 8),
        ("Cube222.vti", 27, 8),
        ("coneAscii.vtp", 7, 10),
        ("Cube222.vts", 27, 8),
        ("Square22.vts", 9, 4),
        ("Cube111.vtr", 8, 1),
        ("Cube222.vtu", 27, 8),
        ("coneAscii.stl", 7, 10),
        ("coneBinary.stl", 7, 10),
    ]
    for file, nbNodes, nbElements in filesToTest:
        reader = VtkReader()
        print(f"Reading test file {file=} ...")
        reader.SetFileName(GetTestDataPath() + file)
        mesh = reader.Read(fileName=GetTestDataPath() + file)
        np.testing.assert_equal(mesh.GetNumberOfNodes(), nbNodes)
        np.testing.assert_equal(mesh.GetNumberOfElements(), nbElements)

    return "ok"


if __name__ == "__main__":
    print(CheckIntegrity())  # pragma: no cover