Coverage for gpaw/lcao/analyse_basis.py: 93%

1import os 

2import sys 

3from optparse import OptionParser 

4 

5 

6def build_parser(): 

7 usage = '%prog [OPTION] [BASIS]...' 

8 parser = OptionParser(usage=usage, version='%prog 1.0') 

9 parser.add_option('-f', '--files', action='store_true', 

10 dest='actual_filenames', 

11 help='Read from specified filenames rather than ' 

12 'searching GPAW setup directories') 

13 parser.add_option('-s', '--save-figs', action='store_true', dest='save', 

14 help='Save figures to disk rather than showing plots') 

15 parser.add_option('-l', '--literal', action='store_true', 

16 help='Do not pre-multiply wave functions by r in plots') 

17 parser.add_option('-n', '--normalize', action='store_true', 

18 help='Plot normalized wave functions') 

19 parser.add_option('-x', '--ext', default='png', 

20 help='Image format [default: %default]') 

21 return parser 

22 

23 

24def main(argv=sys.argv[1:]): 

25 parser = build_parser() 

26 opts, files = parser.parse_args(argv) 

27 

28 import matplotlib.pyplot as plt 

29 from gpaw.basis_data import Basis, BasisPlotter 

30 

31 plotter = BasisPlotter(premultiply=not opts.literal, 

32 normalize=opts.normalize, 

33 show=False, 

34 save=opts.save, 

35 ext=opts.ext) 

36 

37 for path in files: 

38 dir, filename = os.path.split(path) 

39 splitfilename = filename.split('.') 

40 symbol = splitfilename[0] 

41 name = '.'.join(splitfilename[1:-1]) 

42 if opts.actual_filenames: 

43 basis = Basis.read_path(symbol, name, path) 

44 else: # search GPAW setup dirs 

45 basis = Basis.find(symbol, name) 

46 plotter.plot(basis) 

47 

48 if not opts.save: 

49 plt.show()