-
Notifications
You must be signed in to change notification settings - Fork 0
/
cubetools.py
254 lines (222 loc) · 8.6 KB
/
cubetools.py
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
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
# Cube data format functions, adapted from here:
# https://gist.github.com/aditya95sriram/8d1fccbb91dae93c4edf31cd6a22510f
import io
class CubeFile(object):
"""
Object which mimics a cube file opened as a file object
by returning output in the correct format, matching the
metadata of the source cube file and replacing volumetric
data with static data provided as arg to the constructor.
Doesn't copy atoms metadata, retains number of atoms, but
returns dummy atoms
Mimics file object's readline method.
params:
srcname: source file to copy metadata from
const: numeric value to return instead of volumetric data
returns: CubeFile object
"""
def __init__(self, srcname, const=1):
self.cursor = 0
self.const = const
self.src = src = open(srcname)
src.readline(); src.readline(); # comments
_debug(srcname)
self.lines = [" Cubefile created by cubetools.py\n",
" source: {0}\n".format(srcname)]
self.lines.append(src.readline()) # read natm and origin
self.natm = int(self.lines[-1].strip().split()[0])
# read cube dim and vectors along 3 axes
self.lines.extend(src.readline() for i in range(3))
self.src.close()
self.nx, self.ny, self.nz = [int(l.strip().split()[0]) for l in self.lines[3:6]]
self.remvals = self.nz
self.remrows = self.nx*self.ny
for i in range(self.natm):
self.lines.append("{0:^ 8d}".format(1) + "{0:< 12.6f}".format(0)*4 + '\n')
def __del__(self):
self.src.close()
def readline(self):
""" Mimic readline method of file object with cube file opened """
try:
retval = self.lines[self.cursor]
except IndexError:
if not self.remrows:
return ""
if self.remvals <= 6:
nval = min(6,self.remvals)
self.remrows -= 1
self.remvals = self.nz
else:
nval = 6
self.remvals -= nval
return " {0: .5E}".format(self.const)*nval + "\n"
else:
self.cursor += 1
return retval
def _getline(cube):
"""
Read a line from cube file where first field is an int
and the remaining fields are floats.
params:
cube: file object of the cube file
returns: (int, list<float>)
"""
l = cube.readline().strip().split()
return int(l[0]), map(float, l[1:])
def _putline(*args):
"""
Generate a line to be written to a cube file where
the first field is an int and the remaining fields are floats.
params:
*args: first arg is formatted as int and remaining as floats
returns: formatted string to be written to file with trailing newline
"""
s = "{0:^ 8d}".format(args[0])
s += "".join("{0:< 12.6f}".format(arg) for arg in args[1:])
return s + "\n"
def read_cube(fname):
"""
Read cube file into numpy array
params:
fname: filename of cube file
returns: (data: np.array, metadata: dict)
"""
meta = {}
with open(fname, 'r') as cube:
cube.readline(); cube.readline() # ignore comments
natm, meta['org'] = _getline(cube)
nx, meta['xvec'] = _getline(cube)
ny, meta['yvec'] = _getline(cube)
nz, meta['zvec'] = _getline(cube)
meta['atoms'] = [_getline(cube) for i in range(natm)]
data = np.zeros((nx*ny*nz))
idx = 0
for line in cube:
for val in line.strip().split():
data[idx] = float(val)
idx += 1
data = np.reshape(data, (nx, ny, nz))
return data, meta
def read_cube_string(string):
"""
Read cube file into numpy array
params:
fname: filename of cube file
returns: (data: np.array, metadata: dict)
"""
meta = {}
cube = io.StringIO()
cube.write(string)
cube.seek(0)
cube.readline(); cube.readline() # ignore comments
natm, meta['org'] = _getline(cube)
nx, meta['xvec'] = _getline(cube)
ny, meta['yvec'] = _getline(cube)
nz, meta['zvec'] = _getline(cube)
meta['atoms'] = [_getline(cube) for i in range(natm)]
data = np.zeros((nx*ny*nz))
idx = 0
for line in cube:
for val in line.strip().split():
data[idx] = float(val)
idx += 1
data = np.reshape(data, (nx, ny, nz))
return data, meta
def read_imcube(rfname, ifname = ""):
"""
Convenience function to read in two cube files at once,
where one contains the real part and the other contains the
imag part. If only one filename given, other filename is inferred.
params:
rfname: filename of cube file of real part
ifname: optional, filename of cube file of imag part
returns: np.array (real part + j*imag part)
"""
ifname = ifname or rfname.replace('real', 'imag')
_debug("reading from files", rfname, "and", ifname)
re, im = read_cube(rfname), read_cube(ifname)
fin = np.zeros(re[0].shape, dtype='complex128')
if re[1] != im[1]:
_debug("warning: meta data mismatch, real part metadata retained")
fin += re[0]
fin += 1j*im[0]
return fin, re[1]
def write_cube(data, meta, fname):
"""
Write volumetric data to cube file along
params:
data: volumetric data consisting real values
meta: dict containing metadata with following keys
atoms: list of atoms in the form (mass, [position])
org: origin
xvec,yvec,zvec: lattice vector basis
fname: filename of cubefile (existing files overwritten)
returns: None
"""
with open(fname, "w") as cube:
# first two lines are comments
cube.write(" Cubefile created by cubetools.py\n source: none\n")
natm = len(meta['atoms'])
nx, ny, nz = data.shape
cube.write(_putline(natm, *meta['org'])) # 3rd line #atoms and origin
cube.write(_putline(nx, *meta['xvec']))
cube.write(_putline(ny, *meta['yvec']))
cube.write(_putline(nz, *meta['zvec']))
for atom_mass, atom_pos in meta['atoms']:
cube.write(_putline(atom_mass, *atom_pos)) #skip the newline
for i in range(nx):
for j in range(ny):
for k in range(nz):
if (i or j or k) and k%6==0:
cube.write("\n")
cube.write(" {0: .5E}".format(data[i,j,k]))
def write_cube_string(data, meta):
"""
Write volumetric data to cube file along
params:
data: volumetric data consisting real values
meta: dict containing metadata with following keys
atoms: list of atoms in the form (mass, [position])
org: origin
xvec,yvec,zvec: lattice vector basis
fname: filename of cubefile (existing files overwritten)
returns: None
"""
cube = io.StringIO()
# first two lines are comments
cube.write(" Cubefile created by cubetools.py\n source: none\n")
natm = len(meta['atoms'])
nx, ny, nz = data.shape
cube.write(_putline(natm, *meta['org'])) # 3rd line #atoms and origin
cube.write(_putline(nx, *meta['xvec']))
cube.write(_putline(ny, *meta['yvec']))
cube.write(_putline(nz, *meta['zvec']))
for atom_mass, atom_pos in meta['atoms']:
cube.write(_putline(atom_mass, *atom_pos)) #skip the newline
for i in range(nx):
for j in range(ny):
for k in range(nz):
if (i or j or k) and k%6==0:
cube.write("\n")
cube.write(" {0: .5E}".format(data[i,j,k]))
return cube.getvalue()
def write_imcube(data, meta, rfname, ifname=""):
"""
Convenience function to write two cube files from complex valued
volumetric data, one for the real part and one for the imaginary part.
Data about atoms, origin and lattice vectors are kept same for both.
If only one filename given, other filename is inferred.
params:
data: volumetric data consisting complex values
meta: dict containing metadata with following keys
atoms: list of atoms in the form (mass, [position])
org: origin
xvec,yvec,zvec: lattice vector basis
rfname: filename of cube file containing real part
ifname: optional, filename of cube file containing imag part
returns: None
"""
ifname = ifname or rfname.replace('real', 'imag')
_debug("writing data to files", rfname, "and", ifname)
write_cube(data.real, meta, rfname)
write_cube(data.imag, meta, ifname)