1import numpy as np 2 3 4fast = False 5 6 7def write_vti(filename, atoms, data=None): 8 from vtk import vtkStructuredPoints, vtkDoubleArray, vtkXMLImageDataWriter 9 10 # if isinstance(fileobj, str): 11 # fileobj = paropen(fileobj, 'w') 12 13 if isinstance(atoms, list): 14 if len(atoms) > 1: 15 raise ValueError('Can only write one configuration to a VTI file!') 16 atoms = atoms[0] 17 18 if data is None: 19 raise ValueError('VTK XML Image Data (VTI) format requires data!') 20 21 data = np.asarray(data) 22 23 if data.dtype == complex: 24 data = np.abs(data) 25 26 cell = atoms.get_cell() 27 28 if not np.all(cell == np.diag(np.diag(cell))): 29 raise ValueError('Unit cell must be orthogonal') 30 31 bbox = np.array(list(zip(np.zeros(3), cell.diagonal()))).ravel() 32 33 # Create a VTK grid of structured points 34 spts = vtkStructuredPoints() 35 spts.SetWholeBoundingBox(bbox) 36 spts.SetDimensions(data.shape) 37 spts.SetSpacing(cell.diagonal() / data.shape) 38 # spts.SetSpacing(paw.gd.h_c * Bohr) 39 40 # print('paw.gd.h_c * Bohr=',paw.gd.h_c * Bohr) 41 # print('atoms.cell.diagonal() / data.shape=', cell.diagonal()/data.shape) 42 # assert np.all(paw.gd.h_c * Bohr==cell.diagonal()/data.shape) 43 44 # s = paw.wfs.kpt_u[0].psit_nG[0].copy() 45 # data = paw.get_pseudo_wave_function(band=0, kpt=0, spin=0, pad=False) 46 # spts.point_data.scalars = data.swapaxes(0,2).flatten() 47 # spts.point_data.scalars.name = 'scalars' 48 49 # Allocate a VTK array of type double and copy data 50 da = vtkDoubleArray() 51 da.SetName('scalars') 52 da.SetNumberOfComponents(1) 53 da.SetNumberOfTuples(np.prod(data.shape)) 54 55 for i, d in enumerate(data.swapaxes(0, 2).flatten()): 56 da.SetTuple1(i, d) 57 58 # Assign the VTK array as point data of the grid 59 spd = spts.GetPointData() # type(spd) is vtkPointData 60 spd.SetScalars(da) 61 62 """ 63 from vtk.util.vtkImageImportFromArray import vtkImageImportFromArray 64 iia = vtkImageImportFromArray() 65 #iia.SetArray(Numeric_asarray(data.swapaxes(0,2).flatten())) 66 iia.SetArray(Numeric_asarray(data)) 67 ida = iia.GetOutput() 68 ipd = ida.GetPointData() 69 ipd.SetName('scalars') 70 spd.SetScalars(ipd.GetScalars()) 71 """ 72 73 # Save the ImageData dataset to a VTK XML file. 74 w = vtkXMLImageDataWriter() 75 76 if fast: 77 w.SetDataModeToAppend() 78 w.EncodeAppendedDataOff() 79 else: 80 w.SetDataModeToAscii() 81 82 w.SetFileName(filename) 83 w.SetInput(spts) 84 w.Write() 85 86 87def write_vtu(filename, atoms, data=None): 88 from vtk import (VTK_MAJOR_VERSION, vtkUnstructuredGrid, vtkPoints, 89 vtkXMLUnstructuredGridWriter) 90 from vtk.util.numpy_support import numpy_to_vtk 91 92 if isinstance(atoms, list): 93 if len(atoms) > 1: 94 raise ValueError('Can only write one configuration to a VTI file!') 95 atoms = atoms[0] 96 97 # Create a VTK grid of structured points 98 ugd = vtkUnstructuredGrid() 99 100 # add atoms as vtk Points 101 p = vtkPoints() 102 p.SetNumberOfPoints(len(atoms)) 103 p.SetDataTypeToDouble() 104 for i, pos in enumerate(atoms.get_positions()): 105 p.InsertPoint(i, *pos) 106 ugd.SetPoints(p) 107 108 # add atomic numbers 109 numbers = numpy_to_vtk(atoms.get_atomic_numbers(), deep=1) 110 ugd.GetPointData().AddArray(numbers) 111 numbers.SetName("atomic numbers") 112 113 # add tags 114 tags = numpy_to_vtk(atoms.get_tags(), deep=1) 115 ugd.GetPointData().AddArray(tags) 116 tags.SetName("tags") 117 118 # add covalent radii 119 from ase.data import covalent_radii 120 radii = numpy_to_vtk(covalent_radii[atoms.numbers], deep=1) 121 ugd.GetPointData().AddArray(radii) 122 radii.SetName("radii") 123 124 # Save the UnstructuredGrid dataset to a VTK XML file. 125 w = vtkXMLUnstructuredGridWriter() 126 127 if fast: 128 w.SetDataModeToAppend() 129 w.EncodeAppendedDataOff() 130 else: 131 w.GetCompressor().SetCompressionLevel(0) 132 w.SetDataModeToAscii() 133 134 if isinstance(filename, str): 135 w.SetFileName(filename) 136 else: 137 w.SetFileName(filename.name) 138 if VTK_MAJOR_VERSION <= 5: 139 w.SetInput(ugd) 140 else: 141 w.SetInputData(ugd) 142 w.Write() 143