Coverage for eminus/domains.py: 95.56%

1# SPDX-FileCopyrightText: 2021 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""Functions to restrict real-space fields to domains.""" 

4 

5import numbers 

6 

7import numpy as np 

8from scipy.linalg import norm 

9 

10from .logger import log 

11from .tools import center_of_mass 

12 

13 

14def domain_cuboid(obj, length, centers=None): 

15 """Generate a mask for a cuboidal real-space domain. 

16 

17 Args: 

18 obj: Atoms or SCF object. 

19 length: Side length or lengths of the cuboid. 

20 

21 Keyword Args: 

22 centers: Center of the cuboid. 

23 

24 Defaults to the geometric center of mass of the system. For multiple coordinates, 

25 multiple domains will be merged. 

26 

27 Returns: 

28 Boolean mask. 

29 """ 

30 atoms = obj._atoms 

31 

32 if isinstance(length, numbers.Real): 

33 length = length * np.ones(3) 

34 if centers is None: 

35 centers = center_of_mass(atoms.pos) 

36 centers = np.asarray(centers) 

37 # Handle each dimension separately and add them together 

38 if centers.ndim == 1: 

39 mask1 = np.abs(centers[0] - atoms.r[:, 0]) < length[0] 

40 mask2 = np.abs(centers[1] - atoms.r[:, 1]) < length[1] 

41 mask3 = np.abs(centers[2] - atoms.r[:, 2]) < length[2] 

42 mask = mask1 & mask2 & mask3 

43 else: 

44 mask = np.zeros(atoms.Ns, dtype=bool) 

45 for center in centers: 

46 mask1 = np.abs(center[0] - atoms.r[:, 0]) < length[0] 

47 mask2 = np.abs(center[1] - atoms.r[:, 1]) < length[1] 

48 mask3 = np.abs(center[2] - atoms.r[:, 2]) < length[2] 

49 mask = mask | (mask1 & mask2 & mask3) 

50 return mask 

51 

52 

53def domain_isovalue(field, isovalue): 

54 """Generate a mask for an isovalue real-space domain. 

55 

56 Args: 

57 field: Real-space field data. 

58 isovalue: Isovalue for the truncation. 

59 

60 Returns: 

61 Boolean mask. 

62 """ 

63 if field is None: 

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

65 return None 

66 return np.abs(field) > isovalue 

67 

68 

69def domain_sphere(obj, radius, centers=None): 

70 """Generate a mask for a spherical real-space domain. 

71 

72 Args: 

73 obj: Atoms or SCF object. 

74 radius: Radius of the sphere. 

75 

76 Keyword Args: 

77 centers: Center of the sphere. 

78 

79 Defaults to the geometric center of mass of the system. For multiple coordinates, 

80 multiple domains will be merged. 

81 

82 Returns: 

83 Boolean mask. 

84 """ 

85 atoms = obj._atoms 

86 

87 if centers is None: 

88 centers = center_of_mass(atoms.pos) 

89 centers = np.asarray(centers) 

90 if centers.ndim == 1: 

91 mask = norm(centers - atoms.r, axis=1) < radius 

92 else: 

93 mask = np.zeros(atoms.Ns, dtype=bool) 

94 for center in centers: 

95 mask_tmp = norm(center - atoms.r, axis=1) < radius 

96 mask = mask | mask_tmp 

97 return mask 

98 

99 

100def truncate(field, mask): 

101 """Truncate field data for a given mask. 

102 

103 This will not return a smaller array but set all truncated values to zero. 

104 

105 Args: 

106 field: Real-space field data. 

107 mask: Boolean mask. 

108 

109 Returns: 

110 Truncated field. 

111 """ 

112 field_trunc = np.copy(field) 

113 field_trunc[~mask] = 0 

114 return field_trunc