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
|
import numpy as np
fast = False
def write_vti(filename, atoms, data=None):
from vtk import vtkStructuredPoints, vtkDoubleArray, 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()
assert np.all(cell==np.diag(cell.diagonal())), '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 vtkUnstructuredGrid, vtkPoints, vtkXMLUnstructuredGridWriter
#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(list(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()
|
