Statistiques
| Révision :

root / ase / calculators / vasp.py @ 20

Historique | Voir | Annoter | Télécharger (53,06 ko)

1
# Copyright (C) 2008 CSC - Scientific Computing Ltd.
2
"""This module defines an ASE interface to VASP.
3

4
Developed on the basis of modules by Jussi Enkovaara and John
5
Kitchin.  The path of the directory containing the pseudopotential 
6
directories (potpaw,potpaw_GGA, potpaw_PBE, ...) should be set 
7
by the environmental flag $VASP_PP_PATH. 
8

9
The user should also set the environmental flag $VASP_SCRIPT pointing
10
to a python script looking something like::
11

12
   import os
13
   exitcode = os.system('vasp')
14

15
Alternatively, user can set the environmental flag $VASP_COMMAND pointing
16
to the command use the launch vasp e.g. 'vasp' or 'mpirun -n 16 vasp'
17

18
http://cms.mpi.univie.ac.at/vasp/
19

20
-Jonas Bjork j.bjork@liverpool.ac.uk
21
"""
22

    
23
import os
24
import sys
25
from general import Calculator
26
from os.path import join, isfile, islink
27

    
28
import numpy as np
29

    
30
import ase
31

    
32
# Parameters that can be set in INCAR. The values which are None
33
# are not written and default parameters of VASP are used for them.
34

    
35
float_keys = [
36
    'aexx',       # Fraction of exact/DFT exchange
37
    'aggac',      # Fraction of gradient correction to correlation
38
    'aggax',      # Fraction of gradient correction to exchange
39
    'aldac',      # Fraction of LDA correlation energy
40
    'amix',       #
41
    'bmix',       # tags for mixing
42
    'emax',       # energy-range for DOSCAR file
43
    'emin',       #
44
    'enaug',      # Density cutoff
45
    'encut',      # Planewave cutoff
46
    'encutfock',  # FFT grid in the HF related routines 
47
    'hfscreen',   # attribute to change from PBE0 to HSE
48
    'potim',      # time-step for ion-motion (fs)
49
    'nelect',     # total number of electrons
50
    'pomass',     # mass of ions in am
51
    'sigma',      # broadening in eV
52
    'time',       # special control tag
53
    'zval',       # ionic valence
54
    ]
55

    
56
exp_keys = [
57
    'ediff',      # stopping-criterion for electronic upd.
58
    'ediffg',     # stopping-criterion for ionic upd.
59
    'symprec',    # precession in symmetry routines
60
]
61

    
62
string_keys = [
63
    'algo',       # algorithm: Normal (Davidson) | Fast | Very_Fast (RMM-DIIS)
64
    'gga',        # xc-type: PW PB LM or 91
65
    'prec',       # Precission of calculation (Low, Normal, Accurate)
66
    'system',     # name of System
67
    'tebeg',      #
68
    'teend',      # temperature during run
69
]
70

    
71
int_keys = [
72
    'ialgo',      # algorithm: use only 8 (CG) or 48 (RMM-DIIS)
73
    'ibrion',     # ionic relaxation: 0-MD 1-quasi-New 2-CG
74
    'idipol',     # monopol/dipol and quadropole corrections
75
    'iniwav',     # initial electr wf. : 0-lowe 1-rand
76
    'isif',       # calculate stress and what to relax
77
    'ismear',     # part. occupancies: -5 Blochl -4-tet -1-fermi 0-gaus >0 MP
78
    'ispin',      # spin-polarized calculation
79
    'istart',     # startjob: 0-new 1-cont 2-samecut
80
    'isym',       # symmetry: 0-nonsym 1-usesym
81
    'iwavpr',     # prediction of wf.: 0-non 1-charg 2-wave 3-comb
82
    'lmaxmix',    # 
83
    'lorbit',     # create PROOUT
84
    'ngx',        # FFT mesh for wavefunctions, x
85
    'ngxf',       # FFT mesh for charges x
86
    'ngy',        # FFT mesh for wavefunctions, y
87
    'ngyf',       # FFT mesh for charges y
88
    'ngz',        # FFT mesh for wavefunctions, z
89
    'ngzf',       # FFT mesh for charges z
90
    'nbands',     # Number of bands
91
    'nblk',       # blocking for some BLAS calls (Sec. 6.5)
92
    'nbmod',      # specifies mode for partial charge calculation
93
    'nelm',       #
94
    'nelmdl',     # nr. of electronic steps
95
    'nelmin',
96
    'nfree',      # number of steps per DOF when calculting Hessian using finitite differences
97
    'nkred',      # define sub grid of q-points for HF with nkredx=nkredy=nkredz 
98
    'nkredx',      # define sub grid of q-points in x direction for HF 
99
    'nkredy',      # define sub grid of q-points in y direction for HF 
100
    'nkredz',      # define sub grid of q-points in z direction for HF 
101
    'npar',       # parallelization over bands
102
    'nsim',       # evaluate NSIM bands simultaneously if using RMM-DIIS
103
    'nsw',        # number of steps for ionic upd.
104
    'nupdown',    # fix spin moment to specified value
105
    'nwrite',     # verbosity write-flag (how much is written)
106
    'smass',      # Nose mass-parameter (am)
107
    'voskown',    # use Vosko, Wilk, Nusair interpolation
108
]
109

    
110
bool_keys = [
111
    'addgrid',    # finer grid for augmentation charge density
112
    'icharg',     # charge: 1-file 2-atom 10-const
113
    'lasync',     # overlap communcation with calculations
114
    'lcharg',     #
115
    'lcorr',      # Harris-correction to forces
116
    'ldipol',     # potential correction mode
117
    'lelf',       # create ELFCAR
118
    'lhfcalc',    # switch to turn on Hartree Fock calculations
119
    'lpard',      # evaluate partial (band and/or k-point) decomposed charge density
120
    'lplane',     # parallelisation over the FFT grid
121
    'lscalapack', # switch off scaLAPACK
122
    'lscalu',     # switch of LU decomposition
123
    'lsepb',      # write out partial charge of each band seperately?
124
    'lsepk',      # write out partial charge of each k-point seperately?
125
    'lthomas',    # 
126
    'lvtot',      # create WAVECAR/CHGCAR/LOCPOT
127
    'lwave',      #
128
]
129

    
130
list_keys = [
131
    'dipol',      # center of cell for dipol
132
    'eint',       # energy range to calculate partial charge for
133
    'ferwe',      # Fixed band occupation
134
    'iband',      # bands to calculate partial charge for
135
    'magmom',     # initial magnetic moments
136
    'kpuse',      # k-point to calculate partial charge for
137
    'ropt',       # number of grid points for non-local proj in real space
138
    'rwigs',      # Wigner-Seitz radii
139
]
140

    
141
special_keys = [
142
    'lreal',      # non-local projectors in real space
143
]
144

    
145
keys = [    
146
    # 'NBLOCK' and KBLOCK       inner block; outer block
147
    # 'NPACO' and APACO         distance and nr. of slots for P.C.
148
    # 'WEIMIN, EBREAK, DEPER    special control tags
149
]
150

    
151
class Vasp(Calculator):
152
    def __init__(self, restart=None, output_template='vasp', track_output=False, write_input = True, 
153
                 **kwargs):
154
        self.write_input = write_input
155
        self.name = 'Vasp'
156
        self.float_params = {}
157
        self.exp_params = {}
158
        self.string_params = {}
159
        self.int_params = {}
160
        self.bool_params = {}
161
        self.list_params = {}
162
        self.special_params = {}
163
        for key in float_keys:
164
            self.float_params[key] = None
165
        for key in exp_keys:
166
            self.exp_params[key] = None
167
        for key in string_keys:
168
            self.string_params[key] = None
169
        for key in int_keys:
170
            self.int_params[key] = None
171
        for key in bool_keys:
172
            self.bool_params[key] = None
173
        for key in list_keys:
174
            self.list_params[key] = None
175
        for key in special_keys:
176
            self.special_params[key] = None
177
        self.string_params['prec'] = 'Normal'
178

    
179
        self.input_params = {
180
            'xc':     'PW91',  # exchange correlation potential 
181
            'setups': None,    # Special setups (e.g pv, sv, ...)
182
            'txt':    '-',     # Where to send information
183
            'kpts':   (1,1,1), # k-points
184
            'gamma':  False,   # Option to use gamma-sampling instead
185
                               # of Monkhorst-Pack
186
            }
187

    
188
        self.restart = restart
189
        self.track_output = track_output
190
        self.output_template = output_template
191
        if restart:
192
            self.restart_load()
193
            return
194

    
195
        if self.input_params['xc'] not in ['PW91','LDA','PBE']:
196
            raise ValueError(
197
                '%s not supported for xc! use one of: PW91, LDA or PBE.' %
198
                kwargs['xc'])
199
        self.nbands = self.int_params['nbands']
200
        self.atoms = None
201
        self.positions = None
202
        self.run_counts = 0
203
        self.set(**kwargs)
204

    
205
    def set(self, **kwargs):
206
        for key in kwargs:
207
            if self.float_params.has_key(key):
208
                self.float_params[key] = kwargs[key]
209
            elif self.exp_params.has_key(key):
210
                self.exp_params[key] = kwargs[key]
211
            elif self.string_params.has_key(key):
212
                self.string_params[key] = kwargs[key]
213
            elif self.int_params.has_key(key):
214
                self.int_params[key] = kwargs[key]
215
            elif self.bool_params.has_key(key):
216
                self.bool_params[key] = kwargs[key]
217
            elif self.list_params.has_key(key):
218
                self.list_params[key] = kwargs[key]
219
            elif self.special_params.has_key(key):
220
                self.special_params[key] = kwargs[key]
221
            elif self.input_params.has_key(key):
222
                self.input_params[key] = kwargs[key]
223
            else:
224
                raise TypeError('Parameter not defined: ' + key)
225

    
226
    def update(self, atoms):
227
        if self.calculation_required(atoms, ['energy']):
228
            if (self.atoms is None or
229
                self.atoms.positions.shape != atoms.positions.shape):
230
                # Completely new calculation just reusing the same
231
                # calculator, so delete any old VASP files found.
232
                self.clean()
233
            self.calculate(atoms)
234

    
235
    def initialize(self, atoms):
236
        """Initialize a VASP calculation
237

238
        Constructs the POTCAR file. User should specify the PATH
239
        to the pseudopotentials in VASP_PP_PATH environment variable"""
240

    
241
        p = self.input_params
242

    
243
        self.all_symbols = atoms.get_chemical_symbols()
244
        self.natoms = len(atoms)
245
        self.spinpol = atoms.get_initial_magnetic_moments().any()
246
        atomtypes = atoms.get_chemical_symbols()
247

    
248
        # Determine the number of atoms of each atomic species
249
        # sorted after atomic species
250
        special_setups = []
251
        if self.input_params['setups']:
252
            for m in self.input_params['setups']:
253
                try : 
254
                    #special_setup[self.input_params['setups'][m]] = int(m)
255
                    special_setups.append(int(m))
256
                except:
257
                    #print 'setup ' + m + ' is a groups setup'
258
                    continue
259
            #print 'special_setups' , special_setups
260

    
261
        #symbols = {}
262
        #for m,atom in enumerate(atoms):
263
        #    symbol = atom.get_symbol()
264
        #    if m in special_setups:
265
        #        pass
266
        #    else:
267
        #        if not symbols.has_key(symbol):
268
        #            symbols[symbol] = 1
269
        #        else:
270
        #            symbols[symbol] += 1
271
        
272
        symbols = []
273
        for m,atom in enumerate(atoms):
274
            symbol = atom.get_symbol()
275
            if m in special_setups:
276
                pass
277
            else:
278
                ind = 0
279
                bfound = False
280
                for symbolX, iatom in symbols:
281
                    if symbol == symbolX:
282
                        bfound = True
283
                        break
284
                    ind += 1
285
                if bfound:
286
                    symbolX, iatom  = symbols[ind]
287
                    symbols[ind] = symbolX, iatom+1
288
                else:
289
                    symbols.append((symbol, 1))
290
        
291
        # Build the sorting list
292
        self.sort = []
293
        self.sort.extend(special_setups)
294

    
295
        for symbol, iatom in symbols:
296
            for m,atom in enumerate(atoms):
297
                if m in special_setups: 
298
                    pass
299
                else:
300
                    if atom.get_symbol() == symbol:
301
                        self.sort.append(m)
302
        self.resort = range(len(self.sort))
303
        for n in range(len(self.resort)):
304
            self.resort[self.sort[n]] = n
305
        self.atoms_sorted = atoms[self.sort]
306
        
307
        # Check if the necessary POTCAR files exists and
308
        # create a list of their paths.
309
        self.symbol_count = []
310
        for m in special_setups:
311
            for symbol, iatom in symbols:
312
                if symbol == m:
313
                    self.symbol_count.append([atomtypes[m],1])
314
            #self.symbol_count.append([atomtypes[m],1])
315
        for m in symbols:
316
            for symbol, iatom in symbols:
317
                if symbol == m:
318
                    self.symbol_count.append([m,iatom])
319
            #self.symbol_count.append([m,symbols[m]])
320
        sys.stdout.flush()
321
        if self.write_input:
322
            xc = '/'
323
            #print 'p[xc]',p['xc']
324
            if p['xc'] == 'PW91':
325
                xc = '_gga/'
326
            elif p['xc'] == 'PBE':
327
                xc = '_pbe/'
328
            if 'VASP_PP_PATH' in os.environ:
329
                pppaths = os.environ['VASP_PP_PATH'].split(':')
330
            else:
331
                pppaths = []
332
            self.ppp_list = []
333
            # Setting the pseudopotentials, first special setups and 
334
            # then according to symbols
335
            for m in special_setups:
336
                name = 'potpaw'+xc.upper() + p['setups'][str(m)] + '/POTCAR'
337
                found = False
338
                for path in pppaths:
339
                    filename = join(path, name)
340
                    #print 'filename', filename
341
                    if isfile(filename) or islink(filename):
342
                        found = True
343
                        self.ppp_list.append(filename)
344
                        break
345
                    elif isfile(filename + '.Z') or islink(filename + '.Z'):
346
                        found = True
347
                        self.ppp_list.append(filename+'.Z')
348
                        break
349
                if not found:
350
                    raise RuntimeError('No pseudopotential for %s!' % symbol)    
351
            #print 'symbols', symbols 
352
            for symbol in symbols:
353
                try:
354
                    name = 'potpaw'+xc.upper()+symbol + p['setups'][symbol]
355
                except (TypeError, KeyError):
356
                    name = 'potpaw' + xc.upper() + symbol
357
                name += '/POTCAR'
358
                found = False
359
                for path in pppaths:
360
                    filename = join(path, name)
361
                    #print 'filename', filename
362
                    if isfile(filename) or islink(filename):
363
                        found = True
364
                        self.ppp_list.append(filename)
365
                        break
366
                    elif isfile(filename + '.Z') or islink(filename + '.Z'):
367
                        found = True
368
                        self.ppp_list.append(filename+'.Z')
369
                        break
370
                if not found:
371
                    raise RuntimeError('No pseudopotential for %s!' % symbol)
372
        self.converged = None
373
        self.setups_changed = None
374

    
375
    def calculate(self, atoms):
376
        """Generate necessary files in the working directory and run VASP.
377

378
        The method first write VASP input files, then calls the method
379
        which executes VASP. When the VASP run is finished energy, forces, 
380
        etc. are read from the VASP output.
381
        """
382
        
383
        # Write input
384
        from ase.io.vasp import write_vasp
385
        self.initialize(atoms)
386
        write_vasp('POSCAR', self.atoms_sorted, symbol_count = self.symbol_count)
387
        if self.write_input:
388
            self.write_incar(atoms)
389
            self.write_potcar()
390
            self.write_kpoints()
391
        self.write_sort_file()
392
        
393
        # Execute VASP
394
        self.run()
395
        # Read output
396
        atoms_sorted = ase.io.read('CONTCAR', format='vasp')
397
        if self.int_params['ibrion']>-1 and self.int_params['nsw']>0:
398
            # Replace entire atoms object with the one read from CONTCAR.
399
            atoms.__dict__ = atoms_sorted[self.resort].__dict__
400
        self.converged = self.read_convergence()
401
        self.set_results(atoms)
402

    
403
    def set_results(self, atoms):
404
        self.read(atoms)
405
        if self.spinpol:
406
            self.magnetic_moment = self.read_magnetic_moment()
407
            if self.int_params['lorbit']>=10 or (self.int_params['lorbit']!=None and self.list_params['rwigs']):
408
                self.magnetic_moments = self.read_magnetic_moments(atoms)
409
        self.old_float_params = self.float_params.copy()
410
        self.old_exp_params = self.exp_params.copy()
411
        self.old_string_params = self.string_params.copy()
412
        self.old_int_params = self.int_params.copy()
413
        self.old_input_params = self.input_params.copy()
414
        self.old_bool_params = self.bool_params.copy()
415
        self.old_list_params = self.list_params.copy()
416
        self.atoms = atoms.copy()
417
        
418
    def run(self):
419
        """Method which explicitely runs VASP."""
420

    
421
        if self.track_output:
422
            self.out = self.output_template+str(self.run_counts)+'.out'
423
            self.run_counts += 1
424
        else:
425
            self.out = self.output_template+'.out'
426
        stderr = sys.stderr
427
        p=self.input_params
428
        if p['txt'] is None:
429
            sys.stderr = devnull
430
        elif p['txt'] == '-':
431
            pass
432
        elif isinstance(p['txt'], str):
433
            sys.stderr = open(p['txt'], 'w')
434
        if os.environ.has_key('VASP_COMMAND'):
435
            vasp = os.environ['VASP_COMMAND']
436
            exitcode = os.system('%s > %s' % (vasp, self.out))
437
        elif os.environ.has_key('VASP_SCRIPT'):
438
            vasp = os.environ['VASP_SCRIPT']
439
            locals={}
440
            execfile(vasp, {}, locals)
441
            exitcode = locals['exitcode']
442
        else:
443
            raise RuntimeError('Please set either VASP_COMMAND or VASP_SCRIPT environment variable')
444
        sys.stderr = stderr
445
        if exitcode != 0:
446
            raise RuntimeError('Vasp exited with exit code: %d.  ' % exitcode)
447
        
448
    def restart_load(self):
449
        """Method which is called upon restart."""
450
        
451
        # Try to read sorting file
452
        if os.path.isfile('ase-sort.dat'):
453
            self.sort = []
454
            self.resort = []
455
            file = open('ase-sort.dat', 'r')
456
            lines = file.readlines()
457
            file.close()
458
            for line in lines:
459
                data = line.split()
460
                self.sort.append(int(data[0]))
461
                self.resort.append(int(data[1]))
462
            atoms = ase.io.read('CONTCAR', format='vasp')[self.resort]
463
        else:
464
            atoms = ase.io.read('CONTCAR', format='vasp')
465
            self.sort = range(len(atoms))
466
            self.resort = range(len(atoms))
467
        self.atoms = atoms.copy()
468
        self.read_incar()
469
        self.read_outcar()
470
        self.set_results(atoms)
471
        self.read_kpoints()
472
        self.read_potcar()
473
#        self.old_incar_params = self.incar_params.copy()
474
        self.old_input_params = self.input_params.copy()
475
        self.converged = self.read_convergence()
476

    
477
    def clean(self):
478
        """Method which cleans up after a calculation.
479
        
480
        The default files generated by Vasp will be deleted IF this
481
        method is called.
482

483
        """
484
        files = ['CHG', 'CHGCAR', 'POSCAR', 'INCAR', 'CONTCAR', 'DOSCAR',
485
                 'EIGENVAL', 'IBZKPT', 'KPOINTS', 'OSZICAR', 'OUTCAR', 'PCDAT',
486
                 'POTCAR', 'vasprun.xml', 'WAVECAR', 'XDATCAR',
487
                 'PROCAR', 'ase-sort.dat']
488
        for f in files:
489
            try:
490
                os.remove(f)
491
            except OSError:
492
                pass
493

    
494
    def set_atoms(self, atoms):
495
        if (atoms != self.atoms):
496
            self.converged = None
497
        self.atoms = atoms.copy()
498

    
499
    def get_atoms(self):
500
        atoms = self.atoms.copy()
501
        atoms.set_calculator(self)
502
        return atoms
503

    
504
    def get_potential_energy(self, atoms, force_consistent=False):
505
        self.update(atoms)
506
        if force_consistent:
507
            return self.energy_free
508
        else:
509
            return self.energy_zero
510

    
511
    def get_forces(self, atoms):
512
        self.update(atoms)
513
        return self.forces
514

    
515
    def get_stress(self, atoms):
516
        self.update(atoms)
517
        return self.stress
518

    
519
    def read_stress(self):
520
        for line in open('OUTCAR'):
521
            if line.find(' in kB  ') != -1:
522
                stress = -np.array([float(a) for a in line.split()[2:]]) \
523
                         [[0, 1, 2, 4, 5, 3]] \
524
                         * 1e-1 * ase.units.GPa
525
        return stress
526

    
527
    def calculation_required(self, atoms, quantities):
528
        if (self.positions is None or 
529
            (self.atoms != atoms) or
530
            (self.float_params != self.old_float_params) or
531
            (self.exp_params != self.old_exp_params) or
532
            (self.string_params != self.old_string_params) or
533
            (self.int_params != self.old_int_params) or
534
            (self.bool_params != self.old_bool_params) or
535
            (self.list_params != self.old_list_params) or
536
            (self.input_params != self.old_input_params)
537
            or not self.converged):
538
            return True
539
        if 'magmom' in quantities:
540
            return not hasattr(self, 'magnetic_moment')
541
        return False
542

    
543
    def get_number_of_bands(self):
544
        return self.nbands
545

    
546
    def get_k_point_weights(self):
547
        self.update(self.atoms)
548
        return self.read_k_point_weights()
549

    
550
    def get_number_of_spins(self):
551
        return 1 + int(self.spinpol)
552

    
553
    def get_eigenvalues(self, kpt=0, spin=0):
554
        self.update(self.atoms)
555
        return self.read_eigenvalues(kpt, spin)
556

    
557
    def get_fermi_level(self):
558
        return self.fermi
559

    
560
    def get_number_of_grid_points(self):
561
        raise NotImplementedError
562

    
563
    def get_pseudo_density(self):
564
        raise NotImplementedError
565

    
566
    def get_pseudo_wavefunction(self, n=0, k=0, s=0, pad=True):
567
        raise NotImplementedError
568

    
569
    def get_bz_k_points(self):
570
        raise NotImplementedError
571

    
572
    def get_ibz_kpoints(self):
573
        self.update(self.atoms)
574
        return self.read_ibz_kpoints()
575

    
576
    def get_spin_polarized(self):
577
        if not hasattr(self, 'spinpol'):
578
            self.spinpol = self.atoms.get_initial_magnetic_moments().any()
579
        return self.spinpol            
580

    
581
    def get_magnetic_moment(self, atoms):
582
        self.update(atoms)
583
        return self.magnetic_moment
584
        
585
    def get_magnetic_moments(self, atoms):
586
        if self.int_params['lorbit']>=10 or self.list_params['rwigs']:
587
            self.update(atoms)
588
            return self.magnetic_moments
589
        else:
590
            raise RuntimeError(
591
                "The combination %s for lorbit with %s for rwigs not supported to calculate magnetic moments" % (p['lorbit'], p['rwigs']))
592

    
593
    def get_dipole_moment(self, atoms):
594
        """Returns total dipole moment of the system."""
595
        self.update(atoms)
596
        return self.dipole
597

    
598
    def get_xc_functional(self):
599
        return self.input_params['xc']
600

    
601
    def write_incar(self, atoms, **kwargs):
602
        """Writes the INCAR file."""
603
        incar = open('INCAR', 'w')
604
        incar.write('INCAR created by Atomic Simulation Environment\n')
605
        for key, val in self.float_params.items():
606
            if val is not None:
607
                incar.write(' %s = %5.6f\n' % (key.upper(), val))
608
        for key, val in self.exp_params.items():
609
            if val is not None:
610
                incar.write(' %s = %5.2e\n' % (key.upper(), val))
611
        for key, val in self.string_params.items():
612
            if val is not None:
613
                incar.write(' %s = %s\n' % (key.upper(), val))
614
        for key, val in self.int_params.items():
615
            if val is not None:
616
                incar.write(' %s = %d\n' % (key.upper(), val))
617
        for key, val in self.list_params.items():
618
            if val is not None:
619
                incar.write(' %s = ' % key.upper())
620
                if key in ('dipol', 'eint', 'ferwe', 'ropt', 'rwigs'):
621
                    [incar.write('%.4f ' % x) for x in val]
622
                elif key in ('iband', 'kpuse'):
623
                    [incar.write('%i ' % x) for x in val]
624
                elif key == 'magmom':
625
                    list = [[1, val[0]]]
626
                    for n in range(1, len(val)):
627
                        if val[n] == val[n-1]:
628
                            list[-1][0] += 1
629
                        else:
630
                            list.append([1, val[n]])
631
                    [incar.write('%i*%.4f ' % (mom[0], mom[1])) for mom in list]
632
                incar.write('\n')
633
        for key, val in self.bool_params.items():
634
            if val is not None:
635
                incar.write(' %s = ' % key.upper())
636
                if val:
637
                    incar.write('.TRUE.\n')
638
                else:
639
                    incar.write('.FALSE.\n')
640
        for key, val in self.special_params.items():
641
            if val is not None:
642
                incar.write(' %s = ' % key.upper())
643
                if key is 'lreal':
644
                    if type(val)==type('str'):
645
                        incar.write(val+'\n')
646
                    elif type(val)==type(True):
647
                       if val:
648
                           incar.write('.TRUE.\n')
649
                       else:
650
                           incar.write('.FALSE.\n') 
651
        if self.spinpol and not self.int_params['ispin']:
652
            incar.write(' ispin = 2\n'.upper())
653
            # Write out initial magnetic moments
654
            magmom = atoms.get_initial_magnetic_moments()[self.sort]
655
            list = [[1, magmom[0]]]
656
            for n in range(1, len(magmom)):
657
                if magmom[n] == magmom[n-1]:
658
                    list[-1][0] += 1
659
                else:
660
                    list.append([1, magmom[n]])
661
            incar.write(' magmom = '.upper())
662
            [incar.write('%i*%.4f ' % (mom[0], mom[1])) for mom in list]
663
            incar.write('\n')
664
        incar.close()
665

    
666
    def write_kpoints(self, **kwargs):
667
        """Writes the KPOINTS file."""
668
        p = self.input_params
669
        kpoints = open('KPOINTS', 'w')
670
        kpoints.write('KPOINTS created by Atomic Simulation Environemnt\n')
671
        shape=np.array(p['kpts']).shape
672
        if len(shape)==1:
673
            kpoints.write('0\n')
674
            if p['gamma']:
675
                kpoints.write('Gamma\n')
676
            else:
677
                kpoints.write('Monkhorst-Pack\n')
678
            [kpoints.write('%i ' % kpt) for kpt in p['kpts']]
679
            kpoints.write('\n0 0 0')
680
        elif len(shape)==2:
681
            kpoints.write('%i \n' % (len(p['kpts'])))
682
            kpoints.write('Cartesian\n')
683
            for n in range(len(p['kpts'])):
684
                [kpoints.write('%f ' % kpt) for kpt in p['kpts'][n]]
685
                if shape[1]==4:
686
                    kpoints.write('\n')
687
                elif shape[1]==3:
688
                    kpoints.write('1.0 \n')
689
        kpoints.close()
690

    
691
    def write_potcar(self,suffix = ""):
692
        """Writes the POTCAR file."""
693
        import tempfile
694
        potfile = open('POTCAR'+suffix,'w')
695
        for filename in self.ppp_list:
696
            if filename.endswith('R'):
697
                for line in open(filename, 'r'):
698
                    potfile.write(line)
699
            elif filename.endswith('.Z'):
700
                file_tmp = tempfile.NamedTemporaryFile()
701
                os.system('gunzip -c %s > %s' % (filename, file_tmp.name))
702
                for line in file_tmp.readlines():
703
                    potfile.write(line)
704
                file_tmp.close()
705
        potfile.close()
706

    
707
    def write_sort_file(self):
708
        """Writes a sortings file.
709

710
        This file contains information about how the atoms are sorted in
711
        the first column and how they should be resorted in the second
712
        column. It is used for restart purposes to get sorting right
713
        when reading in an old calculation to ASE."""
714

    
715
        file = open('ase-sort.dat', 'w')
716
        for n in range(len(self.sort)):
717
            file.write('%5i %5i \n' % (self.sort[n], self.resort[n]))
718

    
719
    # Methods for reading information from OUTCAR files:
720
    def read_energy(self, all=None):
721
        [energy_free, energy_zero]=[0, 0]
722
        if all:
723
            energy_free = []
724
            energy_zero = []
725
        for line in open('OUTCAR', 'r'):
726
            # Free energy
727
            if line.startswith('  free  energy   toten'):
728
                if all:
729
                    energy_free.append(float(line.split()[-2]))
730
                else:
731
                    energy_free = float(line.split()[-2])
732
            # Extrapolated zero point energy
733
            if line.startswith('  energy  without entropy'):
734
                if all:
735
                    energy_zero.append(float(line.split()[-1]))
736
                else:
737
                    energy_zero = float(line.split()[-1])
738
        return [energy_free, energy_zero]
739

    
740
    def read_forces(self, atoms, all=False):
741
        """Method that reads forces from OUTCAR file.
742

743
        If 'all' is switched on, the forces for all ionic steps
744
        in the OUTCAR file be returned, in other case only the
745
        forces for the last ionic configuration is returned."""
746

    
747
        file = open('OUTCAR','r')
748
        lines = file.readlines()
749
        file.close()
750
        n=0
751
        if all:
752
            all_forces = []
753
        for line in lines:
754
            if line.rfind('TOTAL-FORCE') > -1:
755
                forces=[]
756
                for i in range(len(atoms)):
757
                    forces.append(np.array([float(f) for f in lines[n+2+i].split()[3:6]]))
758
                if all:
759
                    all_forces.append(np.array(forces)[self.resort])
760
            n+=1
761
        if all:
762
            return np.array(all_forces)
763
        else:
764
            return np.array(forces)[self.resort]
765

    
766
    def read_fermi(self):
767
        """Method that reads Fermi energy from OUTCAR file"""
768
        E_f=None
769
        for line in open('OUTCAR', 'r'):
770
            if line.rfind('E-fermi') > -1:
771
                E_f=float(line.split()[2])
772
        return E_f
773

    
774
    def read_dipole(self):
775
        dipolemoment=np.zeros([1,3])
776
        for line in open('OUTCAR', 'r'):
777
            if line.rfind('dipolmoment') > -1:
778
                dipolemoment=np.array([float(f) for f in line.split()[1:4]])
779
        return dipolemoment
780

    
781
    def read_magnetic_moments(self, atoms):
782
        magnetic_moments = np.zeros(len(atoms))
783
        n = 0
784
        lines = open('OUTCAR', 'r').readlines()
785
        for line in lines:
786
            if line.rfind('magnetization (x)') > -1:
787
                for m in range(len(atoms)):
788
                    magnetic_moments[m] = float(lines[n + m + 4].split()[4])
789
            n += 1
790
        return np.array(magnetic_moments)[self.resort]
791

    
792
    def read_magnetic_moment(self):
793
        n=0
794
        for line in open('OUTCAR','r'):
795
            if line.rfind('number of electron  ') > -1:
796
                magnetic_moment=float(line.split()[-1])
797
            n+=1
798
        return magnetic_moment
799

    
800
    def read_nbands(self):
801
        for line in open('OUTCAR', 'r'):
802
            if line.rfind('NBANDS') > -1:
803
                return int(line.split()[-1])
804

    
805
    def read_convergence(self):
806
        """Method that checks whether a calculation has converged."""
807
        converged = None
808
        # First check electronic convergence
809
        for line in open('OUTCAR', 'r'):
810
            if line.rfind('EDIFF  ') > -1:
811
                ediff = float(line.split()[2])
812
            if line.rfind('total energy-change')>-1:
813
                split = line.split(':')
814
                a = float(split[1].split('(')[0])
815
                b = float(split[1].split('(')[1][0:-2])
816
                if [abs(a), abs(b)] < [ediff, ediff]:
817
                    converged = True
818
                else:
819
                    converged = False
820
                    continue
821
        # Then if ibrion > 0, check whether ionic relaxation condition been fulfilled
822
        if self.int_params['ibrion'] > 0:
823
            if not self.read_relaxed():
824
                converged = False
825
            else:
826
                converged = True
827
        return converged
828

    
829
    def read_ibz_kpoints(self):
830
        lines = open('OUTCAR', 'r').readlines()
831
        ibz_kpts = []
832
        n = 0
833
        i = 0
834
        for line in lines:
835
            if line.rfind('Following cartesian coordinates')>-1:
836
                m = n+2
837
                while i==0:
838
                    ibz_kpts.append([float(lines[m].split()[p]) for p in range(3)])
839
                    m += 1
840
                    if lines[m]==' \n':
841
                        i = 1
842
            if i == 1:
843
                continue
844
            n += 1
845
        ibz_kpts = np.array(ibz_kpts)
846
        return np.array(ibz_kpts)
847

    
848
    def read_k_point_weights(self):
849
        file = open('IBZKPT')
850
        lines = file.readlines()
851
        file.close()
852
        kpt_weights = []
853
        for n in range(3, len(lines)):
854
            kpt_weights.append(float(lines[n].split()[3]))
855
        kpt_weights = np.array(kpt_weights)
856
        kpt_weights /= np.sum(kpt_weights)
857
        return kpt_weights
858

    
859
    def read_eigenvalues(self, kpt=0, spin=0):
860
        file = open('EIGENVAL', 'r')
861
        lines = file.readlines()
862
        file.close()
863
        eigs = []
864
        for n in range(8+kpt*(self.nbands+2), 8+kpt*(self.nbands+2)+self.nbands):
865
            eigs.append(float(lines[n].split()[spin+1]))
866
        return np.array(eigs)
867

    
868
    def read_relaxed(self):
869
        for line in open('OUTCAR', 'r'):
870
            if line.rfind('reached required accuracy') > -1:
871
                return True
872
        return False
873

    
874
# The below functions are used to restart a calculation and are under early constructions
875

    
876
    def read_incar(self, filename='INCAR'):
877
        """Method that imports settings from INCAR file."""
878

    
879
        self.spinpol = False
880
        file=open(filename, 'r')
881
        file.readline()
882
        lines=file.readlines()
883
        for line in lines:
884
            try:
885
                data = line.split()
886
                # Skip empty and commented lines. 
887
                if len(data) == 0:
888
                    continue
889
                elif data[0][0] in ['#', '!']:
890
                    continue
891
                key = data[0].lower()
892
                if key in float_keys:
893
                    self.float_params[key] = float(data[2])
894
                elif key in exp_keys:
895
                    self.exp_params[key] = float(data[2])
896
                elif key in string_keys:
897
                    self.string_params[key] = str(data[2])
898
                elif key in int_keys:
899
                    if key == 'ispin':
900
                        if int(data[2]) == 2:
901
                            self.spinpol = True
902
                    else:
903
                        self.int_params[key] = int(data[2])
904
                elif key in bool_keys:
905
                    if 'true' in data[2].lower():
906
                        self.bool_params[key] = True
907
                    elif 'false' in data[2].lower():
908
                        self.bool_params[key] = False
909
                elif key in list_keys:
910
                    list = []
911
                    if key in ('dipol', 'eint', 'ferwe', 'ropt', 'rwigs'):
912
                        for a in data[2:]:
913
                            list.append(float(a))
914
                    elif key in ('iband', 'kpuse'):
915
                        for a in data[2:]:
916
                            list.append(int(a))
917
                    self.list_params[key] = list
918
                    if key == 'magmom':
919
                        done = False
920
                        for a in data[2:]:
921
                            if '!' in a or done:
922
                                done = True
923
                            elif '*' in a: 
924
                                a = a.split('*')
925
                                for b in range(int(a[0])):
926
                                    list.append(float(a[1]))
927
                            else:
928
                                list.append(float(a))
929
                        list = np.array(list)
930
                        self.atoms.set_initial_magnetic_moments(list[self.resort])
931
            except KeyError:
932
                raise IOError('Keyword "%s" in INCAR is not known by calculator.' % key)
933
            except IndexError:
934
                raise IOError('Value missing for keyword "%s".' % key)
935

    
936
    def read_outcar(self):
937
        # Spin polarized calculation?
938
        file = open('OUTCAR', 'r')
939
        lines = file.readlines()
940
        file.close()
941
        for line in lines:
942
            if line.rfind('ISPIN') > -1:
943
                if int(line.split()[2])==2:
944
                    self.spinpol = True
945
                else:
946
                    self.spinpol = None
947
        self.energy_free, self.energy_zero = self.read_energy()
948
        self.forces = self.read_forces(self.atoms)
949
        self.dipole = self.read_dipole()
950
        self.fermi = self.read_fermi()
951
        self.stress = self.read_stress()
952
        self.nbands = self.read_nbands()
953
        p=self.int_params
954
        q=self.list_params
955
        if self.spinpol:
956
            self.magnetic_moment = self.read_magnetic_moment()
957
            if p['lorbit']>=10 or (p['lorbit']!=None and q['rwigs']):
958
                self.magnetic_moments = self.read_magnetic_moments(self.atoms)
959
        self.set(nbands=self.nbands)
960

    
961
    def read_kpoints(self, filename='KPOINTS'):
962
        file = open(filename, 'r')
963
        lines = file.readlines()
964
        file.close()
965
        type = lines[2].split()[0].lower()[0]
966
        if type in ['g', 'm']:
967
            if type=='g':
968
                self.set(gamma=True)
969
            kpts = np.array([int(lines[3].split()[i]) for i in range(3)])
970
            self.set(kpts=kpts)
971
        elif type in ['c', 'k']:
972
            raise NotImplementedError('Only Monkhorst-Pack and gamma centered grid supported for restart.')
973
        else:
974
            raise NotImplementedError('Only Monkhorst-Pack and gamma centered grid supported for restart.')
975
    
976
    def read_potcar(self):
977
        """ Method that reads the Exchange Correlation functional from POTCAR file.
978
        """
979
        file = open('POTCAR', 'r')
980
        lines = file.readlines()
981
        file.close()
982

    
983
        # Search for key 'LEXCH' in POTCAR
984
        xc_flag = None
985
        for line in lines:
986
            key = line.split()[0].upper()
987
            if key == 'LEXCH':
988
                xc_flag = line.split()[-1].upper()
989
                break
990

    
991
        if xc_flag is None:
992
            raise ValueError('LEXCH flag not found in POTCAR file.')
993

    
994
        # Values of parameter LEXCH and corresponding XC-functional
995
        xc_dict = {'PE':'PBE', '91':'PW91', 'CA':'LDA'}
996

    
997
        if xc_flag not in xc_dict.keys():
998
            raise ValueError(
999
                'Unknown xc-functional flag found in POTCAR, LEXCH=%s' % xc_flag)
1000

    
1001
        self.input_params['xc'] = xc_dict[xc_flag]
1002

    
1003

    
1004
class VaspChargeDensity(object):
1005
    """Class for representing VASP charge density"""
1006

    
1007
    def __init__(self, filename='CHG'):
1008
        # Instance variables
1009
        self.atoms = []   # List of Atoms objects
1010
        self.chg = []     # Charge density
1011
        self.chgdiff = [] # Charge density difference, if spin polarized
1012
        self.aug = ''     # Augmentation charges, not parsed just a big string
1013
        self.augdiff = '' # Augmentation charge differece, is spin polarized
1014
        
1015
        # Note that the augmentation charge is not a list, since they
1016
        # are needed only for CHGCAR files which store only a single
1017
        # image.
1018
        if filename != None:
1019
            self.read(filename)
1020

    
1021
    def is_spin_polarized(self):
1022
        if len(self.chgdiff) > 0:
1023
            return True
1024
        return False
1025

    
1026
    def _read_chg(self, fobj, chg, volume):
1027
        """Read charge from file object
1028

1029
        Utility method for reading the actual charge density (or
1030
        charge density difference) from a file object. On input, the
1031
        file object must be at the beginning of the charge block, on
1032
        output the file position will be left at the end of the
1033
        block. The chg array must be of the correct dimensions.
1034

1035
        """
1036
        # VASP writes charge density as
1037
        # WRITE(IU,FORM) (((C(NX,NY,NZ),NX=1,NGXC),NY=1,NGYZ),NZ=1,NGZC)
1038
        # Fortran nested implied do loops; innermost index fastest
1039
        # First, just read it in
1040
        for zz in range(chg.shape[2]):
1041
            for yy in range(chg.shape[1]):
1042
                chg[:, yy, zz] = np.fromfile(fobj, count = chg.shape[0],
1043
                                             sep=' ')
1044
        chg /= volume
1045

    
1046
    def read(self, filename='CHG'):
1047
        """Read CHG or CHGCAR file.
1048

1049
        If CHG contains charge density from multiple steps all the
1050
        steps are read and stored in the object. By default VASP
1051
        writes out the charge density every 10 steps.
1052

1053
        chgdiff is the difference between the spin up charge density
1054
        and the spin down charge density and is thus only read for a
1055
        spin-polarized calculation.
1056

1057
        aug is the PAW augmentation charges found in CHGCAR. These are
1058
        not parsed, they are just stored as a string so that they can
1059
        be written again to a CHGCAR format file.
1060

1061
        """
1062
        import ase.io.vasp as aiv
1063
        f = open(filename)
1064
        self.atoms = []
1065
        self.chg = []
1066
        self.chgdiff = []
1067
        self.aug = ''
1068
        self.augdiff = ''
1069
        while True:
1070
            try:
1071
                atoms = aiv.read_vasp(f)
1072
            except ValueError, e:
1073
                # Probably an empty line, or we tried to read the 
1074
                # augmentation occupancies in CHGCAR
1075
                break 
1076
            f.readline()
1077
            ngr = f.readline().split()
1078
            ng = (int(ngr[0]), int(ngr[1]), int(ngr[2]))
1079
            chg = np.empty(ng)
1080
            self._read_chg(f, chg, atoms.get_volume())
1081
            self.chg.append(chg)
1082
            self.atoms.append(atoms)
1083
            # Check if the file has a spin-polarized charge density part, and
1084
            # if so, read it in.
1085
            fl = f.tell()
1086
            # First check if the file has an augmentation charge part (CHGCAR file.)
1087
            line1 = f.readline()
1088
            if line1=='':
1089
                break
1090
            elif line1.find('augmentation') != -1:
1091
                augs = [line1]
1092
                while True:
1093
                    line2 = f.readline()
1094
                    if line2.split() == ngr:
1095
                        self.aug = ''.join(augs)
1096
                        augs = []
1097
                        chgdiff = np.empty(ng)
1098
                        self._read_chg(f, chgdiff, atoms.get_volume())
1099
                        self.chgdiff.append(chgdiff)
1100
                    elif line2 == '':
1101
                        break
1102
                    else:
1103
                        augs.append(line2)
1104
                if len(self.aug) == 0:
1105
                    self.aug = ''.join(augs)
1106
                    augs = []
1107
                else:
1108
                    self.augdiff = ''.join(augs)
1109
                    augs = []
1110
            elif line1.split() == ngr:
1111
                chgdiff = np.empty(ng)
1112
                self._read_chg(f, chgdiff, atoms.get_volume())
1113
                self.chgdiff.append(chgdiff)
1114
            else:
1115
                f.seek(fl)
1116
        f.close()
1117

    
1118
    def _write_chg(self, fobj, chg, volume, format='chg'):
1119
        """Write charge density
1120

1121
        Utility function similar to _read_chg but for writing.
1122

1123
        """
1124
        # Make a 1D copy of chg, must take transpose to get ordering right
1125
        chgtmp=chg.T.ravel()
1126
        # Multiply by volume
1127
        chgtmp=chgtmp*volume
1128
        # Must be a tuple to pass to string conversion
1129
        chgtmp=tuple(chgtmp)
1130
        # CHG format - 10 columns
1131
        if format.lower() == 'chg':
1132
            # Write all but the last row
1133
            for ii in range((len(chgtmp)-1)/10):
1134
                fobj.write(' %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G\
1135
 %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G\n' % chgtmp[ii*10:(ii+1)*10]
1136
                           )
1137
            # If the last row contains 10 values then write them without a newline
1138
            if len(chgtmp)%10==0:
1139
                fobj.write(' %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G\
1140
 %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G' % chgtmp[len(chgtmp)-10:len(chgtmp)])
1141
            # Otherwise write fewer columns without a newline
1142
            else:
1143
                for ii in range(len(chgtmp)%10):
1144
                    fobj.write((' %#11.5G') % chgtmp[len(chgtmp)-len(chgtmp)%10+ii])
1145
        # Other formats - 5 columns
1146
        else:
1147
            # Write all but the last row
1148
            for ii in range((len(chgtmp)-1)/5):
1149
                fobj.write(' %17.10E %17.10E %17.10E %17.10E %17.10E\n' % chgtmp[ii*5:(ii+1)*5])
1150
            # If the last row contains 5 values then write them without a newline
1151
            if len(chgtmp)%5==0:
1152
                fobj.write(' %17.10E %17.10E %17.10E %17.10E %17.10E' % chgtmp[len(chgtmp)-5:len(chgtmp)])
1153
            # Otherwise write fewer columns without a newline
1154
            else:
1155
                for ii in range(len(chgtmp)%5):
1156
                    fobj.write((' %17.10E') % chgtmp[len(chgtmp)-len(chgtmp)%5+ii])
1157
        # Write a newline whatever format it is
1158
        fobj.write('\n')
1159
        # Clean up
1160
        del chgtmp
1161

    
1162
    def write(self, filename='CHG', format=None):
1163
        """Write VASP charge density in CHG format.
1164

1165
        filename: str
1166
            Name of file to write to.
1167
        format: str
1168
            String specifying whether to write in CHGCAR or CHG
1169
            format.
1170

1171
        """
1172
        import ase.io.vasp as aiv
1173
        if format == None:
1174
            if filename.lower().find('chgcar') != -1:
1175
                format = 'chgcar'
1176
            elif filename.lower().find('chg') != -1:
1177
                format = 'chg'
1178
            elif len(self.chg) == 1:
1179
                format = 'chgcar'
1180
            else:
1181
                format = 'chg'
1182
        f = open(filename, 'w')
1183
        for ii, chg in enumerate(self.chg):
1184
            if format == 'chgcar' and ii != len(self.chg) - 1:
1185
                continue # Write only the last image for CHGCAR
1186
            aiv.write_vasp(f, self.atoms[ii], direct=True)
1187
            f.write('\n')
1188
            for dim in chg.shape:
1189
                f.write(' %4i' % dim)
1190
            f.write('\n')
1191
            vol = self.atoms[ii].get_volume()
1192
            self._write_chg(f, chg, vol, format)
1193
            if format == 'chgcar':
1194
                f.write(self.aug)
1195
            if self.is_spin_polarized():
1196
                if format == 'chg':
1197
                    f.write('\n')
1198
                for dim in chg.shape:
1199
                    f.write(' %4i' % dim)
1200
                self._write_chg(f, self.chgdiff[ii], vol, format)
1201
                if format == 'chgcar':
1202
                    f.write('\n')
1203
                    f.write(self.augdiff)
1204
            if format == 'chg' and len(self.chg) > 1:
1205
                f.write('\n')
1206
        f.close()
1207

    
1208

    
1209
class VaspDos(object):
1210
    """Class for representing density-of-states produced by VASP
1211

1212
    The energies are in property self.energy
1213

1214
    Site-projected DOS is accesible via the self.site_dos method.
1215

1216
    Total and integrated DOS is accessible as numpy.ndarray's in the
1217
    properties self.dos and self.integrated_dos. If the calculation is
1218
    spin polarized, the arrays will be of shape (2, NDOS), else (1,
1219
    NDOS).
1220

1221
    The self.efermi property contains the currently set Fermi
1222
    level. Changing this value shifts the energies.
1223
    
1224
    """
1225

    
1226
    def __init__(self, doscar='DOSCAR', efermi=0.0):
1227
        """Initialize"""
1228
        self._efermi = 0.0
1229
        self.read_doscar(doscar)
1230
        self.efermi = efermi
1231

    
1232
    def _set_efermi(self, efermi):
1233
        """Set the Fermi level."""
1234
        ef = efermi - self._efermi
1235
        self._efermi = efermi
1236
        self._total_dos[0, :] = self._total_dos[0, :] - ef
1237
        try:
1238
            self._site_dos[:, 0, :] = self._site_dos[:, 0, :] - ef
1239
        except IndexError:
1240
            pass
1241

    
1242
    def _get_efermi(self):
1243
        return self._efermi
1244

    
1245
    efermi = property(_get_efermi, _set_efermi, None, "Fermi energy.")
1246

    
1247
    def _get_energy(self):
1248
        """Return the array with the energies."""
1249
        return self._total_dos[0, :]
1250
    energy = property(_get_energy, None, None, "Array of energies")
1251

    
1252
    def site_dos(self, atom, orbital):
1253
        """Return an NDOSx1 array with dos for the chosen atom and orbital.
1254

1255
        atom: int
1256
            Atom index
1257
        orbital: int or str
1258
            Which orbital to plot
1259

1260
        If the orbital is given as an integer:
1261
        If spin-unpolarized calculation, no phase factors:
1262
        s = 0, p = 1, d = 2
1263
        Spin-polarized, no phase factors:
1264
        s-up = 0, s-down = 1, p-up = 2, p-down = 3, d-up = 4, d-down = 5
1265
        If phase factors have been calculated, orbitals are
1266
        s, py, pz, px, dxy, dyz, dz2, dxz, dx2
1267
        double in the above fashion if spin polarized.
1268

1269
        """
1270
        # Integer indexing for orbitals starts from 1 in the _site_dos array
1271
        # since the 0th column contains the energies
1272
        if isinstance(orbital, int):
1273
            return self._site_dos[atom, orbital + 1, :]
1274
        n = self._site_dos.shape[1]
1275
        if n == 4:
1276
            norb = {'s':1, 'p':2, 'd':3}
1277
        elif n == 7:
1278
            norb = {'s+':1, 's-up':1, 's-':2, 's-down':2,
1279
                    'p+':3, 'p-up':3, 'p-':4, 'p-down':4,
1280
                    'd+':5, 'd-up':5, 'd-':6, 'd-down':6}
1281
        elif n == 10:
1282
            norb = {'s':1, 'py':2, 'pz':3, 'px':4,
1283
                    'dxy':5, 'dyz':6, 'dz2':7, 'dxz':8,
1284
                    'dx2':9}
1285
        elif n == 19:
1286
            norb = {'s+':1, 's-up':1, 's-':2, 's-down':2,
1287
                    'py+':3, 'py-up':3, 'py-':4, 'py-down':4,
1288
                    'pz+':5, 'pz-up':5, 'pz-':6, 'pz-down':6,
1289
                    'px+':7, 'px-up':7, 'px-':8, 'px-down':8,
1290
                    'dxy+':9, 'dxy-up':9, 'dxy-':10, 'dxy-down':10,
1291
                    'dyz+':11, 'dyz-up':11, 'dyz-':12, 'dyz-down':12,
1292
                    'dz2+':13, 'dz2-up':13, 'dz2-':14, 'dz2-down':14,
1293
                    'dxz+':15, 'dxz-up':15, 'dxz-':16, 'dxz-down':16,
1294
                    'dx2+':17, 'dx2-up':17, 'dx2-':18, 'dx2-down':18}
1295
        return self._site_dos[atom, norb[orbital.lower()], :]
1296

    
1297
    def _get_dos(self):
1298
        if self._total_dos.shape[0] == 3:
1299
            return self._total_dos[1, :]
1300
        elif self._total_dos.shape[0] == 5:
1301
            return self._total_dos[1:3, :]
1302
    dos = property(_get_dos, None, None, 'Average DOS in cell')
1303

    
1304
    def _get_integrated_dos(self):
1305
        if self._total_dos.shape[0] == 3:
1306
            return self._total_dos[2, :]
1307
        elif self._total_dos.shape[0] == 5:
1308
            return self._total_dos[3:5, :]
1309
    integrated_dos = property(_get_integrated_dos, None, None,
1310
                              'Integrated average DOS in cell')
1311

    
1312
    def read_doscar(self, fname="DOSCAR"):
1313
        """Read a VASP DOSCAR file"""
1314
        f = open(fname)
1315
        natoms = int(f.readline().split()[0])
1316
        [f.readline() for nn in range(4)]  # Skip next 4 lines.
1317
        # First we have a block with total and total integrated DOS
1318
        ndos = int(f.readline().split()[2])
1319
        dos = []
1320
        for nd in xrange(ndos):
1321
            dos.append(np.array([float(x) for x in f.readline().split()]))
1322
        self._total_dos = np.array(dos).T
1323
        # Next we have one block per atom, if INCAR contains the stuff
1324
        # necessary for generating site-projected DOS
1325
        dos = []
1326
        for na in xrange(natoms):
1327
            line = f.readline()
1328
            if line == '':
1329
                # No site-projected DOS
1330
                break
1331
            ndos = int(line.split()[2])
1332
            line = f.readline().split()
1333
            cdos = np.empty((ndos, len(line)))
1334
            cdos[0] = np.array(line)
1335
            for nd in xrange(1, ndos):
1336
                line = f.readline().split()
1337
                cdos[nd] = np.array([float(x) for x in line])
1338
            dos.append(cdos.T)
1339
        self._site_dos = np.array(dos)
1340

    
1341

    
1342
import pickle
1343

    
1344
class xdat2traj:
1345
    def __init__(self, trajectory=None, atoms=None, poscar=None, 
1346
                 xdatcar=None, sort=None, calc=None):
1347
        if not poscar:
1348
            self.poscar = 'POSCAR'
1349
        else:
1350
            self.poscar = poscar
1351
        if not atoms:
1352
            self.atoms = ase.io.read(self.poscar, format='vasp')
1353
        else:
1354
            self.atoms = atoms
1355
        if not xdatcar:
1356
            self.xdatcar = 'XDATCAR'
1357
        else:
1358
            self.xdatcar = xdatcar
1359
        if not trajectory:
1360
            self.trajectory = 'out.traj'
1361
        else:
1362
            self.trajectory = trajectory
1363
        if not calc:
1364
            self.calc = Vasp()
1365
        else:
1366
            self.calc = calc
1367
        if not sort: 
1368
            if not hasattr(self.calc, 'sort'):
1369
                self.calc.sort = range(len(self.atoms))
1370
        else:
1371
            self.calc.sort = sort
1372
        self.calc.resort = range(len(self.calc.sort))
1373
        for n in range(len(self.calc.resort)):
1374
            self.calc.resort[self.calc.sort[n]] = n
1375
        self.out = ase.io.trajectory.PickleTrajectory(self.trajectory, mode='w')
1376
        self.energies = self.calc.read_energy(all=True)[1]
1377
        self.forces = self.calc.read_forces(self.atoms, all=True)
1378

    
1379
    def convert(self):
1380
        lines = open(self.xdatcar).readlines()
1381
        if len(lines[5].split())==0:
1382
            del(lines[0:6])
1383
        elif len(lines[4].split())==0:
1384
            del(lines[0:5])
1385
        step = 0
1386
        iatom = 0
1387
        scaled_pos = []
1388
        for line in lines:
1389
            if iatom == len(self.atoms):
1390
                if step == 0:
1391
                    self.out.write_header(self.atoms[self.calc.resort])
1392
                scaled_pos = np.array(scaled_pos)
1393
                self.atoms.set_scaled_positions(scaled_pos)
1394
                d = {'positions': self.atoms.get_positions()[self.calc.resort],
1395
                     'cell': self.atoms.get_cell(),
1396
                     'momenta': None,
1397
                     'energy': self.energies[step],
1398
                     'forces': self.forces[step],
1399
                     'stress': None}
1400
                pickle.dump(d, self.out.fd, protocol=-1)
1401
                scaled_pos = []
1402
                iatom = 0
1403
                step += 1
1404
            else:
1405
                
1406
                iatom += 1
1407
                scaled_pos.append([float(line.split()[n]) for n in range(3)])
1408

    
1409
        # Write also the last image
1410
        # I'm sure there is also more clever fix...
1411
        scaled_pos = np.array(scaled_pos)
1412
        self.atoms.set_scaled_positions(scaled_pos)
1413
        d = {'positions': self.atoms.get_positions()[self.calc.resort],
1414
             'cell': self.atoms.get_cell(),
1415
             'momenta': None,
1416
             'energy': self.energies[step],
1417
             'forces': self.forces[step],
1418
             'stress': None}
1419
        pickle.dump(d, self.out.fd, protocol=-1)
1420

    
1421
        self.out.fd.close()