index | modules | gitlab

# Source code for ase.io.vtkxml

    import numpy as np

    fast = False

    def write_vti(filename, atoms, data=None):
        from vtk import vtkDoubleArray, vtkStructuredPoints, vtkXMLImageDataWriter

        # if isinstance(fileobj, str):
        #     fileobj = paropen(fileobj, 'w')

        if isinstance(atoms, list):
            if len(atoms) > 1:
                raise ValueError('Can only write one configuration to a VTI file!')
            atoms = atoms[0]

        if data is None:
            raise ValueError('VTK XML Image Data (VTI) format requires data!')

        data = np.asarray(data)

        if data.dtype == complex:
            data = np.abs(data)

        cell = atoms.get_cell()

        if not np.all(cell == np.diag(np.diag(cell))):
            raise ValueError('Unit cell must be orthogonal')

        bbox = np.array(list(zip(np.zeros(3), cell.diagonal()))).ravel()

        # Create a VTK grid of structured points
        spts = vtkStructuredPoints()
        spts.SetWholeBoundingBox(bbox)
        spts.SetDimensions(data.shape)
        spts.SetSpacing(cell.diagonal() / data.shape)
        # spts.SetSpacing(paw.gd.h_c * Bohr)

        # print('paw.gd.h_c * Bohr=',paw.gd.h_c * Bohr)
        # print('atoms.cell.diagonal() / data.shape=', cell.diagonal()/data.shape)
        # assert np.all(paw.gd.h_c * Bohr==cell.diagonal()/data.shape)

        # s = paw.wfs.kpt_u[0].psit_nG[0].copy()
        # data = paw.get_pseudo_wave_function(band=0, kpt=0, spin=0, pad=False)
        # spts.point_data.scalars = data.swapaxes(0,2).flatten()
        # spts.point_data.scalars.name = 'scalars'

        # Allocate a VTK array of type double and copy data
        da = vtkDoubleArray()
        da.SetName('scalars')
        da.SetNumberOfComponents(1)
        da.SetNumberOfTuples(np.prod(data.shape))

        for i, d in enumerate(data.swapaxes(0, 2).flatten()):
            da.SetTuple1(i, d)

        # Assign the VTK array as point data of the grid
        spd = spts.GetPointData()  # type(spd) is vtkPointData
        spd.SetScalars(da)

        """
        from vtk.util.vtkImageImportFromArray import vtkImageImportFromArray
        iia = vtkImageImportFromArray()
        #iia.SetArray(Numeric_asarray(data.swapaxes(0,2).flatten()))
        iia.SetArray(Numeric_asarray(data))
        ida = iia.GetOutput()
        ipd = ida.GetPointData()
        ipd.SetName('scalars')
        spd.SetScalars(ipd.GetScalars())
        """

        # Save the ImageData dataset to a VTK XML file.
        w = vtkXMLImageDataWriter()

        if fast:
            w.SetDataModeToAppend()
            w.EncodeAppendedDataOff()
        else:
            w.SetDataModeToAscii()

        w.SetFileName(filename)
        w.SetInput(spts)
        w.Write()

    def write_vtu(filename, atoms, data=None):
        from vtk import (
            VTK_MAJOR_VERSION,
            vtkPoints,
            vtkUnstructuredGrid,
            vtkXMLUnstructuredGridWriter,
        )
        from vtk.util.numpy_support import numpy_to_vtk

        if isinstance(atoms, list):
            if len(atoms) > 1:
                raise ValueError('Can only write one configuration to a VTI file!')
            atoms = atoms[0]

        # Create a VTK grid of structured points
        ugd = vtkUnstructuredGrid()

        # add atoms as vtk Points
        p = vtkPoints()
        p.SetNumberOfPoints(len(atoms))
        p.SetDataTypeToDouble()
        for i, pos in enumerate(atoms.get_positions()):
            p.InsertPoint(i, *pos)
        ugd.SetPoints(p)

        # add atomic numbers
        numbers = numpy_to_vtk(atoms.get_atomic_numbers(), deep=1)
        ugd.GetPointData().AddArray(numbers)
        numbers.SetName("atomic numbers")

        # add tags
        tags = numpy_to_vtk(atoms.get_tags(), deep=1)
        ugd.GetPointData().AddArray(tags)
        tags.SetName("tags")

        # add covalent radii
        from ase.data import covalent_radii
        radii = numpy_to_vtk(covalent_radii[atoms.numbers], deep=1)
        ugd.GetPointData().AddArray(radii)
        radii.SetName("radii")

        # Save the UnstructuredGrid dataset to a VTK XML file.
        w = vtkXMLUnstructuredGridWriter()

        if fast:
            w.SetDataModeToAppend()
            w.EncodeAppendedDataOff()
        else:
            w.GetCompressor().SetCompressionLevel(0)
            w.SetDataModeToAscii()

        if isinstance(filename, str):
            w.SetFileName(filename)
        else:
            w.SetFileName(filename.name)
        if VTK_MAJOR_VERSION <= 5:
            w.SetInput(ugd)
        else:
            w.SetInputData(ugd)
        w.Write()