gpr.py 8.89 KB
Newer Older
Malthe Kjær Bisbo's avatar
Malthe Kjær Bisbo committed
1
import numpy as np
Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
2
3
from scipy.linalg import cholesky, cho_solve, solve_triangular
from scipy.optimize import minimize
Malthe Kjær Bisbo's avatar
Malthe Kjær Bisbo committed
4

Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
5
from kernel import gauss_kernel, double_gauss_kernel
6
7
8
from descriptor.fingerprint import Fingerprint
from prior.prior import repulsive_prior

Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
9
import warnings
10
11
12
13
14
15

class gpr_memory():
    def __init__(self, descriptor, prior):
        self.descriptor = descriptor
        self.prior = prior

Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
16
        self.initialize_data()
17

Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
    def initialize_data(self):
        self.energies = None
        self.features = None
        self.prior_values = None
        
    def get_data(self):
        return self.energies, self.features, self.prior_values
        
    def save_data(self, atoms_list, add_data=True):
        if not add_data:
            self.initialize_data()
        
        self.save_energies(atoms_list)
        self.save_features(atoms_list)
        self.save_prior_values(atoms_list)
33

Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
34
35
36
    def save_energies(self, atoms_list):
        energies_save = np.array([a.get_potential_energy() for a in atoms_list])
        if self.energies is None:
37
38
39
40
            self.energies = energies_save
        else:
            self.energies = np.r_[self.energies, energies_save]
    
Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
41
42
43
    def save_features(self, atoms_list):
        features_save = self.descriptor.get_featureMat(atoms_list)
        if self.features is None:
44
            self.features = features_save
Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
45
        else:
46
47
            self.features = np.r_[self.features, features_save]

Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
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
    def save_prior_values(self, atoms_list):
        if self.prior is not None:
            prior_values_save = np.array([self.prior.energy(a) for a in atoms_list])
            if self.prior_values is None:
                self.prior_values = prior_values_save
            else:
                self.prior_values = np.r_[self.prior_values, prior_values_save]
        else:
            self.prior_values = 0

class gpr_memory_test():
    def __init__(self, descriptor, prior):
        self.initialize_data()

    def initialize_data(self):
        self.E = None
        self.X = None
        self.prior_values = None
        
    def get_data(self):
        return self.E, self.X, self.prior_values
        
    def save_data(self, atoms_list, add_data=True):
        if not add_data:
            self.initialize_data()
        
        self.save_energies(atoms_list)
        self.save_features(atoms_list)
        self.save_prior_values(atoms_list)

    def save_energies(self, atoms_list):
        E_save = np.array([a.get_potential_energy() for a in atoms_list])
        if self.E is None:
            self.E = E_save
        else:
            self.E = np.r_[self.E, E_save]
    
    def save_features(self, atoms_list):
        X_save = self.descriptor.get_featureMat(atoms_list)
        if self.X is None:
            self.X = X_save
        else:
            self.X = np.r_[self.X, X_save]

    def save_prior_values(self, atoms_list):
        if self.prior is not None:
            prior_values_save = np.array([self.prior.energy(a) for a in atoms_list])
            if self.prior_values is None:
                self.prior_values = prior_values_save
            else:
                self.prior_values = np.r_[self.prior_values, prior_values_save]
        else:
            self.prior_values = 0
101
102


Malthe Kjær Bisbo's avatar
Malthe Kjær Bisbo committed
103
104
105
106
107
108
109
110
111
112
113
class gpr():
    """Gaussian Process Regression
    
    Parameters:
    
    descriptor:
    
    kernel:
    
    prior:
    """
Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
114
    def __init__(self, descriptor=None, kernel='double', prior=None, n_restarts_optimizer=1):
115
116
117
118
119
120
121
122
123
        if descriptor is None:
            self.descriptor = Fingerprint()
        else:
            self.descriptor = descriptor

        if kernel is 'single':
            self.kernel = gauss_kernel()
        elif kernel is 'double':
            self.kernel = double_gauss_kernel()
Malthe Kjær Bisbo's avatar
Malthe Kjær Bisbo committed
124
125
        else:
            self.kernel = kernel
126
127
128
129
130
131

        if prior is None:
            self.prior = repulsive_prior()
        else:
            self.prior = prior

Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
132
133
        self.n_restarts_optimizer = n_restarts_optimizer

134
        self.memory = gpr_memory(self.descriptor, self.prior)
Malthe Kjær Bisbo's avatar
Malthe Kjær Bisbo committed
135

Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
136
    def predict_energy(self, a, eval_std=False):
Malthe Kjær Bisbo's avatar
Malthe Kjær Bisbo committed
137
138
139
        x = self.descriptor.get_feature(a)
        k = self.kernel.kernel_vector(x, self.X)

Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
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
        E = np.dot(k,self.alpha) + self.bias + self.prior.energy(a)

        if eval_std:
            # Lines 5 and 6 in GPML
            vk = np.dot(self.K_inv, k)
            K0 = self.kernel.kernel(x,x)
            E_std = np.sqrt(K0 - np.dot(k, vk))
            return E, E_std
        else:
            return E

    def predict_forces(self, a, eval_std=False):

        # Calculate descriptor and its gradient
        x = self.descriptor.get_feature(a)
        x_ddr = self.descriptor.get_featureGradient(a).T

        # Calculate kernel and its derivative
        k_ddx = self.kernel.kernel_jacobian(x, self.X)
        k_ddr = np.dot(k_ddx, x_ddr)

        F = -np.dot(k_ddr.T, self.alpha) + self.prior.forces(a)

        if eval_std:
            k = self.kernel.kernel_vector(x, self.X)
            vk = np.dot(self.K_inv, k)
            K0 = self.kernel.kernel_value(x,x)
            g = K0 - np.dot(k.T, vk)
            assert g >= 0
            F_std = -1/np.sqrt(g) * np.dot(k_ddr.T, vk)
            return F, F_std
        else:
            return F
173
174

    def set_bias(self):
Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
175
176
177
178
179
180
181
182
        self.bias = np.mean(self.memory.energies - self.memory.prior_values)

    def train(self, atoms_list=None):
        if atoms_list is not None:
            self.memory.save_data(atoms_list)
            self.set_bias()
        self.E, self.X, self.prior_values = self.memory.get_data()
        self.Y = self.E - self.prior_values - self.bias
183
        
Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
184
185
        K = self.kernel(self.X)
        L = cholesky(K, lower=True)
186
        
Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
187
188
189
190
191
192
        self.alpha = cho_solve((L, True), self.Y)
        self.K_inv = cho_solve((L, True), np.eye(K.shape[0]))
    
    def optimize_hyperparameters(self, atoms_list=None, comm=None):
        if self.n_restarts_optimizer == 0:
            self.train(atoms_list)
Malthe Kjær Bisbo's avatar
Malthe Kjær Bisbo committed
193

Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
194
195
196
197
198
        if atoms_list is not None:
            self.memory.save_data(atoms_list)
            self.set_bias()
        self.E, self.X, self.prior_values = self.memory.get_data()
        self.Y = self.E - self.prior_values - self.bias
Malthe Kjær Bisbo's avatar
Malthe Kjær Bisbo committed
199

Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
        results = []
        for i in range(self.n_restarts_optimizer):
            #if i == 0 and comm.rank == 0:
            #    theta_initial = self.kernel.theta
            #else
            theta_initial = np.random.uniform(self.kernel.theta_bounds[:, 0],
                                              self.kernel.theta_bounds[:, 1])
            res = self.constrained_optimization(theta_initial)
            results.append(res)

        index_min = np.argmin(np.array([r[1] for r in results]))
        res_min = results[index_min]
        if comm is not None:
            pass
        self.kernel.theta = res_min[0]
        self.lml = -res_min[1]
Malthe Kjær Bisbo's avatar
Malthe Kjær Bisbo committed
216

Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
217
        self.train()
Malthe Kjær Bisbo's avatar
Malthe Kjær Bisbo committed
218
    
Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
219
220
221
    def neg_log_marginal_likelihood(self, theta=None, eval_gradient=True):
        if theta is not None:
            self.kernel.theta = theta
Malthe Kjær Bisbo's avatar
Malthe Kjær Bisbo committed
222

Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
223
224
225
226
        if eval_gradient:
            K, K_gradient = self.kernel(self.X, eval_gradient)
        else:
            K = self.kernel(self.X)
Malthe Kjær Bisbo's avatar
Malthe Kjær Bisbo committed
227

Malthe Kjær Bisbo's avatar
update    
Malthe Kjær Bisbo committed
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
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
        L = cholesky(K, lower=True)
        alpha = cho_solve((L, True), self.Y)

        #log_likelihood = -0.5 * np.einsum("i,i", Y, alpha)
        lml = -0.5 * np.dot(self.Y, alpha)
        lml -= np.sum(np.log(np.diag(L)))
        lml -= K.shape[0]/2 * np.log(2*np.pi)

        if eval_gradient:
            # Equation (5.9) in GPML
            K_inv = cho_solve((L, True), np.eye(K.shape[0]))
            tmp = np.einsum("i,j->ij", alpha, alpha) - K_inv

            lml_gradient = 0.5*np.einsum("ij,kij->k", tmp, K_gradient)
            return -lml, -lml_gradient
        else:
            return -lml

    def constrained_optimization(self, theta_initial):
        result = minimize(self.neg_log_marginal_likelihood,
                          x0=theta_initial,
                          bounds=self.kernel.theta_bounds,
                          method='L-BFGS-B',
                          jac=True,
                          options={'gtol': 1e-2, 'ftol': 1e-4})
        if not result.success:
            warnings.warn(f"L_BFGS_B terminated with state: {result.message}")
        return result.x, result.fun

    def numerical_neg_lml(self, dx=1e-4):
        N_data = self.X.shape[0]
        theta = np.copy(self.kernel.theta)
        N_hyper = len(theta)
        lml_ddTheta = np.zeros((N_hyper))
        for i in range(N_hyper):
            theta_up = np.copy(theta)
            theta_down = np.copy(theta)
            theta_up[i] += 0.5*dx
            theta_down[i] -= 0.5*dx

            d_theta = np.exp(theta_up[i])-np.exp(theta_down[i])
            
            lml_up = self.neg_log_marginal_likelihood(theta_up, eval_gradient=False)
            lml_down = self.neg_log_marginal_likelihood(theta_down, eval_gradient=False)
            #lml_ddTheta[i] = (lml_up - lml_down)/dx
            lml_ddTheta[i] = (lml_up - lml_down)/d_theta
        return lml_ddTheta
    
    def get_calculator():
        pass
Malthe Kjær Bisbo's avatar
Malthe Kjær Bisbo committed
278