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
|
# fmt: off
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()
|
