Statistiques
| Révision :

root / ase / calculators / vasp.py @ 5

Historique | Voir | Annoter | Télécharger (52 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.bool_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
        symbols = {}
252
        if self.input_params['setups']:
253
            for m in self.input_params['setups']:
254
                try : 
255
                    #special_setup[self.input_params['setups'][m]] = int(m)
256
                    special_setups.append(int(m))
257
                except:
258
                    #print 'setup ' + m + ' is a groups setup'
259
                    continue
260
            #print 'special_setups' , special_setups
261

    
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
        # Build the sorting list
273
        self.sort = []
274
        self.sort.extend(special_setups)
275

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

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

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

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

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

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

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

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

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

    
474
    def get_atoms(self):
475
        atoms = self.atoms.copy()
476
        atoms.set_calculator(self)
477
        return atoms
478

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

    
486
    def get_forces(self, atoms):
487
        self.update(atoms)
488
        return self.forces
489

    
490
    def get_stress(self, atoms):
491
        self.update(atoms)
492
        return self.stress
493

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

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

    
518
    def get_number_of_bands(self):
519
        return self.nbands
520

    
521
    def get_k_point_weights(self):
522
        self.update(self.atoms)
523
        return self.read_k_point_weights()
524

    
525
    def get_number_of_spins(self):
526
        return 1 + int(self.spinpol)
527

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

    
532
    def get_fermi_level(self):
533
        return self.fermi
534

    
535
    def get_number_of_grid_points(self):
536
        raise NotImplementedError
537

    
538
    def get_pseudo_density(self):
539
        raise NotImplementedError
540

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

    
544
    def get_bz_k_points(self):
545
        raise NotImplementedError
546

    
547
    def get_ibz_kpoints(self):
548
        self.update(self.atoms)
549
        return self.read_ibz_kpoints()
550

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

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

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

    
573
    def get_xc_functional(self):
574
        return self.input_params['xc']
575

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

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

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

    
682
    def write_sort_file(self):
683
        """Writes a sortings file.
684

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    
849
# The below functions are used to restart a calculation and are under early constructions
850

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

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

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

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

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

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

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

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

    
976
        self.input_params['xc'] = xc_dict[xc_flag]
977

    
978

    
979
class VaspChargeDensity(object):
980
    """Class for representing VASP charge density"""
981

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

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

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

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

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

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

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

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

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

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

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

1096
        Utility function similar to _read_chg but for writing.
1097

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

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

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

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

    
1183

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

1187
    The energies are in property self.energy
1188

1189
    Site-projected DOS is accesible via the self.site_dos method.
1190

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

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

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

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

    
1217
    def _get_efermi(self):
1218
        return self._efermi
1219

    
1220
    efermi = property(_get_efermi, _set_efermi, None, "Fermi energy.")
1221

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

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

1230
        atom: int
1231
            Atom index
1232
        orbital: int or str
1233
            Which orbital to plot
1234

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

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

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

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

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

    
1316

    
1317
import pickle
1318

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

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

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

    
1396
        self.out.fd.close()