Coverage for gpaw/sphere/rshe.py: 100%
88 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-14 00:18 +0000
1import numpy as np
3from gpaw.atom.radialgd import RadialGridDescriptor
4from gpaw.sphere.integrate import integrate_lebedev
7class RealSphericalHarmonicsExpansion:
8 """Expansion in real spherical harmonics of a function f(r)."""
10 def __init__(self,
11 rgd: RadialGridDescriptor,
12 f_gM, Y_nL, L_M=None):
13 """Construct the expansion
15 Parameters
16 ----------
17 f_gM : np.array
18 f as a function of radial index g and reduced spherical harmonic
19 index M.
20 Y_nL : np.array
21 Real spherical harmonics on the angular Lebedev quadrature as a
22 function of the composite spherical harmonics index L=(l,m).
23 L_M : np.array
24 L index for every reduced expansion index M.
25 """
26 self.rgd = rgd
27 self.f_gM = f_gM
28 self.Y_nL = Y_nL
30 if L_M is None:
31 # Assume that all the composite indices L=(l,m) are represented
32 assert f_gM.shape[1] == self.nL
33 L_M = np.arange(self.nL)
34 self.L_M = L_M
36 @classmethod
37 def from_spherical_grid(cls, rgd, f_ng, Y_nL):
38 r"""Expand the function f(r) in real spherical harmonics.
40 / ^ ^ ^
41 f (r) = |dr Y (r) f(rr)
42 lm / lm
44 Note that the Lebedev quadrature, which is used to perform the angular
45 integral above, is exact up to polynomial order l=11. This implies that
46 expansion coefficients up to l=5 are exact.
48 Parameters
49 ----------
50 f_ng : np.array
51 f as a function of angular index n (on the Lebedev quadrature) and
52 radial index g.
53 Y_nL : np.array
54 Real spherical harmonics on the angular Lebedev quadrature as a
55 function of the composite spherical harmonics index L=(l,m).
56 """
57 # Include coefficients up to l = 5, where nL = (l + 1)**2
58 nL = min(Y_nL.shape[1], 36)
60 # Integrate Y_lm(r) * f(r) on the angular grid
61 f_gL = integrate_lebedev(
62 Y_nL[:, np.newaxis, :nL] * f_ng[..., np.newaxis])
64 return cls(rgd, f_gL, Y_nL)
66 def reduce_expansion(self, L_M):
67 """
68 Produce a new expansion with only the spherical harmonic indices L_M.
69 """
70 # Translate requested indices L_M to the internal index M
71 M_M = []
72 for L in L_M:
73 lookup = np.where(self.L_M == L)[0]
74 assert len(lookup) == 1
75 M_M.append(lookup[0])
77 return RealSphericalHarmonicsExpansion(
78 self.rgd, self.f_gM[:, M_M], self.Y_nL, L_M=L_M)
80 @property
81 def nL(self):
82 return self.Y_nL.shape[1]
84 @property
85 def nM(self):
86 return len(self.L_M)
88 @property
89 def lmax(self):
90 flmax = np.sqrt(self.nL)
91 lmax = int(flmax)
92 assert abs(flmax - lmax) < 1e-8
93 return lmax
95 @property
96 def l_L(self):
97 l_L = []
98 for l in range(self.lmax + 1):
99 l_L += [l] * (2 * l + 1)
100 return l_L
102 @property
103 def l_M(self):
104 return [self.l_L[L] for L in self.L_M]
106 def evaluate_on_quadrature(self):
107 """Evaluate the function f(r) on the angular Lebedev quadrature."""
108 Y_nM = self.Y_nL[:, self.L_M]
109 return Y_nM @ self.f_gM.T
112def calculate_reduced_rshe(rgd, f_ng, Y_nL, lmax=-1, wmin=None):
113 """Expand a function f(r) in real spherical harmonics with a reduced number
114 of expansion coefficients."""
115 rshe = RealSphericalHarmonicsExpansion.from_spherical_grid(rgd, f_ng, Y_nL)
116 L_M, info_string = assess_rshe_reduction(f_ng, rshe, lmax=lmax, wmin=wmin)
117 rshe = rshe.reduce_expansion(L_M)
118 return rshe, info_string
121def assess_rshe_reduction(f_ng, rshe, lmax=-1, wmin=None):
122 """Assess how to reduce the number of expansion coefficients.
124 The composite index L=(l,m) is reduced to an index M, which iterates the
125 expansion coefficients which contribute with a weight larger than wmin to
126 the surface norm square of the function f(r) on average. The M index is
127 further restricted to include coefficients only up to lmax.
128 """
129 # We do not expand beyond l=5
130 if lmax == -1:
131 lmax = 5
132 assert lmax in range(6)
134 # We assume to start with a full expansion
135 assert rshe.nM == rshe.nL
136 f_gL = rshe.f_gM
138 # Filter away (l,m)-coefficients based on their average weight in
139 # completing the surface norm square f(r)
140 fsns_g = integrate_lebedev(f_ng ** 2) # surface norm square
141 mask_g = fsns_g > 1e-12 # Base filter on finite surface norm squares only
142 fw_gL = f_gL[mask_g] ** 2 / fsns_g[mask_g, np.newaxis] # weight of each L
143 rshew_L = np.average(fw_gL, axis=0) # Average over the radial grid
145 # Take rshe coefficients up to l <= lmax (<= 5) which contribute with
146 # at least wmin to the surface norm square on average
147 nL = min(rshe.nL, (lmax + 1)**2)
148 L_L = np.arange(nL)
149 if wmin is not None:
150 assert isinstance(wmin, float) and wmin > 0.
151 L_M = np.where(rshew_L[L_L] >= wmin)[0]
152 else:
153 L_M = L_L
155 info_string = get_reduction_info_string(L_M, fw_gL, rshew_L)
157 return L_M, info_string
160def get_reduction_info_string(L_M, fw_gL, rshew_L):
161 """Construct info string about the reduced expansion."""
162 isl = []
163 isl.append('{:6} {:10} {:10} {:8}'.format('(l,m)', 'max weight',
164 'avg weight', 'included'))
165 for L, (fw_g, rshew) in enumerate(zip(fw_gL.T, rshew_L)):
166 included = L in L_M
167 isl.append('\n' + get_rshe_coefficient_info_string(
168 L, included, rshew, fw_g))
170 avg_cov = np.average(np.sum(fw_gL[:, L_M], axis=1))
171 isl.append(f'\nIn total: {avg_cov} of the surface norm square is '
172 'covered on average')
174 tot_avg_cov = np.average(np.sum(fw_gL, axis=1))
175 isl.append(f'\nIn total: {tot_avg_cov} of the surface norm '
176 'square could be covered on average')
178 return ''.join(isl)
181def get_rshe_coefficient_info_string(L, included, rshew, fw_g):
182 """Construct info string about the weight of a given coefficient."""
183 l = int(np.sqrt(L))
184 m = L - l * (l + 1)
185 included = 'yes' if included else 'no'
186 info_string = '{:6} {:1.8f} {:1.8f} {:8}'.format(f'({l},{m})',
187 np.max(fw_g),
188 rshew, included)
189 return info_string