Coverage for gpaw/point_groups/cli.py: 87%

1import argparse 

2from typing import List, Union 

3 

4import numpy as np 

5from ase import Atoms 

6from gpaw.new.ase_interface import ASECalculator, GPAW 

7from gpaw.point_groups import SymmetryChecker, point_group_names 

8from gpaw.typing import Array1D, Array3D 

9 

10 

11class CubeCalc: 

12 """Wrap cube-file in a calculator.""" 

13 def __init__(self, function: Array3D, atoms: Atoms): 

14 self.function = function 

15 self.atoms = atoms 

16 

17 def get_pseudo_wave_function(self, 

18 band: int, 

19 spin: int) -> Array3D: 

20 return self.function 

21 

22 def get_eigenvalues(self, spin: int) -> Array1D: 

23 return np.zeros(1) 

24 

25 def get_number_of_spins(self): 

26 return 1 

27 

28 

29def main(argv: List[str] = None) -> None: 

30 parser = argparse.ArgumentParser( 

31 prog='python3 -m gpaw.point_groups', 

32 description='Analyse point-group of atoms and wave-functions.') 

33 add = parser.add_argument 

34 add('pg', metavar='point-group', choices=point_group_names, 

35 help='Name of point-group: C2, C2v, C3v, D2d, D3h, D5, D5h, ' 

36 'Ico, Ih, Oh, Td or Th.') 

37 add('file', metavar='input-file', 

38 help='Cube-file, gpw-file or something else with atoms in it.') 

39 add('-c', '--center', help='Center specified as one or more atoms. ' 

40 'Use chemical symbols or sequence numbers.') 

41 add('-r', '--radius', default=2.5, 

42 help='Cutoff radius (in Å) used for wave function overlaps.') 

43 add('-b', '--bands', default=':', metavar='N1:N2', 

44 help='Band range.') 

45 add('-a', '--axes', default='', 

46 help='Example: "-a z=x,x=-y".') 

47 args = parser.parse_intermixed_args(argv) 

48 

49 calc: Union[None, ASECalculator, CubeCalc] 

50 

51 if args.file.endswith('.gpw'): 

52 calc = GPAW(args.file) 

53 atoms = calc.atoms 

54 n1, n2 = (int(x) if x else 0 for x in args.bands.split(':')) 

55 elif args.file.endswith('.cube'): 

56 from ase.io.cube import read_cube 

57 with open(args.file) as fd: 

58 dct = read_cube(fd) 

59 calc = CubeCalc(dct['data'], dct['atoms']) 

60 atoms = dct['atoms'] 

61 n1 = 0 

62 n2 = 1 

63 else: 

64 from ase.io import read 

65 atoms_maybe = read(args.file) 

66 assert isinstance(atoms_maybe, Atoms) 

67 atoms = atoms_maybe 

68 calc = None 

69 

70 if args.center: 

71 symbols = set(args.center.split(',')) 

72 center = np.zeros(3) 

73 n = 0 

74 for a, (symbol, position) in enumerate(zip(atoms.symbols, 

75 atoms.positions)): 

76 if symbol in symbols or str(a) in symbols: 

77 center += position 

78 n += 1 

79 center /= n 

80 else: 

81 center = atoms.cell.sum(0) / 2 

82 print('Center:', center, f'(atoms: {n})') 

83 

84 radius = float(args.radius) 

85 

86 kwargs = {} 

87 for axis in args.axes.split(',') if args.axes else []: 

88 axis1, axis2 = axis.split('=') 

89 kwargs[axis1] = axis2 

90 

91 checker = SymmetryChecker(args.pg, center, radius, **kwargs) 

92 

93 ok = checker.check_atoms(atoms) 

94 print(f'{args.pg}-symmetry:', 'Yes' if ok else 'No') 

95 

96 if calc: 

97 nspins = calc.get_number_of_spins() 

98 for spin in range(nspins): 

99 if nspins == 2: 

100 print('Spin', ['up', 'down'][spin]) 

101 checker.check_calculation(calc, n1, n2, spin=spin)