Statistiques
| Révision :

root / ase / calculators / jacapo / jacapo.py @ 19

Historique | Voir | Annoter | Télécharger (140,74 ko)

1
'''
2
python module for ASE2-free and Numeric-free dacapo
3

4
U{John Kitchin<mailto:jkitchin@andrew.cmu.edu>} December 25, 2008
5

6
This module supports numpy directly.
7

8
* ScientificPython2.8 is required
9

10
 - this is the first version to use numpy by default.
11

12
see https://wiki.fysik.dtu.dk/stuff/nc/ for dacapo netcdf variable
13
documentation
14
'''
15

    
16
__docformat__ = 'restructuredtext'
17

    
18
import exceptions, glob, os, pickle, string
19
from Scientific.IO.NetCDF import NetCDFFile as netCDF
20
import numpy as np
21
import subprocess as sp
22

    
23
import validate
24
import changed 
25

    
26
import logging
27
log = logging.getLogger('Jacapo')
28

    
29
handler = logging.StreamHandler()
30
formatter = logging.Formatter('''\
31
%(levelname)-10s function: %(funcName)s lineno: %(lineno)-4d \
32
%(message)s''')
33
handler.setFormatter(formatter)
34
log.addHandler(handler)
35

    
36

    
37
class DacapoRunning(exceptions.Exception):
38
    """Raised when ncfile.status = 'running'"""
39
    pass
40

    
41
class DacapoAborted(exceptions.Exception):
42
    """Raised when ncfile.status = 'aborted'"""
43
    pass
44

    
45
class DacapoInput(exceptions.Exception):
46
    ''' raised for bad input variables'''
47
    pass
48

    
49
class DacapoAbnormalTermination(exceptions.Exception):
50
    """Raised when text file does not end correctly"""
51
    pass
52

    
53
def read(ncfile):
54
    '''return atoms and calculator from ncfile
55

56
    >>> atoms, calc = read('co.nc')
57
    '''
58
    calc = Jacapo(ncfile)
59
    atoms = calc.get_atoms() #this returns a copy
60
    return (atoms, calc)
61

    
62
class Jacapo:
63
    '''
64
    Python interface to the Fortran DACAPO code
65
    '''
66
    
67
    __name__ = 'Jacapo'
68
    __version__ = 0.4
69
    
70
    #dictionary of valid input variables and default settings
71
    default_input = {'atoms':None,
72
                     'pw':350,
73
                     'dw':350,
74
                     'xc':'PW91',
75
                     'nbands':None,
76
                     'ft':0.1,
77
                     'kpts':(1,1,1),
78
                     'spinpol':False,
79
                     'fixmagmom':None,
80
                     'symmetry':False,
81
                     'calculate_stress':False,
82
                     'dipole':{'status':False,
83
                               'mixpar':0.2,
84
                               'initval':0.0,
85
                               'adddipfield':0.0,
86
                               'position':None},                         
87
                     'status':'new',
88
                     'pseudopotentials':None,
89
                     'extracharge':None,
90
                     'extpot':None,
91
                     'fftgrid':None,
92
                     'ascii_debug':'Off',
93
                     'ncoutput':{'wf':'Yes',
94
                                 'cd':'Yes',
95
                                 'efp':'Yes',
96
                                 'esp':'Yes'},
97
                     'ados':None,
98
                     'decoupling':None,
99
                     'external_dipole':None,
100
                     'convergence':{'energy':0.00001,
101
                                    'density':0.0001,
102
                                    'occupation':0.001,
103
                                    'maxsteps':None,
104
                                    'maxtime':None},
105
                     'charge_mixing':{'method':'Pulay',
106
                                      'mixinghistory':10,
107
                                      'mixingcoeff':0.1,
108
                                      'precondition':'No',
109
                                      'updatecharge':'Yes'},
110
                     'electronic_minimization':{'method':'eigsolve',
111
                                                'diagsperband':2},
112
                     'occupationstatistics':'FermiDirac',
113
                     'fftgrid':{'soft':None,
114
                                'hard':None},
115
                     'mdos':None,
116
                     'psp':None
117
                   }
118
    
119
    def __init__(self,
120
                 nc='out.nc',
121
                 outnc=None,
122
                 debug=logging.WARN,
123
                 stay_alive=False,
124
                 **kwargs):
125
        '''
126
        Initialize the Jacapo calculator
127

128
        :Parameters:
129

130
          nc : string
131
           output netcdf file, or input file if nc already exists
132

133
          outnc : string
134
           output file. by default equal to nc
135

136
          debug : integer
137
           logging debug level.
138
           
139
        Valid kwargs:
140

141
          atoms : ASE.Atoms instance
142
            atoms is an ase.Atoms object that will be attached
143
            to this calculator.
144
        
145
          pw : integer
146
            sets planewave cutoff
147

148
          dw : integer
149
            sets density cutoff
150

151
          kpts : iterable
152
            set chadi-cohen, monkhorst-pack kpt grid,
153
            e.g. kpts = (2,2,1) or explicit list of kpts
154
            
155
          spinpol : Boolean
156
            sets whether spin-polarization is used or not.
157

158
          fixmagmom : float
159
            set the magnetic moment of the unit cell. only used
160
            in spin polarize calculations
161

162
          ft : float
163
            set the Fermi temperature used in occupation smearing
164

165
          xc : string
166
            set the exchange-correlation functional.
167
            one of ['PZ','VWN','PW91','PBE','RPBE','revPBE'],
168

169
          dipole
170
            boolean
171
            turn the dipole correction on (True) or off (False)
172

173
            or:
174
            dictionary of parameters to fine-tune behavior
175
            {'status':False,
176
            'mixpar':0.2,
177
            'initval':0.0,
178
            'adddipfield':0.0,
179
            'position':None}
180
            
181
          nbands : integer
182
            set the number of bands
183

184
          symmetry : Boolean
185
            Turn symmetry reduction on (True) or off (False)
186

187
          stress : Boolean
188
            Turn stress calculation on (True) or off (False)
189

190
          debug : level for logging
191
            could be something like
192
            logging.DEBUG or an integer 0-50. The higher the integer,
193
            the less information you see set debug level (0 = off, 10 =
194
            extreme)
195
        
196
        Modification of the nc file only occurs at calculate time if needed
197

198
        >>> calc = Jacapo('CO.nc')
199
        
200
        reads the calculator from CO.nc if it exists or
201
        minimally initializes CO.nc with dimensions if it does not exist. 
202

203
        >>> calc = Jacapo('CO.nc', pw=300)
204

205
        reads the calculator from CO.nc or initializes it if
206
        it does not exist and changes the planewave cutoff energy to
207
        300eV
208

209
         >>> atoms = Jacapo.read_atoms('CO.nc')
210

211
        returns the atoms in the netcdffile CO.nc, with the calculator
212
        attached to it.
213

214
        >>> atoms, calc = read('CO.nc')
215
        
216
        '''
217
        self.debug = debug
218
        log.setLevel(debug)
219

    
220
        self.pars = Jacapo.default_input.copy()
221
        self.pars_uptodate = {}
222

    
223
        log.debug(self.pars)
224
        
225
        for key in self.pars:
226
            self.pars_uptodate[key] = False
227
            
228
        self.kwargs = kwargs
229
        self.set_psp_database()
230

    
231
        self.set_nc(nc)
232
        #assume not ready at init, rely on code later to change this 
233
        self.ready = False  
234

    
235
        # need to set a default value for stay_alive
236
        self.stay_alive = stay_alive
237
        
238
        # Jacapo('out.nc') should return a calculator with atoms in
239
        # out.nc attached or initialize out.nc
240
        if os.path.exists(nc):
241

    
242
            # for correct updating, we need to set the correct frame number
243
            # before setting atoms or calculator
244

    
245
            self._set_frame_number()
246

    
247
            self.atoms = self.read_only_atoms(nc)
248

    
249
            #if atoms object is passed to
250
            #__init__ we assume the user wants the atoms object
251
            # updated to the current state in the file.
252
            if 'atoms' in kwargs:
253
                log.debug('Updating the atoms in kwargs')
254

    
255
                atoms = kwargs['atoms']
256
                atoms.set_cell(self.atoms.get_cell())
257
                atoms.set_positions(self.atoms.get_positions())
258
                atoms.calc = self
259
                
260
            #update the parameter list from the ncfile
261
            self.update_input_parameters()
262
    
263
            self.ready = True
264
                   
265
        #change output file if needed
266
        if outnc:
267
            self.set_nc(outnc)          
268
            
269
        if len(kwargs) > 0:
270

    
271
            if 'stress' in kwargs:
272
                raise DacapoInput, '''\
273
                stress keyword is deprecated.
274
                you must use calculate_stress instead'''
275
            
276
            #make sure to set calculator on atoms if it was in kwargs
277
            #and do this first, since some parameters need info from atoms
278
            if 'atoms' in kwargs:
279
                #we need to set_atoms here so the atoms are written to
280
                #the ncfile
281
                self.set_atoms(kwargs['atoms'])
282
                kwargs['atoms'].calc = self
283
                del kwargs['atoms'] #so we don't call it in the next
284
                                    #line. we don't want to do that
285
                                    #because it will update the _frame
286
                                    #counter, and that should not be
287
                                    #done here.
288
                                
289
            self.set(**kwargs) #if nothing changes, nothing will be done
290
                    
291
    def set(self, **kwargs):
292
        '''set a parameter
293

294
        parameter is stored in dictionary that is processed later if a
295
        calculation is need.
296
        '''
297
        for key in kwargs:
298
            if key not in self.default_input:
299
                raise DacapoInput, '%s is not valid input' % key
300

    
301
            if kwargs[key] is None:
302
                continue
303
            
304
            #now check for valid input
305
            validf = 'validate.valid_%s' % key
306
            valid = eval('%s(kwargs[key])' % validf)
307
            if not valid:
308
                s = 'Warning invalid input detected for key "%s" %s'
309
                log.warn(s % (key,
310
                              kwargs[key]))
311
                raise DacapoInput, s % (key, kwargs[key])
312
                                       
313
            #now see if key has changed
314
            if key in self.pars:
315
                changef = 'changed.%s_changed' % key
316
                if os.path.exists(self.get_nc()):
317
                    notchanged = not eval('%s(self,kwargs[key])' % changef)
318
                else:
319
                    notchanged = False
320
                log.debug('%s notchanged = %s' % (key, notchanged))
321

    
322
                if notchanged:
323
                    continue
324

    
325
            log.debug('setting: %s. self.ready = False ' % key)
326
            
327
            self.pars[key] = kwargs[key]
328
            self.pars_uptodate[key] = False
329
            self.ready = False
330
            log.debug('exiting set function')
331

    
332
    def write_input(self):
333
        '''write out input parameters as needed
334

335
        you must define a self._set_keyword function that does all the
336
        actual writing.
337
        '''
338
        log.debug('Writing input variables out')
339
        log.debug(self.pars)
340
        
341
        if 'DACAPO_READONLY' in os.environ:
342
            raise Exception, 'DACAPO_READONLY set and you tried to write!'
343
        
344
        if self.ready:
345
            return
346
        # Only write out changed parameters. this function does not do
347
        # the writing, that is done for each variable in private
348
        # functions.
349
        for key in self.pars:
350
            if self.pars_uptodate[key] is False:
351
                setf = 'set_%s' % key
352

    
353
                if self.pars[key] is None:
354
                    continue
355
                
356
                log.debug('trying to call: %s' % setf)
357
                log.debug('self.%s(self.pars[key])' % setf)
358
                log.debug('key = %s' % str(self.pars[key]))
359

    
360
                if isinstance(self.pars[key], dict):
361
                    eval('self.%s(**self.pars[key])' % setf)
362
                else:
363
                    eval('self.%s(self.pars[key])' % setf)
364
                
365
                self.pars_uptodate[key] = True #update the changed flag
366

    
367
                log.debug('wrote %s: %s' % (key, str(self.pars[key])))
368

    
369
        #set Jacapo version
370
        ncf = netCDF(self.get_nc(), 'a')
371
        ncf.Jacapo_version = Jacapo.__version__
372
        ncf.sync()
373
        ncf.close()
374

    
375
    def update_input_parameters(self):
376
        '''read in all the input parameters from the netcdfile'''
377

    
378
        log.debug('Updating parameters')
379
        
380
        for key in self.default_input:
381
            getf = 'self.get_%s()' % key
382
            log.debug('getting key: %s' % key)
383
            self.pars[key] = eval(getf)
384
            self.pars_uptodate[key] = True
385
        return self.pars
386
    
387
    def initnc(self, ncfile=None):
388
        '''create an ncfile with minimal dimensions in it
389

390
        this makes sure the dimensions needed for other set functions
391
        exist when needed.'''
392
        
393
        if ncfile is None:
394
            ncfile = self.get_nc()
395
        else:
396
            self.set_nc(ncfile)
397
            
398
        log.debug('initializing %s' % ncfile)
399
                
400
        base = os.path.split(ncfile)[0]
401
        if base is not '' and not os.path.isdir(base):
402
            os.makedirs(base)
403
        
404
        ncf = netCDF(ncfile, 'w')
405
        #first, we define some dimensions we always need
406
        #unlimited
407
        ncf.createDimension('number_ionic_steps', None)
408
        ncf.createDimension('dim1', 1)
409
        ncf.createDimension('dim2', 2)
410
        ncf.createDimension('dim3', 3)
411
        ncf.createDimension('dim4', 4)
412
        ncf.createDimension('dim5', 5)
413
        ncf.createDimension('dim6', 6)
414
        ncf.createDimension('dim7', 7)
415
        ncf.createDimension('dim20', 20) #for longer strings
416
        ncf.status  = 'new'
417
        ncf.history = 'Dacapo'
418
        ncf.jacapo_version = Jacapo.__version__
419
        ncf.close()
420
        
421
        self.ready = False
422
        self._frame = 0
423
        
424
    def __del__(self):
425
        '''If calculator is deleted try to stop dacapo program
426
        '''
427
        
428
        if hasattr(self, '_dacapo'):
429
            if self._dacapo.poll()==None:
430
                self.execute_external_dynamics(stopprogram=True)
431
        #and clean up after Dacapo
432
        if os.path.exists('stop'):
433
            os.remove('stop')
434
        #remove slave files
435
        txt = self.get_txt()
436
        if txt is not None:
437
            slv = txt + '.slave*'
438
            for slvf in glob.glob(slv):
439
                os.remove(slvf)
440
        
441
    def __str__(self):
442
        '''
443
        pretty-print the calculator and atoms.
444

445
        we read everything directly from the ncfile to prevent
446
        triggering any calculations
447
        '''
448
        s = []
449
        if self.nc is None:
450
            return 'No netcdf file attached to this calculator'
451
        if not os.path.exists(self.nc):
452
            return 'ncfile does not exist yet'
453
        
454
        nc = netCDF(self.nc, 'r')
455
        s.append('  ---------------------------------')
456
        s.append('  Dacapo calculation from %s' % self.nc)
457
        if hasattr(nc, 'status'):
458
            s.append('  status = %s' % nc.status)
459
        if hasattr(nc, 'version'):
460
            s.append('  version = %s' % nc.version)
461
        if hasattr(nc, 'Jacapo_version'):
462
            s.append('  Jacapo version = %s' % nc.Jacapo_version[0])
463
        
464
        energy = nc.variables.get('TotalEnergy', None)
465
        
466
        if energy and energy[:][-1] < 1E36: # missing values get
467
                                              # returned at 9.3E36
468
            s.append('  Energy = %1.6f eV' % energy[:][-1])
469
        else:
470
            s.append('  Energy = None')
471
            
472
        s.append('')
473
        
474
        atoms = self.get_atoms()
475
        
476
        if atoms is None:
477
            s.append('  no atoms defined')
478
        else:
479
            uc = atoms.get_cell()
480
            #a, b, c = uc
481
            s.append("  Unit Cell vectors (angstroms)")
482
            s.append("         x        y       z   length")
483

    
484
            for i, v in enumerate(uc):
485
                L = (np.sum(v**2))**0.5 #vector length
486
                s.append("  a%i [% 1.4f % 1.4f % 1.4f] %1.2f" % (i,
487
                                                                 v[0],
488
                                                                 v[1],
489
                                                                 v[2],
490
                                                                 L))
491
                                                                 
492
            stress = nc.variables.get('TotalStress', None)
493
            if stress is not None:
494
                stress = np.take(stress[:].ravel(), [0, 4, 8, 5, 2, 1])
495
                s.append('  Stress: xx,   yy,    zz,    yz,    xz,    xy')
496
                s1 = '       % 1.3f % 1.3f % 1.3f % 1.3f % 1.3f % 1.3f'
497
                s.append(s1 % tuple(stress))
498
            else:
499
                s.append('  No stress calculated.')
500
            s.append('  Volume = %1.2f A^3' % atoms.get_volume())
501
            s.append('')
502

    
503
            z = "  Atom,  sym, position (in x,y,z),     tag, rmsForce and psp"
504
            s.append(z)
505

    
506
            #this is just the ncvariable
507
            forces = nc.variables.get('DynamicAtomForces', None)
508
           
509
            for i, atom in enumerate(atoms):
510
                sym = atom.get_symbol()
511
                pos = atom.get_position()
512
                tag = atom.get_tag()
513
                if forces is not None and (forces[:][-1][i] < 1E36).all():
514
                    f = forces[:][-1][i]
515
                    # Lars Grabow: this seems to work right for some
516
                    # reason, but I would expect this to be the right
517
                    # index order f=forces[-1][i][:]
518
                    # frame,atom,direction
519
                    rmsforce = (np.sum(f**2))**0.5
520
                else:
521
                    rmsforce = None
522
                    
523
                st = "  %2i   %3.12s  " % (i, sym)
524
                st += "[% 7.3f%7.3f% 7.3f] " % tuple(pos)
525
                st += " %2s  " % tag
526
                if rmsforce is not None:
527
                    st += " %4.3f " % rmsforce
528
                else:
529
                    st += ' None '
530
                st += " %s"  % (self.get_psp(sym))        
531
                s.append(st)
532
                
533
        s.append('')
534
        s.append('  Details:')
535
        xc = self.get_xc()
536
        if xc is not None:
537
            s.append('  XCfunctional        = %s' % self.get_xc())
538
        else:
539
            s.append('  XCfunctional        = Not defined')
540
        s.append('  Planewavecutoff     = %i eV' % int(self.get_pw()))
541
        dw = self.get_dw()
542
        if dw:
543
            s.append('  Densitywavecutoff   = %i eV' % int(self.get_dw()))
544
        else:
545
            s.append('  Densitywavecutoff   = None')
546
        ft = self.get_ft()
547
        if ft is not None:
548
            s.append('  FermiTemperature    = %f kT' % ft)
549
        else:
550
            s.append('  FermiTemperature    = not defined')
551
        nelectrons = self.get_valence()
552
        if nelectrons is not None:
553
            s.append('  Number of electrons = %1.1f'  % nelectrons)
554
        else:
555
            s.append('  Number of electrons = None')
556
        s.append('  Number of bands     = %s'  % self.get_nbands())
557
        s.append('  Kpoint grid         = %s' % str(self.get_kpts_type()))
558
        s.append('  Spin-polarized      = %s' % self.get_spin_polarized())
559
        s.append('  Dipole correction   = %s' % self.get_dipole())
560
        s.append('  Symmetry            = %s' % self.get_symmetry())
561
        s.append('  Constraints         = %s' % str(atoms._get_constraints()))
562
        s.append('  ---------------------------------')
563
        nc.close()
564
        return string.join(s, '\n')            
565

    
566
    #todo figure out other xc psp databases
567
    def set_psp_database(self, xc=None):
568
        '''
569
        get the xc-dependent psp database
570

571
        :Parameters:
572

573
         xc : string
574
           one of 'PW91', 'PBE', 'revPBE', 'RPBE', 'PZ'
575

576
        
577
        not all the databases are complete, and that means
578
        some psp do not exist.
579

580
        note: this function is not supported fully. only pw91 is
581
        imported now. Changing the xc at this point results in loading
582
        a nearly empty database, and I have not thought about how to
583
        resolve that
584
        '''
585
        
586
        if xc == 'PW91' or xc is None:
587
            from pw91_psp import defaultpseudopotentials
588
        else:
589
            log.warn('PW91 pseudopotentials are being used!')
590
            #todo build other xc psp databases
591
            from pw91_psp import defaultpseudopotentials
592

    
593
        self.psp = defaultpseudopotentials
594

    
595
    def _set_frame_number(self, frame=None):
596
        'set framenumber in the netcdf file'
597
        
598
        if frame is None:
599
            nc = netCDF(self.nc, 'r')
600
            if 'TotalEnergy' in nc.variables:
601
                frame = nc.variables['TotalEnergy'].shape[0]
602
                # make sure the last energy is reasonable. Sometime
603
                # the field is empty if the calculation ran out of
604
                # walltime for example. Empty values get returned as
605
                # 9.6E36.  Dacapos energies should always be negative,
606
                # so if the energy is > 1E36, there is definitely
607
                # something wrong and a restart is required.
608
                if nc.variables.get('TotalEnergy', None)[-1] > 1E36:
609
                    log.warn("Total energy > 1E36. NC file is incomplete. \
610
                    calc.restart required")
611
                    self.restart()
612
            else:
613
                frame = 1
614
            nc.close()
615
            log.info("Current frame number is: %i" % (frame-1))
616
        self._frame = frame-1  #netCDF starts counting with 1
617

    
618
    def _increment_frame(self):
619
        'increment the framenumber'
620
        
621
        log.debug('incrementing frame')
622
        self._frame += 1
623

    
624
    def set_pw(self, pw):
625
        '''set the planewave cutoff.
626

627
        :Parameters:
628

629
         pw : integer
630
           the planewave cutoff in eV
631
           
632
        this function checks to make sure the density wave cutoff is
633
        greater than or equal to the planewave cutoff.'''
634
        
635
        nc = netCDF(self.nc, 'a')
636
        if 'PlaneWaveCutoff' in nc.variables:
637
            vpw = nc.variables['PlaneWaveCutoff']
638
            vpw.assignValue(pw)
639
        else:
640
            vpw = nc.createVariable('PlaneWaveCutoff', 'd', ('dim1',))
641
            vpw.assignValue(pw)
642

    
643
        if 'Density_WaveCutoff' in nc.variables:
644
            vdw = nc.variables['Density_WaveCutoff']
645
            dw = vdw.getValue()
646
            if pw > dw:
647
                vdw.assignValue(pw) #make them equal
648
        else:
649
            vdw = nc.createVariable('Density_WaveCutoff', 'd', ('dim1',))
650
            vdw.assignValue(pw) 
651
        nc.close()
652
        self.restart() #nc dimension change for number_plane_Wave dimension
653
        self.set_status('new')
654
        self.ready = False
655

    
656
    def set_dw(self, dw):
657
        '''set the density wave cutoff energy.
658

659
        :Parameters:
660

661
          dw : integer
662
            the density wave cutoff
663

664
        The function checks to make sure it is not less than the
665
        planewave cutoff.
666

667
        Density_WaveCutoff describes the kinetic energy neccesary to
668
        represent a wavefunction associated with the total density,
669
        i.e. G-vectors for which $\vert G\vert^2$ $<$
670
        4*Density_WaveCutoff will be used to describe the total
671
        density (including augmentation charge and partial core
672
        density). If Density_WaveCutoff is equal to PlaneWaveCutoff
673
        this implies that the total density is as soft as the
674
        wavefunctions described by the kinetic energy cutoff
675
        PlaneWaveCutoff. If a value of Density_WaveCutoff is specified
676
        (must be larger than or equal to PlaneWaveCutoff) the program
677
        will run using two grids, one for representing the
678
        wavefunction density (softgrid_dim) and one representing the
679
        total density (hardgrid_dim). If the density can be
680
        reprensented on the same grid as the wavefunction density
681
        Density_WaveCutoff can be chosen equal to PlaneWaveCutoff
682
        (default).
683
        '''
684

    
685
        pw = self.get_pw()
686
        if pw > dw:
687
            log.warn('Planewave cutoff %i is greater \
688
than density cutoff %i' % (pw, dw))
689
        
690
        ncf = netCDF(self.nc, 'a')
691
        if 'Density_WaveCutoff' in ncf.variables:
692
            vdw = ncf.variables['Density_WaveCutoff']
693
            vdw.assignValue(dw)
694
        else:
695
            vdw = ncf.createVariable('Density_WaveCutoff', 'i', ('dim1',))
696
            vdw.assignValue(dw)
697
        ncf.close()
698
        self.restart() #nc dimension change
699
        self.set_status('new')
700
        self.ready = False
701

    
702
    def set_xc(self, xc):
703
        '''Set the self-consistent exchange-correlation functional
704

705
        :Parameters:
706

707
         xc : string
708
           Must be one of 'PZ', 'VWN', 'PW91', 'PBE', 'revPBE', 'RPBE'
709

710
        Selects which density functional to use for
711
        exchange-correlation when performing electronic minimization
712
        (the electronic energy is minimized with respect to this
713
        selected functional) Notice that the electronic energy is also
714
        evaluated non-selfconsistently by DACAPO for other
715
        exchange-correlation functionals Recognized options :
716

717
        * "PZ" (Perdew Zunger LDA-parametrization)
718
        * "VWN" (Vosko Wilk Nusair LDA-parametrization)
719
        * "PW91" (Perdew Wang 91 GGA-parametrization)
720
        * "PBE" (Perdew Burke Ernzerhof GGA-parametrization)
721
        * "revPBE" (revised PBE/1 GGA-parametrization)
722
        * "RPBE" (revised PBE/2 GGA-parametrization)
723

724
        option "PZ" is not allowed for spin polarized
725
        calculation; use "VWN" instead.
726
        '''
727
        nc = netCDF(self.nc, 'a')
728
        v = 'ExcFunctional'
729
        if v in nc.variables:
730
            nc.variables[v][:] = np.array('%7s' % xc, 'c') 
731
        else:
732
            vxc = nc.createVariable('ExcFunctional', 'c', ('dim7',))
733
            vxc[:] = np.array('%7s' % xc, 'c')
734
        nc.close()
735
        self.set_status('new')
736
        self.ready = False    
737

    
738
    def set_nbands(self, nbands=None):
739
        '''Set the number of bands. a few unoccupied bands are
740
        recommended.
741

742
        :Parameters:
743

744
          nbands : integer
745
            the number of bands.
746
            
747
        if nbands = None the function returns with nothing done. At
748
        calculate time, if there are still no bands, they will be set
749
        by:
750

751
        the number of bands is calculated as
752
        $nbands=nvalence*0.65 + 4$
753
        '''
754
        if nbands is None:
755
            return
756

    
757
        self.delete_ncattdimvar(self.nc,
758
                                ncdims=['number_of_bands'],
759
                                ncvars=[])
760
                       
761
        nc = netCDF(self.nc, 'a')
762
        v = 'ElectronicBands'
763
        if v in nc.variables:
764
            vnb = nc.variables[v]
765
        else:
766
            vnb = nc.createVariable('ElectronicBands', 'c', ('dim1',))
767
            
768
        vnb.NumberOfBands = nbands
769
        nc.sync()
770
        nc.close()
771
        self.set_status('new')
772
        self.ready = False
773

    
774
    def set_kpts(self, kpts):
775
        '''
776
        set the kpt grid.
777

778
        Parameters:
779

780
        kpts: (n1,n2,n3) or [k1,k2,k3,...] or one of these
781
        chadi-cohen sets:
782
         
783
        * cc6_1x1
784
        * cc12_2x3
785
        * cc18_sq3xsq3
786
        * cc18_1x1
787
        * cc54_sq3xsq3
788
        * cc54_1x1
789
        * cc162_sq3xsq3
790
        * cc162_1x1
791
        
792
        (n1,n2,n3) creates an n1 x n2 x n3 monkhorst-pack grid,
793
        [k1,k2,k3,...] creates a kpt-grid based on the kpoints
794
        defined in k1,k2,k3,...
795
        
796
        There is also a possibility to have Dacapo (fortran) create
797
        the Kpoints in chadi-cohen or monkhorst-pack form. To do this
798
        you need to set the KpointSetup.gridtype attribute, and
799
        KpointSetup.
800

801
        KpointSetup = [3,0,0]
802
        KpointSetup.gridtype = 'ChadiCohen'
803

804
        KpointSetup(1)         Chadi-Cohen k-point set
805
        1         6 k-points 1x1
806
        2         18-kpoints sqrt(3)*sqrt(3)
807
        3         18-kpoints 1x1
808
        4         54-kpoints sqrt(3)*sqrt(3)
809
        5         54-kpoints 1x1
810
        6         162-kpoints 1x1
811
        7         12-kpoints 2x3
812
        8         162-kpoints 3xsqrt 3
813

814
        or
815
        KpointSetup = [4,4,4]
816
        KpointSetup.gridtype = 'MonkhorstPack'
817
        we do not use this functionality.
818
        '''
819

    
820
        #chadi-cohen
821
        if isinstance(kpts, str):
822
            exec('from ase.dft.kpoints import %s' % kpts)
823
            listofkpts = eval(kpts)
824
            gridtype = kpts #stored in ncfile
825
            #uc = self.get_atoms().get_cell()
826
            #listofkpts = np.dot(ccgrid,np.linalg.inv(uc.T))
827

    
828
        #monkhorst-pack grid
829
        if np.array(kpts).shape == (3,):
830
            from ase.dft.kpoints import monkhorst_pack
831
            N1, N2, N3 = kpts
832
            listofkpts = monkhorst_pack((N1, N2, N3))
833
            gridtype = 'Monkhorst-Pack %s' % str(tuple(kpts))
834

    
835
        #user-defined list is provided
836
        if len(np.array(kpts).shape) == 2:
837
            listofkpts = kpts
838
            gridtype = 'user_defined_%i_kpts' % len(kpts)  #stored in ncfile
839
                    
840
        nbzkpts = len(listofkpts)
841

    
842
        #we need to get dimensions stored temporarily so
843
        #we can delete all dimensions and variables associated with
844
        #kpoints before we save them back out.
845
        nc2 = netCDF(self.nc, 'r')
846
        ncdims = nc2.dimensions
847
        nc2.close()
848

    
849
        if 'number_BZ_kpoints' in ncdims:
850
            self.delete_ncattdimvar(self.nc,
851
                                    ncdims=['number_plane_waves',
852
                                            'number_BZ_kpoints',
853
                                            'number_IBZ_kpoints'])
854

    
855
        # now define dim and var
856
        nc = netCDF(self.nc, 'a')
857
        nc.createDimension('number_BZ_kpoints', nbzkpts)
858
        bv = nc.createVariable('BZKpoints', 'd', ('number_BZ_kpoints',
859
                                                  'dim3'))
860

    
861
        bv[:] = listofkpts
862
        bv.gridtype = gridtype
863
        nc.sync()
864
        nc.close()
865

    
866
        log.debug('kpts = %s' % str(self.get_kpts()))
867
        
868
        self.set_status('new')
869
        self.ready = False
870

    
871
    def atoms_are_equal(self, atoms):
872
        '''
873
        comparison of atoms to self.atoms using tolerances to account
874
        for float/double differences and float math.
875
        '''
876
        
877
        TOL = 1.0e-6 #angstroms
878

    
879
        a = self.atoms.arrays
880
        b = atoms.arrays
881

    
882
        #match number of atoms in cell
883
        lenmatch = len(atoms) == len(self.atoms)
884
        if lenmatch is not True:
885
            return False #the next two comparisons fail in this case.
886
        
887
        #match positions in cell
888
        posmatch = (abs(a['positions'] - b['positions']) <= TOL).all()
889
        #match cell
890
        cellmatch = (abs(self.atoms.get_cell()
891
                         - atoms.get_cell()) <= TOL).all()
892
        
893
        if lenmatch and posmatch and cellmatch:
894
            return True
895
        else:
896
            return False
897
                
898
    def set_atoms(self, atoms):
899
        '''attach an atoms to the calculator and update the ncfile
900

901
        :Parameters:
902

903
          atoms
904
            ASE.Atoms instance
905
          
906
        '''
907

    
908
        log.debug('setting atoms to: %s' % str(atoms))
909
        
910
        if hasattr(self, 'atoms') and self.atoms is not None:
911
            #return if the atoms are the same. no change needs to be made
912
            if self.atoms_are_equal(atoms):
913
                log.debug('No change to atoms in set_atoms, returning')
914
                return
915
            
916
            # some atoms already exist. Test if new atoms are
917
            # different from old atoms.
918
            # this is redundant
919
            if atoms != self.atoms:
920
                # the new atoms are different from the old ones. Start
921
                # a new frame.
922
                log.debug('atoms != self.atoms, incrementing')
923
                self._increment_frame()
924
                
925
        self.atoms = atoms.copy()
926
        self.ready = False
927
        log.debug('self.atoms = %s' % str(self.atoms))
928

    
929
    def set_ft(self, ft):
930
        '''set the Fermi temperature for occupation smearing
931

932
        :Parameters:
933

934
          ft : float
935
            Fermi temperature in kT (eV)
936

937
        Electronic temperature, corresponding to gaussian occupation
938
        statistics. Device used to stabilize the convergence towards
939
        the electronic ground state. Higher values stabilizes the
940
        convergence. Values in the range 0.1-1.0 eV are recommended,
941
        depending on the complexity of the Fermi surface (low values
942
        for d-metals and narrow gap semiconducters, higher for free
943
        electron-like metals).
944
        '''
945
        
946
        nc = netCDF(self.nc, 'a')
947
        v = 'ElectronicBands'
948
        if v in nc.variables:
949
            vnb = nc.variables[v]
950
        else:
951
            vnb = nc.createVariable('ElectronicBands', 'c', ('dim1',))
952

    
953
        vnb.OccupationStatistics_FermiTemperature = ft
954
        nc.sync()
955
        nc.close()
956
        self.set_status('new')
957
        self.ready = False
958

    
959
    def set_status(self, status):
960
        '''set the status flag in the netcdf file
961

962
        :Parameters:
963

964
          status : string
965
            status flag, e.g. 'new', 'finished'
966
        '''
967
        
968
        nc = netCDF(self.nc, 'a')
969
        nc.status = status
970
        nc.sync()
971
        nc.close()
972
        log.debug('set status to %s' % status)
973

    
974
    def get_spinpol(self):
975
        'Returns the spin polarization setting, either True or False'
976
        
977
        nc = netCDF(self.nc, 'r')
978
        v = 'ElectronicBands'
979
        if v in nc.variables:
980
            vnb = nc.variables[v]
981
            if hasattr(vnb, 'SpinPolarization'):
982
                spinpol = vnb.SpinPolarization
983
            else:
984
                spinpol = 1
985
        else:
986
            spinpol = 1
987

    
988
        nc.close()
989
        if spinpol == 1:
990
            return False
991
        else:
992
            return True
993

    
994
    def set_spinpol(self, spinpol=False):
995
        '''set Spin polarization.
996

997
        :Parameters:
998

999
         spinpol : Boolean
1000
           set_spinpol(True)  spin-polarized.
1001
           set_spinpol(False) no spin polarization, default
1002

1003
        Specify whether to perform a spin polarized or unpolarized
1004
        calculation.
1005
        '''
1006
        
1007
        nc = netCDF(self.nc, 'a')
1008
        v = 'ElectronicBands'
1009
        if v in nc.variables:
1010
            vnb = nc.variables[v]
1011
        else:
1012
            vnb = nc.createVariable('ElectronicBands', 'c', ('dim1',))
1013

    
1014
        if spinpol is True:
1015
            vnb.SpinPolarization = 2
1016
        else:
1017
            vnb.SpinPolarization = 1
1018

    
1019
        nc.sync()
1020
        nc.close()
1021
        self.set_status('new')
1022
        self.ready = False
1023

    
1024
    def set_fixmagmom(self, fixmagmom=None):
1025
        '''set a fixed magnetic moment for a spin polarized calculation
1026

1027
        :Parameters:
1028

1029
          fixmagmom : float
1030
            the magnetic moment of the cell in Bohr magnetons
1031
        '''
1032
        
1033
        if fixmagmom is None:
1034
            return
1035
        
1036
        nc = netCDF(self.nc,'a')
1037
        v = 'ElectronicBands'
1038
        if v in nc.variables:
1039
            vnb = nc.variables[v]
1040
        else:
1041
            vnb = nc.createVariable('ElectronicBands', 'c', ('dim1',))
1042

    
1043
        vnb.SpinPolarization = 2 #You must want spin-polarized
1044
        vnb.FixedMagneticMoment = fixmagmom
1045
        nc.sync()
1046
        nc.close()
1047
        self.set_status('new')
1048
        self.ready = False
1049

    
1050
    def get_fixmagmom(self):
1051
        'returns the value of FixedMagneticMoment'
1052
        
1053
        nc = netCDF(self.nc,'r')
1054
        if 'ElectronicBands' in nc.variables:
1055
            v = nc.variables['ElectronicBands']
1056
            if hasattr(v,'FixedMagneticMoment'):
1057
                fixmagmom = v.FixedMagneticMoment
1058
            else:
1059
                fixmagmom = None
1060
        else:
1061
            fixmagmom = None
1062
        nc.close()
1063
        return fixmagmom
1064
            
1065
    def set_calculate_stress(self, stress=True):
1066
        '''Turn on stress calculation
1067

1068
        :Parameters:
1069

1070
          stress : boolean
1071
            set_calculate_stress(True) calculates stress
1072
            set_calculate_stress(False) do not calculate stress
1073
        '''
1074
        
1075
        nc = netCDF(self.get_nc(),'a')
1076
        vs = 'NetCDFOutputControl'
1077
        if vs in nc.variables:
1078
            v = nc.variables[vs]
1079
        else:
1080
            v = nc.createVariable('NetCDFOutputControl', 'c', ('dim1',))
1081

    
1082
        if stress is True:
1083
            v.PrintTotalStress = 'Yes'
1084
        else:
1085
            v.PrintTotalStress = 'No'
1086
        nc.sync()
1087
        nc.close()
1088
        self.set_status('new')
1089
        self.ready = False
1090

    
1091
    def set_nc(self, nc='out.nc'):
1092
        '''
1093
        set filename for the netcdf and text output for this calculation
1094

1095
        :Parameters:
1096

1097
          nc : string
1098
            filename for netcdf file
1099
                
1100
        if the ncfile attached to the calculator is changing, the old
1101
        file will be copied to the new file if it doesn not exist so
1102
        that all the calculator details are preserved. Otherwise, the 
1103

1104
        if the ncfile does not exist, it will get initialized.
1105

1106
        the text file will have the same basename as the ncfile, but
1107
        with a .txt extension.
1108
        '''
1109
        
1110
        #the first time this is called, there may be no self.nc defined
1111
        if not hasattr(self, 'nc'):
1112
            self.nc = nc    
1113
            
1114
        #check if the name is changing and if so, copy the old ncfile
1115
        #to the new one.  This is necessary to ensure all the
1116
        #calculator details are copied over. if the file already
1117
        #exists we use the contents of the existing file
1118
        if nc != self.nc and not os.path.exists(nc):
1119
            log.debug('copying %s to %s' % (self.nc, nc))
1120
            #import shutil
1121
            #shutil.copy(self.nc,nc)
1122
            base = os.path.split(nc)[0]
1123
            if not os.path.isdir(base) and base is not '':
1124
                os.makedirs(base)
1125
            status = os.system('cp %s %s' % (self.nc, nc))
1126
            if status != 0:
1127
                raise Exception, 'Copying ncfile failed.'
1128
            self.nc = nc
1129
        
1130
        elif os.path.exists(nc):
1131
            self._set_frame_number()
1132
            self.set_psp_database()
1133
            self.atoms = self.read_only_atoms(nc)
1134
            self.nc = nc
1135
            self.update_input_parameters()
1136
                
1137
        #I always want the text file set based on the ncfile
1138
        #and I never want to set this myself.
1139
        base = os.path.splitext(self.nc)[0]
1140
        self.txt = base + '.txt'
1141
                 
1142
    def set_psp(self,
1143
                sym=None,
1144
                z=None,
1145
                psp=None):
1146
        '''
1147
        set the pseudopotential file for a species or an atomic number.
1148

1149
        :Parameters:
1150

1151
         sym : string
1152
           chemical symbol of the species
1153

1154
          z : integer
1155
            the atomic number of the species
1156

1157
          psp : string
1158
            filename of the pseudopotential
1159

1160
        
1161
        you can only set sym or z.
1162

1163
        examples::
1164
        
1165
          set_psp('N',psp='pspfile')
1166
          set_psp(z=6,psp='pspfile')
1167
        '''
1168
        log.debug(str([sym, z, psp]))
1169
        if (sym, z, psp) == (None, None, None):
1170
            return
1171
        
1172
        if (sym is None and z is not None):
1173
            from ase.data import chemical_symbols
1174
            sym = chemical_symbols[z]
1175
        elif (sym is not None and z is None):
1176
            pass
1177
        else:
1178
            raise Exception, 'You can only specify Z or sym!'
1179

    
1180
        if not hasattr(self, 'psp'):
1181
            self.set_psp_database()
1182

    
1183
        #only make change if needed
1184
        if sym not in self.psp:
1185
            self.psp[sym] = psp
1186
            self.ready = False
1187
            self.set_status('new')
1188
        elif self.psp[sym] != psp:
1189
            self.psp[sym] = psp
1190
            self.ready = False
1191
            self.set_status('new')
1192

    
1193
        if not self.ready:
1194
            #now we update the netcdf file
1195
            ncf = netCDF(self.nc, 'a')
1196
            vn = 'AtomProperty_%s' % sym
1197
            if vn not in ncf.variables:
1198
                p = ncf.createVariable(vn, 'c', ('dim20',))
1199
            else:
1200
                p = ncf.variables[vn]
1201

    
1202
            ppath = self.get_psp(sym=sym)
1203
            p.PspotFile = ppath
1204
            ncf.close()
1205

    
1206
    def get_pseudopotentials(self):
1207
        'get pseudopotentials set for atoms attached to calculator'
1208
        
1209
        if self.atoms is None:
1210
            return None
1211
        
1212
        psp = {}
1213
        for atom in self.atoms:
1214
            psp[atom.symbol] = self.psp[atom.symbol]
1215
        return psp
1216

    
1217
    def get_symmetry(self):
1218
        '''return the type of symmetry used'''
1219
        
1220
        nc = netCDF(self.nc, 'r')
1221
        if 'UseSymmetry' in nc.variables:
1222
            sym = string.join(nc.variables['UseSymmetry'][:],'').strip()
1223
        else:
1224
            sym = None      
1225
        nc.close()
1226
        if sym in ['Off', None]:
1227
            return False
1228
        elif sym == 'Maximum':
1229
            return True
1230
        else:
1231
            raise Exception, 'Type of symmetry not recognized: %s' % sym
1232
       
1233
    def set_symmetry(self, val=False):
1234
        '''set how symmetry is used to reduce k-points
1235

1236
        :Parameters:
1237

1238
         val : Boolean
1239
           set_sym(True) Maximum symmetry is used
1240
           set_sym(False) No symmetry is used
1241

1242
        This variable controls the if and how DACAPO should attempt
1243
        using symmetry in the calculation. Imposing symmetry generally
1244
        speeds up the calculation and reduces numerical noise to some
1245
        extent. Symmetry should always be applied to the maximum
1246
        extent, when ions are not moved. When relaxing ions, however,
1247
        the symmetry of the equilibrium state may be lower than the
1248
        initial state. Such an equilibrium state with lower symmetry
1249
        is missed, if symmetry is imposed. Molecular dynamics-like
1250
        algorithms for ionic propagation will generally not break the
1251
        symmetry of the initial state, but some algorithms, like the
1252
        BFGS may break the symmetry of the initial state. Recognized
1253
        options:
1254

1255
        "Off": No symmetry will be imposed, apart from time inversion
1256
        symmetry in recipical space. This is utilized to reduce the
1257
        k-point sampling set for Brillouin zone integration and has no
1258
        influence on the ionic forces/motion.
1259

1260
        "Maximum": DACAPO will look for symmetry in the supplied
1261
        atomic structure and extract the highest possible symmetry
1262
        group. During the calculation, DACAPO will impose the found
1263
        spatial symmetry on ionic forces and electronic structure,
1264
        i.e. the symmetry will be conserved during the calculation.
1265
        '''
1266
        
1267
        if val:
1268
            symval = 'Maximum'
1269
        else:
1270
            symval = 'Off'
1271
        
1272
        ncf = netCDF(self.get_nc(), 'a')
1273
        if 'UseSymmetry' not in ncf.variables:
1274
            sym = ncf.createVariable('UseSymmetry', 'c', ('dim7',))
1275
        else:
1276
            sym = ncf.variables['UseSymmetry']
1277
            
1278
        sym[:] = np.array('%7s' % symval, 'c')
1279
        ncf.sync()
1280
        ncf.close()
1281
        self.set_status('new')
1282
        self.ready = False
1283

    
1284
    def set_extracharge(self, val):
1285
        '''add extra charge to unit cell
1286

1287
        :Parameters:
1288

1289
          val : float
1290
            extra electrons to add or subtract from the unit cell
1291

1292
        Fixed extra charge in the unit cell (i.e. deviation from
1293
        charge neutrality). This assumes a compensating, positive
1294
        constant backgound charge (jellium) to forge overall charge
1295
        neutrality.
1296
        '''
1297
        
1298
        nc = netCDF(self.get_nc(), 'a')
1299
        if 'ExtraCharge' in nc.variables:
1300
            v = nc.variables['ExtraCharge']
1301
        else:
1302
            v = nc.createVariable('ExtraCharge', 'd', ('dim1',))
1303

    
1304
        v.assignValue(val)
1305
        nc.sync()
1306
        nc.close()
1307

    
1308
    def get_extracharge(self):
1309
        'Return the extra charge set in teh calculator'
1310
        
1311
        nc = netCDF(self.get_nc(), 'r')
1312
        if 'ExtraCharge' in nc.variables:
1313
            v = nc.variables['ExtraCharge']
1314
            exchg = v.getValue()
1315
        else:
1316
            exchg = None
1317
        nc.close()
1318
        return exchg
1319

    
1320
    def get_extpot(self):
1321
        'return the external potential set in teh calculator'
1322
        
1323
        nc = netCDF(self.get_nc(), 'a')
1324
        if 'ExternalPotential' in nc.variables:
1325
            v = nc.variables['ExternalPotential']
1326
            extpot = v[:]
1327
        else:
1328
            extpot = None
1329

    
1330
        nc.close()
1331
        return extpot
1332
    
1333
    def set_extpot(self, potgrid):
1334
        '''add external potential of value
1335

1336
        see this link before using this
1337
        https://listserv.fysik.dtu.dk/pipermail/campos/2003-August/000657.html
1338
        
1339
        :Parameters:
1340

1341
          potgrid : np.array with shape (nx,ny,nz)
1342
            the shape must be the same as the fft soft grid
1343
            the value of the potential to add
1344

1345
        
1346
        you have to know both of the fft grid dimensions ahead of time!
1347
        if you know what you are doing, you can set the fft_grid you want
1348
        before hand with:
1349
        calc.set_fftgrid((n1,n2,n3))
1350
        '''
1351
        
1352
        nc = netCDF(self.get_nc(), 'a')
1353
        if 'ExternalPotential' in nc.variables:
1354
            v = nc.variables['ExternalPotential']
1355
        else:
1356
            # I assume here you have the dimensions of potgrid correct
1357
            # and that the soft and hard grids are the same. 
1358
            # if softgrid is defined, Dacapo requires hardgrid to be
1359
            # defined too.
1360
            s1, s2, s3 = potgrid.shape
1361
            if 'softgrid_dim1' not in nc.dimensions:
1362
                nc.createDimension('softgrid_dim1', s1)
1363
                nc.createDimension('softgrid_dim2', s2)
1364
                nc.createDimension('softgrid_dim3', s3)
1365
                nc.createDimension('hardgrid_dim1', s1)
1366
                nc.createDimension('hardgrid_dim2', s2)
1367
                nc.createDimension('hardgrid_dim3', s3)
1368
                
1369
            v = nc.createVariable('ExternalPotential',
1370
                                  'd',
1371
                                  ('softgrid_dim1',
1372
                                   'softgrid_dim2',
1373
                                   'softgrid_dim3',))
1374
        v[:] = potgrid
1375
        nc.sync()
1376
        nc.close()
1377
        self.set_status('new')
1378
        self.ready = False
1379
        
1380
    def set_fftgrid(self, soft=None, hard=None):
1381
        '''
1382
        sets the dimensions of the FFT grid to be used
1383

1384
        :Parameters:
1385

1386
          soft : (n1,n2,n3) integers
1387
            make a n1 x n2 x n3 grid
1388

1389
          hard : (n1,n2,n3) integers
1390
            make a n1 x n2 x n3 grid
1391

1392
        
1393
        >>> calc.set_fftgrid(soft=[42,44,46])
1394
        sets the soft and hard grid dimensions to 42,44,46
1395

1396
        >>> calc.set_fftgrid(soft=[42,44,46],hard=[80,84,88])
1397
        sets the soft grid dimensions to 42,44,46 and the hard
1398
        grid dimensions to 80,84,88
1399
        
1400
        These are the fast FFt grid numbers listed in fftdimensions.F
1401
        
1402
        data list_of_fft /2,  4,  6,  8, 10, 12, 14, 16, 18, 20, &
1403
        22,24, 28, 30,32, 36, 40, 42, 44, 48, &
1404
        56,60, 64, 66, 70, 72, 80, 84, 88, 90, &
1405
        96,108,110,112,120,126,128,132,140,144,154, &
1406
        160,168,176,180,192,198,200, &
1407
        216,240,264,270,280,288,324,352,360,378,384,400,432, &
1408
        450,480,540,576,640/
1409
        
1410
        otherwise you will get some errors from mis-dimensioned variables.
1411
        
1412
        this is usually automatically set by Dacapo.
1413
        '''
1414
        
1415
        if soft is not None:
1416
            self.delete_ncattdimvar(self.nc,
1417
                                    ncdims=['softgrid_dim1',
1418
                                            'softgrid_dim2',
1419
                                            'softgrid_dim3'
1420
                                            ],
1421
                                    ncvars=[])
1422
        
1423
            
1424
            nc = netCDF(self.get_nc(), 'a')
1425
            nc.createDimension('softgrid_dim1', soft[0])
1426
            nc.createDimension('softgrid_dim2', soft[1])
1427
            nc.createDimension('softgrid_dim3', soft[2])
1428
            nc.sync()
1429
            nc.close()
1430

    
1431
            if hard is None:
1432
                hard = soft
1433

    
1434
        if hard is not None:
1435
            self.delete_ncattdimvar(self.nc,
1436
                                    ncdims=['hardgrid_dim1',
1437
                                            'hardgrid_dim2',
1438
                                            'hardgrid_dim3'
1439
                                            ],
1440
                                    ncvars=[])
1441
            nc = netCDF(self.get_nc(),'a')
1442
            nc.createDimension('hardgrid_dim1', hard[0])
1443
            nc.createDimension('hardgrid_dim2', hard[1])
1444
            nc.createDimension('hardgrid_dim3', hard[2])
1445
            nc.sync()
1446
            nc.close()
1447

    
1448
        self.set_status('new')
1449
        self.ready = False
1450

    
1451
    def get_ascii_debug(self):
1452
        'Return the debug settings in Dacapo'
1453
        
1454
        nc = netCDF(self.get_nc(), 'r')
1455
        if 'PrintDebugInfo' in nc.variables:
1456
            v = nc.variables['PrintDebugInfo']
1457
            debug = string.join(v[:], '')
1458
        else:
1459
            debug = None
1460
        nc.close()
1461
        return debug
1462
    
1463
    def set_ascii_debug(self, level):
1464
        '''set the debug level for Dacapo
1465

1466
        :Parameters:
1467
        
1468
          level : string
1469
            one of 'Off', 'MediumLevel', 'HighLevel'
1470
        '''
1471
        
1472
        nc = netCDF(self.get_nc(), 'a')
1473
        if 'PrintDebugInfo' in nc.variables:
1474
            v = nc.variables['PrintDebugInfo']
1475
        else:
1476
            if 'dim20' not in nc.dimensions:
1477
                nc.createDimension('dim20', 20)
1478
            v = nc.createVariable('PrintDebugInfo', 'c', ('dim20',))
1479

    
1480
        v[:] = np.array('%20s' % level, dtype='c')
1481
        nc.sync()
1482
        nc.close()
1483
        self.set_status('new')
1484
        self.ready = False
1485

    
1486
    def get_ncoutput(self):
1487
        'returns the control variables for the ncfile'
1488
        
1489
        nc = netCDF(self.get_nc(), 'a')
1490
        if 'NetCDFOutputControl' in nc.variables:
1491
            v = nc.variables['NetCDFOutputControl']
1492
            ncoutput = {}
1493
            if hasattr(v, 'PrintWaveFunction'):
1494
                ncoutput['wf'] = v.PrintWaveFunction
1495
            if hasattr(v, 'PrintChargeDensity'):
1496
                ncoutput['cd'] = v.PrintChargeDensity
1497
            if hasattr(v, 'PrintEffPotential'):
1498
                ncoutput['efp'] = v.PrintEffPotential
1499
            if hasattr(v, 'PrintElsPotential'):
1500
                ncoutput['esp'] = v.PrintElsPotential
1501
        else:
1502
            ncoutput = None
1503
        nc.close()
1504
        return ncoutput
1505
        
1506
    def set_ncoutput(self,
1507
                     wf=None,
1508
                     cd=None,
1509
                     efp=None,
1510
                     esp=None):
1511
        '''set the output of large variables in the netcdf output file
1512

1513
        :Parameters:
1514

1515
          wf : string
1516
            controls output of wavefunction. values can
1517
            be 'Yes' or 'No'
1518

1519
          cd : string
1520
            controls output of charge density. values can
1521
            be 'Yes' or 'No'
1522

1523
          efp : string
1524
            controls output of effective potential. values can
1525
            be 'Yes' or 'No'
1526

1527
          esp : string
1528
            controls output of electrostatic potential. values can
1529
            be 'Yes' or 'No'
1530
        '''
1531
        nc = netCDF(self.get_nc(), 'a')
1532
        if 'NetCDFOutputControl' in nc.variables:
1533
            v = nc.variables['NetCDFOutputControl']
1534
        else:
1535
            v = nc.createVariable('NetCDFOutputControl', 'c', ())
1536

    
1537
        if wf is not None:
1538
            v.PrintWaveFunction = wf
1539
        if cd is not None:
1540
            v.PrintChargeDensity = cd
1541
        if efp is not None:
1542
            v.PrintEffPotential = efp
1543
        if esp is not None:
1544
            v.PrintElsPotential = esp
1545

    
1546
        nc.sync()
1547
        nc.close()
1548
        self.set_status('new')
1549
        self.ready = False
1550

    
1551
    def get_ados(self, **kwargs):
1552
        '''
1553
        attempt at maintaining backward compatibility with get_ados
1554
        returning data
1555

1556
        Now when we call calc.get_ados() it will return settings,
1557

1558
        and calc.get_ados(atoms=[],...) should return data
1559

1560
        '''
1561
        
1562
        if len(kwargs) != 0:
1563
            return self.get_ados_data(**kwargs)
1564
        
1565
        nc = netCDF(self.get_nc(),'r')
1566
        if 'PrintAtomProjectedDOS' in nc.variables:
1567
            v = nc.variables['PrintAtomProjectedDOS']
1568
            ados = {}
1569
            if hasattr(v, 'EnergyWindow'):
1570
                ados['energywindow'] = v.EnergyWindow
1571
            if hasattr(v, 'EnergyWidth'):
1572
                ados['energywidth'] = v.EnergyWidth
1573
            if hasattr(v, 'NumberEnergyPoints'):
1574
                ados['npoints'] = v.NumberEnergyPoints
1575
            if hasattr(v, 'CutoffRadius'):
1576
                ados['cutoff'] = v.CutoffRadius
1577
        else:
1578
            ados = None
1579

    
1580
        nc.close()
1581
        return ados
1582
        
1583
    def set_ados(self,
1584
                 energywindow=(-15,5),
1585
                 energywidth=0.2,
1586
                 npoints=250,
1587
                 cutoff=1.0):
1588
        '''
1589
        setup calculation of atom-projected density of states
1590

1591
        :Parameters:
1592

1593
          energywindow : (float, float)
1594
            sets (emin,emax) in eV referenced to the Fermi level
1595

1596
          energywidth : float
1597
            the gaussian used in smearing
1598

1599
          npoints : integer
1600
            the number of points to sample the DOS at
1601

1602
          cutoff : float
1603
            the cutoff radius in angstroms for the integration.
1604
        '''
1605
        
1606
        nc = netCDF(self.get_nc(), 'a')
1607
        if 'PrintAtomProjectedDOS' in nc.variables:
1608
            v = nc.variables['PrintAtomProjectedDOS']
1609
        else:
1610
            v = nc.createVariable('PrintAtomProjectedDOS', 'c', ())
1611

    
1612
        v.EnergyWindow = energywindow
1613
        v.EnergyWidth  = energywidth
1614
        v.NumberEnergyPoints = npoints
1615
        v.CutoffRadius = cutoff
1616

    
1617
        nc.sync()
1618
        nc.close()
1619
        self.set_status('new')
1620
        self.ready = False
1621

    
1622
    def get_mdos(self):
1623
        'return multicentered projected dos parameters'
1624
        nc = netCDF(self.get_nc(),'r')
1625

    
1626
        mdos = {}
1627
        
1628
        if 'MultiCenterProjectedDOS' in nc.variables:
1629
            v = nc.variables['MultiCenterProjectedDOS']
1630
            mdos['energywindow'] = v.EnergyWindow
1631
            mdos['energywidth'] = v.EnergyWidth
1632
            mdos['numberenergypoints'] = v.NumberEnergyPoints
1633
            mdos['cutoffradius'] = v.CutoffRadius
1634
            mdos['mcenters'] = eval(v.mcenters)
1635
                
1636
        nc.close()
1637

    
1638
        return mdos
1639

    
1640
    def get_mdos_data(self,
1641
                      spin=0,
1642
                      cutoffradius='infinite'):
1643
        '''returns data from multicentered projection
1644

1645

1646
        returns (mdos, rotmat)
1647

1648
        the rotation matrices are retrieved from the text file. I am
1649
        not sure what you would do with these, but there was a note
1650
        about them in the old documentation so I put the code to
1651
        retrieve them here. the syntax for the return value is:
1652
        rotmat[atom#][label] returns the rotation matrix for the
1653
        center on the atom# for label. I do not not know what the
1654
        label refers to.
1655
        '''
1656
        
1657
        if self.calculation_required():
1658
            self.calculate()
1659
        
1660
        nc = netCDF(self.get_nc(),'r')
1661
        icut = 1 #short
1662
        if cutoffradius == "infinite":
1663
            icut = 0
1664
            
1665
        #var = nc.variables['MultiCenterProjectedDOS']
1666
        integrated = nc.variables['MultiCenterProjectedDOS_IntegratedDOS'][:]
1667
        tz = 'MultiCenterProjectedDOS_EnergyResolvedDOS'
1668
        energyresolved = nc.variables[tz][:]
1669
        energygrid = nc.variables['MultiCenterProjectedDOS_EnergyGrid'][:]
1670

    
1671
        number_of_multicenters  = integrated.shape[0]
1672
        #number_of_cutoff = integrated.shape[1]
1673
        #number_of_spin = integrated.shape[2]
1674

    
1675
        multicenterprojections = []
1676
        for multicenter in range(number_of_multicenters): 
1677
            #orbitals = var[multicenter]
1678
            energyresolveddata = energyresolved[multicenter, icut, spin, :]
1679
            #integrateddata     = integrated[multicenter, icut, spin]
1680
            multicenterprojections.append([energygrid, energyresolveddata])
1681

    
1682
        log.info('Found %d multicenters' % len(multicenterprojections))
1683
        nc.close()
1684

    
1685
        #now parse the text file for the rotation matrices
1686
        rot_mat_lines = []
1687
        txt = self.get_txt()
1688
        if os.path.exists(txt):
1689
            f = open(txt,'r')
1690
            for line in f:
1691
                if 'MUL: Rmatrix' in line:
1692
                    rot_mat_lines.append(line)
1693
            f.close()
1694

    
1695
            rotmat = []
1696
            for line in rot_mat_lines:
1697
                fields = line.split()
1698
                novl = int(fields[2])
1699
                ncen = int(fields[3])
1700
                row = [float(x) for x in fields[4:]]
1701

    
1702
                try:
1703
                    rotmat[novl-1][ncen-1].append(row)
1704
                except IndexError:
1705
                    try:
1706
                        rotmat[novl-1].append([])
1707
                        rotmat[novl-1][ncen-1].append(row)
1708
                    except IndexError:
1709
                        rotmat.append([])
1710
                        rotmat[novl-1].append([])
1711
                    rotmat[novl-1][ncen-1].append(row)
1712
        else:
1713
            rotmat = None
1714
                    
1715
        return (multicenterprojections, rotmat)
1716
            
1717
    def set_mdos(self,
1718
                 mcenters=None,
1719
                 energywindow=(-15,5),
1720
                 energywidth=0.2,
1721
                 numberenergypoints=250,
1722
                 cutoffradius=1.0):
1723
        '''Setup multicentered projected DOS.
1724

1725
        mcenters
1726
           a list of tuples containing (atom#,l,m,weight)
1727
           (0,0,0,1.0) specifies (atom 0, l=0, m=0, weight=1.0) an s orbital
1728
           on atom 0
1729
           
1730
           (1,1,1,1.0) specifies (atom 1, l=1, m=1, weight=1.0) a p orbital
1731
           with m = +1 on atom 0
1732

1733
           l=0 s-orbital
1734
           l=1 p-orbital
1735
           l=2 d-orbital
1736

1737
           m in range of ( -l ... 0 ... +l )
1738

1739
           The direction cosines for which the spherical harmonics are
1740
           set up are using the next different atom in the list
1741
           (cyclic) as direction pointer, so the z-direction is chosen
1742
           along the direction to this next atom. At the moment the
1743
           rotation matrices is only given in the text file, you can
1744
           use grep'MUL: Rmatrix' out_o2.txt to get this information.
1745
           
1746
        adapated from old MultiCenterProjectedDOS.py
1747
        '''
1748
        if mcenters is None:
1749
            return
1750
        
1751
        nc = netCDF(self.get_nc(), 'a')
1752
        
1753
        _listofmcenters_ = mcenters
1754
        
1755
        # get number of multi centers
1756
        ncenters = len(_listofmcenters_)
1757
        # get max number of orbitals any center 
1758
        max_orbitals = max(map(len, _listofmcenters_))
1759

    
1760
        mmatrix = np.zeros([ncenters, max_orbitals, 4], np.float)
1761
        ncenter = 0
1762
        for multicenter in _listofmcenters_: 
1763
            norbital = 0
1764
            for orbital in multicenter: 
1765
                mmatrix[ncenter, norbital] = orbital 
1766
                norbital = norbital + 1 
1767

    
1768
            # signal that this multicenter contains less than
1769
            # max_orbital orbitals
1770
            if len(multicenter) < max_orbitals: 
1771
                mmatrix[ncenter, len(multicenter):max_orbitals] = (-1.0, 0,
1772
                                                                   0, 0)
1773

    
1774
            ncenter = ncenter + 1
1775

    
1776
        nc.createDimension('max_orbitals', max_orbitals)
1777
        nc.createDimension('number_of_multicenters', ncenters)
1778

    
1779
        if 'MultiCenterProjectedDOS' in nc.variables:
1780
            v = nc.variables['MultiCenterProjectedDOS']
1781
        else:
1782
            v = nc.createVariable('MultiCenterProjectedDOS',
1783
                                  'd',
1784
                                  ('number_of_multicenters',
1785
                                   'max_orbitals',
1786
                                   'dim4'))
1787

    
1788
        v.EnergyWindow = energywindow
1789
        v.EnergyWidth = energywidth
1790
        v.NumberEnergyPoints = numberenergypoints
1791
        v.CutoffRadius = cutoffradius
1792

    
1793
        #this is kind of hacky, but it is needed for get_mdos so you
1794
        #can tell if the input is changed.
1795
        v.mcenters = str(mcenters)
1796

    
1797
        v[:] = mmatrix
1798

    
1799
        nc.sync()
1800
        nc.close()        
1801

    
1802
    def set_debug(self, debug):
1803
        '''
1804
        set debug level for python logging
1805

1806
        debug should be an integer from 0-100 or one of the logging
1807
        constants like logging.DEBUG, logging.WARN, etc...
1808

1809
        '''
1810
        
1811
        self.debug = debug
1812
        log.setLevel(debug)
1813

    
1814
    def get_debug(self):
1815
        'Return the python logging level'
1816
        
1817
        return self.debug
1818

    
1819
    def get_decoupling(self):
1820
        'return the electrostatic decoupling parameters'
1821
        
1822
        nc = netCDF(self.get_nc(), 'r')
1823
        if 'Decoupling' in nc.variables:
1824
            v = nc.variables['Decoupling']
1825
            decoupling = {}
1826
            if hasattr(v,'NumberOfGaussians'):
1827
                decoupling['ngaussians'] = v.NumberOfGaussians
1828
            if hasattr(v,'ECutoff'):
1829
                decoupling['ecutoff'] = v.ECutoff
1830
            if hasattr(v,'WidthOfGaussian'):
1831
                decoupling['gausswidth'] = v.WidthOfGaussian
1832
        else:
1833
            decoupling = None
1834
        nc.close()
1835
        return decoupling
1836
        
1837
    def set_decoupling(self,
1838
                       ngaussians=3,
1839
                       ecutoff=100,
1840
                       gausswidth=0.35):
1841
        '''
1842
        Decoupling activates the three dimensional electrostatic
1843
        decoupling. Based on paper by Peter E. Bloechl: JCP 103
1844
        page7422 (1995).
1845

1846
        :Parameters:
1847

1848
          ngaussians : int
1849
            The number of gaussian functions per atom
1850
            used for constructing the model charge of the system
1851

1852
          ecutoff : int
1853
            The cut off energy (eV) of system charge density in
1854
            g-space used when mapping constructing the model change of
1855
            the system, i.e. only charge density components below
1856
            ECutoff enters when constructing the model change.
1857

1858
          gausswidth : float
1859
            The width of the Gaussians defined by
1860
            $widthofgaussian*1.5^(n-1)$  $n$=(1 to numberofgaussians)
1861
            
1862
        '''
1863

    
1864
        nc = netCDF(self.get_nc(), 'a')
1865
        if 'Decoupling' in nc.variables:
1866
            v = nc.variables['Decoupling']
1867
        else:
1868
            v = nc.createVariable('Decoupling', 'c', ())
1869

    
1870
        v.NumberOfGaussians = ngaussians
1871
        v.ECutoff = ecutoff
1872
        v.WidthOfGaussian = gausswidth
1873

    
1874
        nc.sync()
1875
        nc.close()
1876
        self.set_status('new')
1877
        self.ready = False
1878

    
1879
    def set_external_dipole(self,
1880
                            value,
1881
                            position=None):
1882
        '''
1883
        Externally imposed dipole potential. This option overwrites
1884
        DipoleCorrection if set. 
1885

1886
        :Parameters:
1887

1888
          value : float
1889
            units of volts
1890

1891
          position : float
1892
            scaled coordinates along third unit cell direction.
1893
            if None, the compensation dipole layer plane in the
1894
            vacuum position farthest from any other atoms on both
1895
            sides of the slab. Do not set to 0.0.
1896
        '''
1897
        
1898
        var = 'ExternalDipolePotential'
1899
        nc = netCDF(self.get_nc(), 'a')
1900
        if var in nc.variables:
1901
            v = nc.variables[var]
1902
        else:
1903
            v = nc.createVariable('ExternalDipolePotential', 'd', ())
1904
        
1905
        v.assignValue(value)
1906
        if position is not None:
1907
            v.DipoleLayerPosition = position
1908

    
1909
        nc.sync()
1910
        nc.close()
1911
        self.set_status('new')
1912
        self.ready = False
1913

    
1914
    def get_external_dipole(self):
1915
        'return the External dipole settings'
1916

    
1917
        var = 'ExternalDipolePotential'
1918
        nc = netCDF(self.get_nc(),'r')
1919
        if var in nc.variables:
1920
            v = nc.variables[var]
1921
            value = v.getValue()
1922
            if hasattr(v, 'DipoleLayerPosition'):
1923
                position = v.DipoleLayerPosition
1924
            else:
1925
                position = None
1926

    
1927
            ed = {'value':value, 'position':position}
1928
        else:
1929
            ed = None
1930
        nc.close()
1931
        return ed
1932
            
1933
    def set_dipole(self,
1934
                   status=True,
1935
                   mixpar=0.2,
1936
                   initval=0.0,
1937
                   adddipfield=0.0,
1938
                   position=None):
1939
        '''turn on and set dipole correction scheme
1940

1941
        :Parameters:
1942

1943
          status : Boolean
1944
            True turns dipole on. False turns Dipole off
1945

1946
          mixpar : float
1947
            Mixing Parameter for the the dipole correction field
1948
            during the electronic minimization process. If instabilities
1949
            occur during electronic minimization, this value may be
1950
            decreased.
1951

1952
          initval : float
1953
            initial value to start at
1954

1955
          adddipfield : float
1956
            additional dipole field to add
1957
            units : V/ang
1958
            External additive, constant electrostatic field along
1959
            third unit cell vector, corresponding to an external
1960
            dipole layer. The field discontinuity follows the position
1961
            of the dynamical dipole correction, i.e. if
1962
            DipoleCorrection:DipoleLayerPosition is set, the field
1963
            discontinuity is at this value, otherwise it is at the
1964
            vacuum position farthest from any other atoms on both
1965
            sides of the slab.
1966

1967
          position : float
1968
            scaled coordinates along third unit cell direction.
1969
            If this attribute is set, DACAPO will position the
1970
            compensation dipole layer plane in at the provided value.
1971
            If this attribute is not set, DACAPO will put the compensation
1972
            dipole layer plane in the vacuum position farthest from any
1973
            other atoms on both sides of the slab. Do not set this to
1974
            0.0
1975

1976
        
1977
        calling set_dipole() sets all default values.
1978
            
1979
        '''
1980
        if status == False:
1981
            self.delete_ncattdimvar(self.nc, ncvars=['DipoleCorrection'])
1982
            return
1983
        
1984
        ncf = netCDF(self.get_nc(), 'a')
1985
        if 'DipoleCorrection' not in ncf.variables:
1986
            dip = ncf.createVariable('DipoleCorrection', 'c', ())
1987
        else:
1988
            dip = ncf.variables['DipoleCorrection']
1989
        dip.MixingParameter = mixpar
1990
        dip.InitialValue = initval
1991
        dip.AdditiveDipoleField = adddipfield
1992

    
1993
        if position is not None:
1994
            dip.DipoleLayerPosition = position
1995
            
1996
        ncf.sync()
1997
        ncf.close()
1998
        self.set_status('new')
1999
        self.ready = False
2000
       
2001
    def set_stay_alive(self, value):
2002
        'set the stay alive setting'
2003
        
2004
        self.delete_ncattdimvar(self.nc,
2005
                                ncvars=['Dynamics'])
2006

    
2007
        if value in [True, False]:
2008
            self.stay_alive = value
2009
            #self._dacapo_is_running = False
2010
        else:
2011
            log.debug("stay_alive must be boolean. Value was not changed.")
2012

    
2013
    def get_stay_alive(self):
2014
        'return the stay alive settings'
2015
        
2016
        return self.stay_alive
2017

    
2018
    def get_fftgrid(self):
2019
        'return soft and hard fft grids'
2020
        
2021
        nc = netCDF(self.nc, 'r')
2022
        soft = []
2023
        hard = []
2024
        for d in [1, 2, 3]:
2025
            sd = 'softgrid_dim%i' % d
2026
            hd = 'hardgrid_dim%i' % d
2027
            if sd in nc.dimensions:
2028
                soft.append(nc.dimensions[sd])
2029
                hard.append(nc.dimensions[hd])
2030
        nc.close()
2031
        return ({'soft':soft,
2032
                 'hard':hard})
2033

    
2034
    def get_kpts_type(self):
2035
        'return the kpt grid type'
2036
        
2037
        nc = netCDF(self.nc, 'r')
2038

    
2039
        if 'BZKpoints' in nc.variables:
2040
            bv = nc.variables['BZKpoints']
2041
            if hasattr(bv, 'gridtype'):
2042
                kpts_type = bv.gridtype #string saved in jacapo
2043
            else:
2044
                #no grid attribute, this ncfile was created pre-jacapo
2045
                kpts_type = 'pre-Jacapo: %i kpts' % len(bv[:])
2046
        else:
2047
            kpts_type = 'BZKpoints not defined. [[0,0,0]] used by default.'
2048

    
2049
        nc.close()
2050
        return kpts_type
2051
    
2052
    def get_kpts(self):
2053
        'return the BZ kpts'
2054
        nc = netCDF(self.nc, 'r')
2055

    
2056
        if 'BZKpoints' in nc.variables:
2057
            bv = nc.variables['BZKpoints']
2058
            kpts = bv[:]
2059
        else:
2060
            kpts = ([0, 0, 0]) #default Gamma point used in Dacapo when
2061
                             #BZKpoints not defined
2062

    
2063
        nc.close()
2064
        return kpts
2065
        
2066
    def get_nbands(self):
2067
        'return the number of bands used in the calculation'
2068
        nc = netCDF(self.nc, 'r')
2069

    
2070
        if 'ElectronicBands' in nc.variables:
2071
            v = nc.variables['ElectronicBands']
2072
            if hasattr(v, 'NumberOfBands'):
2073
                nbands = v.NumberOfBands[0]
2074
            else:
2075
                nbands = None
2076
        else:
2077
            nbands = None
2078
            
2079
        nc.close()
2080
        return nbands
2081
    
2082
    def get_ft(self):
2083
        'return the FermiTemperature used in the calculation'
2084
        nc = netCDF(self.nc, 'r')
2085

    
2086
        if 'ElectronicBands' in nc.variables:
2087
            v = nc.variables['ElectronicBands']
2088
            if hasattr(v, 'OccupationStatistics_FermiTemperature'):
2089
                ft = v.OccupationStatistics_FermiTemperature
2090
            else:
2091
                ft = None
2092
        else:
2093
            ft = None
2094
        nc.close()
2095
        return ft
2096
    
2097
    def get_dipole(self):
2098
        'return dictionary of parameters if the DipoleCorrection was used'
2099
        
2100
        nc = netCDF(self.get_nc(), 'r')
2101
        pars = {}
2102
        if 'DipoleCorrection' in nc.variables:
2103
            v = nc.variables['DipoleCorrection']
2104
            pars['status'] = True
2105
            if hasattr(v, 'MixingParameter'):
2106
                pars['mixpar'] = v.MixingParameter
2107
            if hasattr(v, 'InitialValue'):
2108
                pars['initval'] = v.InitialValue
2109
            if hasattr(v, 'AdditiveDipoleField'):
2110
                pars['adddipfield'] = v.AdditiveDipoleField
2111
            if hasattr(v, 'DipoleLayerPosition'):
2112
                pars['position'] = v.DipoleLayerPosition
2113
            
2114
        else:
2115
            pars = False
2116
        nc.close()
2117
        return pars
2118
        
2119
    def get_pw(self):
2120
        'return the planewave cutoff used'
2121
        
2122
        ncf = netCDF(self.nc, 'r')
2123
        if 'PlaneWaveCutoff' in ncf.variables:
2124
            pw = ncf.variables['PlaneWaveCutoff'].getValue()
2125
        else:
2126
            pw = None
2127
        ncf.close()
2128
        
2129
        if isinstance(pw, int) or isinstance(pw, float):
2130
            return pw
2131
        elif pw is None:
2132
            return None
2133
        else:
2134
            return pw[0]
2135

    
2136
    def get_dw(self):
2137
        'return the density wave cutoff'
2138
        
2139
        ncf = netCDF(self.nc, 'r')
2140
        if 'Density_WaveCutoff' in ncf.variables:
2141
            dw = ncf.variables['Density_WaveCutoff'].getValue()
2142
        else:
2143
            dw = None
2144
        ncf.close()
2145

    
2146
        #some old calculations apparently store ints, while newer ones
2147
        #are lists
2148
        if isinstance(dw, int) or isinstance(dw, float):
2149
            return dw
2150
        else:
2151
            if dw is None:
2152
                return None
2153
            else:
2154
                return dw[0]
2155
    
2156
    def get_xc(self):
2157
        '''return the self-consistent exchange-correlation functional used
2158

2159
        returns a string'''
2160
        
2161
        nc = netCDF(self.nc, 'r')
2162
        v = 'ExcFunctional'
2163
        if v in nc.variables:
2164
            xc = nc.variables[v][:].tostring().strip()
2165
        else:
2166
            xc = None
2167

    
2168
        nc.close()
2169
        return xc
2170

    
2171
    def get_potential_energy(self,
2172
                             atoms=None,
2173
                             force_consistent=False):
2174
        '''
2175
        return the potential energy.
2176
        '''
2177
        
2178
        if self.calculation_required(atoms):
2179
            log.debug('calculation required for energy')
2180
            self.calculate()
2181
        else:
2182
            log.debug('no calculation required for energy')
2183
                        
2184
        nc = netCDF(self.get_nc(), 'r')
2185
        try:
2186
            if force_consistent:
2187
                e = nc.variables['TotalFreeEnergy'][-1]
2188
            else:
2189
                e = nc.variables['TotalEnergy'][-1]
2190
            nc.close()
2191
            return e 
2192
        except (TypeError, KeyError):
2193
            raise RuntimeError('Error in calculating the total energy\n' +
2194
                               'Check ascii out file for error messages')
2195

    
2196
    def get_forces(self, atoms=None):
2197
        """Calculate atomic forces"""
2198
        
2199
        if atoms is None:
2200
            atoms = self.atoms
2201
        if self.calculation_required(atoms):
2202
            self.calculate()
2203
        nc = netCDF(self.get_nc(), 'r')
2204
        forces = nc.variables['DynamicAtomForces'][-1]
2205
        nc.close()
2206
        return forces
2207

    
2208
    def get_atoms(self):
2209
        'return the atoms attached to a calculator()'
2210
        
2211
        if hasattr(self, 'atoms'):
2212
            if self.atoms is None:
2213
                return None
2214
            atoms = self.atoms.copy()
2215
            #it is not obvious the copy of atoms should have teh same
2216
            #calculator
2217
            atoms.set_calculator(self)
2218
        else:
2219
            atoms = None
2220
        return atoms
2221

    
2222
    def get_nc(self):
2223
        'return the ncfile used for output'
2224
        
2225
        return self.nc
2226

    
2227
    def get_txt(self):
2228
        'return the txt file used for output'
2229
        
2230
        if hasattr(self,'txt'):
2231
            return self.txt
2232
        else:
2233
            return None
2234
        
2235
    def get_psp(self, sym=None, z=None):
2236
        '''get the pseudopotential filename from the psp database
2237

2238
        :Parameters:
2239

2240
          sym : string
2241
            the chemical symbol of the species
2242

2243
          z : integer
2244
            the atomic number of the species
2245

2246
        
2247
        you can only specify sym or z. Returns the pseudopotential
2248
        filename, not the full path.
2249
        '''
2250
        
2251
        if (sym is None and z is not None):
2252
            from ase.data import chemical_symbols
2253
            sym = chemical_symbols[z]
2254
        elif (sym is not None and z is None):
2255
            pass
2256
        else:
2257
            raise Exception, 'You can only specify Z or sym!'
2258
        psp = self.psp[sym]
2259
        return psp
2260
    
2261
    def get_spin_polarized(self):
2262
        'Return True if calculate is spin-polarized or False if not'
2263
        
2264
        #self.calculate() #causes recursion error with get_magnetic_moments
2265
        nc = netCDF(self.nc, 'r')
2266
        if 'ElectronicBands' in nc.variables:
2267
            v = nc.variables['ElectronicBands']
2268
            if hasattr(v, 'SpinPolarization'):
2269
                if v.SpinPolarization == 1:
2270
                    spinpol = False
2271
                elif v.SpinPolarization == 2:
2272
                    spinpol = True
2273
            else:
2274
                spinpol = False
2275
        else:
2276
            spinpol = 'Not defined'
2277

    
2278
        nc.close()
2279
        return spinpol
2280

    
2281
    def get_magnetic_moments(self, atoms=None):
2282
        '''return magnetic moments on each atom after the calculation is
2283
        run'''
2284

    
2285
        if self.calculation_required(atoms):
2286
            self.calculate()
2287
        nc = netCDF(self.nc, 'r')
2288
        if 'InitialAtomicMagneticMoment' in nc.variables:
2289
            mom = nc.variables['InitialAtomicMagneticMoment'][:]
2290
        else:
2291
            mom = [0.0]*len(self.atoms)
2292

    
2293
        nc.close()
2294
        return mom
2295

    
2296
    def get_status(self):
2297
        '''get status of calculation from ncfile. usually one of:
2298
        'new',
2299
        'aborted'
2300
        'running'
2301
        'finished'
2302
        None
2303
        '''
2304
        
2305
        nc = netCDF(self.nc, 'r')
2306
        if hasattr(nc, 'status'):
2307
            status = nc.status
2308
        else:
2309
            status = None
2310
        nc.close()
2311
        return status
2312

    
2313
    def get_calculate_stress(self):
2314
        'return whether stress is calculated or not'
2315
        
2316
        nc = netCDF(self.get_nc(), 'r')
2317
        if 'TotalStress' in nc.variables:
2318
            calcstress = True
2319
        else:
2320
            calcstress = False
2321
        nc.close()
2322
        return calcstress
2323
        
2324
    def get_stress(self, atoms=None):
2325
        '''get stress on the atoms.
2326

2327
        you should have set up the calculation
2328
        to calculate stress first.
2329

2330
        returns [sxx, syy, szz, syz, sxz, sxy]'''
2331
        
2332
        if self.calculation_required(atoms):
2333
            self.calculate()
2334

    
2335
        nc = netCDF(self.get_nc(), 'r')
2336
        if 'TotalStress' in nc.variables:
2337
            stress = nc.variables['TotalStress'][:]
2338
            #ase expects the 6-element form
2339
            stress = np.take(stress.ravel(), [0, 4, 8, 5, 2, 1])
2340
        else:
2341
            #stress will not be here if you did not set it up by
2342
            #calling set_stress() or in the __init__
2343
            stress = None
2344
        
2345
        nc.close()
2346
        
2347
        return stress
2348

    
2349
    def get_psp_valence(self, psp):
2350
        '''
2351
        get the psp valence charge on an atom from the pspfile.
2352
        '''
2353
        
2354
        from struct import unpack
2355
        dacapopath = os.environ.get('DACAPOPATH')
2356

    
2357
        if os.path.exists(psp):
2358
            #the pspfile may be in the current directory
2359
            #or defined by an absolute path
2360
            fullpsp = psp
2361
        else:
2362
            #or, it is in the default psp path
2363
            fullpsp = os.path.join(dacapopath, psp)
2364

    
2365
        if os.path.exists(fullpsp.strip()):
2366
            f = open(fullpsp)
2367
            # read past version numbers and text information
2368
            buf = f.read(64)
2369
            # read number valence electrons
2370
            buf = f.read(8)
2371
            fmt = ">d"
2372
            nvalence = unpack(fmt, buf)[0]
2373
            f.close()
2374

    
2375
        else:
2376
            raise Exception, "%s does not exist" % fullpsp
2377

    
2378
        return nvalence
2379

    
2380
    def get_psp_nuclear_charge(self, psp):
2381
        '''
2382
        get the nuclear charge of the atom from teh psp-file.
2383

2384
        This is not the same as the atomic number, nor is it
2385
        necessarily the negative of the number of valence electrons,
2386
        since a psp may be an ion. this function is needed to compute
2387
        centers of ion charge for the dipole moment calculation.
2388

2389
        We read in the valence ion configuration from the psp file and
2390
        add up the charges in each shell.
2391
        '''
2392
        
2393
        from struct import unpack
2394
        dacapopath = os.environ.get('DACAPOPATH')
2395

    
2396
        if os.path.exists(psp):
2397
            #the pspfile may be in the current directory
2398
            #or defined by an absolute path
2399
            fullpsp = psp
2400

    
2401
        else:
2402
            #or, it is in the default psp path
2403
            fullpsp = os.path.join(dacapopath, psp)
2404

    
2405
        if os.path.exists(fullpsp.strip()):
2406
            f = open(fullpsp)
2407
            unpack('>i', f.read(4))[0]
2408
            for i in range(3):
2409
                f.read(4)
2410
            for i in range(3):
2411
                f.read(4)
2412
            f.read(8)
2413
            f.read(20)
2414
            f.read(8)
2415
            f.read(8)
2416
            f.read(8)
2417
            nvalps = unpack('>i', f.read(4))[0]
2418
            f.read(4)
2419
            f.read(8)
2420
            f.read(8)
2421
            wwnlps = []
2422
            for i in range(nvalps):
2423
                f.read(4)
2424
                wwnlps.append(unpack('>d', f.read(8))[0])
2425
                f.read(8)
2426
            f.close()
2427

    
2428
        else:
2429
            raise Exception, "%s does not exist" % fullpsp
2430

    
2431
        return np.array(wwnlps).sum()
2432
       
2433
    def get_valence(self, atoms=None):
2434
        '''return the total number of valence electrons for the
2435
        atoms. valence electrons are read directly from the
2436
        pseudopotentials.
2437

2438
        the psp filenames are stored in the ncfile. They may be just
2439
        the name of the file, in which case the psp may exist in the
2440
        same directory as the ncfile, or in $DACAPOPATH, or the psp
2441
        may be defined by an absolute or relative path. This function
2442
        deals with all these possibilities.
2443
        '''
2444
        
2445
        from struct import unpack
2446
        
2447
        #do not use get_atoms() or recursion occurs
2448
        if atoms is None:
2449
            if hasattr(self, 'atoms'):
2450
                atoms = self.atoms
2451
            else:
2452
                return None
2453

    
2454
        dacapopath = os.environ.get('DACAPOPATH')
2455
        totval = 0.0
2456
        for sym in atoms.get_chemical_symbols():
2457
            psp = self.get_psp(sym)
2458
            
2459
            if os.path.exists(psp):
2460
                #the pspfile may be in the current directory
2461
                #or defined by an absolute path
2462
                fullpsp = psp
2463

    
2464
            #let's also see if we can construct an absolute path to a
2465
            #local or relative path psp.
2466
            abs_path_to_nc = os.path.abspath(self.get_nc())
2467
            base = os.path.split(abs_path_to_nc)[0]
2468
            possible_path_to_psp = os.path.join(base, psp)
2469
            if os.path.exists(possible_path_to_psp):
2470
                fullpsp = possible_path_to_psp
2471
                
2472
            else:
2473
                #or, it is in the default psp path
2474
                fullpsp = os.path.join(dacapopath, psp)
2475
            if os.path.exists(fullpsp.strip()):
2476
                f = open(fullpsp)
2477
                # read past version numbers and text information
2478
                buf = f.read(64)
2479
                # read number valence electrons
2480
                buf = f.read(8)
2481
                fmt = ">d"
2482
                nvalence = unpack(fmt, buf)[0]
2483
                f.close()
2484
                totval += float(nvalence)
2485
            else:
2486
                raise Exception, "%s does not exist" % fullpsp
2487
            
2488
        return totval 
2489

    
2490
    def calculation_required(self, atoms=None, quantities=None):
2491
        '''
2492
        determines if a calculation is needed.
2493

2494
        return True if a calculation is needed to get up to date data.
2495
        return False if no calculation is needed.
2496

2497
        quantities is here because of the ase interface.
2498
        '''
2499
        
2500
        # first, compare if the atoms is the same as the stored atoms
2501
        # if anything has changed, we need to run a calculation
2502
        log.debug('running calculation_required')
2503

    
2504
        if self.nc is None:
2505
            raise Exception, 'No output ncfile specified!'
2506
                
2507
        if atoms is not None:
2508
            if not self.atoms_are_equal(atoms):
2509
                log.debug('found that atoms != self.atoms')
2510
                tol = 1.0e-6 #tolerance that the unit cell is the same
2511
                new = atoms.get_cell()
2512
                old = self.atoms.get_cell()
2513
                #float comparison of equality
2514
                if not np.all(abs(old-new) < tol): 
2515
                    #this often changes the number of planewaves
2516
                    #which requires a complete restart
2517
                    log.debug('restart required! because cell changed')
2518
                    self.restart()
2519
                else:
2520
                    log.debug('Unitcells apparently the same')
2521
                    
2522
                self.set_atoms(atoms) #we have to update the atoms in any case
2523
                return True
2524
            
2525
        #if we make it past the atoms check, we look in the
2526
        #nc file. if parameters have been changed the status
2527
        #will tell us if a calculation is needed
2528

    
2529
        #past this point, atoms was None or equal, so there is nothing to
2530
        #update in the calculator
2531

    
2532
        log.debug('atoms tested equal')
2533
        if os.path.exists(self.nc):
2534
            nc = netCDF(self.nc, 'r')
2535
            if hasattr(nc, 'status'):
2536
                if nc.status == 'finished' and self.ready:
2537
                    nc.close()
2538
                    return False
2539
                elif nc.status == 'running':
2540
                    nc.close()
2541
                    raise DacapoRunning('Dacapo is Running')
2542
                elif nc.status == 'aborted':
2543
                    nc.close()
2544
                    raise DacapoAborted('Dacapo aborted. see txt file!')
2545
                else:
2546
                    log.debug('ncfile exists, but is not ready')
2547
                    nc.close()
2548
                    return True
2549
            else:
2550
                #legacy calculations do not have a status flag in them.
2551
                #let us guess that if the TotalEnergy is there
2552
                #no calculation needs to be run?
2553
                if 'TotalEnergy' in nc.variables:
2554
                    runflag = False
2555
                else:
2556
                    runflag = True
2557
                nc.close()
2558
                log.debug('Legacy calculation')
2559
                return runflag #if no status run calculation
2560
            nc.close()
2561
            
2562
        #default, a calculation is required
2563
        return True
2564

    
2565
    def get_scratch(self):
2566
        '''finds an appropriate scratch directory for the calculation'''
2567
        
2568
        import getpass
2569
        username = getpass.getuser()
2570

    
2571
        scratch_dirs = []
2572
        if os.environ.has_key('SCRATCH'):
2573
            scratch_dirs.append(os.environ['SCRATCH'])
2574
        if os.environ.has_key('SCR'):
2575
            scratch_dirs.append(os.environ['SCR'])
2576
        scratch_dirs.append('/scratch/'+username)
2577
        scratch_dirs.append('/scratch/')
2578
        scratch_dirs.append(os.curdir)
2579
        for scratch_dir in scratch_dirs:
2580
            if os.access(scratch_dir, os.W_OK):
2581
                return scratch_dir
2582
        raise IOError, "No suitable scratch directory and no write access \
2583
        to current dir."
2584

    
2585
    def calculate(self):
2586
        '''run a calculation.
2587

2588
        you have to be a little careful with code in here. Use the
2589
        calculation_required function to tell if a calculation is
2590
        required. It is assumed here that if you call this, you mean
2591
        it.'''
2592

    
2593
        #provide a way to make no calculation get run
2594
        if os.environ.get('DACAPO_DRYRUN', None) is not None:
2595
            raise Exception, '$DACAPO_DRYRUN detected, and a calculation \
2596
            attempted'
2597
    
2598
        if not self.ready:
2599
            log.debug('Calculator is not ready.')
2600
            if not os.path.exists(self.get_nc()):
2601
                self.initnc()
2602

    
2603
            log.debug('writing atoms out')
2604
            log.debug(self.atoms)
2605
            
2606
            self.write_nc() #write atoms to ncfile
2607

    
2608
            log.debug('writing input out')
2609
            self.write_input() #make sure input is uptodate
2610
            
2611

    
2612
            #check that the bands get set
2613
            if self.get_nbands() is None:
2614
                nelectrons = self.get_valence()
2615
                nbands = int(nelectrons * 0.65 + 4)
2616
                self.set_nbands(nbands) 
2617
            
2618
        log.debug('running a calculation')
2619

    
2620
        nc = self.get_nc()
2621
        txt = self.get_txt()
2622
        scratch = self.get_scratch()
2623

    
2624
        if self.stay_alive:
2625
            self.execute_external_dynamics(nc, txt)
2626
            self.ready = True
2627
            self.set_status('finished')
2628
        else:
2629
            cmd = 'dacapo.run  %(innc)s -out %(txt)s -scratch %(scratch)s'
2630
            cmd = cmd % {'innc':nc,
2631
                         'txt':txt,
2632
                         'scratch':scratch}
2633

    
2634
            log.debug(cmd)
2635
            # using subprocess instead of commands subprocess is more
2636
            # flexible and works better for stay_alive
2637
            self._dacapo = sp.Popen(cmd,
2638
                                    stdout=sp.PIPE,
2639
                                    stderr=sp.PIPE,
2640
                                    shell=True)
2641
            status = self._dacapo.wait()
2642
            [stdout, stderr] = self._dacapo.communicate()
2643
            output = stdout+stderr
2644
            
2645
            if status is 0: #that means it ended fine!
2646
                self.ready = True
2647
                self.set_status('finished')
2648
            else:
2649
                log.debug('Status was not 0')
2650
                log.debug(output)
2651
                self.ready = False
2652
            # directory cleanup has been moved to self.__del__()
2653
            del self._dacapo
2654

    
2655
            #Sometimes dacapo dies or is killed abnormally, and in this
2656
            #case an exception should be raised to prevent a geometry
2657
            #optimization from continuing for example. The best way to
2658
            #detect this right now is actually to check the end of the
2659
            #text file to make sure it ends with the right line. The
2660
            #line differs if the job was run in parallel or in serial.
2661
            f = open(txt, 'r')
2662
            lines = f.readlines()
2663
            f.close()
2664

    
2665
            if 'PAR: msexit halting Master' in lines[-1]:
2666
                pass #standard parallel end
2667
            elif ('TIM' in lines[-2]
2668
                  and 'clexit: exiting the program' in lines[-1]):
2669
                pass #standard serial end
2670
            else:
2671
                # text file does not end as expected, print the last
2672
                # 10 lines and raise exception
2673
                log.debug(string.join(lines[-10:-1], ''))
2674
                s = 'Dacapo output txtfile (%s) did not end normally.' 
2675
                raise DacapoAbnormalTermination(s % txt)
2676

    
2677
    def execute_external_dynamics(self,
2678
                                  nc=None,
2679
                                  txt=None,
2680
                                  stoppfile='stop',
2681
                                  stopprogram=None):
2682
        '''
2683
        Implementation of the stay alive functionality with socket
2684
        communication between dacapo and python.  Known limitations:
2685
        It is not possible to start 2 independent Dacapo calculators
2686
        from the same python process, since the python PID is used as
2687
        identifier for the script[PID].py file.
2688
        '''
2689
        
2690
        from socket import socket, AF_INET, SOCK_STREAM, timeout
2691
        import tempfile
2692
        
2693
        if hasattr(self, "_dacapo"):
2694
            msg = "Starting External Dynamics while Dacapo is runnning: %s"
2695
            msg = msg % str(self._dacapo.poll())
2696
            log.debug(msg)
2697
        else:
2698
            log.debug("No dacapo instance has been started yet")
2699
        log.debug("Stopprogram: %s" % stopprogram)
2700

    
2701
        if not nc:
2702
            nc = self.get_nc()
2703
        if not txt:
2704
            txt = self.get_txt()
2705
        tempfile.tempdir = os.curdir
2706

    
2707
        if stopprogram:
2708
            # write stop file
2709
            stfile = open(stoppfile, 'w')
2710
            stfile.write('1 \n')
2711
            stfile.close()
2712

    
2713
            # signal to dacapo that positions are ready
2714
            # let dacapo continue, it is up to the python mainloop 
2715
            # to allow dacapo enough time to finish properly.
2716
            self._client.send('ok too proceed')
2717

    
2718
            # Wait for dacapo to acknowledge that netcdf file has
2719
            # been updated, and analysis part of the code has been
2720
            # terminated. Dacapo sends a signal at the end of call
2721
            # clexit().
2722
            log.info("waiting for dacapo to exit...")
2723
            self.s.settimeout(1200.0) # if dacapo exits with an
2724
                                       # error, self.s.accept()
2725
                                       # should time out,
2726
                                      # but we need to give it
2727
                                      # enough time to write the
2728
                                      # wave function to the nc
2729
                                      # file.
2730
            try:
2731
                self._client, self._addr = self.s.accept() # Last
2732
                                                          # mumble
2733
                                                          # before
2734
                                                          # Dacapo
2735
                                                          # dies.
2736
                os.system("sleep 5") # 5 seconds of silence
2737
                                                         # mourning
2738
                                                         # dacapo.
2739
            except timeout:
2740
                print '''Socket connection timed out.'''
2741
                print '''This usually means Dacapo crashed.'''
2742
                
2743
            # close the socket s
2744
            self.s.close()
2745
            self._client.close()
2746

    
2747
            # remove the script???? file
2748
            ncfile = netCDF(nc, 'r')
2749
            vdyn = ncfile.variables['Dynamics']
2750
            os.system('rm -f '+vdyn.ExternalIonMotion_script)
2751
            ncfile.close()
2752
            os.system('rm -f '+stoppfile)
2753

    
2754
            if self._dacapo.poll()==None:  # dacapo is still not dead!
2755
                # but this should do it!
2756
                sp.Popen("kill -9 "+str(self._dacapo.pid), shell=True)
2757
                #if Dacapo dies for example because of too few
2758
                #bands, subprocess never returns an exitcode.
2759
                #very strange, but at least the program is
2760
                #terminated.  print self._dacapo.returncode
2761
            del self._dacapo
2762
            return
2763

    
2764
        if hasattr(self, '_dacapo') and self._dacapo.poll()==None: 
2765
            # returns  None if  dacapo is running self._dacapo_is_running:
2766

    
2767
            # calculation_required already updated the positions in
2768
            # the nc file
2769
            self._client.send('ok too proceed')
2770

    
2771
        else:
2772

    
2773
            # get process pid that will be used as communication
2774
            # channel
2775
            pid = os.getpid()
2776

    
2777
            # setup communication channel to dacapo
2778
            from sys    import version
2779
            from string import split
2780
            effpid = (pid)%(2**16-1025)+1025 # This translate pid
2781
                                               # [0;99999] to a number
2782
                                               # in [1025;65535] (the
2783
                                               # allowed socket
2784
                                               # numbers)
2785

    
2786
            self.s = socket(AF_INET, SOCK_STREAM)
2787
            foundafreesocket = 0
2788
            while not foundafreesocket:
2789
                try:
2790
                    if split(version)[0] > "2":     # new interface
2791
                        self.s.bind(("", effpid))
2792
                    else:                           # old interface
2793
                        self.s.bind("", effpid)
2794
                    foundafreesocket = 1
2795
                except:
2796
                    effpid = effpid + 1
2797

    
2798
            # write script file that will be used by dacapo
2799
            scriptname = 'script%s.py' % str(pid)
2800
            scriptfile = open(scriptname, 'w')
2801
            scriptfile.write(
2802
"""#!/usr/bin/env python
2803
from socket import *
2804
from sys    import version
2805
from string import split  
2806
s = socket(AF_INET,SOCK_STREAM)
2807
# tell python that dacapo has finished
2808
if split(version)[0] > "2":     # new interface 
2809
     s.connect(("",%(effpid)s))
2810
else:                           # old interface
2811
     s.connect("",%(effpid)s)
2812
# wait for python main loop
2813
s.recv(14)
2814
""" % {'effpid':str(effpid)})
2815
            scriptfile.close()
2816
            os.system('chmod +x ' + scriptname)
2817

    
2818
            # setup dynamics as external and set the script name
2819
            ncfile = netCDF(nc, 'a')
2820
            if 'Dynamics' not in ncfile.variables:
2821
                vdyn = ncfile.createVariable('Dynamics', 'c', ())
2822
            else:
2823
                vdyn = ncfile.variables['Dynamics']
2824
            vdyn.Type = "ExternalIonMotion" 
2825
            vdyn.ExternalIonMotion_script = './'+ scriptname
2826
            ncfile.close()
2827

    
2828
            # dacapo is not running start dacapo non blocking
2829
            scratch_in_nc = tempfile.mktemp()
2830
            os.system('mv '+nc+' '+scratch_in_nc)
2831
            os.system('rm -f '+stoppfile)
2832
            scratch = self.get_scratch()
2833
            cmd = 'dacapo.run'
2834
            cmd += ' %(innc)s %(outnc)s -out %(txt)s -scratch %(scratch)s'
2835
            cmd = cmd % {'innc':scratch_in_nc,
2836
                         'outnc':nc,
2837
                         'txt':txt,
2838
                         'scratch':scratch}
2839

    
2840
            log.debug(cmd)
2841
            self._dacapo = sp.Popen(cmd,
2842
                                    stdout=sp.PIPE,
2843
                                    stderr=sp.PIPE,
2844
                                    shell=True)
2845

    
2846
            self.s.listen(1)
2847

    
2848
        # wait for dacapo  
2849
        self._client, self._addr = self.s.accept()
2850

    
2851
    def write_nc(self, nc=None, atoms=None):
2852
        '''
2853
        write out atoms to a netcdffile.
2854

2855
        This does not write out the calculation parameters!
2856

2857
        :Parameters:
2858

2859
          nc : string
2860
            ncfilename to write to. this file will get clobbered
2861
            if it already exists.
2862

2863
          atoms : ASE.Atoms
2864
            atoms to write. if None use the attached atoms
2865
            if no atoms are attached only the calculator is
2866
            written out. 
2867

2868
        the ncfile is always opened in 'a' mode.
2869

2870
        note: it is good practice to use the atoms argument to make
2871
        sure that the geometry you mean gets written! Otherwise, the
2872
        atoms in the calculator is used, which may be different than
2873
        the external copy of the atoms.
2874

2875
        '''
2876

    
2877
        log.debug('writing atoms to ncfile')
2878
        #no filename was provided to function, use the current ncfile
2879
        if nc is None: 
2880
            nc = self.get_nc()
2881
        
2882
        if nc != self.nc:
2883
            #this means we are writing a new file, and we should copy
2884
            #the old file to it first.  this makes sure the old
2885
            #calculator settings are preserved
2886
            new = nc
2887
            old = self.nc
2888
            log.debug('Copying old ncfile to new ncfile')
2889
            log.debug('cp %s %s' % (old, new))
2890
            os.system('cp %s %s' % (old, new))
2891
              
2892
        if atoms is None:
2893
            atoms = self.get_atoms()
2894

    
2895
        log.debug('self.atoms = %s' % str(self.atoms))
2896
        log.debug('atoms = %s' % str(atoms))
2897

    
2898
        if atoms is not None: #there may still be no atoms attached
2899
            log.debug('about to write to %s' % nc)
2900
            ncf = netCDF(nc, 'a')
2901

    
2902
            if 'number_of_dynamic_atoms' not in ncf.dimensions:
2903
                ncf.createDimension('number_of_dynamic_atoms',
2904
                                    len(atoms))
2905
            else:
2906
                # number of atoms is already a dimension, but we might
2907
                # be setting new atoms here
2908
                # check for same atom symbols (implicitly includes
2909
                # a length check)
2910
                symbols = np.array(['%2s' % s for s in
2911
                                    atoms.get_chemical_symbols()], dtype='c')
2912
                ncsym = ncf.variables['DynamicAtomSpecies'][:]
2913
                if (symbols.size != ncsym.size) or (np.any(ncsym != symbols)):
2914
                    # the number of atoms or their order has changed.
2915
                    # Treat this as a new calculation and reset
2916
                    # number_of_ionic_steps and
2917
                    # number_of_dynamic_atoms.
2918
                    ncf.close() #nc file must be closed for
2919
                                 #delete_ncattdimvar to work correctly
2920
                    self.delete_ncattdimvar(nc, ncattrs=[],
2921
                                            ncdims=['number_of_dynamic_atoms',
2922
                                                    'number_ionic_steps'])
2923
                    ncf = netCDF(nc, 'a')
2924
                    ncf.createDimension('number_of_dynamic_atoms',
2925
                                                  len(atoms)) 
2926
                    ncf.createDimension('number_ionic_steps', None)
2927
                    self._set_frame_number(0)
2928
                    ncf.close() #nc file must be closed for restart to
2929
                                #work correctly
2930
                    self.restart()
2931
                    ncf = netCDF(nc, 'a')
2932

    
2933
            #now, create variables
2934
            if 'DynamicAtomSpecies' not in ncf.variables:
2935
                sym = ncf.createVariable('DynamicAtomSpecies',
2936
                                         'c',
2937
                                         ('number_of_dynamic_atoms',
2938
                                          'dim2',))
2939
            else:
2940
                sym = ncf.variables['DynamicAtomSpecies']
2941

    
2942
            #note explicit array casting was required here
2943
            symbols = atoms.get_chemical_symbols()
2944
            sym[:] = np.array(['%2s' % s for s in symbols], dtype='c')
2945

    
2946
            if 'DynamicAtomPositions' not in ncf.variables:
2947
                pos = ncf.createVariable('DynamicAtomPositions',
2948
                                         'd',
2949
                                         ('number_ionic_steps',
2950
                                          'number_of_dynamic_atoms',
2951
                                          'dim3'))
2952
            else:
2953
                pos = ncf.variables['DynamicAtomPositions']
2954

    
2955
            spos = atoms.get_scaled_positions()
2956
            if pos.typecode() == 'f':
2957
                spos = np.array(spos, dtype=np.float32)
2958
            pos[self._frame, :] = spos
2959

    
2960
            if 'UnitCell' not in ncf.variables:
2961
                uc = ncf.createVariable('UnitCell', 'd',
2962
                                        ('number_ionic_steps',
2963
                                         'dim3', 'dim3'))
2964
            else:
2965
                uc = ncf.variables['UnitCell']
2966

    
2967
            cell = atoms.get_cell()
2968
            if uc.typecode() == 'f':
2969
                cell = np.array(cell, dtype=np.float32)
2970
                
2971
            uc[self._frame, :] = cell
2972

    
2973
            if 'AtomTags' not in ncf.variables:
2974
                tags = ncf.createVariable('AtomTags', 'i',
2975
                                          ('number_of_dynamic_atoms',))
2976
            else:
2977
                tags = ncf.variables['AtomTags']
2978

    
2979
            tags[:] = np.array(atoms.get_tags(), np.int32)
2980

    
2981
            if 'InitialAtomicMagneticMoment' not in ncf.variables:
2982
                mom = ncf.createVariable('InitialAtomicMagneticMoment',
2983
                                         'd',
2984
                                         ('number_of_dynamic_atoms',))
2985
            else:
2986
                mom = ncf.variables['InitialAtomicMagneticMoment']
2987

    
2988
            #explain why we have to use get_initial_magnetic_moments()
2989
            moms = atoms.get_initial_magnetic_moments()
2990
            if mom.typecode() == 'f':
2991
                moms = np.array(moms, dtype=np.float32)
2992
            mom[:] = moms
2993
            
2994
            #finally the atom pseudopotentials
2995
            for sym in atoms.get_chemical_symbols():
2996
                vn = 'AtomProperty_%s' % sym
2997
                if vn not in ncf.variables:
2998
                    p = ncf.createVariable(vn, 'c', ('dim20',))
2999
                else:
3000
                    p = ncf.variables[vn]
3001

    
3002
                ppath = self.get_psp(sym=sym)
3003
                p.PspotFile = ppath
3004

    
3005
            ncf.sync()
3006
            ncf.close()
3007
               
3008
            #store constraints if they exist
3009
            constraints = atoms._get_constraints()
3010
            if constraints != []:
3011
                nc = netCDF(self.get_nc(), 'a')
3012
                if 'constraints' not in nc.variables:
3013
                    if 'dim1' not in nc.dimensions:
3014
                        nc.createDimension('dim1', 1)
3015
                    c = nc.createVariable('constraints', 'c', ('dim1',))
3016
                else:
3017
                    c = nc.variables['constraints']
3018
                #we store the pickle string as an attribute of a
3019
                #netcdf variable because that way we do not have to
3020
                #know how long the string is.  with a character
3021
                #variable you have to specify the dimension of the
3022
                #string ahead of time.
3023
                c.data = pickle.dumps(constraints)
3024
                nc.close()
3025
            else:
3026
                # getting here means there where no constraints on the
3027
                # atoms just written we should check if there are any
3028
                # old constraints left in the ncfile
3029
                # from a previous atoms, and delete them if so
3030
                delete_constraints = False
3031
                nc = netCDF(self.get_nc())
3032
                if 'constraints' in nc.variables:
3033
                    delete_constraints = True
3034
                nc.close()
3035

    
3036
                if delete_constraints:
3037
                    log.debug('deleting old constraints')
3038
                    self.delete_ncattdimvar(self.nc,
3039
                                            ncvars=['constraints'])
3040
        
3041
    def read_atoms(filename):
3042
        '''read atoms and calculator from an existing netcdf file.
3043

3044
        :Parameters:
3045

3046
          filename : string
3047
            name of file to read from.
3048

3049
        static method
3050

3051
        example::
3052
        
3053
          >>> atoms = Jacapo.read_atoms(ncfile)
3054
          >>> calc = atoms.get_calculator()
3055

3056
        this method is here for legacy purposes. I used to use it alot.
3057
        '''
3058
        
3059
        calc = Jacapo(filename)
3060
        atoms = calc.get_atoms()
3061
        return atoms
3062
    
3063
    read_atoms = staticmethod(read_atoms)
3064

    
3065
    def read_only_atoms(self, ncfile):
3066
        '''read only the atoms from an existing netcdf file. Used to
3067
        initialize a calculator from a ncfilename.
3068

3069
        :Parameters:
3070

3071
          ncfile : string
3072
            name of file to read from.
3073

3074
        return ASE.Atoms with no calculator attached or None if no
3075
        atoms found
3076
        '''
3077
        
3078
        from ase import Atoms
3079
        
3080
        nc = netCDF(ncfile, 'r')
3081
        #some ncfiles do not have atoms in them
3082
        if 'UnitCell' not in nc.variables:
3083
            log.debug('no unit cell found in ncfile')
3084
            nc.close()
3085
            return None
3086
        
3087
        cell = nc.variables['UnitCell'][:][-1]
3088
        sym = nc.variables['DynamicAtomSpecies'][:]
3089
        symbols = [x.tostring().strip() for x in sym]
3090
        spos = nc.variables['DynamicAtomPositions'][:][-1]
3091

    
3092
        pos = np.dot(spos, cell)
3093
        
3094
        atoms = Atoms(symbols=symbols,
3095
                      positions=pos,
3096
                      cell=cell)
3097

    
3098
        if 'AtomTags' in nc.variables:
3099
            tags = nc.variables['AtomTags'][:]
3100
            atoms.set_tags(tags)
3101

    
3102
        if 'InitialAtomicMagneticMoment' in nc.variables:
3103
            mom = nc.variables['InitialAtomicMagneticMoment'][:]
3104
            atoms.set_initial_magnetic_moments(mom)
3105

    
3106
        #update psp database
3107
        for sym in symbols:
3108
            vn = 'AtomProperty_%s' % sym
3109
            if vn in nc.variables:
3110
                var = nc.variables[vn]
3111
                pspfile = var.PspotFile
3112
                self.psp[sym] = pspfile
3113

    
3114
        #get constraints if they exist
3115
        c = nc.variables.get('constraints', None)
3116
        if c is not None:
3117
            constraints = pickle.loads(c.data)
3118
            atoms.set_constraint(constraints)
3119
                    
3120
        nc.close()
3121
        
3122
        return atoms
3123
        
3124
    def delete_ncattdimvar(self, ncf, ncattrs=None, ncdims=None, ncvars=None):
3125
        '''
3126
        helper function to delete attributes,
3127
        dimensions and variables in a netcdffile
3128

3129
        this functionality is not implemented for some reason in
3130
        netcdf, so the only way to do this is to copy all the
3131
        attributes, dimensions, and variables to a new file, excluding
3132
        the ones you want to delete and then rename the new file.
3133

3134
        if you delete a dimension, all variables with that dimension
3135
        are also deleted.
3136
        '''
3137

    
3138
        if ncattrs is None:
3139
            ncattrs = []
3140
        if ncdims is None:
3141
            ncdims = []
3142
        if ncvars is None:
3143
            ncvars = []
3144
        
3145
        log.debug('beginning: going to delete dims: %s' % str(ncdims))
3146
        log.debug('beginning: going to delete vars: %s' % str(ncvars))
3147
            
3148
        oldnc = netCDF(ncf, 'r')
3149

    
3150
        #h,tempnc = tempfile.mkstemp(dir='.',suffix='.nc')
3151
        tempnc = ncf+'.temp'
3152
        
3153
        newnc = netCDF(tempnc, 'w')
3154

    
3155
        for attr in dir(oldnc):
3156
            if attr in ['close', 'createDimension',
3157
                        'createVariable', 'flush', 'sync']:
3158
                continue
3159
            if attr in ncattrs:
3160
                continue #do not copy this attribute
3161
            setattr(newnc, attr, getattr(oldnc, attr))
3162
           
3163
        #copy dimensions
3164
        for dim in oldnc.dimensions:
3165
            if dim in ncdims:
3166
                log.debug('deleting %s of %s' % (dim, str(ncdims)))
3167
                continue #do not copy this dimension
3168
            size = oldnc.dimensions[dim]
3169
            
3170
            newnc.createDimension(dim, size)
3171

    
3172
        # we need to delete all variables that depended on a deleted dimension
3173
        for v in oldnc.variables:
3174
            dims1 = oldnc.variables[v].dimensions
3175
            for dim in ncdims:
3176
                if dim in dims1:
3177
                    s = 'deleting "%s" because it depends on dim "%s"'
3178
                    log.debug(s %(v, dim))
3179
                    ncvars.append(v)
3180

    
3181
        #copy variables, except the ones to delete
3182
        for v in oldnc.variables:
3183
            if v in ncvars:
3184
                log.debug('vars to delete: %s ' % ncvars)
3185
                log.debug('deleting ncvar: %s' % v)
3186
                continue #we do not copy this v over
3187

    
3188
            ncvar = oldnc.variables[v]
3189
            tcode = ncvar.typecode()
3190
            #char typecodes do not come out right apparently
3191
            if tcode == " ":
3192
                tcode = 'c'
3193
            
3194
            ncvar2 = newnc.createVariable(v, tcode, ncvar.dimensions)
3195
            try:
3196
                ncvar2[:] = ncvar[:]
3197
            except TypeError:
3198
                #this exception occurs for scalar variables
3199
                #use getValue and assignValue instead
3200
                ncvar2.assignValue(ncvar.getValue())
3201

    
3202
            #and variable attributes
3203
            #print dir(ncvar)
3204
            for att in dir(ncvar):
3205
                if att in ['assignValue', 'getValue', 'typecode']:
3206
                    continue
3207
                setattr(ncvar2, att, getattr(ncvar, att))
3208

    
3209
        oldnc.close()
3210
        newnc.close()
3211

    
3212
        s = 'looking for .nfs files before copying: %s'
3213
        log.debug(s % glob.glob('.nfs*'))
3214
        
3215
        #ack!!! this makes .nfsxxx files!!!
3216
        #os.close(h) #this avoids the stupid .nfsxxx file
3217
        #import shutil
3218
        #shutil.move(tempnc,ncf)
3219

    
3220
        #this seems to avoid making the .nfs files 
3221
        os.system('cp %s %s' % (tempnc, ncf))
3222
        os.system('rm %s' % tempnc)
3223

    
3224
        s = 'looking for .nfs files after copying: %s'
3225
        log.debug(s %  glob.glob('.nfs*'))
3226
        
3227
    def restart(self):
3228
        '''
3229
        Restart the calculator by deleting nc dimensions that will
3230
        be rewritten on the next calculation. This is sometimes required
3231
        when certain dimensions change related to unitcell size changes
3232
        planewave/densitywave cutoffs and kpt changes. These can cause
3233
        fortran netcdf errors if the data does not match the pre-defined
3234
        dimension sizes.
3235

3236
        also delete all the output from previous calculation.
3237
        '''
3238
        
3239
        log.debug('restarting!')
3240
        
3241
        ncdims = ['number_plane_waves',
3242
                  'number_IBZ_kpoints',
3243
                  'softgrid_dim1',
3244
                  'softgrid_dim2',
3245
                  'softgrid_dim3',
3246
                  'hardgrid_dim1',
3247
                  'hardgrid_dim2',
3248
                  'hardgrid_dim3',
3249
                  'max_projectors_per_atom',
3250
                  'atomdos_energygrid_size',
3251
                  'atomdos_angular_channels',
3252
                  'atomdos_radial_orbs']
3253

    
3254
        ncvars = ['TotalEnergy',
3255
                  'TotalFreeEnergy',
3256
                  'EvaluateTotalEnergy',
3257
                  'DynamicAtomForces',
3258
                  'FermiLevel',
3259
                  'EnsembleXCEnergies',
3260
                  'AtomProjectedDOS_IntegratedDOS',
3261
                  'AtomProjectedDOS_OrdinalMap',
3262
                  'NumberPlaneWavesKpoint',
3263
                  'AtomProjectedDOS_EnergyResolvedDOS',
3264
                  'AtomProjectedDOS_EnergyGrid',
3265
                  'EvaluateCorrelationEnergy',
3266
                  'DynamicAtomVelocities',
3267
                  'KpointWeight',
3268
                  'EvaluateExchangeEnergy',
3269
                  'EffectivePotential',
3270
                  'TotalStress',
3271
                  'ChargeDensity',
3272
                  'WaveFunction',
3273
                  'WaveFunctionFFTindex',
3274
                  'NumberOfNLProjectors',
3275
                  'NLProjectorPsi',
3276
                  'TypeNLProjector1',
3277
                  'NumberofNLProjectors',
3278
                  'PartialCoreDensity',
3279
                  'ChargeDensity',
3280
                  'ElectrostaticPotential',
3281
                  'StructureFactor',
3282
                  'EigenValues',
3283
                  'OccupationNumbers']
3284
                  
3285
        self.delete_ncattdimvar(self.nc,
3286
                                ncattrs=[],
3287
                                ncdims=ncdims,
3288
                                ncvars=ncvars)
3289

    
3290
        self.set_status('new')
3291
        self.ready = False
3292

    
3293
    def get_convergence(self):
3294
        'return convergence settings for Dacapo'
3295
        
3296
        nc = netCDF(self.get_nc(), 'r')
3297
        vname = 'ConvergenceControl'
3298
        if vname in nc.variables:
3299
            v = nc.variables[vname]
3300
            convergence = {}
3301
            if hasattr(v, 'AbsoluteEnergyConvergence'):
3302
                convergence['energy'] = v.AbsoluteEnergyConvergence[0]
3303
            if hasattr(v, 'DensityConvergence'):
3304
                convergence['density'] = v.DensityConvergence[0]
3305
            if hasattr(v, 'OccupationConvergence'):
3306
                convergence['occupation'] = v.OccupationConvergence[0]
3307
            if hasattr(v, 'MaxNumberOfSteps'):
3308
                convergence['maxsteps'] = v.MaxNumberOfSteps[0]
3309
            if hasattr(v, 'CPUTimeLimit'):
3310
                convergence['cputime'] = v.CPUTimeLimit[0]
3311
        else:
3312
            convergence = None
3313

    
3314
        nc.close()
3315
        return convergence
3316

    
3317
    def set_convergence(self,
3318
                        energy=0.00001,
3319
                        density=0.0001,
3320
                        occupation=0.001,
3321
                        maxsteps=None,
3322
                        maxtime=None
3323
                        ):
3324
        '''set convergence criteria for stopping the dacapo calculator.
3325

3326
        :Parameters:
3327

3328
          energy : float
3329
            set total energy change (eV) required for stopping
3330

3331
          density : float
3332
            set density change required for stopping
3333

3334
          occupation : float
3335
            set occupation change required for stopping
3336

3337
          maxsteps : integer
3338
            specify maximum number of steps to take
3339

3340
          maxtime : integer
3341
            specify maximum number of hours to run.
3342

3343
        Autopilot not supported here.
3344
        '''
3345
        
3346
        nc = netCDF(self.get_nc(), 'a')
3347
        vname = 'ConvergenceControl'
3348
        if vname in nc.variables:
3349
            v = nc.variables[vname]
3350
        else:
3351
            v = nc.createVariable(vname, 'c', ('dim1',))
3352

    
3353
        if energy is not None:
3354
            v.AbsoluteEnergyConvergence = energy
3355
        if density is not None:
3356
            v.DensityConvergence = density
3357
        if occupation is not None:
3358
            v.OccupationConvergence = occupation
3359
        if maxsteps is not None:
3360
            v.MaxNumberOfSteps = maxsteps
3361
        if maxtime is not None:
3362
            v.CPUTimeLimit = maxtime
3363

    
3364
        nc.sync()
3365
        nc.close()
3366

    
3367
    def get_charge_mixing(self):
3368
        'return charge mixing parameters'
3369
        
3370
        nc = netCDF(self.get_nc(), 'r')
3371
        vname = 'ChargeMixing'
3372
        if vname in nc.variables:
3373
            v = nc.variables[vname]
3374
            charge_mixing = {}
3375
            if hasattr(v, 'Method'):
3376
                charge_mixing['method'] = v.Method
3377
            if hasattr(v, 'UpdateCharge'):
3378
                charge_mixing['updatecharge'] = v.UpdateCharge
3379
            if hasattr(v, 'Pulay_MixingHistory'):
3380
                charge_mixing['mixinghistory'] = v.Pulay_MixingHistory[0]
3381
            if hasattr(v, 'Pulay_DensityMixingCoeff'):
3382
                charge_mixing['mixingcoeff'] = v.Pulay_DensityMixingCoeff[0]
3383
            if hasattr(v, 'Pulay_KerkerPrecondition'):
3384
                charge_mixing['precondition'] = v.Pulay_KerkerPrecondition
3385
        else:
3386
            charge_mixing = None
3387

    
3388
        nc.close()
3389
        return charge_mixing
3390

    
3391
    def set_charge_mixing(self,
3392
                          method='Pulay',
3393
                          mixinghistory=10,
3394
                          mixingcoeff=0.1,
3395
                          precondition='No',
3396
                          updatecharge='Yes'):
3397
        '''set density mixing method and parameters
3398

3399
        :Parameters:
3400

3401
          method : string
3402
            'Pulay' for Pulay mixing. only one supported now
3403

3404
          mixinghistory : integer
3405
            number of iterations to mix
3406
            Number of charge residual vectors stored for generating
3407
            the Pulay estimate on the self-consistent charge density,
3408
            see Sec. 4.2 in Kresse/Furthmuller:
3409
            Comp. Mat. Sci. 6 (1996) p34ff
3410

3411
          mixingcoeff : float
3412
            Mixing coefficient for Pulay charge mixing, corresponding
3413
            to A in G$^1$ in Sec. 4.2 in Kresse/Furthmuller:
3414
            Comp. Mat. Sci. 6 (1996) p34ff
3415
                        
3416
          precondition : string
3417
            'Yes' or 'No'
3418
            
3419
            * "Yes" : Kerker preconditiong is used,
3420
               i.e. q$_0$ is different from zero, see eq. 82
3421
               in Kresse/Furthmuller: Comp. Mat. Sci. 6 (1996).
3422
               The value of q$_0$ is fix to give a damping of 20
3423
               of the lowest q vector.
3424
            
3425
            * "No" : q$_0$ is zero and mixing is linear (default).
3426

3427
          updatecharge : string
3428
            'Yes' or 'No'
3429
            
3430
            * "Yes" : Perform charge mixing according to
3431
               ChargeMixing:Method setting
3432
              
3433
            * "No" : Freeze charge to initial value.
3434
               This setting is useful when evaluating the Harris-Foulkes
3435
               density functional
3436
              
3437
        '''
3438
        
3439
        if method == 'Pulay':
3440
            nc = netCDF(self.get_nc(), 'a')
3441
            vname = 'ChargeMixing'
3442
            if vname in nc.variables:
3443
                v = nc.variables[vname]
3444
            else:
3445
                v = nc.createVariable(vname, 'c', ('dim1',))
3446

    
3447
            v.Method = 'Pulay'
3448
            v.UpdateCharge = updatecharge
3449
            v.Pulay_MixingHistory = mixinghistory
3450
            v.Pulay_DensityMixingCoeff = mixingcoeff
3451
            v.Pulay_KerkerPrecondition = precondition
3452

    
3453
            nc.sync()
3454
            nc.close()
3455

    
3456
        self.ready = False
3457

    
3458
    def set_electronic_minimization(self,
3459
                                    method='eigsolve',
3460
                                    diagsperband=2):
3461
        '''set the eigensolver method
3462

3463
        Selector for which subroutine to use for electronic
3464
        minimization
3465

3466
        Recognized options : "resmin", "eigsolve" and "rmm-diis".
3467

3468
        * "resmin" : Power method (Lennart Bengtson), can only handle
3469
           k-point parallization.
3470

3471
        * "eigsolve : Block Davidson algorithm
3472
           (Claus Bendtsen et al).
3473

3474
        * "rmm-diis : Residual minimization
3475
           method (RMM), using DIIS (direct inversion in the iterate
3476
           subspace) The implementaion follows closely the algorithm
3477
           outlined in Kresse and Furthmuller, Comp. Mat. Sci, III.G/III.H
3478
        
3479
        :Parameters:
3480

3481
          method : string
3482
            should be 'resmin', 'eigsolve' or 'rmm-diis'
3483

3484
          diagsperband : int
3485
            The number of diagonalizations per band for
3486
            electronic minimization algorithms (maps onto internal
3487
            variable ndiapb). Applies for both
3488
            ElectronicMinimization:Method = "resmin" and "eigsolve".
3489
            default value = 2
3490
        '''
3491
        
3492
        nc = netCDF(self.get_nc(), 'a')
3493
        
3494
        vname = 'ElectronicMinimization'
3495
        if vname in nc.variables:
3496
            v = nc.variables[vname]
3497
        else:
3498
            log.debug('Creating ElectronicMinimization')
3499
            v = nc.createVariable(vname, 'c', ('dim1',))
3500

    
3501
        log.debug('setting method for ElectronicMinimization: % s' % method)
3502
        v.Method = method
3503
        log.debug('setting DiagonalizationsBand for ElectronicMinimization')
3504
        if diagsperband is not None:
3505
            v.DiagonalizationsPerBand = diagsperband
3506

    
3507
        log.debug('synchronizing ncfile')
3508
        nc.sync()
3509
        
3510
        nc.close()
3511

    
3512
    def get_electronic_minimization(self):
3513
        '''get method and diagonalizations per band for electronic
3514
        minimization algorithms'''
3515

    
3516
        log.debug('getting electronic minimization parameters')
3517
        
3518
        nc = netCDF(self.get_nc(), 'a')
3519
        vname = 'ElectronicMinimization'
3520
        if vname in nc.variables:
3521
            v = nc.variables[vname]
3522
            method = v.Method
3523
            if hasattr(v, 'DiagonalizationsPerBand'):
3524
                diagsperband = v.DiagonalizationsPerBand[0]
3525
            else:
3526
                diagsperband = None
3527
        else:
3528
            method = None
3529
            diagsperband = None
3530
        nc.close()
3531
        return {'method':method,
3532
                'diagsperband':diagsperband}
3533

    
3534
    def get_occupationstatistics(self):
3535
        'return occupation statistics method'
3536
        
3537
        nc = netCDF(self.get_nc(), 'r')
3538
        if 'ElectronicBands' in nc.variables:
3539
            v = nc.variables['ElectronicBands']
3540
            if hasattr(v, 'OccupationStatistics'):
3541
                occstat = v.OccupationStatistics
3542
            else:
3543
                occstat = None
3544
        else:
3545
            occstat = None
3546
        nc.close()
3547
        return occstat
3548

    
3549
    def set_occupationstatistics(self, method):
3550
        '''
3551
        set the method used for smearing the occupations.
3552

3553
        :Parameters:
3554

3555
          method : string
3556
            one of 'FermiDirac' or 'MethfesselPaxton'
3557
            Currently, the Methfessel-Paxton scheme (PRB 40, 3616 (1989).)
3558
            is implemented to 1th order (which is recommemded by most authors).
3559
            'FermiDirac' is the default
3560
        '''
3561
            
3562
        nc = netCDF(self.get_nc(), 'a')
3563
        if 'ElectronicBands' in nc.variables:
3564
            v = nc.variables['ElectronicBands']
3565
            v.OccupationStatistics = method
3566
        
3567
        nc.sync()
3568
        nc.close()
3569

    
3570
    def get_fermi_level(self):
3571
        'return Fermi level'
3572
        
3573
        if self.calculation_required():
3574
            self.calculate()
3575
        nc = netCDF(self.get_nc(), 'r')
3576
        ef = nc.variables['FermiLevel'][-1]
3577
        nc.close()
3578
        return ef
3579

    
3580
    def get_occupation_numbers(self, kpt=0, spin=0):
3581
        '''return occupancies of eigenstates for a kpt and spin
3582

3583
        :Parameters:
3584

3585
          kpt : integer
3586
            index of the IBZ kpoint you want the occupation of
3587

3588
          spin : integer
3589
            0 or 1
3590
        '''
3591
        
3592
        if self.calculation_required():
3593
            self.calculate()
3594
        nc = netCDF(self.get_nc(), 'r')
3595
        occ = nc.variables['OccupationNumbers'][:][-1][kpt, spin]
3596
        nc.close()
3597
        return occ
3598

    
3599
    def get_xc_energies(self, *functional):
3600
        """
3601
        Get energies for different functionals self-consistent and
3602
        non-self-consistent.
3603

3604
        :Parameters:
3605

3606
          functional : strings
3607
            some set of 'PZ','VWN','PW91','PBE','revPBE', 'RPBE'
3608

3609
        This function returns the self-consistent energy and/or
3610
        energies associated with various functionals. 
3611
        The functionals are currently PZ,VWN,PW91,PBE,revPBE, RPBE.
3612
        The different energies may be useful for calculating improved
3613
        adsorption energies as in B. Hammer, L.B. Hansen and
3614
        J.K. Norskov, Phys. Rev. B 59,7413. 
3615
        Examples: 
3616
        get_xcenergies() #returns all the energies
3617
        get_xcenergies('PBE') # returns the PBE total energy
3618
        get_xcenergies('PW91','PBE','revPBE') # returns a
3619
        # list of energies in the order asked for
3620
        """
3621
        
3622
        if self.calculation_required():
3623
            self.calculate()
3624

    
3625
        nc = netCDF(self.get_nc(), 'r')
3626

    
3627
        funcenergies = nc.variables['EvaluateTotalEnergy'][:][-1]
3628
        xcfuncs = nc.variables['EvalFunctionalOfDensity_XC'][:]
3629

    
3630
        nc.close()
3631
        
3632
        xcfuncs = [xc.tostring().strip() for xc in xcfuncs]
3633
        edict = dict(zip(xcfuncs, funcenergies))
3634

    
3635
        if len(functional) == 0:
3636
            #get all energies by default
3637
            functional = xcfuncs
3638

    
3639
        return [edict[xc] for xc in functional]
3640

    
3641
    # break of compatibility
3642
    def get_ados_data(self,
3643
                      atoms,
3644
                      orbitals,
3645
                      cutoff,
3646
                      spin):
3647
        '''get atom projected data
3648

3649
        :Parameters:
3650

3651
          atoms  
3652
              list of atom indices (integers)
3653

3654
          orbitals
3655
              list of strings
3656
              ['s','p','d'],
3657
              ['px','py','pz']
3658
              ['d_zz', 'dxx-yy', 'd_xy', 'd_xz', 'd_yz']
3659

3660
          cutoff : string
3661
            cutoff radius you want the results for 'short' or 'infinite'
3662

3663
          spin
3664
            : list of integers
3665
            spin you want the results for
3666
            [0] or [1] or [0,1] for both
3667

3668
        returns (egrid, ados)
3669
        egrid has the fermi level at 0 eV
3670
        '''
3671
              
3672
        if self.calculation_required():
3673
            self.calculate()
3674
        nc = netCDF(self.get_nc(), 'r')
3675
        omapvar = nc.variables['AtomProjectedDOS_OrdinalMap']
3676
        omap = omapvar[:] #indices
3677
        c = omapvar.AngularChannels
3678
        channels = [x.strip() for x in c.split(',')] #channel names
3679
        #this has dimensions(nprojections, nspins, npoints)
3680
        ados = nc.variables['AtomProjectedDOS_EnergyResolvedDOS'][:]
3681
        #this is the energy grid for all the atoms
3682
        egrid = nc.variables['AtomProjectedDOS_EnergyGrid'][:]
3683
        nc.close()
3684

    
3685
        #it is apparently not necessary to normalize the egrid to
3686
        #the Fermi level. the data is already for ef = 0.
3687

    
3688
        #get list of orbitals, replace 'p' and 'd' in needed
3689
        orbs = []
3690
        for o in orbitals:
3691
            if o == 'p':
3692
                orbs += ['p_x', 'p_y', 'p_z']
3693
            elif o == 'd':
3694
                orbs += ['d_zz', 'dxx-yy', 'd_xy', 'd_xz', 'd_yz']
3695
            else:
3696
                orbs += [o]
3697

    
3698
        orbinds = [channels.index(x) for x in orbs]
3699

    
3700
        cutdict = {'infinite':0,
3701
                   'short':1}
3702

    
3703
        icut = cutdict[cutoff]
3704

    
3705
        ydata = np.zeros(len(egrid), np.float)
3706
        
3707
        for atomind in atoms:
3708
            for oi in orbinds:
3709
                ind = omap[atomind, icut, oi]
3710

    
3711
                for si in spin:
3712
                    ydata += ados[ind, si]
3713

    
3714
        return (egrid, ydata)
3715

    
3716
    def get_all_eigenvalues(self, spin=0):
3717
        '''return all the eigenvalues at all the kpoints for a spin.
3718

3719
        :Parameters:
3720

3721
          spin : integer
3722
            which spin the eigenvalues are for'''
3723
        
3724
        if self.calculation_required():
3725
            self.calculate()
3726
        nc = netCDF(self.get_nc(), 'r')
3727
        ev = nc.variables['EigenValues'][:][-1][:, spin]
3728
        nc.close()
3729
        return ev
3730
    
3731
    def get_eigenvalues(self, kpt=0, spin=0):
3732
        '''return the eigenvalues for a kpt and spin
3733

3734
        :Parameters:
3735

3736
          kpt : integer
3737
            index of the IBZ kpoint
3738

3739
          spin : integer
3740
            which spin the eigenvalues are for'''
3741
        
3742
        if self.calculation_required():
3743
            self.calculate()
3744
        nc = netCDF(self.get_nc(), 'r')
3745
        ev = nc.variables['EigenValues'][:][-1][kpt, spin]
3746
        nc.close()
3747
        return ev
3748
    
3749
    def get_k_point_weights(self):
3750
        'return the weights on the IBZ kpoints'
3751
        
3752
        if self.calculation_required():
3753
            self.calculate()
3754
        nc = netCDF(self.get_nc(), 'r')
3755
        kw = nc.variables['KpointWeight'][:]
3756
        nc.close()
3757
        return kw
3758

    
3759
    def get_magnetic_moment(self):
3760
        'calculates the magnetic moment (Bohr-magnetons) of the supercell'
3761

    
3762
        if not self.get_spin_polarized():
3763
            return None
3764
        
3765
        if self.calculation_required():
3766
            self.calculate()
3767

    
3768
        nibzk = len(self.get_ibz_kpoints())
3769
        ibzkw = self.get_k_point_weights()
3770
        spinup, spindn = 0.0, 0.0
3771

    
3772
        for k in range(nibzk):
3773

    
3774
            spinup += self.get_occupation_numbers(k, 0).sum()*ibzkw[k]
3775
            spindn += self.get_occupation_numbers(k, 1).sum()*ibzkw[k]
3776

    
3777
        return (spinup - spindn)
3778

    
3779
    def get_number_of_spins(self):
3780
        'if spin-polarized returns 2, if not returns 1'
3781
        
3782
        if self.calculation_required():
3783
            self.calculate()
3784
        nc = netCDF(self.get_nc(), 'r')
3785
        spv = nc.variables['ElectronicBands']
3786
        nc.close()
3787

    
3788
        if hasattr(spv, 'SpinPolarization'):
3789
            return spv.SpinPolarization
3790
        else:
3791
            return 1
3792

    
3793
    def get_ibz_kpoints(self):
3794
        'return list of kpoints in the irreducible brillouin zone'
3795
        
3796
        if self.calculation_required():
3797
            self.calculate()
3798
        nc = netCDF(self.get_nc(), 'r')
3799
        ibz = nc.variables['IBZKpoints'][:]
3800
        nc.close()
3801
        return ibz
3802

    
3803
    get_ibz_k_points = get_ibz_kpoints
3804

    
3805
    def get_bz_k_points(self):
3806
        'return list of kpoints in the Brillouin zone'
3807
        
3808
        nc = netCDF(self.get_nc(), 'r')
3809
        if 'BZKpoints' in nc.variables:
3810
            bz = nc.variables['BZKpoints'][:]
3811
        else:
3812
            bz = None
3813
        nc.close()
3814
        return bz
3815
    
3816
    def get_effective_potential(self, spin=1):
3817
        '''
3818
        returns the realspace local effective potential for the spin.
3819
        the units of the potential are eV
3820

3821
        :Parameters:
3822

3823
          spin : integer
3824
             specify which spin you want, 0 or 1
3825
          
3826
        '''
3827
        
3828
        if self.calculation_required():
3829
            self.calculate()
3830
            
3831
        nc = netCDF(self.get_nc(), 'r')
3832
        efp = np.transpose(nc.variables['EffectivePotential'][:][spin])
3833
        nc.close()
3834
        fftgrids = self.get_fftgrid()
3835
        hardgrid = fftgrids['hard']
3836
        x, y, z = self.get_ucgrid(hardgrid)
3837
        return (x, y, z, efp)
3838
        
3839
    def get_electrostatic_potential(self, spin=0):
3840
        '''get electrostatic potential
3841

3842
        Netcdf documentation::
3843
        
3844
          double ElectrostaticPotential(number_of_spin,
3845
                                        hardgrid_dim3,
3846
                                        hardgrid_dim2,
3847
                                        hardgrid_dim1) ;
3848
                 ElectrostaticPotential:
3849
                     Description = "realspace local effective potential" ;
3850
                     unit = "eV" ;
3851
                     
3852
        '''
3853
        
3854
        if self.calculation_required():
3855
            self.calculate()
3856
            
3857
        nc = netCDF(self.get_nc(), 'r')
3858
        esp = np.transpose(nc.variables['ElectrostaticPotential'][:][spin])
3859
        nc.close()
3860
        fftgrids = self.get_fftgrid()
3861

    
3862
        x, y, z = self.get_ucgrid(fftgrids['hard'])
3863
        
3864
        return (x, y, z, esp)
3865
    
3866
    def get_charge_density(self, spin=0):
3867
        '''
3868
        return x,y,z,charge density data
3869
        
3870
        x,y,z are grids sampling the unit cell
3871
        cd is the charge density data
3872

3873
        netcdf documentation::
3874
        
3875
          ChargeDensity(number_of_spin,
3876
                        hardgrid_dim3,
3877
                        hardgrid_dim2,
3878
                        hardgrid_dim1)
3879
          ChargeDensity:Description = "realspace charge density" ;
3880
                  ChargeDensity:unit = "-e/A^3" ;
3881

3882
        '''
3883
        
3884
        if self.calculation_required():
3885
            self.calculate()
3886
            
3887
        nc = netCDF(self.get_nc(), 'r')
3888
     
3889
        cd = np.transpose(nc.variables['ChargeDensity'][:][spin])
3890

    
3891
        #I am not completely sure why this has to be done
3892
        #it does give units of electrons/ang**3
3893
        vol = self.get_atoms().get_volume()
3894
        cd /= vol
3895
        nc.close()
3896
        grids = self.get_fftgrid()
3897

    
3898
        x, y, z = self.get_ucgrid(grids['hard'])
3899
        return x, y, z, cd
3900

    
3901
    def get_ucgrid(self, dims):
3902
        '''Return X,Y,Z grids for uniform sampling of the unit cell
3903

3904
        dims = (n0,n1,n2)
3905

3906
        n0 points along unitcell vector 0
3907
        n1 points along unitcell vector 1
3908
        n2 points along unitcell vector 2
3909
        '''
3910
        
3911
        n0, n1, n2 = dims
3912
        
3913
        s0 = 1.0/n0
3914
        s1 = 1.0/n1
3915
        s2 = 1.0/n2
3916

    
3917
        X, Y, Z = np.mgrid[0.0:1.0:s0,
3918
                          0.0:1.0:s1,
3919
                          0.0:1.0:s2]
3920
        
3921
        C = np.column_stack([X.ravel(),
3922
                             Y.ravel(),
3923
                             Z.ravel()])
3924
        
3925
        atoms = self.get_atoms()
3926
        uc = atoms.get_cell()
3927
        real = np.dot(C, uc)
3928

    
3929
        #now convert arrays back to unitcell shape
3930
        RX = np.reshape(real[:, 0], (n0, n1, n2))
3931
        RY = np.reshape(real[:, 1], (n0, n1, n2))
3932
        RZ = np.reshape(real[:, 2], (n0, n1, n2))
3933
        return (RX, RY, RZ)
3934

    
3935
    def get_number_of_grid_points(self):
3936
        'return soft fft grid'
3937
        
3938
        # needed by ase.dft.wannier
3939
        fftgrids = self.get_fftgrid()
3940
        return np.array(fftgrids['soft'])
3941
    
3942
    def get_wannier_localization_matrix(self, nbands, dirG, kpoint,
3943
                                        nextkpoint, G_I, spin):
3944
        'return wannier localization  matrix'
3945

    
3946
        if self.calculation_required():
3947
            self.calculate()
3948

    
3949
        if not hasattr(self, 'wannier'):
3950
            from utils.wannier import Wannier
3951
            self.wannier = Wannier(self)
3952
            self.wannier.set_bands(nbands)
3953
            self.wannier.set_spin(spin)
3954
        locmat = self.wannier.get_zi_bloch_matrix(dirG,
3955
                                                  kpoint,
3956
                                                  nextkpoint,
3957
                                                  G_I)
3958
        return locmat
3959

    
3960
    def initial_wannier(self,
3961
                        initialwannier,
3962
                        kpointgrid,
3963
                        fixedstates,
3964
                        edf,
3965
                        spin):
3966
        'return initial wannier'
3967

    
3968
        if self.calculation_required():
3969
            self.calculate()
3970

    
3971
        if not hasattr(self, 'wannier'):
3972
            from utils.wannier import Wannier
3973
            self.wannier = Wannier(self)
3974

    
3975
        self.wannier.set_data(initialwannier)
3976
        self.wannier.set_k_point_grid(kpointgrid)
3977
        self.wannier.set_spin(spin)
3978

    
3979
        waves = [[self.get_reciprocal_bloch_function(band=band,
3980
                                                     kpt=kpt,
3981
                                                     spin=spin)
3982
                  for band in range(self.get_nbands())]
3983
                  for kpt in range(len(self.get_ibz_k_points()))]
3984

    
3985
        self.wannier.setup_m_matrix(waves, self.get_bz_k_points())
3986

    
3987
        #lfn is too keep line length below 78 characters
3988
        lfn = self.wannier.get_list_of_coefficients_and_rotation_matrices
3989
        c, U = lfn((self.get_nbands(), fixedstates, edf))
3990

    
3991
        U = np.array(U)
3992
        for k in range(len(c)):
3993
            c[k] = np.array(c[k])
3994
        return c, U
3995

    
3996
    def get_dipole_moment(self,atoms=None):
3997
        '''
3998
        return dipole moment of unit cell
3999

4000
        Defined by the vector connecting the center of electron charge
4001
        density to the center of nuclear charge density.
4002

4003
        Units = eV*angstrom
4004

4005
        1 Debye = 0.208194 eV*angstrom
4006

4007
        '''
4008
        if self.calculation_required():
4009
            self.calculate()
4010

    
4011
        if atoms is None:
4012
            atoms = self.get_atoms()
4013
        
4014
        #center of electron charge density
4015
        x, y, z, cd = self.get_charge_density()
4016

    
4017
        n1, n2, n3 = cd.shape
4018
        nelements = n1*n2*n3
4019
        voxel_volume = atoms.get_volume()/nelements
4020
        total_electron_charge = -cd.sum()*voxel_volume
4021

    
4022
        
4023
        electron_density_center = np.array([(cd*x).sum(),
4024
                                            (cd*y).sum(),
4025
                                            (cd*z).sum()])
4026
        electron_density_center *= voxel_volume
4027
        electron_density_center /= total_electron_charge
4028
       
4029
        electron_dipole_moment = electron_density_center*total_electron_charge
4030
        electron_dipole_moment *= -1.0 #we need the - here so the two
4031
                                        #negatives don't cancel
4032
        # now the ion charge center
4033
        psps = self.get_pseudopotentials()
4034
        ion_charge_center = np.array([0.0, 0.0, 0.0])
4035
        total_ion_charge = 0.0
4036
        for atom in atoms:
4037
            Z = self.get_psp_nuclear_charge(psps[atom.symbol])
4038
            total_ion_charge += Z
4039
            pos = atom.get_position()
4040
            ion_charge_center += Z*pos
4041

    
4042
        ion_charge_center /= total_ion_charge
4043
        ion_dipole_moment = ion_charge_center*total_ion_charge
4044

    
4045
        dipole_vector = (ion_dipole_moment + electron_dipole_moment)
4046
        return dipole_vector
4047

    
4048
    
4049
    def get_reciprocal_bloch_function(self, band=0, kpt=0, spin=0):
4050
        '''return the reciprocal bloch function. Need for Jacapo
4051
        Wannier class.'''
4052
        
4053
        if self.calculation_required():
4054
            self.calculate()
4055

    
4056
        nc = netCDF(self.get_nc(), 'r')
4057

    
4058
        # read reciprocal bloch function
4059
        npw = nc.variables['NumberPlaneWavesKpoint'][:]
4060
        bf = nc.variables['WaveFunction'][kpt, spin, band]
4061
        wflist = np.zeros(npw[kpt], np.complex)
4062
        wflist.real = bf[0:npw[kpt], 1]
4063
        wflist.imag = bf[0:npw[kpt], 0]
4064

    
4065
        nc.close()
4066

    
4067
        return wflist
4068

    
4069
    def get_reciprocal_fft_index(self, kpt=0):
4070
        '''return the Wave Function FFT Index'''
4071
        
4072
        nc = netCDF(self.get_nc(), 'r')
4073
        recind = nc.variables['WaveFunctionFFTindex'][kpt, :, :]
4074
        nc.close()
4075
        return recind
4076

    
4077
    def get_ensemble_coefficients(self):
4078
        'returns exchange correlation ensemble coefficients'
4079
        
4080
        # adapted from ASE/dacapo.py
4081
        # def GetEnsembleCoefficients(self):
4082
        #     self.Calculate()
4083
        #     E = self.GetPotentialEnergy()
4084
        #     xc = self.GetNetCDFEntry('EnsembleXCEnergies')
4085
        #     Exc = xc[0]
4086
        #     exc_c = self.GetNetCDFEntry('EvaluateCorrelationEnergy')
4087
        #     exc_e =  self.GetNetCDFEntry('EvaluateExchangeEnergy')
4088
        #     exc = exc_c + exc_e
4089
        #     if self.GetXCFunctional() == 'RPBE':
4090
        #             Exc = exc[-1][-1]
4091
        # 
4092
        #     E0 = xc[1]       # Fx = 0
4093
        # 
4094
        #     diff0 = xc[2] # - Exc
4095
        #     diff1 = xc[3] # - Exc
4096
        #     diff2 = xc[4] # - Exc
4097
        #     coefs = (E + E0 - Exc,diff0-E0 ,diff1-E0,diff2-E0)
4098
        #     print 'ensemble: (%.9f, %.9f, %.9f, %.9f)'% coefs
4099
        #     return num.array(coefs)
4100
        if self.calculation_required():
4101
            self.calculate()
4102

    
4103
        E = self.get_potential_energy()
4104
        nc = netCDF(self.get_nc(), 'r')
4105
        if 'EnsembleXCEnergies' in nc.variables:
4106
            v = nc.variables['EnsembleXCEnergies']
4107
            xc = v[:]
4108

    
4109
        EXC = xc[0]
4110

    
4111
        if 'EvaluateCorrelationEnergy' in nc.variables:
4112
            v = nc.variables['EvaluateCorrelationEnergy']
4113
            exc_c = v[:]
4114
            
4115
        if 'EvaluateExchangeEnergy' in nc.variables:
4116
            v = nc.variables['EvaluateExchangeEnergy']
4117
            exc_e = v[:]
4118

    
4119
        exc = exc_c + exc_e
4120

    
4121
        if self.get_xc == 'RPBE':
4122
            EXC = exc[-1][-1]
4123
            
4124
        E0 = xc[1]    # Fx = 0
4125

    
4126
        diff0 = xc[2] # - Exc
4127
        diff1 = xc[3] # - Exc
4128
        diff2 = xc[4] # - Exc
4129
        coefs = (E + E0 - EXC, diff0-E0, diff1-E0, diff2-E0)
4130
        log.info('ensemble: (%.9f, %.9f, %.9f, %.9f)'% coefs)
4131
        return np.array(coefs)
4132

    
4133
    def get_pseudo_wave_function(self, band=0, kpt=0, spin=0, pad=True):
4134

    
4135
        '''return the pseudo wavefunction'''
4136
        
4137
        # pad=True does nothing here.
4138
        if self.calculation_required():
4139
            self.calculate()
4140

    
4141
        ibz = self.get_ibz_kpoints()
4142

    
4143
        #get the reciprocal bloch function
4144
        wflist = self.get_reciprocal_bloch_function(band=band,
4145
                                                    kpt=kpt,
4146
                                                    spin=spin)
4147
        # wflist == Reciprocal Bloch Function
4148
 
4149
        recind = self. get_reciprocal_fft_index(kpt)
4150
        grids = self.get_fftgrid()
4151
        softgrid = grids['soft']
4152
        
4153
        # GetReciprocalBlochFunctionGrid
4154
        wfrec = np.zeros((softgrid), np.complex) 
4155

    
4156
        for i in xrange(len(wflist)):
4157
            wfrec[recind[0, i]-1,
4158
                  recind[1, i]-1,
4159
                  recind[2, i]-1] = wflist[i]
4160

    
4161
        # calculate Bloch Function
4162
        wf = wfrec.copy() 
4163
        dim = wf.shape
4164
        for i in range(len(dim)):
4165
            wf = np.fft.fft(wf, dim[i], axis=i)
4166

    
4167
        #now the phase function to get the bloch phase
4168
        basis = self.get_atoms().get_cell()
4169
        kpoint = np.dot(ibz[kpt], basis) #coordinates of relevant
4170
                                         #kpoint in cartesian
4171
                                         #coordinates
4172
        def phasefunction(coor):
4173
            'return phasefunction'
4174
            pf = np.exp(1.0j*np.dot(kpoint, coor))
4175
            return pf
4176

    
4177
        # Calculating the Bloch phase at the origin (0,0,0) of the grid
4178
        origin = np.array([0., 0., 0.])
4179
        blochphase = phasefunction(origin)
4180
        spatialshape = wf.shape[-len(basis):]
4181
        gridunitvectors = np.array(map(lambda unitvector,
4182
                                       shape:unitvector/shape,
4183
                                       basis,
4184
                                       spatialshape))
4185

    
4186
        for dim in range(len(spatialshape)):
4187
            # Multiplying with the phase at the origin
4188
            deltaphase = phasefunction(gridunitvectors[dim])
4189
            # and calculating phase difference between each point
4190
            newphase = np.fromfunction(lambda i, phase=deltaphase:phase**i,
4191
                                     (spatialshape[dim],))
4192
            blochphase = np.multiply.outer(blochphase, newphase)
4193

    
4194
        return blochphase*wf
4195

    
4196
    def get_wave_function(self, band=0, kpt=0, spin=0):
4197
        '''return the wave function. This is the pseudo wave function
4198
        divided by volume.'''
4199
        
4200
        pwf = self.get_pseudo_wave_function(band=band,
4201
                                            kpt=kpt,
4202
                                            spin=spin,
4203
                                            pad=True)
4204
        vol = self.get_atoms().get_volume()
4205
        fftgrids = self.get_fftgrid()
4206
        softgrid = fftgrids['soft']
4207
    
4208
        x, y, z = self.get_ucgrid((softgrid))
4209

    
4210
        return x, y, z, pwf/np.sqrt(vol)
4211

    
4212
    def strip(self):
4213
        '''remove all large memory nc variables not needed for
4214
        anything I use very often. 
4215
        '''
4216
        self.delete_ncattdimvar(self.nc,
4217
                                ncdims=['max_projectors_per_atom'],
4218
                                ncvars=['WaveFunction',
4219
                                        'WaveFunctionFFTindex',
4220
                                        'NumberOfNLProjectors',
4221
                                        'NLProjectorPsi',
4222
                                        'TypeNLProjector1',
4223
                                        'NumberofNLProjectors',
4224
                                        'PartialCoreDensity',
4225
                                        'ChargeDensity',
4226
                                        'ElectrostaticPotential',
4227
                                        'StructureFactor'])
4228

    
4229
# shortcut function names
4230
Jacapo.get_cd = Jacapo.get_charge_density
4231
Jacapo.get_wf = Jacapo.get_wave_function
4232
Jacapo.get_esp = Jacapo.get_electrostatic_potential
4233
Jacapo.get_occ = Jacapo.get_occupation_numbers
4234
Jacapo.get_ef = Jacapo.get_fermi_level
4235
Jacapo.get_number_of_bands = Jacapo.get_nbands
4236
Jacapo.set_pseudopotentials = Jacapo.set_psp