Coverage for eminus/domains.py: 95.56%

1# SPDX-FileCopyrightText: 2022 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

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

4 

5import copy 

6import numbers 

7 

8from . import backend as xp 

9from .logger import log 

10from .tools import center_of_mass 

11 

12 

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

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

15 

16 Args: 

17 obj: Atoms or SCF object. 

18 length: Side length or lengths of the cuboid. 

19 

20 Keyword Args: 

21 centers: Center of the cuboid. 

22 

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

24 multiple domains will be merged. 

25 

26 Returns: 

27 Boolean mask. 

28 """ 

29 atoms = obj._atoms 

30 

31 if isinstance(length, numbers.Real): 

32 length = length * xp.ones(3) 

33 if centers is None: 

34 centers = center_of_mass(atoms.pos) 

35 centers = xp.asarray(centers) 

36 # Handle each dimension separately and add them together 

37 if centers.ndim == 1: 

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

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

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

41 mask = mask1 & mask2 & mask3 

42 else: 

43 mask = xp.zeros(atoms.Ns, dtype=bool) 

44 for center in centers: 

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

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

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

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

49 return mask 

50 

51 

52def domain_isovalue(field, isovalue): 

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

54 

55 Args: 

56 field: Real-space field data. 

57 isovalue: Isovalue for the truncation. 

58 

59 Returns: 

60 Boolean mask. 

61 """ 

62 if field is None: 

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

64 return None 

65 return xp.abs(field) > isovalue 

66 

67 

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

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

70 

71 Args: 

72 obj: Atoms or SCF object. 

73 radius: Radius of the sphere. 

74 

75 Keyword Args: 

76 centers: Center of the sphere. 

77 

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

79 multiple domains will be merged. 

80 

81 Returns: 

82 Boolean mask. 

83 """ 

84 atoms = obj._atoms 

85 

86 if centers is None: 

87 centers = center_of_mass(atoms.pos) 

88 centers = xp.asarray(centers) 

89 if centers.ndim == 1: 

90 mask = xp.linalg.norm(centers - atoms.r, axis=1) < radius 

91 else: 

92 mask = xp.zeros(atoms.Ns, dtype=bool) 

93 for center in centers: 

94 mask_tmp = xp.linalg.norm(center - atoms.r, axis=1) < radius 

95 mask = mask | mask_tmp 

96 return mask 

97 

98 

99def truncate(field, mask): 

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

101 

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

103 

104 Args: 

105 field: Real-space field data. 

106 mask: Boolean mask. 

107 

108 Returns: 

109 Truncated field. 

110 """ 

111 field_trunc = copy.deepcopy(field) 

112 field_trunc[~mask] = 0 

113 return field_trunc