prior.pyx 7.96 KB
Newer Older
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
import numpy as np
from itertools import product
from ase.ga.utilities import closest_distances_generator
cimport numpy as np

from libc.math cimport *

from cymem.cymem cimport Pool
cimport cython

# Custom functions
ctypedef struct Point:
    double coord[3]

cdef Point subtract(Point p1, Point p2):
    cdef Point p
    p.coord[0] = p1.coord[0] - p2.coord[0]
    p.coord[1] = p1.coord[1] - p2.coord[1]
    p.coord[2] = p1.coord[2] - p2.coord[2]
    return p

cdef Point add(Point p1, Point p2):
    cdef Point p
    p.coord[0] = p1.coord[0] + p2.coord[0]
    p.coord[1] = p1.coord[1] + p2.coord[1]
    p.coord[2] = p1.coord[2] + p2.coord[2]
    return p

cdef double norm(Point p):
    return sqrt(p.coord[0]*p.coord[0] + p.coord[1]*p.coord[1] + p.coord[2]*p.coord[2])

cdef double euclidean(Point p1, Point p2):
    return norm(subtract(p1,p2))

def convert_atom_types(num):
    cdef int Natoms = len(num)
    cdef list atomic_types = sorted(list(set(num)))
    cdef int Ntypes = len(atomic_types)
    cdef list num_converted = [0]*Natoms
    cdef int i, j
    for i in range(Natoms):
        for j in range(Ntypes):
            if num[i] == atomic_types[j]:
                num_converted[i] = j
    return num_converted

47
class RepulsivePrior():
Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
48
    def __init__(self, rcut=6, ratio_of_covalent_radii=0.7):
49
        self.rcut = rcut
50
51
        self.ratio_of_covalent_radii = ratio_of_covalent_radii
        self.initialized = False
52

53
    def initialize_from_atoms(self, atoms):
54
55
56
57
58
59
60
61
62
63
        self.pbc = atoms.get_pbc()
        self.cell = atoms.get_cell()

        self.cell_displacements = self.__get_neighbour_cells_displacement(self.pbc, self.cell)
        self.Ncells = len(self.cell_displacements)

        num = atoms.get_atomic_numbers()
        atomic_types = sorted(list(set(num)))
        self.Ntypes = len(atomic_types)

64
        blmin_dict = closest_distances_generator(num, ratio_of_covalent_radii=self.ratio_of_covalent_radii)
65
66
67
68
69
70
71

        self.blmin = -np.ones((self.Ntypes, self.Ntypes))
        for i1, type1 in enumerate(atomic_types):
            for i2, type2 in enumerate(atomic_types):
                self.blmin[i1,i2] = blmin_dict[(type1, type2)]
        self.blmin = self.blmin.reshape(-1).tolist()

72
73
        self.initialized = True

74
    def energy(self, a):
75
76
        if not self.initialized:
            self.initialize_from_atoms(a)
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

        # Memory allocation pool
        cdef Pool mem
        mem = Pool()

        cdef int Natoms = a.get_number_of_atoms()
        cdef double rcut = self.rcut

        cdef list x_np = a.get_positions().tolist()
        cdef Point *x
        x = <Point*>mem.alloc(Natoms, sizeof(Point))
        cdef int m
        for m in range(Natoms):
            x[m].coord[0] = x_np[m][0]
            x[m].coord[1] = x_np[m][1]
            x[m].coord[2] = x_np[m][2]

        # Get neighbourcells and convert to Point-struct
        cdef int Ncells = self.Ncells
        cdef list cell_displacements_old = self.cell_displacements
        cdef Point *cell_displacements
        cell_displacements = <Point*>mem.alloc(Ncells, sizeof(Point))
        for m in range(Ncells):
            cell_displacements[m].coord[0] = cell_displacements_old[m][0]
            cell_displacements[m].coord[1] = cell_displacements_old[m][1]
            cell_displacements[m].coord[2] = cell_displacements_old[m][2]

        # Convert 2body bondtype list into c-array
        cdef int Ntypes = self.Ntypes
        cdef list blmin_old = self.blmin
        cdef double *blmin
        blmin = <double*>mem.alloc(Ntypes*Ntypes, sizeof(double))
        for m in range(Ntypes*Ntypes):
            blmin[m] = blmin_old[m]

        # Get converted atom Ntypes
        cdef list num_converted_old = convert_atom_types(a.get_atomic_numbers())
        cdef int *num_converted
        num_converted = <int*>mem.alloc(Natoms, sizeof(int))
        for m in range(Natoms):
            num_converted[m] = num_converted_old[m]

        cdef double E=0, rmin, radd
        cdef int i, j
        for i in range(Natoms):
            xi = x[i]
            for cell_index in range(Ncells):
                displacement = cell_displacements[cell_index]
                for j in range(Natoms):
                    rmin = blmin[num_converted[i]*Ntypes + num_converted[j]]
                    radd = 1 - rmin
                    xj = add(x[j], displacement)
                    r = euclidean(xi, xj)
                    if r > 1e-5 and r < rcut:
                        r_scaled = r + radd
                        E += 1/pow(r_scaled,12)
        return E/2

    def forces(self, a):
136
137
138
        if not self.initialized:
            self.initialize_from_atoms(a)

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
        # Memory allocation pool
        cdef Pool mem
        mem = Pool()

        cdef int Natoms = a.get_number_of_atoms()
        cdef int dim = 3
        cdef double rcut = self.rcut

        cdef list x_np = a.get_positions().tolist()
        cdef Point *x
        x = <Point*>mem.alloc(Natoms, sizeof(Point))
        cdef int m
        for m in range(Natoms):
            x[m].coord[0] = x_np[m][0]
            x[m].coord[1] = x_np[m][1]
            x[m].coord[2] = x_np[m][2]

        # Get neighbourcells and convert to Point-struct
        cdef int Ncells = self.Ncells
        cdef list cell_displacements_old = self.cell_displacements
        cdef Point *cell_displacements
        cell_displacements = <Point*>mem.alloc(Ncells, sizeof(Point))
        for m in range(Ncells):
            cell_displacements[m].coord[0] = cell_displacements_old[m][0]
            cell_displacements[m].coord[1] = cell_displacements_old[m][1]
            cell_displacements[m].coord[2] = cell_displacements_old[m][2]

        # Convert 2body bondtype list into c-array
        cdef int Ntypes = self.Ntypes
        cdef list blmin_old = self.blmin
        cdef double *blmin
        blmin = <double*>mem.alloc(Ntypes*Ntypes, sizeof(double))
        for m in range(Ntypes*Ntypes):
            blmin[m] = blmin_old[m]

        # Get converted atom Ntypes
        cdef list num_converted_old = convert_atom_types(a.get_atomic_numbers())
        cdef int *num_converted
        num_converted = <int*>mem.alloc(Natoms, sizeof(int))
        for m in range(Natoms):
            num_converted[m] = num_converted_old[m]

        #cdef double rmin = 0.7 * self.cov_dist
        #cdef double radd = 1 - rmin

        # Initialize Force object
        cdef double *dE
        dE = <double*>mem.alloc(Natoms * dim, sizeof(double))

        cdef int i, j, k
        for i in range(Natoms):
            xi = x[i]
            for cell_index in range(Ncells):
                displacement = cell_displacements[cell_index]
                for j in range(Natoms):
                    rmin = blmin[num_converted[i]*Ntypes + num_converted[j]]
                    radd = 1 - rmin
                    xj = add(x[j], displacement)
                    r = euclidean(xi,xj)
                    if r > 1e-5 and r < rcut:
                        r_scaled = r + radd
                        rijVec = subtract(xi,xj)

                        for k in range(dim):
                            dE[dim*i + k] += -12*rijVec.coord[k] / (pow(r_scaled,13)*r)
                            dE[dim*j + k] += 12*rijVec.coord[k] / (pow(r_scaled,13)*r)

        dE_np = np.zeros((Natoms, dim))
        for i in range(Natoms):
            for k in range(dim):
                dE_np[i,k] = dE[dim*i + k]

        return - dE_np.reshape(-1)/2

    def __get_neighbour_cells_displacement(self, pbc, cell):
        # Calculate neighbour cells
        cell_vec_norms = np.linalg.norm(cell, axis=0)
        neighbours = []
        for i in range(3):
            if pbc[i]:
                ncellmax = int(np.ceil(abs(self.rcut/cell_vec_norms[i]))) + 1  # +1 because atoms can be outside unitcell.
                neighbours.append(range(-ncellmax,ncellmax+1))
            else:
                neighbours.append([0])

        neighbourcells_disp = []
        for x,y,z in product(*neighbours):
            xyz = (x,y,z)
            displacement = np.dot(xyz, cell)
            neighbourcells_disp.append(list(displacement))

        return neighbourcells_disp