PySLM VTK Exporter

vtkExporter.py

from typing import  Dict, List, Tuple,

import struct
import numpy as np

import pyslm
import pyslm.geometry

from pyslm import hatching as hatching
from pyslm import geometry as slm


def writeVTKLayers(layers: List[slm.Layer], filename: str, scaleFactor: float = 1e3):
    """

    :param layers: List of layers to export to VTK
    :param filename: The filename of the VTK .vtp file to write to
    :param scaleFactor: The scale factor to use for the Z coordinates
    """

    fileScanVectors = []
    scanData = []

    appendData = []
    appendDataOffsets = []

    i = 0
    with open(filename, 'w') as fp:

        fp.write('<VTKFile type="PolyData" version="1.0" byte_order="LittleEndian" header_type="UInt32">')

        for layer in layers:

            i += 1

            scanVectors = []
            for geom in layer.geometry:

                if isinstance(geom, slm.HatchGeometry):
                    coords = geom.coords.reshape(-1, 2, 2)
                elif isinstance(geom, slm.ContourGeometry):
                    coords = np.hstack([geom.coords, np.roll(geom.coords, -1, axis=0)])[:-1, :].reshape(-1, 2, 2)
                elif isinstance(geom, slm.PointsGeometry):
                    # Note that we duplicate the coordinates to represent emulate hatch vectors
                    coords = np.tile(geom.coords.reshape(-1, 2), (1, 2)).reshape(-1, 2, 2)

                scanVectors.append(coords)

            if len(scanVectors) == 0:
                continue

            scanVectors = np.vstack(scanVectors).reshape(-1, 2)

            # Append the Layer's Z coordinate to the scan vectors
            scanVectors = np.hstack([scanVectors, np.ones((len(scanVectors), 1)) * layer.z / scaleFactor])

            # Append the current layer's scan vectors to the accumulated list of scan vectors
            fileScanVectors.append(scanVectors)

            #pointsGeom = layer.getPointsGeometry()
            #if len(pointsGeom) > 0:
            #    coords = np.vstack([geom.coords for geom in pointsGeom])

            """
            Plot the sequential index of the hatch vector and generating the colourmap by using the cumulative distance
            across all the scan vectors in order to normalise the length based effectively on the distance
            """
            scanVectors2 = scanVectors[:, :2].reshape(-1, 2, 2)
            delta = scanVectors2[:, 1, :] - scanVectors2[:, 0, :]
            dist = np.sqrt(delta[:, 0] * delta[:, 0] + delta[:, 1] * delta[:, 1])
            cumDist = np.cumsum(dist)

            scanData.append(cumDist)

        fileScanVectors = np.vstack(fileScanVectors)

        # Add all the line collections to the figure
        fp.write('\t<PolyData>')
        fp.write('\t\t<Piece NumberOfPoints="{:d}" NumberOfVerts="0" NumberOfLines="{:d}" NumberOfStrips="0" NumberOfPolys="0">\n'.format(
                fileScanVectors.shape[0], int(fileScanVectors.shape[0] / 2.0)))

        # Write the Points Array
        fp.write('\t\t\t <Points>\n')
        fp.write('\t\t\t\t<DataArray type="Float32" NumberOfComponents="3" Name="Points" format="appended" offset="0" />')
        fp.write('\t\t\t</Points>\n')

        """
        Pack the coordinate data
        """
        s = fileScanVectors.astype(np.float32)
        b = struct.pack('=%sf' % s.size, *s.flatten())
        appendData.append(b)
        appendDataOffsets.append(len(b) + 4)  # Increment by 4 bytes for uint32 size

        # Write the point data
        fp.write('\t\t\t <PointData Scalars="GlobalNodeID">\n')
        orderId = np.arange(0, len(fileScanVectors) * 2, 1)

        """
        Pack the data
        """
        fp.write('\t\t\t\t<DataArray type="Int32" NumberOfComponents="1" '
                 'Name="GlobalNodeID" format="appended" offset="{:d}" />\n'.format(appendDataOffsets[-1]))
        s = orderId.astype(np.int32)
        b = struct.pack('=%sI' % s.size, *s.flatten())
        appendData.append(b)
        appendDataOffsets.append(appendDataOffsets[-1] + len(b) + 4)

        # p.write('\t\t\t\t</DataArray>\n')
        fp.write('\t\t\t</PointData>\n')

        """
        Write out the individual hatch vectors 
        """
        con = np.arange(0, len(fileScanVectors), 1)

        fp.write('\t\t\t<Lines>\n')
        fp.write('\t\t\t\t<DataArray type="Int32" NumberOfComponents="1" '
                 'Name="connectivity" format="appended" offset="{:d}" />\n'.format(appendDataOffsets[-1]))

        s = con.astype(np.int32)
        b = struct.pack('=%sI' % s.size, *s.flatten())
        appendData.append(b)
        appendDataOffsets.append(appendDataOffsets[-1] + len(b) + 4)

        fp.write('\t\t\t\t<DataArray type="Int32" NumberOfComponents="1" '
                 'Name="offsets" format="appended" offset="{:d}" />\n'.format(appendDataOffsets[-1]))

        offsets = np.arange(2, (len(fileScanVectors) + 1), 2)
        s = offsets.astype(np.int32)
        b = struct.pack('=%sI' % s.size, *s.flatten())

        appendData.append(b)
        appendDataOffsets.append(appendDataOffsets[-1] + len(b) + 4)

        fp.write('\t\t\t</Lines>\n')
        fp.write('\t\t\t</Piece>\n')
        fp.write('\t</PolyData>\n')
        fp.write('<AppendedData encoding="raw">\n')

    """
    The previous file must be opened with binary flag to permit the raw data to be written in the 
    <AppendedData> section. 
    """

    with open(filename, 'ab') as fp:

        # Inside the <AppendedData> section, write the underscore character to indicate the start
        # of the raw data section
        fp.write(bytes("\t_", 'ascii'))

        # Iterate across the stored data sections, pack the data and write both the size (bytes) and the
        # data into the file's Appended
        for data in appendData:

            # Write the size of the data section
            numSize = struct.pack('I', len(data))
            fp.write(numSize)

            # Write the raw data section
            fp.write(data)

        # Write the end of the raw data section and close the XML file
        fp.write(bytes('</AppendedData>\n', 'ascii'))
        fp.write(bytes('</VTKFile>\n', 'ascii'))

"""
Create an example data structure for a build file and export this to VTK
"""
testLayer = slm.Layer()
testLayer.z = 1000
testGeom = slm.HatchGeometry()
testGeom.coords = [[1.0, 2.0],
                   [10.0, 2.0],
                   [5.0, 4.0],
                   [15.0, 4.0]]

testLayer.geometry.append(testGeom)

# Write the layer geometry to VTK file format
writeVTKLayers([testLayer], './testFile.vtp')