Statistiques
| Révision :

root / ase / calculators / vasp.py @ 1

Historique | Voir | Annoter | Télécharger (51,89 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, 
153
                 **kwargs):
154
        self.name = 'Vasp'
155
        self.float_params = {}
156
        self.exp_params = {}
157
        self.string_params = {}
158
        self.int_params = {}
159
        self.bool_params = {}
160
        self.list_params = {}
161
        self.special_params = {}
162
        for key in float_keys:
163
            self.float_params[key] = None
164
        for key in exp_keys:
165
            self.exp_params[key] = None
166
        for key in string_keys:
167
            self.string_params[key] = None
168
        for key in int_keys:
169
            self.int_params[key] = None
170
        for key in bool_keys:
171
            self.bool_params[key] = None
172
        for key in list_keys:
173
            self.list_params[key] = None
174
        for key in special_keys:
175
            self.special_params[key] = None
176
        self.string_params['prec'] = 'Normal'
177

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

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

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

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

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

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

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

    
240
        p = self.input_params
241

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

    
247
        # Determine the number of atoms of each atomic species
248
        # sorted after atomic species
249
        special_setups = []
250
        symbols = {}
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
        for m,atom in enumerate(atoms):
262
            symbol = atom.get_symbol()
263
            if m in special_setups:
264
                pass
265
            else:
266
                if not symbols.has_key(symbol):
267
                    symbols[symbol] = 1
268
                else:
269
                    symbols[symbol] += 1
270
        
271
        # Build the sorting list
272
        self.sort = []
273
        self.sort.extend(special_setups)
274

    
275
        for symbol in symbols:
276
            for m,atom in enumerate(atoms):
277
                if m in special_setups: 
278
                    pass
279
                else:
280
                    if atom.get_symbol() == symbol:
281
                        self.sort.append(m)
282
        self.resort = range(len(self.sort))
283
        for n in range(len(self.resort)):
284
            self.resort[self.sort[n]] = n
285
        self.atoms_sorted = atoms[self.sort]
286

    
287
        # Check if the necessary POTCAR files exists and
288
        # create a list of their paths.
289
        self.symbol_count = []
290
        for m in special_setups:
291
            self.symbol_count.append([atomtypes[m],1])
292
        for m in symbols:
293
            self.symbol_count.append([m,symbols[m]])
294
        #print 'self.symbol_count',self.symbol_count 
295
        sys.stdout.flush()
296
        xc = '/'
297
        #print 'p[xc]',p['xc']
298
        if p['xc'] == 'PW91':
299
            xc = '_gga/'
300
        elif p['xc'] == 'PBE':
301
            xc = '_pbe/'
302
        if 'VASP_PP_PATH' in os.environ:
303
            pppaths = os.environ['VASP_PP_PATH'].split(':')
304
        else:
305
            pppaths = []
306
        self.ppp_list = []
307
        # Setting the pseudopotentials, first special setups and 
308
        # then according to symbols
309
        for m in special_setups:
310
            name = 'potpaw'+xc.upper() + p['setups'][str(m)] + '/POTCAR'
311
            found = False
312
            for path in pppaths:
313
                filename = join(path, name)
314
                #print 'filename', filename
315
                if isfile(filename) or islink(filename):
316
                    found = True
317
                    self.ppp_list.append(filename)
318
                    break
319
                elif isfile(filename + '.Z') or islink(filename + '.Z'):
320
                    found = True
321
                    self.ppp_list.append(filename+'.Z')
322
                    break
323
            if not found:
324
                raise RuntimeError('No pseudopotential for %s!' % symbol)    
325
        #print 'symbols', symbols 
326
        for symbol in symbols:
327
            try:
328
                name = 'potpaw'+xc.upper()+symbol + p['setups'][symbol]
329
            except (TypeError, KeyError):
330
                name = 'potpaw' + xc.upper() + symbol
331
            name += '/POTCAR'
332
            found = False
333
            for path in pppaths:
334
                filename = join(path, name)
335
                #print 'filename', filename
336
                if isfile(filename) or islink(filename):
337
                    found = True
338
                    self.ppp_list.append(filename)
339
                    break
340
                elif isfile(filename + '.Z') or islink(filename + '.Z'):
341
                    found = True
342
                    self.ppp_list.append(filename+'.Z')
343
                    break
344
            if not found:
345
                raise RuntimeError('No pseudopotential for %s!' % symbol)
346
        self.converged = None
347
        self.setups_changed = None
348

    
349
    def calculate(self, atoms):
350
        """Generate necessary files in the working directory and run VASP.
351

352
        The method first write VASP input files, then calls the method
353
        which executes VASP. When the VASP run is finished energy, forces, 
354
        etc. are read from the VASP output.
355
        """
356
        
357
        # Write input
358
        from ase.io.vasp import write_vasp
359
        self.initialize(atoms)
360
        write_vasp('POSCAR', self.atoms_sorted, symbol_count = self.symbol_count)
361
        self.write_incar(atoms)
362
        self.write_potcar()
363
        self.write_kpoints()
364
        self.write_sort_file()
365
        
366
        # Execute VASP
367
        self.run()
368
        # Read output
369
        atoms_sorted = ase.io.read('CONTCAR', format='vasp')
370
        if self.int_params['ibrion']>-1 and self.int_params['nsw']>0:
371
            # Replace entire atoms object with the one read from CONTCAR.
372
            atoms.__dict__ = atoms_sorted[self.resort].__dict__
373
        self.converged = self.read_convergence()
374
        self.set_results(atoms)
375

    
376
    def set_results(self, atoms):
377
        self.read(atoms)
378
        if self.spinpol:
379
            self.magnetic_moment = self.read_magnetic_moment()
380
            if self.int_params['lorbit']>=10 or (self.int_params['lorbit']!=None and self.list_params['rwigs']):
381
                self.magnetic_moments = self.read_magnetic_moments(atoms)
382
        self.old_float_params = self.float_params.copy()
383
        self.old_exp_params = self.exp_params.copy()
384
        self.old_string_params = self.string_params.copy()
385
        self.old_int_params = self.int_params.copy()
386
        self.old_input_params = self.input_params.copy()
387
        self.old_bool_params = self.bool_params.copy()
388
        self.old_list_params = self.list_params.copy()
389
        self.atoms = atoms.copy()
390
        
391
    def run(self):
392
        """Method which explicitely runs VASP."""
393

    
394
        if self.track_output:
395
            self.out = self.output_template+str(self.run_counts)+'.out'
396
            self.run_counts += 1
397
        else:
398
            self.out = self.output_template+'.out'
399
        stderr = sys.stderr
400
        p=self.input_params
401
        if p['txt'] is None:
402
            sys.stderr = devnull
403
        elif p['txt'] == '-':
404
            pass
405
        elif isinstance(p['txt'], str):
406
            sys.stderr = open(p['txt'], 'w')
407
        if os.environ.has_key('VASP_COMMAND'):
408
            vasp = os.environ['VASP_COMMAND']
409
            exitcode = os.system('%s > %s' % (vasp, self.out))
410
        elif os.environ.has_key('VASP_SCRIPT'):
411
            vasp = os.environ['VASP_SCRIPT']
412
            locals={}
413
            execfile(vasp, {}, locals)
414
            exitcode = locals['exitcode']
415
        else:
416
            raise RuntimeError('Please set either VASP_COMMAND or VASP_SCRIPT environment variable')
417
        sys.stderr = stderr
418
        if exitcode != 0:
419
            raise RuntimeError('Vasp exited with exit code: %d.  ' % exitcode)
420
        
421
    def restart_load(self):
422
        """Method which is called upon restart."""
423
        
424
        # Try to read sorting file
425
        if os.path.isfile('ase-sort.dat'):
426
            self.sort = []
427
            self.resort = []
428
            file = open('ase-sort.dat', 'r')
429
            lines = file.readlines()
430
            file.close()
431
            for line in lines:
432
                data = line.split()
433
                self.sort.append(int(data[0]))
434
                self.resort.append(int(data[1]))
435
            atoms = ase.io.read('CONTCAR', format='vasp')[self.resort]
436
        else:
437
            atoms = ase.io.read('CONTCAR', format='vasp')
438
            self.sort = range(len(atoms))
439
            self.resort = range(len(atoms))
440
        self.atoms = atoms.copy()
441
        self.read_incar()
442
        self.read_outcar()
443
        self.set_results(atoms)
444
        self.read_kpoints()
445
        self.read_potcar()
446
#        self.old_incar_params = self.incar_params.copy()
447
        self.old_input_params = self.input_params.copy()
448
        self.converged = self.read_convergence()
449

    
450
    def clean(self):
451
        """Method which cleans up after a calculation.
452
        
453
        The default files generated by Vasp will be deleted IF this
454
        method is called.
455

456
        """
457
        files = ['CHG', 'CHGCAR', 'POSCAR', 'INCAR', 'CONTCAR', 'DOSCAR',
458
                 'EIGENVAL', 'IBZKPT', 'KPOINTS', 'OSZICAR', 'OUTCAR', 'PCDAT',
459
                 'POTCAR', 'vasprun.xml', 'WAVECAR', 'XDATCAR',
460
                 'PROCAR', 'ase-sort.dat']
461
        for f in files:
462
            try:
463
                os.remove(f)
464
            except OSError:
465
                pass
466

    
467
    def set_atoms(self, atoms):
468
        if (atoms != self.atoms):
469
            self.converged = None
470
        self.atoms = atoms.copy()
471

    
472
    def get_atoms(self):
473
        atoms = self.atoms.copy()
474
        atoms.set_calculator(self)
475
        return atoms
476

    
477
    def get_potential_energy(self, atoms, force_consistent=False):
478
        self.update(atoms)
479
        if force_consistent:
480
            return self.energy_free
481
        else:
482
            return self.energy_zero
483

    
484
    def get_forces(self, atoms):
485
        self.update(atoms)
486
        return self.forces
487

    
488
    def get_stress(self, atoms):
489
        self.update(atoms)
490
        return self.stress
491

    
492
    def read_stress(self):
493
        for line in open('OUTCAR'):
494
            if line.find(' in kB  ') != -1:
495
                stress = -np.array([float(a) for a in line.split()[2:]]) \
496
                         [[0, 1, 2, 4, 5, 3]] \
497
                         * 1e-1 * ase.units.GPa
498
        return stress
499

    
500
    def calculation_required(self, atoms, quantities):
501
        if (self.positions is None or 
502
            (self.atoms != atoms) or
503
            (self.float_params != self.old_float_params) or
504
            (self.exp_params != self.old_exp_params) or
505
            (self.string_params != self.old_string_params) or
506
            (self.int_params != self.old_int_params) or
507
            (self.bool_params != self.old_bool_params) or
508
            (self.list_params != self.old_list_params) or
509
            (self.input_params != self.old_input_params)
510
            or not self.converged):
511
            return True
512
        if 'magmom' in quantities:
513
            return not hasattr(self, 'magnetic_moment')
514
        return False
515

    
516
    def get_number_of_bands(self):
517
        return self.nbands
518

    
519
    def get_k_point_weights(self):
520
        self.update(self.atoms)
521
        return self.read_k_point_weights()
522

    
523
    def get_number_of_spins(self):
524
        return 1 + int(self.spinpol)
525

    
526
    def get_eigenvalues(self, kpt=0, spin=0):
527
        self.update(self.atoms)
528
        return self.read_eigenvalues(kpt, spin)
529

    
530
    def get_fermi_level(self):
531
        return self.fermi
532

    
533
    def get_number_of_grid_points(self):
534
        raise NotImplementedError
535

    
536
    def get_pseudo_density(self):
537
        raise NotImplementedError
538

    
539
    def get_pseudo_wavefunction(self, n=0, k=0, s=0, pad=True):
540
        raise NotImplementedError
541

    
542
    def get_bz_k_points(self):
543
        raise NotImplementedError
544

    
545
    def get_ibz_kpoints(self):
546
        self.update(self.atoms)
547
        return self.read_ibz_kpoints()
548

    
549
    def get_spin_polarized(self):
550
        if not hasattr(self, 'spinpol'):
551
            self.spinpol = self.atoms.get_initial_magnetic_moments().any()
552
        return self.spinpol            
553

    
554
    def get_magnetic_moment(self, atoms):
555
        self.update(atoms)
556
        return self.magnetic_moment
557
        
558
    def get_magnetic_moments(self, atoms):
559
        if self.int_params['lorbit']>=10 or self.list_params['rwigs']:
560
            self.update(atoms)
561
            return self.magnetic_moments
562
        else:
563
            raise RuntimeError(
564
                "The combination %s for lorbit with %s for rwigs not supported to calculate magnetic moments" % (p['lorbit'], p['rwigs']))
565

    
566
    def get_dipole_moment(self, atoms):
567
        """Returns total dipole moment of the system."""
568
        self.update(atoms)
569
        return self.dipole
570

    
571
    def get_xc_functional(self):
572
        return self.input_params['xc']
573

    
574
    def write_incar(self, atoms, **kwargs):
575
        """Writes the INCAR file."""
576
        incar = open('INCAR', 'w')
577
        incar.write('INCAR created by Atomic Simulation Environment\n')
578
        for key, val in self.float_params.items():
579
            if val is not None:
580
                incar.write(' %s = %5.6f\n' % (key.upper(), val))
581
        for key, val in self.exp_params.items():
582
            if val is not None:
583
                incar.write(' %s = %5.2e\n' % (key.upper(), val))
584
        for key, val in self.string_params.items():
585
            if val is not None:
586
                incar.write(' %s = %s\n' % (key.upper(), val))
587
        for key, val in self.int_params.items():
588
            if val is not None:
589
                incar.write(' %s = %d\n' % (key.upper(), val))
590
        for key, val in self.list_params.items():
591
            if val is not None:
592
                incar.write(' %s = ' % key.upper())
593
                if key in ('dipol', 'eint', 'ferwe', 'ropt', 'rwigs'):
594
                    [incar.write('%.4f ' % x) for x in val]
595
                elif key in ('iband', 'kpuse'):
596
                    [incar.write('%i ' % x) for x in val]
597
                elif key == 'magmom':
598
                    list = [[1, val[0]]]
599
                    for n in range(1, len(val)):
600
                        if val[n] == val[n-1]:
601
                            list[-1][0] += 1
602
                        else:
603
                            list.append([1, val[n]])
604
                    [incar.write('%i*%.4f ' % (mom[0], mom[1])) for mom in list]
605
                incar.write('\n')
606
        for key, val in self.bool_params.items():
607
            if val is not None:
608
                incar.write(' %s = ' % key.upper())
609
                if val:
610
                    incar.write('.TRUE.\n')
611
                else:
612
                    incar.write('.FALSE.\n')
613
        for key, val in self.special_params.items():
614
            if val is not None:
615
                incar.write(' %s = ' % key.upper())
616
                if key is 'lreal':
617
                    if type(val)==type('str'):
618
                        incar.write(val+'\n')
619
                    elif type(val)==type(True):
620
                       if val:
621
                           incar.write('.TRUE.\n')
622
                       else:
623
                           incar.write('.FALSE.\n') 
624
        if self.spinpol and not self.int_params['ispin']:
625
            incar.write(' ispin = 2\n'.upper())
626
            # Write out initial magnetic moments
627
            magmom = atoms.get_initial_magnetic_moments()[self.sort]
628
            list = [[1, magmom[0]]]
629
            for n in range(1, len(magmom)):
630
                if magmom[n] == magmom[n-1]:
631
                    list[-1][0] += 1
632
                else:
633
                    list.append([1, magmom[n]])
634
            incar.write(' magmom = '.upper())
635
            [incar.write('%i*%.4f ' % (mom[0], mom[1])) for mom in list]
636
            incar.write('\n')
637
        incar.close()
638

    
639
    def write_kpoints(self, **kwargs):
640
        """Writes the KPOINTS file."""
641
        p = self.input_params
642
        kpoints = open('KPOINTS', 'w')
643
        kpoints.write('KPOINTS created by Atomic Simulation Environemnt\n')
644
        shape=np.array(p['kpts']).shape
645
        if len(shape)==1:
646
            kpoints.write('0\n')
647
            if p['gamma']:
648
                kpoints.write('Gamma\n')
649
            else:
650
                kpoints.write('Monkhorst-Pack\n')
651
            [kpoints.write('%i ' % kpt) for kpt in p['kpts']]
652
            kpoints.write('\n0 0 0')
653
        elif len(shape)==2:
654
            kpoints.write('%i \n' % (len(p['kpts'])))
655
            kpoints.write('Cartesian\n')
656
            for n in range(len(p['kpts'])):
657
                [kpoints.write('%f ' % kpt) for kpt in p['kpts'][n]]
658
                if shape[1]==4:
659
                    kpoints.write('\n')
660
                elif shape[1]==3:
661
                    kpoints.write('1.0 \n')
662
        kpoints.close()
663

    
664
    def write_potcar(self,suffix = ""):
665
        """Writes the POTCAR file."""
666
        import tempfile
667
        potfile = open('POTCAR'+suffix,'w')
668
        for filename in self.ppp_list:
669
            if filename.endswith('R'):
670
                for line in open(filename, 'r'):
671
                    potfile.write(line)
672
            elif filename.endswith('.Z'):
673
                file_tmp = tempfile.NamedTemporaryFile()
674
                os.system('gunzip -c %s > %s' % (filename, file_tmp.name))
675
                for line in file_tmp.readlines():
676
                    potfile.write(line)
677
                file_tmp.close()
678
        potfile.close()
679

    
680
    def write_sort_file(self):
681
        """Writes a sortings file.
682

683
        This file contains information about how the atoms are sorted in
684
        the first column and how they should be resorted in the second
685
        column. It is used for restart purposes to get sorting right
686
        when reading in an old calculation to ASE."""
687

    
688
        file = open('ase-sort.dat', 'w')
689
        for n in range(len(self.sort)):
690
            file.write('%5i %5i \n' % (self.sort[n], self.resort[n]))
691

    
692
    # Methods for reading information from OUTCAR files:
693
    def read_energy(self, all=None):
694
        [energy_free, energy_zero]=[0, 0]
695
        if all:
696
            energy_free = []
697
            energy_zero = []
698
        for line in open('OUTCAR', 'r'):
699
            # Free energy
700
            if line.startswith('  free  energy   toten'):
701
                if all:
702
                    energy_free.append(float(line.split()[-2]))
703
                else:
704
                    energy_free = float(line.split()[-2])
705
            # Extrapolated zero point energy
706
            if line.startswith('  energy  without entropy'):
707
                if all:
708
                    energy_zero.append(float(line.split()[-1]))
709
                else:
710
                    energy_zero = float(line.split()[-1])
711
        return [energy_free, energy_zero]
712

    
713
    def read_forces(self, atoms, all=False):
714
        """Method that reads forces from OUTCAR file.
715

716
        If 'all' is switched on, the forces for all ionic steps
717
        in the OUTCAR file be returned, in other case only the
718
        forces for the last ionic configuration is returned."""
719

    
720
        file = open('OUTCAR','r')
721
        lines = file.readlines()
722
        file.close()
723
        n=0
724
        if all:
725
            all_forces = []
726
        for line in lines:
727
            if line.rfind('TOTAL-FORCE') > -1:
728
                forces=[]
729
                for i in range(len(atoms)):
730
                    forces.append(np.array([float(f) for f in lines[n+2+i].split()[3:6]]))
731
                if all:
732
                    all_forces.append(np.array(forces)[self.resort])
733
            n+=1
734
        if all:
735
            return np.array(all_forces)
736
        else:
737
            return np.array(forces)[self.resort]
738

    
739
    def read_fermi(self):
740
        """Method that reads Fermi energy from OUTCAR file"""
741
        E_f=None
742
        for line in open('OUTCAR', 'r'):
743
            if line.rfind('E-fermi') > -1:
744
                E_f=float(line.split()[2])
745
        return E_f
746

    
747
    def read_dipole(self):
748
        dipolemoment=np.zeros([1,3])
749
        for line in open('OUTCAR', 'r'):
750
            if line.rfind('dipolmoment') > -1:
751
                dipolemoment=np.array([float(f) for f in line.split()[1:4]])
752
        return dipolemoment
753

    
754
    def read_magnetic_moments(self, atoms):
755
        magnetic_moments = np.zeros(len(atoms))
756
        n = 0
757
        lines = open('OUTCAR', 'r').readlines()
758
        for line in lines:
759
            if line.rfind('magnetization (x)') > -1:
760
                for m in range(len(atoms)):
761
                    magnetic_moments[m] = float(lines[n + m + 4].split()[4])
762
            n += 1
763
        return np.array(magnetic_moments)[self.resort]
764

    
765
    def read_magnetic_moment(self):
766
        n=0
767
        for line in open('OUTCAR','r'):
768
            if line.rfind('number of electron  ') > -1:
769
                magnetic_moment=float(line.split()[-1])
770
            n+=1
771
        return magnetic_moment
772

    
773
    def read_nbands(self):
774
        for line in open('OUTCAR', 'r'):
775
            if line.rfind('NBANDS') > -1:
776
                return int(line.split()[-1])
777

    
778
    def read_convergence(self):
779
        """Method that checks whether a calculation has converged."""
780
        converged = None
781
        # First check electronic convergence
782
        for line in open('OUTCAR', 'r'):
783
            if line.rfind('EDIFF  ') > -1:
784
                ediff = float(line.split()[2])
785
            if line.rfind('total energy-change')>-1:
786
                split = line.split(':')
787
                a = float(split[1].split('(')[0])
788
                b = float(split[1].split('(')[1][0:-2])
789
                if [abs(a), abs(b)] < [ediff, ediff]:
790
                    converged = True
791
                else:
792
                    converged = False
793
                    continue
794
        # Then if ibrion > 0, check whether ionic relaxation condition been fulfilled
795
        if self.int_params['ibrion'] > 0:
796
            if not self.read_relaxed():
797
                converged = False
798
            else:
799
                converged = True
800
        return converged
801

    
802
    def read_ibz_kpoints(self):
803
        lines = open('OUTCAR', 'r').readlines()
804
        ibz_kpts = []
805
        n = 0
806
        i = 0
807
        for line in lines:
808
            if line.rfind('Following cartesian coordinates')>-1:
809
                m = n+2
810
                while i==0:
811
                    ibz_kpts.append([float(lines[m].split()[p]) for p in range(3)])
812
                    m += 1
813
                    if lines[m]==' \n':
814
                        i = 1
815
            if i == 1:
816
                continue
817
            n += 1
818
        ibz_kpts = np.array(ibz_kpts)
819
        return np.array(ibz_kpts)
820

    
821
    def read_k_point_weights(self):
822
        file = open('IBZKPT')
823
        lines = file.readlines()
824
        file.close()
825
        kpt_weights = []
826
        for n in range(3, len(lines)):
827
            kpt_weights.append(float(lines[n].split()[3]))
828
        kpt_weights = np.array(kpt_weights)
829
        kpt_weights /= np.sum(kpt_weights)
830
        return kpt_weights
831

    
832
    def read_eigenvalues(self, kpt=0, spin=0):
833
        file = open('EIGENVAL', 'r')
834
        lines = file.readlines()
835
        file.close()
836
        eigs = []
837
        for n in range(8+kpt*(self.nbands+2), 8+kpt*(self.nbands+2)+self.nbands):
838
            eigs.append(float(lines[n].split()[spin+1]))
839
        return np.array(eigs)
840

    
841
    def read_relaxed(self):
842
        for line in open('OUTCAR', 'r'):
843
            if line.rfind('reached required accuracy') > -1:
844
                return True
845
        return False
846

    
847
# The below functions are used to restart a calculation and are under early constructions
848

    
849
    def read_incar(self, filename='INCAR'):
850
        """Method that imports settings from INCAR file."""
851

    
852
        self.spinpol = False
853
        file=open(filename, 'r')
854
        file.readline()
855
        lines=file.readlines()
856
        for line in lines:
857
            try:
858
                data = line.split()
859
                # Skip empty and commented lines. 
860
                if len(data) == 0:
861
                    continue
862
                elif data[0][0] in ['#', '!']:
863
                    continue
864
                key = data[0].lower()
865
                if key in float_keys:
866
                    self.float_params[key] = float(data[2])
867
                elif key in exp_keys:
868
                    self.exp_params[key] = float(data[2])
869
                elif key in string_keys:
870
                    self.string_params[key] = str(data[2])
871
                elif key in int_keys:
872
                    if key == 'ispin':
873
                        if int(data[2]) == 2:
874
                            self.spinpol = True
875
                    else:
876
                        self.int_params[key] = int(data[2])
877
                elif key in bool_keys:
878
                    if 'true' in data[2].lower():
879
                        self.bool_params[key] = True
880
                    elif 'false' in data[2].lower():
881
                        self.bool_params[key] = False
882
                elif key in list_keys:
883
                    list = []
884
                    if key in ('dipol', 'eint', 'ferwe', 'ropt', 'rwigs'):
885
                        for a in data[2:]:
886
                            list.append(float(a))
887
                    elif key in ('iband', 'kpuse'):
888
                        for a in data[2:]:
889
                            list.append(int(a))
890
                    self.list_params[key] = list
891
                    if key == 'magmom':
892
                        done = False
893
                        for a in data[2:]:
894
                            if '!' in a or done:
895
                                done = True
896
                            elif '*' in a: 
897
                                a = a.split('*')
898
                                for b in range(int(a[0])):
899
                                    list.append(float(a[1]))
900
                            else:
901
                                list.append(float(a))
902
                        list = np.array(list)
903
                        self.atoms.set_initial_magnetic_moments(list[self.resort])
904
            except KeyError:
905
                raise IOError('Keyword "%s" in INCAR is not known by calculator.' % key)
906
            except IndexError:
907
                raise IOError('Value missing for keyword "%s".' % key)
908

    
909
    def read_outcar(self):
910
        # Spin polarized calculation?
911
        file = open('OUTCAR', 'r')
912
        lines = file.readlines()
913
        file.close()
914
        for line in lines:
915
            if line.rfind('ISPIN') > -1:
916
                if int(line.split()[2])==2:
917
                    self.spinpol = True
918
                else:
919
                    self.spinpol = None
920
        self.energy_free, self.energy_zero = self.read_energy()
921
        self.forces = self.read_forces(self.atoms)
922
        self.dipole = self.read_dipole()
923
        self.fermi = self.read_fermi()
924
        self.stress = self.read_stress()
925
        self.nbands = self.read_nbands()
926
        p=self.int_params
927
        q=self.list_params
928
        if self.spinpol:
929
            self.magnetic_moment = self.read_magnetic_moment()
930
            if p['lorbit']>=10 or (p['lorbit']!=None and q['rwigs']):
931
                self.magnetic_moments = self.read_magnetic_moments(self.atoms)
932
        self.set(nbands=self.nbands)
933

    
934
    def read_kpoints(self, filename='KPOINTS'):
935
        file = open(filename, 'r')
936
        lines = file.readlines()
937
        file.close()
938
        type = lines[2].split()[0].lower()[0]
939
        if type in ['g', 'm']:
940
            if type=='g':
941
                self.set(gamma=True)
942
            kpts = np.array([int(lines[3].split()[i]) for i in range(3)])
943
            self.set(kpts=kpts)
944
        elif type in ['c', 'k']:
945
            raise NotImplementedError('Only Monkhorst-Pack and gamma centered grid supported for restart.')
946
        else:
947
            raise NotImplementedError('Only Monkhorst-Pack and gamma centered grid supported for restart.')
948
    
949
    def read_potcar(self):
950
        """ Method that reads the Exchange Correlation functional from POTCAR file.
951
        """
952
        file = open('POTCAR', 'r')
953
        lines = file.readlines()
954
        file.close()
955

    
956
        # Search for key 'LEXCH' in POTCAR
957
        xc_flag = None
958
        for line in lines:
959
            key = line.split()[0].upper()
960
            if key == 'LEXCH':
961
                xc_flag = line.split()[-1].upper()
962
                break
963

    
964
        if xc_flag is None:
965
            raise ValueError('LEXCH flag not found in POTCAR file.')
966

    
967
        # Values of parameter LEXCH and corresponding XC-functional
968
        xc_dict = {'PE':'PBE', '91':'PW91', 'CA':'LDA'}
969

    
970
        if xc_flag not in xc_dict.keys():
971
            raise ValueError(
972
                'Unknown xc-functional flag found in POTCAR, LEXCH=%s' % xc_flag)
973

    
974
        self.input_params['xc'] = xc_dict[xc_flag]
975

    
976

    
977
class VaspChargeDensity(object):
978
    """Class for representing VASP charge density"""
979

    
980
    def __init__(self, filename='CHG'):
981
        # Instance variables
982
        self.atoms = []   # List of Atoms objects
983
        self.chg = []     # Charge density
984
        self.chgdiff = [] # Charge density difference, if spin polarized
985
        self.aug = ''     # Augmentation charges, not parsed just a big string
986
        self.augdiff = '' # Augmentation charge differece, is spin polarized
987
        
988
        # Note that the augmentation charge is not a list, since they
989
        # are needed only for CHGCAR files which store only a single
990
        # image.
991
        if filename != None:
992
            self.read(filename)
993

    
994
    def is_spin_polarized(self):
995
        if len(self.chgdiff) > 0:
996
            return True
997
        return False
998

    
999
    def _read_chg(self, fobj, chg, volume):
1000
        """Read charge from file object
1001

1002
        Utility method for reading the actual charge density (or
1003
        charge density difference) from a file object. On input, the
1004
        file object must be at the beginning of the charge block, on
1005
        output the file position will be left at the end of the
1006
        block. The chg array must be of the correct dimensions.
1007

1008
        """
1009
        # VASP writes charge density as
1010
        # WRITE(IU,FORM) (((C(NX,NY,NZ),NX=1,NGXC),NY=1,NGYZ),NZ=1,NGZC)
1011
        # Fortran nested implied do loops; innermost index fastest
1012
        # First, just read it in
1013
        for zz in range(chg.shape[2]):
1014
            for yy in range(chg.shape[1]):
1015
                chg[:, yy, zz] = np.fromfile(fobj, count = chg.shape[0],
1016
                                             sep=' ')
1017
        chg /= volume
1018

    
1019
    def read(self, filename='CHG'):
1020
        """Read CHG or CHGCAR file.
1021

1022
        If CHG contains charge density from multiple steps all the
1023
        steps are read and stored in the object. By default VASP
1024
        writes out the charge density every 10 steps.
1025

1026
        chgdiff is the difference between the spin up charge density
1027
        and the spin down charge density and is thus only read for a
1028
        spin-polarized calculation.
1029

1030
        aug is the PAW augmentation charges found in CHGCAR. These are
1031
        not parsed, they are just stored as a string so that they can
1032
        be written again to a CHGCAR format file.
1033

1034
        """
1035
        import ase.io.vasp as aiv
1036
        f = open(filename)
1037
        self.atoms = []
1038
        self.chg = []
1039
        self.chgdiff = []
1040
        self.aug = ''
1041
        self.augdiff = ''
1042
        while True:
1043
            try:
1044
                atoms = aiv.read_vasp(f)
1045
            except ValueError, e:
1046
                # Probably an empty line, or we tried to read the 
1047
                # augmentation occupancies in CHGCAR
1048
                break 
1049
            f.readline()
1050
            ngr = f.readline().split()
1051
            ng = (int(ngr[0]), int(ngr[1]), int(ngr[2]))
1052
            chg = np.empty(ng)
1053
            self._read_chg(f, chg, atoms.get_volume())
1054
            self.chg.append(chg)
1055
            self.atoms.append(atoms)
1056
            # Check if the file has a spin-polarized charge density part, and
1057
            # if so, read it in.
1058
            fl = f.tell()
1059
            # First check if the file has an augmentation charge part (CHGCAR file.)
1060
            line1 = f.readline()
1061
            if line1=='':
1062
                break
1063
            elif line1.find('augmentation') != -1:
1064
                augs = [line1]
1065
                while True:
1066
                    line2 = f.readline()
1067
                    if line2.split() == ngr:
1068
                        self.aug = ''.join(augs)
1069
                        augs = []
1070
                        chgdiff = np.empty(ng)
1071
                        self._read_chg(f, chgdiff, atoms.get_volume())
1072
                        self.chgdiff.append(chgdiff)
1073
                    elif line2 == '':
1074
                        break
1075
                    else:
1076
                        augs.append(line2)
1077
                if len(self.aug) == 0:
1078
                    self.aug = ''.join(augs)
1079
                    augs = []
1080
                else:
1081
                    self.augdiff = ''.join(augs)
1082
                    augs = []
1083
            elif line1.split() == ngr:
1084
                chgdiff = np.empty(ng)
1085
                self._read_chg(f, chgdiff, atoms.get_volume())
1086
                self.chgdiff.append(chgdiff)
1087
            else:
1088
                f.seek(fl)
1089
        f.close()
1090

    
1091
    def _write_chg(self, fobj, chg, volume, format='chg'):
1092
        """Write charge density
1093

1094
        Utility function similar to _read_chg but for writing.
1095

1096
        """
1097
        # Make a 1D copy of chg, must take transpose to get ordering right
1098
        chgtmp=chg.T.ravel()
1099
        # Multiply by volume
1100
        chgtmp=chgtmp*volume
1101
        # Must be a tuple to pass to string conversion
1102
        chgtmp=tuple(chgtmp)
1103
        # CHG format - 10 columns
1104
        if format.lower() == 'chg':
1105
            # Write all but the last row
1106
            for ii in range((len(chgtmp)-1)/10):
1107
                fobj.write(' %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G\
1108
 %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G\n' % chgtmp[ii*10:(ii+1)*10]
1109
                           )
1110
            # If the last row contains 10 values then write them without a newline
1111
            if len(chgtmp)%10==0:
1112
                fobj.write(' %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G\
1113
 %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G' % chgtmp[len(chgtmp)-10:len(chgtmp)])
1114
            # Otherwise write fewer columns without a newline
1115
            else:
1116
                for ii in range(len(chgtmp)%10):
1117
                    fobj.write((' %#11.5G') % chgtmp[len(chgtmp)-len(chgtmp)%10+ii])
1118
        # Other formats - 5 columns
1119
        else:
1120
            # Write all but the last row
1121
            for ii in range((len(chgtmp)-1)/5):
1122
                fobj.write(' %17.10E %17.10E %17.10E %17.10E %17.10E\n' % chgtmp[ii*5:(ii+1)*5])
1123
            # If the last row contains 5 values then write them without a newline
1124
            if len(chgtmp)%5==0:
1125
                fobj.write(' %17.10E %17.10E %17.10E %17.10E %17.10E' % chgtmp[len(chgtmp)-5:len(chgtmp)])
1126
            # Otherwise write fewer columns without a newline
1127
            else:
1128
                for ii in range(len(chgtmp)%5):
1129
                    fobj.write((' %17.10E') % chgtmp[len(chgtmp)-len(chgtmp)%5+ii])
1130
        # Write a newline whatever format it is
1131
        fobj.write('\n')
1132
        # Clean up
1133
        del chgtmp
1134

    
1135
    def write(self, filename='CHG', format=None):
1136
        """Write VASP charge density in CHG format.
1137

1138
        filename: str
1139
            Name of file to write to.
1140
        format: str
1141
            String specifying whether to write in CHGCAR or CHG
1142
            format.
1143

1144
        """
1145
        import ase.io.vasp as aiv
1146
        if format == None:
1147
            if filename.lower().find('chgcar') != -1:
1148
                format = 'chgcar'
1149
            elif filename.lower().find('chg') != -1:
1150
                format = 'chg'
1151
            elif len(self.chg) == 1:
1152
                format = 'chgcar'
1153
            else:
1154
                format = 'chg'
1155
        f = open(filename, 'w')
1156
        for ii, chg in enumerate(self.chg):
1157
            if format == 'chgcar' and ii != len(self.chg) - 1:
1158
                continue # Write only the last image for CHGCAR
1159
            aiv.write_vasp(f, self.atoms[ii], direct=True)
1160
            f.write('\n')
1161
            for dim in chg.shape:
1162
                f.write(' %4i' % dim)
1163
            f.write('\n')
1164
            vol = self.atoms[ii].get_volume()
1165
            self._write_chg(f, chg, vol, format)
1166
            if format == 'chgcar':
1167
                f.write(self.aug)
1168
            if self.is_spin_polarized():
1169
                if format == 'chg':
1170
                    f.write('\n')
1171
                for dim in chg.shape:
1172
                    f.write(' %4i' % dim)
1173
                self._write_chg(f, self.chgdiff[ii], vol, format)
1174
                if format == 'chgcar':
1175
                    f.write('\n')
1176
                    f.write(self.augdiff)
1177
            if format == 'chg' and len(self.chg) > 1:
1178
                f.write('\n')
1179
        f.close()
1180

    
1181

    
1182
class VaspDos(object):
1183
    """Class for representing density-of-states produced by VASP
1184

1185
    The energies are in property self.energy
1186

1187
    Site-projected DOS is accesible via the self.site_dos method.
1188

1189
    Total and integrated DOS is accessible as numpy.ndarray's in the
1190
    properties self.dos and self.integrated_dos. If the calculation is
1191
    spin polarized, the arrays will be of shape (2, NDOS), else (1,
1192
    NDOS).
1193

1194
    The self.efermi property contains the currently set Fermi
1195
    level. Changing this value shifts the energies.
1196
    
1197
    """
1198

    
1199
    def __init__(self, doscar='DOSCAR', efermi=0.0):
1200
        """Initialize"""
1201
        self._efermi = 0.0
1202
        self.read_doscar(doscar)
1203
        self.efermi = efermi
1204

    
1205
    def _set_efermi(self, efermi):
1206
        """Set the Fermi level."""
1207
        ef = efermi - self._efermi
1208
        self._efermi = efermi
1209
        self._total_dos[0, :] = self._total_dos[0, :] - ef
1210
        try:
1211
            self._site_dos[:, 0, :] = self._site_dos[:, 0, :] - ef
1212
        except IndexError:
1213
            pass
1214

    
1215
    def _get_efermi(self):
1216
        return self._efermi
1217

    
1218
    efermi = property(_get_efermi, _set_efermi, None, "Fermi energy.")
1219

    
1220
    def _get_energy(self):
1221
        """Return the array with the energies."""
1222
        return self._total_dos[0, :]
1223
    energy = property(_get_energy, None, None, "Array of energies")
1224

    
1225
    def site_dos(self, atom, orbital):
1226
        """Return an NDOSx1 array with dos for the chosen atom and orbital.
1227

1228
        atom: int
1229
            Atom index
1230
        orbital: int or str
1231
            Which orbital to plot
1232

1233
        If the orbital is given as an integer:
1234
        If spin-unpolarized calculation, no phase factors:
1235
        s = 0, p = 1, d = 2
1236
        Spin-polarized, no phase factors:
1237
        s-up = 0, s-down = 1, p-up = 2, p-down = 3, d-up = 4, d-down = 5
1238
        If phase factors have been calculated, orbitals are
1239
        s, py, pz, px, dxy, dyz, dz2, dxz, dx2
1240
        double in the above fashion if spin polarized.
1241

1242
        """
1243
        # Integer indexing for orbitals starts from 1 in the _site_dos array
1244
        # since the 0th column contains the energies
1245
        if isinstance(orbital, int):
1246
            return self._site_dos[atom, orbital + 1, :]
1247
        n = self._site_dos.shape[1]
1248
        if n == 4:
1249
            norb = {'s':1, 'p':2, 'd':3}
1250
        elif n == 7:
1251
            norb = {'s+':1, 's-up':1, 's-':2, 's-down':2,
1252
                    'p+':3, 'p-up':3, 'p-':4, 'p-down':4,
1253
                    'd+':5, 'd-up':5, 'd-':6, 'd-down':6}
1254
        elif n == 10:
1255
            norb = {'s':1, 'py':2, 'pz':3, 'px':4,
1256
                    'dxy':5, 'dyz':6, 'dz2':7, 'dxz':8,
1257
                    'dx2':9}
1258
        elif n == 19:
1259
            norb = {'s+':1, 's-up':1, 's-':2, 's-down':2,
1260
                    'py+':3, 'py-up':3, 'py-':4, 'py-down':4,
1261
                    'pz+':5, 'pz-up':5, 'pz-':6, 'pz-down':6,
1262
                    'px+':7, 'px-up':7, 'px-':8, 'px-down':8,
1263
                    'dxy+':9, 'dxy-up':9, 'dxy-':10, 'dxy-down':10,
1264
                    'dyz+':11, 'dyz-up':11, 'dyz-':12, 'dyz-down':12,
1265
                    'dz2+':13, 'dz2-up':13, 'dz2-':14, 'dz2-down':14,
1266
                    'dxz+':15, 'dxz-up':15, 'dxz-':16, 'dxz-down':16,
1267
                    'dx2+':17, 'dx2-up':17, 'dx2-':18, 'dx2-down':18}
1268
        return self._site_dos[atom, norb[orbital.lower()], :]
1269

    
1270
    def _get_dos(self):
1271
        if self._total_dos.shape[0] == 3:
1272
            return self._total_dos[1, :]
1273
        elif self._total_dos.shape[0] == 5:
1274
            return self._total_dos[1:3, :]
1275
    dos = property(_get_dos, None, None, 'Average DOS in cell')
1276

    
1277
    def _get_integrated_dos(self):
1278
        if self._total_dos.shape[0] == 3:
1279
            return self._total_dos[2, :]
1280
        elif self._total_dos.shape[0] == 5:
1281
            return self._total_dos[3:5, :]
1282
    integrated_dos = property(_get_integrated_dos, None, None,
1283
                              'Integrated average DOS in cell')
1284

    
1285
    def read_doscar(self, fname="DOSCAR"):
1286
        """Read a VASP DOSCAR file"""
1287
        f = open(fname)
1288
        natoms = int(f.readline().split()[0])
1289
        [f.readline() for nn in range(4)]  # Skip next 4 lines.
1290
        # First we have a block with total and total integrated DOS
1291
        ndos = int(f.readline().split()[2])
1292
        dos = []
1293
        for nd in xrange(ndos):
1294
            dos.append(np.array([float(x) for x in f.readline().split()]))
1295
        self._total_dos = np.array(dos).T
1296
        # Next we have one block per atom, if INCAR contains the stuff
1297
        # necessary for generating site-projected DOS
1298
        dos = []
1299
        for na in xrange(natoms):
1300
            line = f.readline()
1301
            if line == '':
1302
                # No site-projected DOS
1303
                break
1304
            ndos = int(line.split()[2])
1305
            line = f.readline().split()
1306
            cdos = np.empty((ndos, len(line)))
1307
            cdos[0] = np.array(line)
1308
            for nd in xrange(1, ndos):
1309
                line = f.readline().split()
1310
                cdos[nd] = np.array([float(x) for x in line])
1311
            dos.append(cdos.T)
1312
        self._site_dos = np.array(dos)
1313

    
1314

    
1315
import pickle
1316

    
1317
class xdat2traj:
1318
    def __init__(self, trajectory=None, atoms=None, poscar=None, 
1319
                 xdatcar=None, sort=None, calc=None):
1320
        if not poscar:
1321
            self.poscar = 'POSCAR'
1322
        else:
1323
            self.poscar = poscar
1324
        if not atoms:
1325
            self.atoms = ase.io.read(self.poscar, format='vasp')
1326
        else:
1327
            self.atoms = atoms
1328
        if not xdatcar:
1329
            self.xdatcar = 'XDATCAR'
1330
        else:
1331
            self.xdatcar = xdatcar
1332
        if not trajectory:
1333
            self.trajectory = 'out.traj'
1334
        else:
1335
            self.trajectory = trajectory
1336
        if not calc:
1337
            self.calc = Vasp()
1338
        else:
1339
            self.calc = calc
1340
        if not sort: 
1341
            if not hasattr(self.calc, 'sort'):
1342
                self.calc.sort = range(len(self.atoms))
1343
        else:
1344
            self.calc.sort = sort
1345
        self.calc.resort = range(len(self.calc.sort))
1346
        for n in range(len(self.calc.resort)):
1347
            self.calc.resort[self.calc.sort[n]] = n
1348
        self.out = ase.io.trajectory.PickleTrajectory(self.trajectory, mode='w')
1349
        self.energies = self.calc.read_energy(all=True)[1]
1350
        self.forces = self.calc.read_forces(self.atoms, all=True)
1351

    
1352
    def convert(self):
1353
        lines = open(self.xdatcar).readlines()
1354
        if len(lines[5].split())==0:
1355
            del(lines[0:6])
1356
        elif len(lines[4].split())==0:
1357
            del(lines[0:5])
1358
        step = 0
1359
        iatom = 0
1360
        scaled_pos = []
1361
        for line in lines:
1362
            if iatom == len(self.atoms):
1363
                if step == 0:
1364
                    self.out.write_header(self.atoms[self.calc.resort])
1365
                scaled_pos = np.array(scaled_pos)
1366
                self.atoms.set_scaled_positions(scaled_pos)
1367
                d = {'positions': self.atoms.get_positions()[self.calc.resort],
1368
                     'cell': self.atoms.get_cell(),
1369
                     'momenta': None,
1370
                     'energy': self.energies[step],
1371
                     'forces': self.forces[step],
1372
                     'stress': None}
1373
                pickle.dump(d, self.out.fd, protocol=-1)
1374
                scaled_pos = []
1375
                iatom = 0
1376
                step += 1
1377
            else:
1378
                
1379
                iatom += 1
1380
                scaled_pos.append([float(line.split()[n]) for n in range(3)])
1381

    
1382
        # Write also the last image
1383
        # I'm sure there is also more clever fix...
1384
        scaled_pos = np.array(scaled_pos)
1385
        self.atoms.set_scaled_positions(scaled_pos)
1386
        d = {'positions': self.atoms.get_positions()[self.calc.resort],
1387
             'cell': self.atoms.get_cell(),
1388
             'momenta': None,
1389
             'energy': self.energies[step],
1390
             'forces': self.forces[step],
1391
             'stress': None}
1392
        pickle.dump(d, self.out.fd, protocol=-1)
1393

    
1394
        self.out.fd.close()