Coverage for gpaw/atom/basisfromfile.py: 99%

1from __future__ import annotations 

2 

3from argparse import Namespace 

4from collections.abc import Callable 

5from operator import methodcaller 

6from os import PathLike 

7from pathlib import Path 

8from typing import NamedTuple 

9 

10from ..basis_data import Basis, parse_basis_name 

11from ..setup_data import SetupData, read_maybe_unzipping, search_for_file 

12from ..typing import Self 

13from .all_electron import ValenceData 

14from .basis import BasisMaker 

15 

16 

17BasisGetter = Callable[[BasisMaker], Basis] 

18 

19 

20def read_setup_and_generate_basis( 

21 setup: str | PathLike, 

22 getter: BasisGetter, /, 

23 **kwargs) -> tuple[SetupData, Basis]: 

24 setupdata = read_setupdata(setup) 

25 valdata = ValenceData.from_setupdata_onthefly_potentials(setupdata) 

26 return setupdata, getter(BasisMaker(valdata, **kwargs)) 

27 

28 

29class BasisInfo(NamedTuple): 

30 zetacount: int 

31 polcount: int 

32 name: str | None = None 

33 

34 @classmethod 

35 def from_name(cls, name: str) -> Self: 

36 zc, pc = parse_basis_name(name) 

37 return cls(zc, pc, name) 

38 

39 

40def parse_j_values(j: str) -> list[int]: 

41 return [int(value) for value in j.split(',')] 

42 

43 

44def parse_tail_norm(tail: str) -> list[float]: 

45 return [float(value) for value in tail.split(',')] 

46 

47 

48def add_common_args(add: Callable) -> None: 

49 add('-t', '--type', 

50 default='dzp', metavar='<type>', type=BasisInfo.from_name, 

51 help='type of basis (e.g.: sz, dzp, qztp, 4z3p) ' 

52 '[default: %(default)s]') 

53 add('-E', '--energy-shift', 

54 default=.1, metavar='<energy>', type=float, 

55 help='use given energy shift to determine cutoff ' 

56 '[default/eV: %(default)s]') 

57 add('-T', '--tail-norm', 

58 default=[0.16, 0.3, 0.6], dest='tailnorm', 

59 metavar='<norm>[,<norm>[,...]]', type=parse_tail_norm, 

60 help='use the given fractions to define the split-valence cutoffs ' 

61 '[default: %(default)s]') 

62 add('--rcut-max', 

63 default=16., metavar='<rcut>', type=float, 

64 help='max cutoff for confined atomic orbitals. ' 

65 'This option has no effect on orbitals with smaller cutoff ' 

66 '[default/Bohr: %(default)s]') 

67 add('--rcut-pol-rel', default=1.0, metavar='<rcut>', type=float, 

68 help='polarization function cutoff relative to largest ' 

69 'single-zeta cutoff [default: %(default)s]') 

70 add('--rchar-pol-rel', metavar='<rchar>', type=float, 

71 help='characteristic radius of Gaussian when not using interpolation ' 

72 'scheme, relative to rcut') 

73 add('--vconf-amplitude', default=12., metavar='<alpha>', type=float, 

74 help='set proportionality constant of smooth ' 

75 'confinement potential [default: %(default)s]') 

76 add('--vconf-rstart-rel', default=.6, metavar='<ri/rc>', type=float, 

77 help='set inner cutoff for smooth confinement potential ' 

78 'relative to hard cutoff [default: %(default)s]') 

79 add('--vconf-sharp-confinement', action='store_true', 

80 help='use sharp rather than smooth confinement potential') 

81 add('--lpol', metavar='<l>', type=int, 

82 help='angular momentum quantum number of polarization function. ' 

83 'Default behaviour is to take the lowest l which is not ' 

84 'among the valence states') 

85 add('--jvalues', metavar='<j>[,<j>[,...]]', type=parse_j_values, 

86 help='explicitly specify which states to include. ' 

87 'Numbering corresponds to generator\'s valence state ordering. ' 

88 'For example: 0,1,2') 

89 

90 

91def read_setupdata(path: str | PathLike) -> SetupData: 

92 setupdata = SetupData(symbol=None, xcsetupname=None, readxml=False) 

93 setupdata.read_xml(read_maybe_unzipping(Path(path))) 

94 return setupdata 

95 

96 

97def get_basis_getter(args: Namespace) -> BasisGetter: 

98 if args.vconf_sharp_confinement: 

99 vconf_args = None 

100 else: 

101 vconf_args = args.vconf_amplitude, args.vconf_rstart_rel 

102 return methodcaller('generate', args.type.zetacount, args.type.polcount, 

103 tailnorm=args.tailnorm, 

104 energysplit=args.energy_shift, 

105 rcutpol_rel=args.rcut_pol_rel, 

106 rcutmax=args.rcut_max, 

107 rcharpol_rel=args.rchar_pol_rel, 

108 vconf_args=vconf_args, 

109 l_pol=args.lpol, 

110 jvalues=args.jvalues) 

111 

112 

113def main(args: Namespace) -> None: 

114 get_basis = get_basis_getter(args) 

115 tokens = [] 

116 if args.name: 

117 tokens.append(args.name) 

118 tokens += [args.type.name, 'basis'] 

119 for filename in args.file: 

120 print(f'Generating basis set for {filename!r}') 

121 if args.search: 

122 found_filename, _ = search_for_file(filename) 

123 print(f'Search result: {filename!r} -> {found_filename!r}') 

124 filename = found_filename 

125 setupdata, basis = read_setup_and_generate_basis(filename, get_basis) 

126 # Should the setupname be added as part of the name, too? 

127 # Probably not, since we don't include the xcname either. 

128 # But I suppose it depends more on the runtime behaviour when 

129 # GPAW actually picks setups/basis sets for a calculation. 

130 outputfile = '.'.join([setupdata.symbol, *tokens]) 

131 with open(outputfile, 'w') as fd: 

132 basis.write_to(fd) 

133 

134 

135class CLICommand: 

136 """Create basis sets from setup files.""" 

137 

138 @staticmethod 

139 def add_arguments(parser) -> None: 

140 add = parser.add_argument 

141 add('file', metavar='<filename>', nargs='+', help='setup data file') 

142 add('--name', 

143 metavar='<name>', 

144 help='basis name to be included in output filename') 

145 add('-s', '--search', 

146 action='store_true', 

147 help='instead of treating <filename> as paths, ' 

148 'search the installed datasets for them') 

149 add_common_args(add) 

150 

151 run = staticmethod(main)