Source code for ase.io.vtkxml

import numpy as np

fast = False


[docs]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()
[docs]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()