Coverage for eminus/io/cube.py: 95.24%

1# SPDX-FileCopyrightText: 2023 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""CUBE file handling.""" 

4 

5import textwrap 

6import time 

7 

8import numpy as np 

9 

10from eminus import backend as xp 

11from eminus.data import NUMBER2SYMBOL, SYMBOL2NUMBER 

12from eminus.logger import log 

13from eminus.version import __version__ 

14 

15 

16def read_cube(filename): 

17 """Load atom and cell data from CUBE files. 

18 

19 There is no standard for CUBE files. The following format has been used. 

20 File format definition: https://h5cube-spec.readthedocs.io/en/latest/cubeformat.html 

21 

22 Args: 

23 filename: CUBE input file path/name. 

24 

25 Returns: 

26 Species, positions, charges, cell size, sampling, and field array. 

27 """ 

28 if not filename.endswith((".cub", ".cube")): 

29 filename += ".cube" 

30 

31 # Atomic units and a cell that starts at (0,0,0) are assumed. 

32 with open(filename, encoding="utf-8") as fh: 

33 lines = fh.readlines() 

34 

35 # The first and second lines can contain comments, print them if available 

36 comment = f"{lines[0].strip()}\n{lines[1].strip()}" 

37 if comment: 

38 log.info(f'CUBE file comment: "{comment}"') 

39 

40 # Lines 4 to 6 contain the sampling per axis and the cell basis vectors with length a/s 

41 s = np.empty(3, dtype=int) 

42 a = np.empty((3, 3)) 

43 for i, line in enumerate(lines[3:6]): 

44 line_split = line.strip().split() 

45 s[i] = line_split[0] 

46 a[i] = s[i] * np.asarray(line_split[1:], dtype=float) 

47 a, s = xp.asarray(a), xp.asarray(s) 

48 

49 atom = [] 

50 pos = [] 

51 Z = [] 

52 # Following lines contain atom positions with the format: atom-id charge x-pos y-pos z-pos 

53 offset = 0 

54 for line in lines[6:]: 

55 line_split = line.strip().split() 

56 # If the first value is not a (positive) integer, we have reached the field data 

57 if not line_split[0].isdigit(): 

58 break 

59 atom.append(NUMBER2SYMBOL[int(line_split[0])]) 

60 Z.append(float(line_split[1])) 

61 pos.append(np.asarray(line_split[2:5], dtype=float)) 

62 offset += 1 

63 pos = xp.asarray(np.asarray(pos)) 

64 

65 # The rest of the data is the field data 

66 # Split the strings, flatten the lists of lists, and convert to a float numpy array 

67 tmp_list = [l.split() for l in lines[6 + offset :]] 

68 field_list = [item for sublist in tmp_list for item in sublist] 

69 field = xp.asarray(np.asarray(field_list, dtype=float)) 

70 return atom, pos, Z, a, s, field 

71 

72 

73def write_cube(obj, filename, field, fods=None, elec_symbols=("X", "He")): 

74 """Generate CUBE files from given field data. 

75 

76 There is no standard for CUBE files. The following format has been used to work with VESTA. 

77 File format definition: https://h5cube-spec.readthedocs.io/en/latest/cubeformat.html 

78 

79 Args: 

80 obj: Atoms or SCF object. 

81 filename: CUBE output file path/name. 

82 field: Real-space field data. 

83 

84 Keyword Args: 

85 fods: FOD coordinates to write. 

86 elec_symbols: Identifier for up and down FODs. 

87 """ 

88 # Atomic units are assumed, so there is no need for conversion. 

89 atoms = obj._atoms 

90 

91 if not filename.endswith((".cub", ".cube")): 

92 filename += ".cube" 

93 

94 if "He" in atoms.atom and atoms.unrestricted: 

95 log.warning( 

96 'You need to modify "elec_symbols" to write helium with FODs in the spin-' 

97 "polarized case." 

98 ) 

99 

100 # Make sure we have real-valued data in the correct order 

101 if field is None: 

102 log.warning('The provided field is "None".') 

103 return 

104 field = xp.real(field) 

105 

106 with open(filename, "w", encoding="utf-8") as fp: 

107 # The first two lines have to be a comment 

108 # Print information about the file and program, and the file creation time 

109 fp.write(f"File generated with eminus {__version__} on {time.ctime()}\n\n") 

110 # Number of atoms (int), and origin of the coordinate system (float) 

111 # The origin is normally at 0,0,0 but we could move our box, so take the minimum 

112 if fods is None: 

113 fp.write(f"{atoms.Natoms} ") 

114 else: 

115 fp.write(f"{atoms.Natoms + sum(len(i) for i in fods)} ") 

116 fp.write("0.0 0.0 0.0\n") 

117 # Number of points per axis (int), and vector defining the axis (float) 

118 fp.write( 

119 f"{atoms.s[0]} {atoms.a[0, 0] / atoms.s[0]:.6f} {atoms.a[0, 1] / atoms.s[0]:.6f}" 

120 f" {atoms.a[0, 2] / atoms.s[0]:.6f}\n" 

121 f"{atoms.s[1]} {atoms.a[1, 0] / atoms.s[1]:.6f} {atoms.a[1, 1] / atoms.s[1]:.6f}" 

122 f" {atoms.a[1, 2] / atoms.s[1]:.6f}\n" 

123 f"{atoms.s[2]} {atoms.a[2, 0] / atoms.s[2]:.6f} {atoms.a[2, 1] / atoms.s[2]:.6f}" 

124 f" {atoms.a[2, 2] / atoms.s[2]:.6f}\n" 

125 ) 

126 # Atomic number (int), atomic charge (float), and atom position (floats) for every atom 

127 fp.writelines( 

128 f"{SYMBOL2NUMBER[atoms.atom[ia]]} {atoms.Z[ia]:.3f} " 

129 f"{atoms.pos[ia, 0]: .6f} {atoms.pos[ia, 1]: .6f} {atoms.pos[ia, 2]: .6f}\n" 

130 for ia in range(atoms.Natoms) 

131 ) 

132 if fods is not None: 

133 for s in range(len(fods)): 

134 fp.writelines( 

135 f"{SYMBOL2NUMBER[elec_symbols[s]]} 0.000 " 

136 f"{ie[0]: .6f} {ie[1]: .6f} {ie[2]: .6f}\n" 

137 for ie in fods[s] 

138 ) 

139 # Field data (float) with scientific formatting 

140 # We have s[0]*s[1] chunks values with a length of s[2] 

141 for i in range(atoms.s[0] * atoms.s[1]): 

142 # Print every round of values, so we can add empty lines between them 

143 data_str = "%+1.6e " * atoms.s[2] % tuple(field[i * atoms.s[2] : (i + 1) * atoms.s[2]]) 

144 # Print a maximum of 6 values per row 

145 # Max width for this formatting is 90, since 6*len("+1.00000e-000 ")=90 

146 # Setting break_on_hyphens to False greatly improves the textwrap.fill performance 

147 fp.write(f"{textwrap.fill(data_str, width=90, break_on_hyphens=False)}\n\n")