Coverage for gpaw/domain.py: 55%
110 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-14 00:18 +0000
1# Copyright (C) 2003 CAMP
2# Please see the accompanying LICENSE file for further information.
4"""Domain object.
6This module contains the definition of the ``Domain`` class and some
7helper functins for parallel domain decomposition. """
9import numpy as np
11UNIFORM = False # distribute grid points uniformly
14class Domain:
15 """Domain class.
17 A ``Domain`` object (in domain.py) holds informaion on the unit
18 cell and the boundary conditions"""
20 def __init__(self, cell, pbc, comm, parsize_c, N_c):
21 """Create Domain object from a unit cell and boundary conditions.
23 The arguments are the lengths of the three axes, followed by a
24 tuple of three periodic-boundary condition flags (``bool``s).
26 Parallel stuff:
27 =============== ==================================================
28 ``sdisp_cd`` Scaled displacement in direction ``d`` along axis
29 ``c``.
30 ``comm`` MPI-communicator for this domain.
31 ``neighbor_cd`` Rank of neighbor CPU in direction ``d`` along axis
32 ``c``.
33 ``parpos_c`` Position of this CPU in the 3D grid of all CPUs.
34 ``parsize_c`` Domain decomposition along the three axes.
35 ``stride_c`` Strides.
36 =============== ==================================================
37 """
38 cell_c = np.array(cell, float)
39 if cell_c.ndim == 1:
40 self.cell_cv = np.diag(cell_c)
41 else:
42 self.cell_cv = cell_c
43 cell_c = (self.cell_cv**2).sum(1)**0.5
44 self.icell_cv = np.linalg.inv(self.cell_cv).T
45 self.xxxucell_cv = np.array([self.cell_cv[x] / cell_c[x]
46 for x in range(3)])
47 self.xxxiucell_cv = np.linalg.inv(self.xxxucell_cv.T) # Jacobian
49 self.pbc_c = np.asarray(pbc, bool)
51 if isinstance(parsize_c, int):
52 assert parsize_c == comm.size
53 parsize_c = None
55 self.set_decomposition(comm, parsize_c, N_c)
57 def set_decomposition(self, comm, parsize_c=None, N_c=None):
58 """Set MPI-communicator and do domain decomposition.
60 With ``parsize_c=(a, b, c)``, the domin will be divided in
61 ``a*b*c`` sub-domains - one for each CPU in ``comm``. If
62 ``parsize_c`` is not given, the number of grid points will be
63 used to suggest a good domain decomposition."""
65 self.comm = comm
67 if parsize_c is None:
68 # XXX This should rather be:
69 # parsize_c = decompose_domain(N_c - not self.pbc_c, comm.size)
70 parsize_c = decompose_domain(N_c, comm.size)
71 self.parsize_c = np.array(parsize_c)
73 self.stride_c = np.array([parsize_c[1] * parsize_c[2],
74 parsize_c[2],
75 1])
77 if np.prod(self.parsize_c) != self.comm.size:
78 raise RuntimeError('Bad domain decomposition! '
79 'CPU counts %s do not multiply to '
80 'communicator size %d' % (self.parsize_c,
81 self.comm.size))
83 self.parpos_c = self.get_processor_position_from_rank()
84 self.find_neighbor_processors()
86 def get_ranks_from_positions(self, spos_ac):
87 """Calculate rank of domain containing scaled position."""
88 rnk_ac = np.floor(spos_ac * self.parsize_c).astype(int)
89 if (rnk_ac < 0).any() or (rnk_ac >= self.parsize_c).any():
90 raise ValueError('Some atom is too close to the zero-boundary!')
91 return np.dot(rnk_ac, self.stride_c)
93 def get_rank_from_position(self, spos_c):
94 return self.get_ranks_from_positions(np.array([spos_c]))[0]
96 def get_rank_from_processor_position(self, parpos_c):
97 return parpos_c[2] + self.parsize_c[2] * \
98 (parpos_c[1] + self.parsize_c[1] * parpos_c[0])
100 def get_processor_position_from_rank(self, rank=None):
101 """Calculate position of a domain in the 3D grid of all domains."""
102 if rank is None:
103 rank = self.comm.rank
104 parpos_c = np.array([rank // self.stride_c[0],
105 (rank % self.stride_c[0]) // self.stride_c[1],
106 rank % self.stride_c[1]])
107 assert np.dot(parpos_c, self.stride_c) == rank
108 return parpos_c
110 def find_neighbor_processors(self):
111 """Find neighbor processors - surprise!
113 With ``i`` and ``d`` being the indices for the axis (x, y or
114 z) and direction + or - (0 or 1), two attributes are
115 calculated:
117 * ``neighbor_cd``: Rank of neighbor.
118 * ``sdisp_cd``: Scaled displacement for neighbor.
119 """
121 self.neighbor_cd = np.zeros((3, 2), int)
122 self.sdisp_cd = np.zeros((3, 2), int)
123 for c in range(3):
124 p = self.parpos_c[c]
125 for d in range(2):
126 sdisp = 2 * d - 1
127 pd = p + sdisp
128 pd0 = pd % self.parsize_c[c]
129 self.neighbor_cd[c, d] = (self.comm.rank +
130 (pd0 - p) * self.stride_c[c])
131 if pd0 != pd:
132 # Wrap around the box?
133 if self.pbc_c[c]:
134 # Yes:
135 self.sdisp_cd[c, d] = -sdisp
136 else:
137 # No:
138 self.neighbor_cd[c, d] = -1
141def decompose_domain(ngpts_c, ncores):
142 """Determine parsize based on grid size and number of processors.
144 If global variable UNIFORM is True, the returned parsize will
145 result in a uniform distribution of gridpoints amongst processors.
147 If UNIFORM is False, the returned parsize may result in a variable
148 number of points on each cpu.
149 """
150 if ncores == 1:
151 return (1, 1, 1)
152 n1, n2, n3 = ngpts_c
153 plist = prims(ncores)
154 pdict = {}
155 for n in plist:
156 pdict[n] = 0
157 for n in plist:
158 pdict[n] += 1
159 candidates = factorizations(list(pdict.items()))
160 mincost = 10000000000.0
161 best = None
162 for p1, p2, p3 in candidates:
163 if UNIFORM and (n1 % p1 != 0 or n2 % p2 != 0 or n3 % p3 != 0):
164 continue
165 elif UNIFORM:
166 m1 = n1 // p1
167 m2 = n2 // p2
168 m3 = n3 // p3
169 else:
170 m1 = float(n1) / p1
171 m2 = float(n2) / p2
172 m3 = float(n3) / p3
173 cost = abs(m1 - m2) + abs(m2 - m3) + abs(m3 - m1)
174 # A long z-axis (unit stride) is best:
175 if m1 <= m2 <= m3:
176 cost -= 0.1
177 if cost < mincost:
178 mincost = cost
179 best = (p1, p2, p3)
180 if best is None:
181 raise RuntimeError("Can't decompose a %dx%dx%d grid on %d CPUs!" %
182 (n1, n2, n3, ncores))
183 return best
186def factorizations(f):
187 p, n = f[0]
188 facs0 = []
189 for n1 in range(n + 1):
190 for n2 in range(n - n1 + 1):
191 n3 = n - n1 - n2
192 facs0.append((p**n1, p**n2, p**n3))
193 if len(f) == 1:
194 return facs0
195 else:
196 facs = factorizations(f[1:])
197 facs1 = []
198 for p1, p2, p3 in facs0:
199 facs1 += [(p1 * q1, p2 * q2, p3 * q3) for q1, q2, q3 in facs]
200 return facs1
203primes = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53,
204 59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113,
205 127, 131, 137, 139, 149, 151, 157, 163, 167, 173, 179, 181,
206 191, 193, 197, 199]
209def prims(p):
210 if p == 1:
211 return []
212 for n in primes:
213 if p % n == 0:
214 return prims(p / n) + [n]
215 raise RuntimeError