Coverage for gpaw/response/mpa_sampling.py: 98%
54 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-08 00:17 +0000
1from __future__ import annotations
2import numpy as np
3from gpaw.typing import Array1D
6def sampling_branches(w_dist: Array1D,
7 parallel_lines: int = 2,
8 varpi: float = 1,
9 eta0: float = 1e-5,
10 eta_rest: float = 0.1) -> Array1D:
11 """
12 w_dist Array of points in the positive real axis.
13 parallel_lines How many (1-2) parallel lines to the real frequency axis
14 the sampling has.
15 varpi Distance of the second line to the real axis.
16 eta0 Imaginary part of the first point of the first line.
17 eta_rest Imaginary part of the rest of the points of the first
18 line.
19 """
20 if parallel_lines not in [1, 2]:
21 raise ValueError('parallel_lines must be either 1 or 2.')
23 if len(w_dist) == 1:
24 assert eta0 >= 0
25 assert parallel_lines == 2
26 w_grid = np.concatenate([w_dist + 1j * eta0, w_dist + 1j * varpi])
27 return w_grid
29 if parallel_lines == 1: # only one branch
30 assert varpi >= 0
31 w_grid = w_dist + 1j * varpi
32 return w_grid
34 # parallel lines == 2
35 assert eta0 >= 0
36 assert eta_rest >= 0
37 assert varpi > eta0 and varpi > eta_rest
38 w_grid = np.concatenate((np.array([w_dist[0] + 1j * eta0]),
39 w_dist[1:] + 1j * eta_rest, w_dist + 1j * varpi))
40 assert len(w_grid.shape) == 1
41 return w_grid
44def frequency_distribution(npoles: int,
45 wrange: tuple[float, float] | Array1D,
46 alpha: float = 1) -> Array1D:
48 if npoles == 1:
49 w_grid = np.array([0.])
50 return w_grid
52 assert wrange[0].real >= 0
53 if not wrange[1].real > wrange[0].real:
54 raise ValueError('Frequency range inverted.')
56 if alpha == 0: # homogeneous distribution
57 w_grid = np.linspace(wrange[0], wrange[1], 2 * npoles)
58 return w_grid
60 ws = wrange[1] - wrange[0]
61 w_grid = wrange[0] + ws * semi_homogenous_partition(npoles)**alpha
62 return w_grid
65def mpa_frequency_sampling(npoles: int,
66 wrange: tuple[float, float],
67 varpi: float,
68 eta0: float,
69 eta_rest: float,
70 parallel_lines: int = 2,
71 alpha: float = 1) -> Array1D:
72 """
73 This function creates a frequency grid in the complex plane.
74 The grid can have 1 or 2 branches with points non homogeneously
75 distributed along the real frequency axis.
76 See Fig. 1 and Eq. (18) of Ref. [1], and Eq. (11) of Ref. [2]
77 [1] DA. Leon et al, PRB 104, 115157 (2021)
78 [2] DA. Leon et al, PRB 107, 155130 (2023)
80 Parameters
81 ----------
82 npoles : numper of poles (half the number of frequency points)
83 wrange : [w_start, w_end]. An array of two complex numbers defining
84 the sampling range
85 varpi : immaginary part of the second branch
86 eta0 : damping factor for the first point of the first branch
87 eta_rest : damping for the rest of the points of the first branch
88 parallel_lines : Either 1 or 2, how many lines there are parallel
89 to the real axis.
90 alpha : exponent of the distribution of points along the real axis
91 ______________________________________________________________________
92 Example: double parallel sampling with 9 poles
93 ----------------------------------------------------------------------
94 complex frequency plane
95 ----------------------------------------------------------------------
96 | (wrange[0], varpi) ... . . . . . . (wrange[1], varpi) |
97 | |
98 | .. . . . . . . (wrange[1], eta_rest) |
99 | (wrange[0], eta0) . |
100 |____________________________________________________________________|
101 """
102 _wrange = np.array(wrange)
103 grid_p = frequency_distribution(npoles, _wrange, alpha)
104 grid_w = sampling_branches(grid_p, parallel_lines=parallel_lines,
105 varpi=varpi, eta0=eta0, eta_rest=eta_rest)
106 return grid_w
109def semi_homogenous_partition(npoles: int) -> Array1D:
110 """
111 Returns a semi-homogenous partition with n-poles between 0 and 1
112 according to Eq. (18) of Ref. [1], and Eq. (11) of Ref. [2]
113 [1] DA. Leon et al, PRB 104, 115157 (2021)
114 [2] DA. Leon et al, PRB 107, 155130 (2023)
115 """
116 small_cases = {1: np.array([0.0]),
117 2: np.array([0.0, 1.0]),
118 3: np.array([0.0, 0.5, 1.0])}
119 if npoles < 4:
120 return small_cases[npoles]
121 # Calculate the grid spacing
122 # Round up to the next power of 2. This will determine the minimum spacing
123 dw = 1 / 2**np.ceil(np.log2(npoles))
125 # Get the previous power of two, by searching down,
126 # e.g. lp(4) = 2, lp(7)=4, lp(8)=4, lp(9)=8
127 lp = 2**int(np.log2(npoles - 1))
129 # There are usually 2 kinds of intervals in a semi homogenous grid,
130 # they are always in order such that smallest intervals are closer to zero.
131 # The interval sizes are dw, 2*dw.
132 # There is an exception to this rule when npoles == power of two + 1,
133 # because then there would be only one type of interval with the rule
134 # below. To keep the grid inhomogenous, one adds the third interval 4 * dw.
135 if npoles == lp + 1:
136 np1 = 2
137 np3 = 1
138 else:
139 np1 = (npoles - (lp + 1)) * 2
140 np3 = 0
141 # The number of intervals is always one less, than the number of points
142 # in the grid. Therefore, We can deduce np2 from np1 and np3.
143 np2 = npoles - 1 - np1 - np3
145 # Build up the intervals onto an array
146 dw_n = np.repeat([0, 1, 2, 4], [1, np1, np2, np3])
148 # Sum over the intervals to build the point grid
149 w_grid = np.cumsum(dw_n) * dw
150 return w_grid