/usr/share/pyshared/ase/io/vtkxml.py is in python-ase 3.6.0.2515-1.
This file is owned by root:root, with mode 0o644.
The actual contents of the file can be viewed below.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 | import numpy as np
#from Numeric import asarray as Numeric_asarray
#from ase.units import Bohr
#from ase.parallel import paropen
fast = False
# -------------------------------------------------------------------
from vtk import vtkStructuredPoints, vtkDoubleArray, vtkXMLImageDataWriter
def write_vti(filename, atoms, data):
#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()
assert np.all(cell==np.diag(cell.diagonal())), 'Unit cell must be orthogonal'
bbox = np.array(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()
# -------------------------------------------------------------------
from vtk import vtkStructuredGrid, vtkPoints, vtkXMLStructuredGridWriter
def write_vts(filename, atoms, data=None):
raise NotImplementedError
# -------------------------------------------------------------------
from vtk import vtkUnstructuredGrid, vtkPoints, vtkXMLUnstructuredGridWriter
def write_vtu(filename, atoms, data=None):
#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 Unstructured Grid (VTI) format requires data!')
data = np.asarray(data)
if data.dtype == complex:
data = np.abs(data)
"""
cell = atoms.get_cell()
assert np.all(cell==np.diag(cell.diagonal())), 'Unit cell must be orthogonal' #TODO bounding box with general unit cell?!
bbox = np.array(zip(np.zeros(3),cell.diagonal())).ravel()
# Create a VTK grid of structured points
ugd = vtkUnstructuredGrid()
ugd.SetWholeBoundingBox(bbox)
"""
# Allocate a VTK array of type double and copy data
da = vtkDoubleArray()
da.SetName('scalars')
da.SetNumberOfComponents(3)
da.SetNumberOfTuples(len(atoms))
for i,pos in enumerate(atoms.get_positions()):
da.SetTuple3(i,pos[0],pos[1],pos[2])
"""
p = vtkPoints()
p.SetNumberOfPoints(len(atoms))
p.SetDataTypeToDouble()
for i,pos in enumerate(atoms.get_positions()):
p.InsertPoint(i,pos[0],pos[1],pos[2])
ugd.SetPoints(p)
# Assign the VTK array as point data of the grid
#upd = ugd.GetPointData() # type(spd) is vtkPointData
#upd.SetScalars(da)
# Save the UnstructuredGrid dataset to a VTK XML file.
w = vtkXMLUnstructuredGridWriter()
if fast:
w.SetDataModeToAppend()
w.EncodeAppendedDataOff()
else:
w.GetCompressor().SetCompressionLevel(0)
w.SetDataModeToAscii()
w.SetFileName(filename)
w.SetInput(ugd)
w.Write()
# -------------------------------------------------------------------
def read_vti(filename):
raise NotImplementedError
def read_vts(filename):
raise NotImplementedError
def read_vtu(filename):
raise NotImplementedError
# -------------------------------------------------------------------
from vtk import vtkXMLFileReadTester
def probe_vtkxml(filename):
"""something..."""
r = vtkXMLFileReadTester()
r.SetFileName(filename)
if r.TestReadFile():
return r.GetFileDataType()
else:
return None
|