kernel.py 21.2 KB
Newer Older
1
2
3
4
5
6
7
8
9
#! /usr/bin/env python
import soap
import soap.tools

import os
import numpy as np
import logging
from momo import osio, endl, flush

10
11
12
13
14
15
16
17
18
# THERE IS SOMETHING WRONG WITH THE GRADIENTS HERE (DIRECTION SEEMS FINE, BUT MAGNITUDE IS NOT?)! REMARK ~ok, error source found
# TODO Unit test for global spectrum                                    REMARK ~ok, automatize
# TODO Check outer kernel derivative                                    REMARK ~ok, automatize
# TODO Check normalization derivative (compute on C++ level already)    REMARK ~ok, -> C++

# TODO PCA + identification of eigenstructures
# TODO Sample Bethe tree

class TrajectoryLogger(object):
Carl Poelking's avatar
Carl Poelking committed
19
20
    def __init__(self, outfile, mode='w'):
        self.ofs = open(outfile, mode)
21
22
23
24
25
26
27
28
29
30
31
    def logFrame(self, structure):
        # Write first frame
        self.ofs.write('%d\n\n' % structure.n_particles)
        for p in structure.particles:
            r = p.pos
            self.ofs.write('%s %+1.7f %+1.7f %+1.7f\n' % (p.type, r[0], r[1], r[2]))
        return
    def close(self):
        self.ofs.close()
        return

Carl Poelking's avatar
Carl Poelking committed
32
class KernelAdaptorGeneric(object):
33
34
    def __init__(self, options):
        return
Carl Poelking's avatar
Carl Poelking committed
35
36
    def getListAtomic(self, spectrum):
        return spectrum
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
    def adapt(self, spectrum):
        # IX = [
        #    X1 ->,
        #    X2 ->,
        #    ...
        # ]
        IX = np.zeros((0,0), dtype='complex128')
        dimX = -1
        # TODO More adaptors:
        # TODO Particle ID mask, center-type mask
        # TODO For type-resolved adaptors, increase dimensionality accordingly
        for atomic_i in spectrum:
            #print atomic_i.getCenter().pos
            #if pid_list:
            #    pid = atomic_i.getCenter().id
            #    if not pid in pid_list:
            #        logging.debug("Skip, adapt atomic, pid = %d" % pid)
            #        continue
            
            Xi_unnorm, Xi_norm = self.adaptScalar(atomic_i)
            dimX = Xi_norm.shape[0]
            #Xi = np.real(atomic_i.getPower("","").array)
            #dimX = Xi.shape[0]*Xi.shape[1]
            #Xi = Xi.reshape((dimX))
            
            #print "dim(X) =", dimX
            if not IX.any():
                IX = np.copy(Xi_norm) # TODO Is this necessary?
                IX.resize((1, dimX))
            else:
                i = IX.shape[0]
                IX.resize((i+1, dimX))
                IX[-1,:] = Xi_norm
            #print IX
        return IX
    def adaptScalar(self, atomic):
        X = np.real(atomic.getPower("","").array)
        dimX = X.shape[0]*X.shape[1]
        X = np.real(X.reshape((dimX)))
        # Normalize
        X_norm = X/np.dot(X,X)**0.5
        return X, X_norm
    def adaptGradients(self, atomic, nb_pid, X):
Carl Poelking's avatar
Carl Poelking committed
80
        # NOTE X must not be normalized (=> X = X')
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
        dxnkl_pid = atomic.getPowerGradGeneric(nb_pid)
        dX_dx = dxnkl_pid.getArrayGradX()
        dX_dy = dxnkl_pid.getArrayGradY()
        dX_dz = dxnkl_pid.getArrayGradZ()        
        dimX = dX_dx.shape[0]*dX_dx.shape[1]
        dX_dx = np.real(dX_dx.reshape((dimX)))
        dX_dy = np.real(dX_dy.reshape((dimX)))
        dX_dz = np.real(dX_dz.reshape((dimX)))
        # Normalize
        mag_X = np.dot(X,X)**0.5
        dX_dx = dX_dx/mag_X - np.dot(X, dX_dx)/mag_X**3 * X
        dX_dy = dX_dy/mag_X - np.dot(X, dX_dy)/mag_X**3 * X
        dX_dz = dX_dz/mag_X - np.dot(X, dX_dz)/mag_X**3 * X
        return dX_dx, dX_dy, dX_dz
        
Carl Poelking's avatar
Carl Poelking committed
96
97
98
class KernelAdaptorGlobalGeneric(object):
    def __init__(self, options):
        return
Carl Poelking's avatar
Carl Poelking committed
99
100
    def getListAtomic(self, spectrum):
        return [ spectrum.getGlobal() ]
Carl Poelking's avatar
Carl Poelking committed
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
    def adapt(self, spectrum):
        # EXTRACT A SET OF CENTER-BASED POWER EXPANSIONS
        # Here: only global
        IX = np.zeros((0,0), dtype='complex128')
        dimX = -1
        atomic_global = spectrum.getGlobal()
        Xi_unnorm, Xi_norm = self.adaptScalar(atomic_global)
        dimX = Xi_norm.shape[0]
        IX = np.copy(Xi_norm)
        IX.resize((1,dimX))
        return IX
    def adaptScalar(self, atomic):
        # EXTRACT POWER EXPANSION FROM ATOMIC SPECTRUM
        # Here: type "":"" (= generic)
        X = np.real(atomic.getPower("","").array)
        dimX = X.shape[0]*X.shape[1]
        X = np.real(X.reshape((dimX)))
        # Normalize
        X_norm = X/np.dot(X,X)**0.5
        return X, X_norm
    def adaptGradients(self, atomic, nb_pid, X):
        # NOTE X is not normalized (=> X = X')
        dxnkl_pid = atomic.getPowerGradGeneric(nb_pid)
        dX_dx = dxnkl_pid.getArrayGradX()
        dX_dy = dxnkl_pid.getArrayGradY()
        dX_dz = dxnkl_pid.getArrayGradZ()        
        dimX = dX_dx.shape[0]*dX_dx.shape[1]
        dX_dx = np.real(dX_dx.reshape((dimX)))
        dX_dy = np.real(dX_dy.reshape((dimX)))
        dX_dz = np.real(dX_dz.reshape((dimX)))
        # Normalize
        mag_X = np.dot(X,X)**0.5
        dX_dx = dX_dx/mag_X - np.dot(X, dX_dx)/mag_X**3 * X
        dX_dy = dX_dy/mag_X - np.dot(X, dX_dy)/mag_X**3 * X
        dX_dz = dX_dz/mag_X - np.dot(X, dX_dz)/mag_X**3 * X
        return dX_dx, dX_dy, dX_dz

138
139
140
141
142
143
144
145
146
147
148
class KernelFunctionDot(object):
    def __init__(self, options):
        self.delta = float(options.get('kernel.delta'))
        self.xi = float(options.get('kernel.xi'))
        return
    def computeDot(self, IX, X, xi, delta):
        return delta**2 * np.dot(IX,X)**xi
    def compute(self, IX, X):
        return self.computeDot(IX, X, self.xi, self.delta)    
    def computeDerivativeOuter(self, IX, X):
        c = self.computeDot(IX, X, self.xi-1, self.delta)
Carl Poelking's avatar
Carl Poelking committed
149
150
151
152
153
154
155
156
        return self.xi*np.diag(c).dot(IX)
    def computeBlock(self, IX, return_distance=False):
        # Order such that: IX[i] => fingerprint observation i
        K = self.delta**2 * IX.dot(IX.T)**self.xi
        if return_distance:
            return (abs(1.-K))**0.5
        else:
            return K
Carl Poelking's avatar
Carl Poelking committed
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
    def computeBlockDot(self, IX, return_distance=False):
        return self.computeBlock(IX, return_distance)

class KernelFunctionDotHarmonic(object):
    """
    C_i = ( d^2 [ IX.X ]^xi - mu_i )^2
    """
    def __init__(self, options):
        self.delta = float(options.get('kernel.delta'))
        self.xi = float(options.get('kernel.xi'))
        self.mu = float(options.get('kernel.mu'))
        self.kfctdot = KernelFunctionDot(options)
        return
    def compute(self, IX, X):
        if type(self.mu) == float:
            mu = np.zeros((IX.shape[0]))
Carl Poelking's avatar
Carl Poelking committed
173
            mu.fill(self.mu)
Carl Poelking's avatar
Carl Poelking committed
174
175
176
177
178
179
180
181
182
183
184
185
186
            self.mu = mu
        C = self.kfctdot.compute(IX, X)
        return (C-self.mu)**2
    def computeDerivativeOuter(self, IX, X):
        if type(self.mu) == float:
            mu = np.zeros((IX.shape[0]))
            mu.fill(self.mu)
            self.mu = mu
        C = self.kfctdot.compute(IX, X)
        IC = self.kfctdot.computeDerivativeOuter(IX, X)
        return 2*(C-self.mu)*IC
    def computeBlockDot(self, IX, return_distance=False):
        return self.kfctdot.computeBlock(IX, return_distance)
Carl Poelking's avatar
Carl Poelking committed
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202

class KernelFunctionDotHarmonicDist(object):
    def __init__(self, options):
        self.d0 = float(options.get('kernel.dotharmdist_d0'))
        self.eps = float(options.get('kernel.dotharmdist_eps'))
        self.kfctdot = KernelFunctionDot(options)
        return
    def compute(self, IX, X):
        C = self.kfctdot.compute(IX, X)
        D = (1.-C+self.eps)**0.5
        return (D-self.d0)**2
    def computeDerivativeOuter(self, IX, X):
        C = self.kfctdot.compute(IX, X)
        D = (1.-C+self.eps)**0.5
        IC = self.kfctdot.computeDerivativeOuter(IX, X)
        return 2*(D-self.d0)*0.5/D*(-1.)*IC
Carl Poelking's avatar
Carl Poelking committed
203
        
Carl Poelking's avatar
Carl Poelking committed
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
class KernelFunctionDotLj(object):
    def __init__(self, options):
        self.sigma = float(options.get('kernel.lj_sigma'))
        self.eps_cap = float(options.get('kernel.lj_eps_cap'))
        self.kfctdot = KernelFunctionDot(options)
    def compute(self, IX, X):
        C = self.kfctdot.compute(IX, X)
        D = (1.-C+self.eps_cap)**0.5
        return (self.sigma/D)**12 - (self.sigma/D)**6
    def computeDerivativeOuter(self, IX, X):
        C = self.kfctdot.compute(IX, X)
        D = (1.-C+self.eps_cap)**0.5
        IC = self.kfctdot.computeDerivativeOuter(IX, X)
        return (-6.*self.sigma**12*(1./D)**7 + 3.*self.sigma**6*(1./D)**4)*IC
    def computeBlockDot(self, IX, return_distance=False):
        return self.kfctdot.computeBlock(IX, return_distance)        
Carl Poelking's avatar
Carl Poelking committed
220

Carl Poelking's avatar
Carl Poelking committed
221
class KernelFunctionDot3Harmonic(object):
Carl Poelking's avatar
Carl Poelking committed
222
223
224
225
226
    def __init__(self, options):
        self.kfctdot = KernelFunctionDot(options)
    def compute(self, IX_A, IX_B, X):
        C_A = self.kfctdot.compute(IX_A, X)
        C_B = self.kfctdot.compute(IX_B, X)
Carl Poelking's avatar
Carl Poelking committed
227
228
229
230
231
        # TODO Generalize this to multi-environment representation:
        assert IX_A.shape[0] == 1
        assert IX_B.shape[0] == 1
        C_AB = self.kfctdot.compute(IX_A, IX_B[0])
        return (C_A - C_B)**2 + (C_A - C_AB)**2 + (C_B - C_AB)**2
Carl Poelking's avatar
Carl Poelking committed
232
    def computeDerivativeOuter(self, IX_A, IX_B, X):
Carl Poelking's avatar
Carl Poelking committed
233
234
235
236
237
238
239
240
241
242
        C_A = self.kfctdot.compute(IX_A, X)
        C_B = self.kfctdot.compute(IX_B, X)
        # TODO Generalize this to multi-environment representation:
        assert IX_A.shape[0] == 1
        assert IX_B.shape[0] == 1
        C_AB = self.kfctdot.compute(IX_A, IX_B[0])
        IC_A = self.kfctdot.computeDerivativeOuter(IX_A, X)
        IC_B = self.kfctdot.computeDerivativeOuter(IX_B, X)
        return 2*(C_A - C_B)*(IC_A - IC_B) + 2*(C_A - C_AB)*IC_A + 2*(C_B - C_AB)*IC_B

Carl Poelking's avatar
Carl Poelking committed
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
class KernelFunctionDot3HarmonicDist(object):
    def __init__(self, options):
        self.kfctdot = KernelFunctionDot(options)
        self.d0 = float(options.get('kernel.dot3harmdist_d0'))
        self.eps = float(options.get('kernel.dot3harmdist_eps'))
    def compute(self, IX_A, IX_B, X):
        C_A = self.kfctdot.compute(IX_A, X)
        C_B = self.kfctdot.compute(IX_B, X)
        # TODO Generalize this to multi-environment representation:
        assert IX_A.shape[0] == 1
        assert IX_B.shape[0] == 1
        C_AB = self.kfctdot.compute(IX_A, IX_B[0])
        D_A = (1. - C_A + self.eps)**0.5
        D_B = (1. - C_B + self.eps)**0.5
        D_AB = (1. - C_AB + self.eps)**0.5
        return (D_A - D_B)**2 + (D_A - D_AB)**2 + (D_B - D_AB)**2 + (D_A - self.d0)**2 + (D_B - self.d0)**2
    def computeDerivativeOuter(self, IX_A, IX_B, X):
        C_A = self.kfctdot.compute(IX_A, X)
        C_B = self.kfctdot.compute(IX_B, X)
        # TODO Generalize this to multi-environment representation:
        assert IX_A.shape[0] == 1
        assert IX_B.shape[0] == 1
        C_AB = self.kfctdot.compute(IX_A, IX_B[0])
        D_A = (1. - C_A + self.eps)**0.5
        D_B = (1. - C_B + self.eps)**0.5
        D_AB = (1. - C_AB + self.eps)**0.5
        IC_A = self.kfctdot.computeDerivativeOuter(IX_A, X)
        IC_B = self.kfctdot.computeDerivativeOuter(IX_B, X)
        return 2*(D_A - D_B)*(0.5/D_A*(-1.)*IC_A - 0.5/D_B*(-1.)*IC_B) + 2*(D_A - D_AB)*0.5/D_A*(-1.)*IC_A + 2*(D_B - D_AB)*0.5/D_B*(-1.)*IC_B + 2*(D_A-self.d0)*1./D_A*(-1.)*IC_A + 2*(D_B-self.d0)*1./D_B*(-1.)*IC_B

Carl Poelking's avatar
Carl Poelking committed
273
274
275
276
KernelAdaptorFactory = {
'generic': KernelAdaptorGeneric, 
'global-generic': KernelAdaptorGlobalGeneric 
}     
Carl Poelking's avatar
Carl Poelking committed
277
278

KernelFunctionFactory = { 
Carl Poelking's avatar
Carl Poelking committed
279
'dot': KernelFunctionDot, 
Carl Poelking's avatar
Carl Poelking committed
280
'dot-harmonic': KernelFunctionDotHarmonic, 
Carl Poelking's avatar
Carl Poelking committed
281
'dot-harmonic-dist': KernelFunctionDotHarmonicDist,
Carl Poelking's avatar
Carl Poelking committed
282
'dot-lj': KernelFunctionDotLj,
Carl Poelking's avatar
Carl Poelking committed
283
284
'dot-3-harmonic': KernelFunctionDot3Harmonic, 
'dot-3-harmonic-dist': KernelFunctionDot3HarmonicDist
Carl Poelking's avatar
Carl Poelking committed
285
}
286

Carl Poelking's avatar
Carl Poelking committed
287
class KernelPotential(object):
288
289
290
291
292
293
294
295
    def __init__(self, options):
        logging.info("Construct kernel potential ...")
        self.basis = soap.Basis(options)
        self.options = options        
        # CORE DATA
        self.IX = None
        self.alpha = None
        self.dimX = None # <- Also used as flag set by first ::acquire call
Carl Poelking's avatar
Carl Poelking committed
296
        self.labels = []
297
298
299
300
301
302
        # KERNEL
        logging.info("Choose kernel function ...")
        self.kernelfct = KernelFunctionFactory[options.get('kernel.type')](options)   
        # ADAPTOR
        logging.info("Choose adaptor ...")
        self.adaptor = KernelAdaptorFactory[options.get('kernel.adaptor')](options)
Carl Poelking's avatar
Carl Poelking committed
303
        self.use_global_spectrum = True if options.get('kernel.adaptor') == 'global-generic' else False
304
305
306
        # INCLUSIONS / EXCLUSIONS
        # -> Already be enforced when computing spectra
        return
Carl Poelking's avatar
Carl Poelking committed
307
308
309
310
311
312
313
314
    def computeKernelMatrix(self, return_distance=False):
        return self.kernelfct.computeBlock(self.IX, return_distance)
    def computeDotKernelMatrix(self, return_distance=False):
        # Even if the, e.g., dot-harmonic kernel is used for optimization, 
        # the dot-product based kernel matrix may still be relevant
        return self.kernelfct.computeBlockDot(self.IX, return_distance)
    def nConfigs(self):
        return self.IX.shape[0]
Carl Poelking's avatar
Carl Poelking committed
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
    def importAcquire(self, IX_acqu, alpha):
        n_acqu = IX_acqu.shape[0]
        dim_acqu = IX_acqu.shape[1]
        alpha_acqu = np.zeros((n_acqu))
        alpha_acqu.fill(alpha)
        if not self.dimX:
            # First time ...
            self.dimX = dim_acqu
            self.IX = IX_acqu
            self.alpha = alpha_acqu
        else:
            # Check and extend ...
            assert self.dimX == dim_acqu # Acquired descr. should match linear dim. of previous descr.'s
            I = self.IX.shape[0]
            self.IX.resize((I+n_acqu, self.dimX))
            self.IX[I:I+n_acqu,:] = IX_acqu
            self.alpha.resize((I+n_acqu))
            self.alpha[I:I+n_acqu] = alpha_acqu
        #print self.alpha
        #print self.IX
        logging.info("Imported %d environments." % n_acqu)
        return
Carl Poelking's avatar
Carl Poelking committed
337
    def acquire(self, structure, alpha, label=None):
338
339
340
341
342
        logging.info("Acquire ...")
        spectrum = soap.Spectrum(structure, self.options, self.basis)
        spectrum.compute()
        spectrum.computePower()
        spectrum.computePowerGradients()
Carl Poelking's avatar
Carl Poelking committed
343
344
        if self.use_global_spectrum:
            spectrum.computeGlobal()
345
346
347
348
349
350
351
352
        # New X's
        logging.info("Adapt spectrum ...")
        IX_acqu = self.adaptor.adapt(spectrum)
        n_acqu = IX_acqu.shape[0]
        dim_acqu = IX_acqu.shape[1]
        # New alpha's
        alpha_acqu = np.zeros((n_acqu))
        alpha_acqu.fill(alpha)
Carl Poelking's avatar
Carl Poelking committed
353
354
355
        # Labels
        labels = [ label for i in range(n_acqu) ]
        self.labels = self.labels + labels
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
        if not self.dimX:
            # First time ...
            self.dimX = dim_acqu
            self.IX = IX_acqu
            self.alpha = alpha_acqu
        else:
            # Check and extend ...
            assert self.dimX == dim_acqu # Acquired descr. should match linear dim. of previous descr.'s
            I = self.IX.shape[0]
            self.IX.resize((I+n_acqu, self.dimX))
            self.IX[I:I+n_acqu,:] = IX_acqu
            self.alpha.resize((I+n_acqu))
            self.alpha[I:I+n_acqu] = alpha_acqu
        #print self.alpha
        #print self.IX
        logging.info("Acquired %d environments." % n_acqu)
        return
Carl Poelking's avatar
Carl Poelking committed
373
    def computeEnergy(self, structure, return_prj_mat=False):
374
375
376
377
378
379
        logging.info("Start energy ...")
        energy = 0.0
        spectrum = soap.Spectrum(structure, self.options, self.basis)
        spectrum.compute()
        spectrum.computePower()
        spectrum.computePowerGradients()
Carl Poelking's avatar
Carl Poelking committed
380
381
382
        # TODO The kernel policy should take care of this:
        if self.use_global_spectrum:
            spectrum.computeGlobal()
383
384
385
386
        IX_acqu = self.adaptor.adapt(spectrum)        
        n_acqu = IX_acqu.shape[0]
        dim_acqu = IX_acqu.shape[1]
        logging.info("Compute energy from %d atomic environments ..." % n_acqu)
Carl Poelking's avatar
Carl Poelking committed
387
        projection_matrix = []
388
389
        for n in range(n_acqu):
            X = IX_acqu[n]
390
            #print "X_MAG", np.dot(X,X)
391
392
            ic = self.kernelfct.compute(self.IX, X)
            energy += self.alpha.dot(ic)
Carl Poelking's avatar
Carl Poelking committed
393
394
395
396
397
398
            #print "Projection", ic
            projection_matrix.append(ic)
        if return_prj_mat:
            return energy, projection_matrix
        else:
            return energy
399
400
401
402
403
404
405
406
407
408
409
    def computeForces(self, structure, verbose=False):
        logging.info("Start forces ...")
        if verbose:
            for p in structure: print p.pos
        forces = [ np.zeros((3)) for i in range(structure.n_particles) ]
        logging.info("Compute forces on %d particles ..." % structure.n_particles)
        # Compute X's, dX's
        spectrum = soap.Spectrum(structure, self.options, self.basis)
        spectrum.compute()
        spectrum.computePower()
        spectrum.computePowerGradients()
Carl Poelking's avatar
Carl Poelking committed
410
411
412
413
414
        if self.use_global_spectrum:
            atomic_global = spectrum.computeGlobal()
            spectrum_iter = [ atomic_global ]
        else:
            spectrum_iter = spectrum
415
        # Extract & compute force
Carl Poelking's avatar
Carl Poelking committed
416
        for atomic in spectrum_iter:
417
            pid = atomic.getCenter().id if not self.use_global_spectrum else -1
418
419
420
421
422
423
424
            #if not pid in self.pid_list_force:
            #    logging.debug("Skip forces derived from environment with pid = %d" % pid)
            #    continue
            #if pid != 1: continue
            nb_pids = atomic.getNeighbourPids()
            logging.info("  Center %d" % (pid))
            # neighbour-pid-independent kernel "prevector" (outer derivative)
425
            X_unnorm, X_norm = self.adaptor.adaptScalar(atomic)
Carl Poelking's avatar
Carl Poelking committed
426
            dIC = self.kernelfct.computeDerivativeOuter(self.IX, X_norm)
427
428
429
430
            alpha_dIC = self.alpha.dot(dIC)
            for nb_pid in nb_pids:
                # Force on neighbour
                logging.info("    -> Nb %d" % (nb_pid))
431
                dX_dx, dX_dy, dX_dz = self.adaptor.adaptGradients(atomic, nb_pid, X_unnorm)
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
                
                force_x = -alpha_dIC.dot(dX_dx)
                force_y = -alpha_dIC.dot(dX_dy)
                force_z = -alpha_dIC.dot(dX_dz)
                
                forces[nb_pid-1][0] += force_x
                forces[nb_pid-1][1] += force_y
                forces[nb_pid-1][2] += force_z
                #print forces
                #raw_input('...')
        return forces
        
def restore_positions(structure, positions):
    idx = 0
    for part in structure:
        part.pos = positions[idx]
        idx += 1
    return [ part.pos for part in structure ]

def perturb_positions(structure, exclude_pid=[]):
    for part in structure:
        if part.id in exclude_pid: continue
        dx = np.random.uniform(-1.,1.)
        dy = np.random.uniform(-1.,1.)
        dz = np.random.uniform(-1.,1.)
        part.pos = part.pos + 0.1*np.array([dx,dy,dz])
    return [ part.pos for part in structure ]
459
    
460
def random_positions(structure, exclude_pid=[], b0=-1., b1=1., x=True, y=True, z=True):
461
462
    for part in structure:
        if part.id in exclude_pid: continue
463
464
465
        dx = np.random.uniform(b0,b1) if x else 0.
        dy = np.random.uniform(b0,b1) if y else 0.
        dz = np.random.uniform(b0,b1) if z else 0.
466
467
        part.pos = np.array([dx,dy,dz])
    return [ part.pos for part in structure ]
468
469
470
471
472
473
474
475
476

def apply_force_step(structure, forces, scale, constrain_particles=[]):
    max_step = 0.05
    max_f = 0.0
    for f in forces:
        df = scale*np.dot(f,f)**0.5
        if df > max_f: max_f = df
    if max_f > max_step: scale = scale*max_step/max_f
    else: pass
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
    #print "Scale =", scale
    idx = -1
    for part in structure:
        idx += 1
        if part.id in constrain_particles: 
            print "Skip force step, pid =", part.id
            continue
        #if np.random.uniform(0.,1.) > 0.5: continue
        part.pos = part.pos + scale*forces[idx]
    return [ part.pos for part in structure ]

def apply_force_norm_step(structure, forces, scale, constrain_particles=[]):
    min_step = 0.5
    max_f = 0.0
    for f in forces:
        df = np.dot(f,f)**0.5
        if df > max_f: max_f = df
    scale = min_step/max_f
495
496
497
498
499
500
501
502
503
504
    idx = -1
    for part in structure:
        idx += 1
        if part.id in constrain_particles: 
            print "Skip force step, pid =", part.id
            continue
        #if np.random.uniform(0.,1.) > 0.5: continue
        part.pos = part.pos + scale*forces[idx]
    return [ part.pos for part in structure ]

Carl Poelking's avatar
Carl Poelking committed
505
def evaluate_energy(positions, structure, kernelpot, opt_pids, verbose=False, ofs=None, average=False):
506
507
508
509
510
511
512
513
514
515
516
    if verbose: print "Energy"
    # Impose positions
    pid_pos = positions.reshape((opt_pids.shape[0],3))
    for pidx, pid in enumerate(opt_pids):
        pos = pid_pos[pidx,:]
        particle = structure.getParticle(pid)        
        particle.pos = pos        
    for part in structure:
        if verbose: print part.id, part.type, part.pos
    # Evaluate energy function
    energy = kernelpot.computeEnergy(structure)
Carl Poelking's avatar
Carl Poelking committed
517
    if average: energy /= kernelpot.nConfigs()
518
519
    # Log
    if ofs: ofs.logFrame(structure)
Carl Poelking's avatar
Carl Poelking committed
520
521
    if verbose: print energy
    print energy
522
523
    return energy

Carl Poelking's avatar
Carl Poelking committed
524
def evaluate_energy_gradient(positions, structure, kernelpot, opt_pids, verbose=False, ofs=None, average=False):
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
    if verbose: print "Forces"
    # Adjust positions
    pid_pos = positions.reshape((opt_pids.shape[0],3))
    for pidx, pid in enumerate(opt_pids):
        pos = pid_pos[pidx,:]
        particle = structure.getParticle(pid)        
        particle.pos = pos        
    for part in structure:
        if verbose: print part.id, part.type, part.pos
    # Evaluate forces
    forces = kernelpot.computeForces(structure)
    gradients = -1.*np.array(forces)
    opt_pidcs = opt_pids-1
    gradients_short = gradients[opt_pidcs]
    gradients_short = gradients_short.flatten()
Carl Poelking's avatar
Carl Poelking committed
540
541
    if average:
        gradients_short /= kernelpot.nConfigs()
542
543
544
545
    if verbose: print gradients_short
    #forces[2] = 0. # CONSTRAIN TO Z=0 PLANE
    return gradients_short