Coverage for eminus/cell.py: 100.00%

1# SPDX-FileCopyrightText: 2023 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""Cell wrapper function.""" 

4 

5from . import backend as xp 

6from .atoms import Atoms 

7from .data import LATTICE_VECTORS 

8from .utils import molecule2list 

9 

10#: Crystal structures with their respective lattice and basis. 

11STRUCTURES = { 

12 "sc": { 

13 "lattice": "sc", 

14 "basis": [[0, 0, 0]], 

15 }, 

16 "fcc": { 

17 "lattice": "fcc", 

18 "basis": [[0, 0, 0]], 

19 }, 

20 "bcc": { 

21 "lattice": "bcc", 

22 "basis": [[0, 0, 0]], 

23 }, 

24 "tetragonal": { 

25 "lattice": "sc", 

26 "basis": [[0, 0, 0]], 

27 }, 

28 "orthorhombic": { 

29 "lattice": "sc", 

30 "basis": [[0, 0, 0]], 

31 }, 

32 "hexagonal": { 

33 "lattice": "hexagonal", 

34 "basis": [[0, 0, 0]], 

35 }, 

36 "diamond": { 

37 "lattice": "fcc", 

38 "basis": [[0, 0, 0], [1 / 4, 1 / 4, 1 / 4]], 

39 }, 

40 "zincblende": { 

41 "lattice": "fcc", 

42 "basis": [[0, 0, 0], [1 / 4, 1 / 4, 1 / 4]], 

43 }, 

44 "rocksalt": { 

45 "lattice": "fcc", 

46 "basis": [[0, 0, 0], [1 / 2, 0, 0]], 

47 }, 

48 "cesiumchloride": { 

49 "lattice": "sc", 

50 "basis": [[0, 0, 0], [1 / 2, 1 / 2, 1 / 2]], 

51 }, 

52 "fluorite": { 

53 "lattice": "fcc", 

54 "basis": [[0, 0, 0], [1 / 4, 1 / 4, 1 / 4], [3 / 4, 3 / 4, 3 / 4]], 

55 }, 

56} 

57 

58 

59def Cell( 

60 atom, 

61 lattice, 

62 ecut, 

63 a, 

64 basis=None, 

65 bands=None, 

66 kmesh=1, 

67 smearing=0, 

68 magnetization=None, 

69 verbose=None, 

70 **kwargs, 

71): 

72 """Wrapper to create Atoms classes for crystal systems. 

73 

74 Args: 

75 atom: Atom symbols. 

76 lattice: Lattice system. 

77 ecut: Cut-off energy. 

78 a: Cell size. 

79 

80 Keyword Args: 

81 basis: Lattice basis. 

82 bands: Number of bands (has to be larger or equal than the occupied states). 

83 kmesh: k-point mesh. 

84 smearing: Smearing width in Hartree. 

85 magnetization: Initial magnetization. 

86 verbose: Level of output. 

87 **kwargs: Keyword arguments to pass to the Atoms object. 

88 

89 Returns: 

90 Atoms object. 

91 """ 

92 # Get the lattice vectors from a string or use them directly 

93 if isinstance(lattice, str): 

94 lattice = lattice.lower() 

95 if basis is None: 

96 basis = STRUCTURES[lattice]["basis"] 

97 lattice = STRUCTURES[lattice]["lattice"] 

98 lattice_vectors = xp.asarray(LATTICE_VECTORS[lattice], dtype=float) 

99 else: 

100 if basis is None: 

101 basis = [[0, 0, 0]] 

102 lattice_vectors = xp.asarray(lattice, dtype=float) 

103 

104 # Only scale the lattice vectors with if a is given 

105 if a is not None: 

106 a = xp.asarray(a, dtype=float) 

107 lattice_vectors = a * lattice_vectors 

108 basis = a * xp.atleast_2d(xp.asarray(basis, dtype=float)) 

109 

110 # Account for different atom and basis sizes 

111 if isinstance(atom, str): 

112 atom_list = molecule2list(atom) 

113 else: 

114 atom_list = atom 

115 if len(atom_list) != len(basis): 

116 atom = [atom] * len(basis) 

117 

118 # Build the atoms object 

119 atoms = Atoms(atom, basis, ecut=ecut, a=lattice_vectors, verbose=verbose, **kwargs) 

120 # Handle k-points and states 

121 atoms.kpts.kmesh = kmesh 

122 if isinstance(lattice, str): 

123 atoms.kpts.lattice = lattice 

124 atoms.occ.smearing = smearing 

125 if bands is not None: 

126 atoms.occ.bands = bands 

127 if magnetization is not None: 

128 atoms.occ.magnetization = magnetization 

129 return atoms