-
Notifications
You must be signed in to change notification settings - Fork 1
/
locpot_import.pyx
204 lines (176 loc) · 8.49 KB
/
locpot_import.pyx
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
import numpy as np
cimport numpy as np
from atomic_symbols_and_numbers import atomic_symbol_to_number, atomic_number_to_symbol
DTYPE_FLOAT = np.float64
ctypedef np.float64_t DTYPE_FLOAT_t
DTYPE_INT = np.int
ctypedef np.int_t DTYPE_INT_t
# Reads the voxel vectors from file and checks if they are orthogonal
def check_grid_orthogonality(filename):
cdef int i
cdef double scaling_factor
cdef np.ndarray[DTYPE_FLOAT_t, ndim=2] cell_vectors = np.zeros((3, 3), dtype=DTYPE_FLOAT)
in_file = open(filename, 'r')
is_orthogonal = False
is_z_orthogonal = False
# read comment line
title = in_file.readline()
# read unit cell
splitline = in_file.readline().split()
scaling_factor = float(splitline[0])
for i in range(3):
splitline = in_file.readline().split()
cell_vectors[i, 0] = scaling_factor * float(splitline[0])
cell_vectors[i, 1] = scaling_factor * float(splitline[1])
cell_vectors[i, 2] = scaling_factor * float(splitline[2])
if (np.dot(cell_vectors[0, :], cell_vectors[1, :]) == 0 and
np.dot(cell_vectors[0, :], cell_vectors[2, :]) == 0 and
np.dot(cell_vectors[1, :], cell_vectors[2, :]) == 0):
is_orthogonal = True
if (cell_vectors[2, 0] == 0 and cell_vectors[2, 1] == 0 and
cell_vectors[0, 2] == 0 and cell_vectors[1, 2] == 0):
is_z_orthogonal = True
return is_orthogonal, is_z_orthogonal
def read_locpot(filename, orthogonal=True, return_unstructured_data=False):
cdef int ia, n_atoms, ix, iy, iz, idata, i_atom_type, ia_per_type
cdef double scaling_factor
cdef np.ndarray[DTYPE_INT_t, ndim=1] n_voxels = np.zeros(3, dtype=DTYPE_INT)
cdef np.ndarray[DTYPE_FLOAT_t, ndim=1] voxel_diffs = np.zeros(3, dtype=DTYPE_FLOAT)
cdef np.ndarray[DTYPE_FLOAT_t, ndim=1] vector_coeffs = np.zeros(3, dtype=DTYPE_FLOAT)
cdef np.ndarray[DTYPE_FLOAT_t, ndim=2] cell_vectors = np.zeros((3, 3), dtype=DTYPE_FLOAT)
cdef np.ndarray[DTYPE_FLOAT_t, ndim=2] voxel_vectors = np.zeros((3, 3), dtype=DTYPE_FLOAT)
cdef np.ndarray[DTYPE_INT_t, ndim=1] n_atoms_per_type
cdef np.ndarray[DTYPE_INT_t, ndim=1] atom_types
cdef np.ndarray[DTYPE_FLOAT_t, ndim=2] atom_pos
cdef np.ndarray[DTYPE_FLOAT_t, ndim=1] xs, ys, zs
cdef np.ndarray[DTYPE_FLOAT_t, ndim=3] grid_data
cdef np.ndarray[DTYPE_FLOAT_t, ndim=2] volume_data
in_file = open(filename, 'r')
# read comment line
title = in_file.readline()
# read unit cell
splitline = in_file.readline().split()
scaling_factor = float(splitline[0])
for ia in range(3):
splitline = in_file.readline().split()
cell_vectors[ia, 0] = scaling_factor * float(splitline[0])
cell_vectors[ia, 1] = scaling_factor * float(splitline[1])
cell_vectors[ia, 2] = scaling_factor * float(splitline[2])
# read atom types and positions
splitline = in_file.readline().split()
atom_symbols = splitline
splitline = in_file.readline().split()
n_atoms_per_type = np.zeros(len(splitline), dtype=DTYPE_INT)
for ia, atoms_per_symbol in enumerate(splitline):
n_atoms_per_type[ia] = int(atoms_per_symbol)
n_atoms = n_atoms_per_type.sum()
atom_types = np.zeros(n_atoms, dtype=DTYPE_INT)
atom_pos = np.zeros((n_atoms, 3), dtype=DTYPE_FLOAT)
splitline = in_file.readline().split()
coordinate_type = splitline[0]
ia = 0
if coordinate_type == 'Direct':
for i_atom_type, atom_symbol in enumerate(atom_symbols):
for ia_per_type in range(n_atoms_per_type[i_atom_type]):
atom_types[ia] = atomic_symbol_to_number(atom_symbol)
splitline = in_file.readline().split()
vector_coeffs[0] = float(splitline[0])
vector_coeffs[1] = float(splitline[1])
vector_coeffs[2] = float(splitline[2])
if orthogonal:
atom_pos[ia, 0] = vector_coeffs[0]*cell_vectors[0, 0]
atom_pos[ia, 1] = vector_coeffs[1]*cell_vectors[1, 1]
atom_pos[ia, 2] = vector_coeffs[2]*cell_vectors[2, 2]
else:
atom_pos[ia, 0] = vector_coeffs[0]*cell_vectors[0, 0] + vector_coeffs[1]*cell_vectors[1, 0] + vector_coeffs[2]*cell_vectors[2, 0]
atom_pos[ia, 1] = vector_coeffs[0]*cell_vectors[0, 1] + vector_coeffs[1]*cell_vectors[1, 1] + vector_coeffs[2]*cell_vectors[2, 1]
atom_pos[ia, 2] = vector_coeffs[0]*cell_vectors[0, 2] + vector_coeffs[1]*cell_vectors[1, 2] + vector_coeffs[2]*cell_vectors[2, 2]
ia = ia + 1
elif coordinate_type == 'Cartesian':
for i_atom_type, atom_symbol in enumerate(atom_symbols):
for ia_per_type in range(n_atoms_per_type[i_atom_type]):
atom_types[ia] = atomic_symbol_to_number(atom_symbol)
splitline = in_file.readline().split()
atom_pos[ia, 0] = scaling_factor * float(splitline[0])
atom_pos[ia, 1] = scaling_factor * float(splitline[1])
atom_pos[ia, 2] = scaling_factor * float(splitline[2])
ia = ia +1
else:
raise Exception('Unknown coordinate type: {}'.format(coordinate_type))
# read volumetric electrostatic potential data
splitline = in_file.readline().split()
if not splitline: # probably always empty line here
splitline = in_file.readline().split()
n_voxels[0] = int(splitline[0])
n_voxels[1] = int(splitline[1])
n_voxels[2] = int(splitline[2])
voxel_vectors[0, :] = cell_vectors[0, :] / n_voxels[0]
voxel_vectors[1, :] = cell_vectors[1, :] / n_voxels[1]
voxel_vectors[2, :] = cell_vectors[2, :] / n_voxels[2]
if orthogonal:
grid_data = np.zeros((n_voxels[0], n_voxels[1], n_voxels[2]), dtype=DTYPE_FLOAT)
idata = 0
ix = 0
iy = 0
iz = 0
while idata < n_voxels[0]*n_voxels[1]*n_voxels[2]:
splitline = in_file.readline().split()
for data_entry in splitline:
grid_data[ix, iy, iz] = float(data_entry)
idata = idata + 1
ix = ix + 1
if ix == n_voxels[0]:
iy = iy + 1
ix = 0
if iy == n_voxels[1]:
iz = iz + 1
iy = 0
# create the grid points
voxel_diffs[0] = voxel_vectors[0, 0]
voxel_diffs[1] = voxel_vectors[1, 1]
voxel_diffs[2] = voxel_vectors[2, 2]
xs = np.arange(n_voxels[0], dtype=DTYPE_FLOAT)*voxel_diffs[0]
ys = np.arange(n_voxels[1], dtype=DTYPE_FLOAT)*voxel_diffs[1]
zs = np.arange(n_voxels[2], dtype=DTYPE_FLOAT)*voxel_diffs[2]
return atom_types, atom_pos, xs, ys, zs, grid_data
elif not return_unstructured_data:
grid_data = np.zeros((n_voxels[0], n_voxels[1], n_voxels[2]), dtype=DTYPE_FLOAT)
idata = 0
ix = 0
iy = 0
iz = 0
while idata < n_voxels[0]*n_voxels[1]*n_voxels[2]:
splitline = in_file.readline().split()
for data_entry in splitline:
grid_data[ix, iy, iz] = float(data_entry)
idata = idata + 1
ix = ix + 1
if ix == n_voxels[0]:
iy = iy + 1
ix = 0
if iy == n_voxels[1]:
iz = iz + 1
iy = 0
return atom_types, atom_pos, n_voxels, voxel_vectors, grid_data
else:
volume_data = np.zeros((n_voxels[0]*n_voxels[1]*n_voxels[2], 4), dtype=DTYPE_FLOAT)
idata = 0
ix = 0
iy = 0
iz = 0
while idata < n_voxels[0]*n_voxels[1]*n_voxels[2]:
splitline = in_file.readline().split()
for data_entry in splitline:
volume_data[idata, 0] = ix*voxel_vectors[0, 0] + iy*voxel_vectors[1, 0] + iz*voxel_vectors[2, 0]
volume_data[idata, 1] = ix*voxel_vectors[0, 1] + iy*voxel_vectors[1, 1] + iz*voxel_vectors[2, 1]
volume_data[idata, 2] = ix*voxel_vectors[0, 2] + iy*voxel_vectors[1, 2] + iz*voxel_vectors[2, 2]
volume_data[idata, 3] = float(data_entry)
idata = idata + 1
ix = ix + 1
if ix == n_voxels[0]:
iy = iy + 1
ix = 0
if iy == n_voxels[1]:
iz = iz + 1
iy = 0
return atom_types, atom_pos, n_voxels, voxel_vectors, volume_data