9b737b8GPU implementation of fmm3dbie's interior Helmholtz impedance solver, with a showcase room problemJeremy Magland 1"""CUDA kernels (compiled at run time with NVRTC through CuPy) for the
2Helmholtz layer potentials on quadrilateral patch surfaces.
4Six kernels are handled together, in the order used by fmm3dbie's
5helm_rpcomb_imp solver:
7 0: S'_k 1: S_kappa 2: S'_kappa 3: D'_k - D'_kappa 4: S_k 5: D_k
9with kappa = i|k|, G_k = exp(ikr)/(4 pi r), S' the normal derivative at the
10target, D the normal derivative at the source.
12Each layer potential at a target x is split by patch. A patch j is *near*
13x when |x - c_j| <= rfac r_j (c_j, r_j the patch centroid and radius); this
14test is the same device function in every kernel, so the split is
15consistent. Far patches are summed with the smooth (oversampled Gauss-
16Legendre) rule in dense_apply; near patches use precomputed weights
17(near_weights) applied in sparse_apply. Because near patches never enter the
18dense sum, the dense sum has no large cancelling terms and can be evaluated
19in single precision.
21Near weights are computed as Legendre moments I_n = int K P_n |J| du dv,
22converted to node weights with the vals-to-coefs matrix U: w_j = sum_n U[n, j] I_n.
24 * self patch (target on the patch): the square is split into four
25 triangles fanned from the target to the edges. On each triangle we use
26 polar-type coordinates (rho along the ray, x the position on the edge)
27 and the substitution x = x_c + h_eff sinh(xi) chosen from the first
28 fundamental form at the target, which removes the near-singular
29 behaviour in the angular direction. Tensor Gauss-Legendre in (rho, xi),
30 with the xi range split so that each piece is short relative to the
31 singularity and to the wavelength.
32 * other near patches: a quadtree on the reference square, refined while the
33 target is within rfac_q times the radius of a leaf or the leaf is large
34 compared with the wavelength, with a tensor Gauss-Legendre rule on each
35 leaf (fmm3dbie's istrat=1 strategy).
37The geometry inside the integrals is either the polynomial interpolant of
38the node data (gmode 0, as in fmm3dbie) or, for the rounded cube, the exact
39superellipsoid parametrization (gmode 1). The exact geometry is much more
40accurate at nodes next to patch edges: the double layer kernel is very
41sensitive to small inconsistencies between interpolated positions and
42normals.
43"""
45import cupy as cp
46import numpy as np
48_SRC = r"""
49#include <cupy/complex.cuh>
50typedef complex<double> cdouble;
52#define NORD1_MAX 13
53#define NPOLS_MAX 169
54#define NT 128
55#define NSUB_MAX 32
56#define OVER4PI 0.07957747154594767
58__device__ inline void legendre(double x, int norder, double* P) {
59 P[0] = 1.0;
60 if (norder >= 1) P[1] = x;
61 for (int n = 1; n < norder; n++)
62 P[n + 1] = ((2 * n + 1) * x * P[n] - n * P[n - 1]) / (n + 1);
63}
65__device__ inline double ipow(double x, int n) {
66 double r = 1.0;
67 for (int i = 0; i < n; i++) r *= x;
68 return r;
69}
71// xyz, r_u, r_v at (u, v) of one patch.
72// gmode 0: from the Legendre coefficients c[comp * npols + n] of the patch.
73// gmode 1: superellipsoid face chart; gp = n(3), e1(3), e2(3), s0, s1, t0, t1.
74__device__ inline void geometry(int gmode, const double* coefs, int norder, const double* gpar,
75 double ga, int gpe, int patch, double u, double v, double* g) {
76 if (gmode == 1) {
77 const double* gp = gpar + 13 * (size_t)patch;
78 double hs = 0.5 * (gp[10] - gp[9]), ht = 0.5 * (gp[12] - gp[11]);
79 double s = gp[9] + (u + 1) * hs, t = gp[11] + (v + 1) * ht;
80 double c[3], S = 0.0;
81 for (int q = 0; q < 3; q++) { c[q] = gp[q] + s * gp[3 + q] + t * gp[6 + q]; S += ipow(c[q], gpe); }
82 double N = pow(S, 1.0 / gpe);
83 double dNu = 0.0, dNv = 0.0;
84 for (int q = 0; q < 3; q++) {
85 double gN = ipow(c[q] / N, gpe - 1);
86 dNu += gN * gp[3 + q]; dNv += gN * gp[6 + q];
87 }
88 for (int q = 0; q < 3; q++) {
89 g[q] = ga * c[q] / N;
90 g[3 + q] = ga * hs * (gp[3 + q] / N - c[q] * dNu / (N * N));
91 g[6 + q] = ga * ht * (gp[6 + q] / N - c[q] * dNv / (N * N));
92 }
93 return;
94 }
95 int n1 = norder + 1, npols = n1 * n1;
96 const double* c = coefs + (size_t)patch * 9 * npols;
97 double Lu[NORD1_MAX], Lv[NORD1_MAX];
98 legendre(u, norder, Lu);
99 legendre(v, norder, Lv);
100 for (int comp = 0; comp < 9; comp++) {
101 double s = 0.0;
102 const double* cc = c + comp * npols;
103 for (int b = 0; b < n1; b++) {
104 double sb = 0.0;
105 for (int a = 0; a < n1; a++) sb += cc[b * n1 + a] * Lu[a];
106 s += sb * Lv[b];
107 }
108 g[comp] = s;
109 }
110}
112// The near test, shared by every kernel so that the near/far split is consistent.
113__device__ inline bool is_near(const double* x, const double* cen, double r2) {
114 double dx = x[0] - cen[0], dy = x[1] - cen[1], dz = x[2] - cen[2];
115 return dx * dx + dy * dy + dz * dz <= r2;
116}
118// The six kernels at one (target, source) pair, times a real weight.
119template <typename T>
120__device__ inline void kernels6(const T* xt, const T* nt, const T* y, const T* ns, T w,
121 complex<T> zk, T zkabs, int selmask, complex<T>* K) {
122 T dx = xt[0] - y[0], dy = xt[1] - y[1], dz = xt[2] - y[2];
123 T r2 = dx * dx + dy * dy + dz * dz;
124 T r = sqrt(r2);
125 T rinv = (T)1 / r;
126 T r2inv = rinv * rinv;
127 T dnt = dx * nt[0] + dy * nt[1] + dz * nt[2];
128 T dns = dx * ns[0] + dy * ns[1] + dz * ns[2];
129 complex<T> izkr = complex<T>(0, 1) * zk * r;
130 complex<T> Ek(0, 0);
131 T Ec = 0;
132 // selmask: 2 bits per kernel m at bits 2m, 2m+1
133 if (selmask & 0xFC3) Ek = exp(izkr) * ((T)(OVER4PI) * rinv * w); // kernels 0, 3, 4, 5
134 if (selmask & 0x0FC) Ec = exp(-zkabs * r) * ((T)(OVER4PI) * rinv * w); // kernels 1, 2, 3
135 complex<T> fk = (T)1 - izkr;
136 T fc = (T)1 + zkabs * r;
137 K[4] = Ek;
138 K[1] = complex<T>(Ec, 0);
139 K[0] = -dnt * fk * Ek * r2inv;
140 K[2] = complex<T>(-dnt * fc * Ec * r2inv, 0);
141 K[5] = dns * fk * Ek * r2inv;
142 if (selmask & 0x0C0) {
143 T nn = ns[0] * nt[0] + ns[1] * nt[1] + ns[2] * nt[2];
144 complex<T> gk = ((T)3 - (T)3 * izkr - zk * zk * r2);
145 T gc = (T)3 + (T)3 * zkabs * r + zkabs * zkabs * r2;
146 K[3] = (nn * fk * r2inv - dns * dnt * gk * r2inv * r2inv) * Ek
147 - complex<T>((nn * fc * r2inv - dns * dnt * gc * r2inv * r2inv) * Ec, 0);
148 } else {
149 K[3] = complex<T>(0, 0);
150 }
151}
153// ---------------------------------------------------------------------------
154// Near pairs: two passes (counts, then fill using row offsets), one thread per target.
155extern "C" __global__ void near_pairs(
156 int ntarg, const double* targ_x, int npatches, const double* cen, const double* r2,
157 int* counts, const int* row_ptr, int* pair_p)
158{
159 int i = blockIdx.x * blockDim.x + threadIdx.x;
160 if (i >= ntarg) return;
161 double x[3] = {targ_x[3 * i], targ_x[3 * i + 1], targ_x[3 * i + 2]};
162 int n = 0;
163 for (int j = 0; j < npatches; j++) {
164 if (is_near(x, cen + 3 * j, r2[j])) {
165 if (pair_p != NULL) pair_p[row_ptr[i] + n] = j;
166 n++;
167 }
168 }
169 counts[i] = n;
170}
172// ---------------------------------------------------------------------------
173// Quadtree leaves for the (target, patch) pairs that are not self pairs.
174// Two passes: counts only (leaves == NULL), then fill using offsets.
175extern "C" __global__ void build_leaves(
176 int npairs, const int* pair_t, const int* pair_p, const double* targ_x,
177 const int* targ_patch, int gmode, const double* coefs, int norder, const double* gpar,
178 double ga, int gpe, double rfac_q, int nlmax, double zkabs, double kR_max,
179 int* counts, const int* offsets, double* leaves)
180{
181 int p = blockIdx.x * blockDim.x + threadIdx.x;
182 if (p >= npairs) return;
183 int it = pair_t[p], ip = pair_p[p];
184 if (targ_patch[it] == ip) { counts[p] = 0; return; }
185 double tx = targ_x[3 * it], ty = targ_x[3 * it + 1], tz = targ_x[3 * it + 2];
186 double su[64], sv[64], sh[64];
187 int sd[64];
188 int sp = 1;
189 su[0] = 0.0; sv[0] = 0.0; sh[0] = 1.0; sd[0] = 0;
190 int n = 0;
191 while (sp > 0) {
192 sp--;
193 double cu = su[sp], cv = sv[sp], h = sh[sp];
194 int d = sd[sp];
195 bool refine = false;
196 if (d < nlmax) {
197 // centroid and radius from a 3x3 sample grid (corners alone underestimate
198 // the extent of strongly curved patches)
199 double xc[9][3], cen[3] = {0, 0, 0}, g[9];
200 for (int k = 0; k < 9; k++) {
201 geometry(gmode, coefs, norder, gpar, ga, gpe, ip, cu + h * (k % 3 - 1), cv + h * (k / 3 - 1), g);
202 for (int q = 0; q < 3; q++) { xc[k][q] = g[q]; cen[q] += g[q] / 9.0; }
203 }
204 double R2 = 0.0;
205 for (int k = 0; k < 9; k++) {
206 double s = 0.0;
207 for (int q = 0; q < 3; q++) s += (xc[k][q] - cen[q]) * (xc[k][q] - cen[q]);
208 R2 = fmax(R2, s);
209 }
210 double d2 = (tx - cen[0]) * (tx - cen[0]) + (ty - cen[1]) * (ty - cen[1]) + (tz - cen[2]) * (tz - cen[2]);
211 // refine while the target is close relative to the leaf, or the leaf is
212 // large compared with the wavelength (the leaf rule must resolve exp(ikr))
213 refine = (d2 < rfac_q * rfac_q * R2) || (zkabs * zkabs * R2 > kR_max * kR_max);
214 }
215 if (refine) {
216 double hh = 0.5 * h;
217 const double cs[4][2] = {{-1, -1}, {1, -1}, {1, 1}, {-1, 1}};
218 for (int k = 0; k < 4; k++) {
219 su[sp] = cu + hh * cs[k][0]; sv[sp] = cv + hh * cs[k][1]; sh[sp] = hh; sd[sp] = d + 1; sp++;
220 }
221 } else {
222 if (leaves != NULL) {
223 size_t o = (size_t)(offsets[p] + n) * 3;
224 leaves[o] = cu; leaves[o + 1] = cv; leaves[o + 2] = h;
225 }
226 n++;
227 }
228 }
229 counts[p] = n;
230}
232// ---------------------------------------------------------------------------
233// Near-field weights: one block per (target, patch) pair.
234struct SelfItem { int edge; double xi0, xi1; };
236extern "C" __global__ void near_weights(
237 int npairs, const int* pair_t, const int* pair_p,
238 const double* targ_x, const double* targ_n, const int* targ_patch, const double* targ_uv,
239 int gmode, const double* coefs, int norder, const double* gpar, double ga, int gpe,
240 const int* leaf_ptr, const double* leaves, int nq, const double* gl_x, const double* gl_w,
241 int nqr, const double* glr_x, const double* glr_w, int nqx, const double* glx_x, const double* glx_w,
242 double sub_len, double osc_len, const double* umat, cdouble zk, double zkabs, cdouble* W)
243{
244 __shared__ double sKr[6][NT], sKi[6][NT];
245 __shared__ double sLu[NORD1_MAX][NT], sLv[NORD1_MAX][NT];
246 __shared__ SelfItem sitems[4 * NSUB_MAX];
247 __shared__ double sEdge[4][3]; // per edge: h, x_c, h_eff
248 __shared__ int snitems;
250 int p = blockIdx.x;
251 if (p >= npairs) return;
252 int t = threadIdx.x;
253 int it = pair_t[p], ip = pair_p[p];
254 int n1 = norder + 1, npols = n1 * n1;
255 double xt[3] = {targ_x[3 * it], targ_x[3 * it + 1], targ_x[3 * it + 2]};
256 double nt[3] = {targ_n[3 * it], targ_n[3 * it + 1], targ_n[3 * it + 2]};
257 bool self = (targ_patch[it] == ip);
258 double u0 = 0, v0 = 0;
260 // per-thread moment accumulators for basis indices n = t, t + NT
261 const int NACC = (NPOLS_MAX + NT - 1) / NT;
262 double Ir[NACC][6], Ii[NACC][6];
263 for (int a = 0; a < NACC; a++) for (int m = 0; m < 6; m++) { Ir[a][m] = 0; Ii[a][m] = 0; }
265 int nitems, nodes_per_item;
266 if (self) {
267 u0 = targ_uv[2 * it]; v0 = targ_uv[2 * it + 1];
268 if (t == 0) {
269 double g[9];
270 geometry(gmode, coefs, norder, gpar, ga, gpe, ip, u0, v0, g);
271 double E = g[3] * g[3] + g[4] * g[4] + g[5] * g[5];
272 double F = g[3] * g[6] + g[4] * g[7] + g[5] * g[8];
273 double G = g[6] * g[6] + g[7] * g[7] + g[8] * g[8];
274 double det = sqrt(fmax(E * G - F * F, 1e-300));
275 int cnt = 0;
276 for (int e = 0; e < 4; e++) {
277 double h, x1, x2, xc, heff;
278 if (e == 0) { h = 1 - u0; } else if (e == 1) { h = -1 - u0; }
279 else if (e == 2) { h = 1 - v0; } else { h = -1 - v0; }
280 if (e < 2) { x1 = -1 - v0; x2 = 1 - v0; xc = -h * F / G; heff = fabs(h) * det / G; }
281 else { x1 = -1 - u0; x2 = 1 - u0; xc = -h * F / E; heff = fabs(h) * det / E; }
282 sEdge[e][0] = h; sEdge[e][1] = xc; sEdge[e][2] = heff;
283 if (fabs(h) < 1e-12) continue;
284 double xi0 = asinh((x1 - xc) / heff), xi1 = asinh((x2 - xc) / heff);
285 // sub-intervals of length <= sub_len in xi (singular structure) whose
286 // extent along the edge is <= dxmax (so that exp(ikr) is resolved)
287 double arc = sqrt(e < 2 ? G : E);
288 double dxmax = zkabs * arc > 0 ? osc_len / (zkabs * arc) : 1e300;
289 double xi = xi0;
290 while (xi < xi1 && cnt < 4 * NSUB_MAX) {
291 double xh = xc + heff * sinh(xi);
292 double xn = fmin(xi + sub_len, asinh((xh + dxmax - xc) / heff));
293 if (xn > xi1 || cnt == 4 * NSUB_MAX - 1) xn = xi1;
294 sitems[cnt].edge = e; sitems[cnt].xi0 = xi; sitems[cnt].xi1 = xn; cnt++;
295 xi = xn;
296 }
297 }
298 snitems = cnt;
299 }
300 __syncthreads();
301 nitems = snitems;
302 nodes_per_item = nqr * nqx;
303 } else {
304 nitems = leaf_ptr[p + 1] - leaf_ptr[p];
305 nodes_per_item = nq * nq;
306 }
308 int ntot = nitems * nodes_per_item;
309 for (int base = 0; base < ntot; base += NT) {
310 int g = base + t;
311 double Lu[NORD1_MAX], Lv[NORD1_MAX];
312 cdouble K[6];
313 bool active = g < ntot;
314 if (active) {
315 int item = g / nodes_per_item, loc = g - item * nodes_per_item;
316 double u, v, wq;
317 if (self) {
318 int e = sitems[item].edge;
319 int ir = loc / nqx, ix = loc - ir * nqx;
320 double h = sEdge[e][0], xc = sEdge[e][1], heff = sEdge[e][2];
321 double rho = glr_x[ir], wr = glr_w[ir];
322 double xi0 = sitems[item].xi0, xi1 = sitems[item].xi1;
323 double xi = xi0 + 0.5 * (xi1 - xi0) * (glx_x[ix] + 1.0);
324 double wxi = 0.5 * (xi1 - xi0) * glx_w[ix];
325 double x = xc + heff * sinh(xi);
326 double dx = heff * cosh(xi) * wxi;
327 if (e < 2) { u = u0 + rho * h; v = v0 + rho * x; }
328 else { u = u0 + rho * x; v = v0 + rho * h; }
329 wq = fabs(h) * rho * wr * dx;
330 } else {
331 int li = leaf_ptr[p] + item;
332 double cu = leaves[3 * li], cv = leaves[3 * li + 1], hf = leaves[3 * li + 2];
333 int a = loc / nq, b = loc - a * nq;
334 u = cu + hf * gl_x[a]; v = cv + hf * gl_x[b];
335 wq = hf * hf * gl_w[a] * gl_w[b];
336 }
337 double gg[9];
338 geometry(gmode, coefs, norder, gpar, ga, gpe, ip, u, v, gg);
339 double ns[3] = {gg[4] * gg[8] - gg[5] * gg[7], gg[5] * gg[6] - gg[3] * gg[8], gg[3] * gg[7] - gg[4] * gg[6]};
340 double jac = sqrt(ns[0] * ns[0] + ns[1] * ns[1] + ns[2] * ns[2]);
341 ns[0] /= jac; ns[1] /= jac; ns[2] /= jac;
342 kernels6<double>(xt, nt, gg, ns, wq * jac, zk, zkabs, 0xFFF, K);
343 legendre(u, norder, Lu);
344 legendre(v, norder, Lv);
345 }
346 for (int m = 0; m < 6; m++) { sKr[m][t] = active ? K[m].real() : 0.0; sKi[m][t] = active ? K[m].imag() : 0.0; }
347 for (int a = 0; a < n1; a++) { sLu[a][t] = active ? Lu[a] : 0.0; sLv[a][t] = active ? Lv[a] : 0.0; }
348 __syncthreads();
349 for (int acc = 0; acc < NACC; acc++) {
350 int n = t + acc * NT;
351 if (n < npols) {
352 int a = n % n1, b = n / n1;
353 for (int i = 0; i < NT; i++) {
354 double P = sLu[a][i] * sLv[b][i];
355 for (int m = 0; m < 6; m++) { Ir[acc][m] += P * sKr[m][i]; Ii[acc][m] += P * sKi[m][i]; }
356 }
357 }
358 }
359 __syncthreads();
360 }
362 // moments -> node weights, w_j = sum_n U[n, j] I_n. Reuse sLu/sLv as I storage.
363 double* sIr = &sLu[0][0]; // NORD1_MAX * NT = 1664 >= 6 * NPOLS_MAX = 1014
364 double* sIi = &sLv[0][0];
365 for (int acc = 0; acc < NACC; acc++) {
366 int n = t + acc * NT;
367 if (n < npols) for (int m = 0; m < 6; m++) { sIr[m * NPOLS_MAX + n] = Ir[acc][m]; sIi[m * NPOLS_MAX + n] = Ii[acc][m]; }
368 }
369 __syncthreads();
370 for (int j = t; j < npols; j += NT) {
371 for (int m = 0; m < 6; m++) {
372 double wr = 0, wi = 0;
373 for (int n = 0; n < npols; n++) {
374 double u = umat[n * npols + j];
375 wr += u * sIr[m * NPOLS_MAX + n];
376 wi += u * sIi[m * NPOLS_MAX + n];
377 }
378 W[((size_t)p * 6 + m) * npols + j] = cdouble(wr, wi);
379 }
380 }
381}
383// ---------------------------------------------------------------------------
384// Far-field pass: out[m][i] = sum over patches j not near x_i of
385// sum_l K_m(x_i, y_l) w_l q_l over the oversampled nodes l of patch j, for the
386// kernels selected in selmask (4 bits per kernel: 0 skip, 1 density qa, 2 density qb).
387// Kernel evaluation in precision T; per-patch partial sums are added into double accumulators.
388template <typename T>
389__global__ void dense_apply(
390 int ntarg, const double* targ_x, const T* targ_n, int npatches, int npols_o,
391 const double* cen, const double* r2, const T* src_x, const T* src_n, const T* src_w,
392 const complex<T>* qa, const complex<T>* qb, complex<T> zk, T zkabs, int selmask, cdouble* out)
393{
394 __shared__ T tile[NT][11];
395 __shared__ double scen[4];
396 int i = blockIdx.x * blockDim.x + threadIdx.x;
397 double xd[3] = {0, 0, 0};
398 T xt[3] = {0, 0, 0}, nt[3] = {0, 0, 0};
399 if (i < ntarg) {
400 for (int q = 0; q < 3; q++) { xd[q] = targ_x[3 * i + q]; xt[q] = (T)xd[q]; nt[q] = targ_n[3 * i + q]; }
401 }
402 int sel[6];
403 int kmask = 0;
404 for (int m = 0; m < 6; m++) { sel[m] = (selmask >> (4 * m)) & 3; if (sel[m]) kmask |= 3 << (2 * m); }
405 cdouble acc[6];
406 for (int m = 0; m < 6; m++) acc[m] = cdouble(0, 0);
407 for (int j = 0; j < npatches; j++) {
408 if (threadIdx.x < 3) scen[threadIdx.x] = cen[3 * j + threadIdx.x];
409 if (threadIdx.x == 3) scen[3] = r2[j];
410 __syncthreads();
411 bool far = (i < ntarg) && !is_near(xd, scen, scen[3]);
412 for (int base = 0; base < npols_o; base += NT) {
413 int l = base + threadIdx.x;
414 if (l < npols_o) {
415 size_t s = (size_t)j * npols_o + l;
416 for (int q = 0; q < 3; q++) { tile[threadIdx.x][q] = src_x[3 * s + q]; tile[threadIdx.x][3 + q] = src_n[3 * s + q]; }
417 tile[threadIdx.x][6] = src_w[s];
418 tile[threadIdx.x][7] = qa[s].real(); tile[threadIdx.x][8] = qa[s].imag();
419 tile[threadIdx.x][9] = qb[s].real(); tile[threadIdx.x][10] = qb[s].imag();
420 }
421 __syncthreads();
422 if (far) {
423 int nl = min(NT, npols_o - base);
424 complex<T> part[6];
425 for (int m = 0; m < 6; m++) part[m] = complex<T>(0, 0);
426 for (int s = 0; s < nl; s++) {
427 complex<T> K[6];
428 kernels6<T>(xt, nt, &tile[s][0], &tile[s][3], tile[s][6], zk, zkabs, kmask, K);
429 complex<T> a(tile[s][7], tile[s][8]), b(tile[s][9], tile[s][10]);
430 for (int m = 0; m < 6; m++) {
431 if (sel[m] == 1) part[m] += K[m] * a;
432 else if (sel[m] == 2) part[m] += K[m] * b;
433 }
434 }
435 for (int m = 0; m < 6; m++) acc[m] += cdouble((double)part[m].real(), (double)part[m].imag());
436 }
437 __syncthreads();
438 }
439 }
440 if (i < ntarg) for (int m = 0; m < 6; m++) out[(size_t)m * ntarg + i] = acc[m];
441}
443// ---------------------------------------------------------------------------
444// Near contributions: out[m][i] += sum_{near pairs p of i} sum_l W[p][m][l] q[d_m][patch nodes l]
445extern "C" __global__ void sparse_apply(
446 int ntarg, const int* row_ptr, const int* pair_p, int npols, int nreq, const int* req_m,
447 const int* req_d, const cdouble* W, const cdouble* q, int npts, cdouble* out)
448{
449 int i = blockIdx.x * blockDim.x + threadIdx.x;
450 if (i >= ntarg) return;
451 cdouble acc[6];
452 for (int r = 0; r < nreq; r++) acc[r] = cdouble(0, 0);
453 for (int p = row_ptr[i]; p < row_ptr[i + 1]; p++) {
454 int j = pair_p[p];
455 for (int r = 0; r < nreq; r++) {
456 int m = req_m[r];
457 const cdouble* w = W + ((size_t)p * 6 + m) * npols;
458 const cdouble* qq = q + (size_t)req_d[r] * npts + (size_t)j * npols;
459 cdouble s(0, 0);
460 for (int l = 0; l < npols; l++) s += w[l] * qq[l];
461 acc[r] += s;
462 }
463 }
464 for (int r = 0; r < nreq; r++) out[(size_t)req_m[r] * ntarg + i] += acc[r];
465}
466"""
468_module = None
471def module():
472 global _module
473 if _module is None:
474 _module = cp.RawModule(code=_SRC, options=("-std=c++17",),
475 name_expressions=["dense_apply<double>", "dense_apply<float>"])
476 _module.compile() # compile (or load from CuPy's cache) now rather than on first use
477 return _module
480def _i32(a):
481 return cp.ascontiguousarray(cp.asarray(a, dtype=cp.int32))
484def _f64(a):
485 return cp.ascontiguousarray(cp.asarray(a, dtype=cp.float64))
488class Geometry:
489 """Device-side description of the patch geometry used inside the integrals."""
491 def __init__(self, surf):
492 npols = surf.npols
493 self.norder = surf.norder
494 self.coefs = cp.asarray(np.ascontiguousarray(
495 surf.srccoefs.reshape(9, surf.npatches, npols).transpose(1, 0, 2)))
496 if getattr(surf, "exact", None) is not None:
497 self.gmode = 1
498 self.gpar = _f64(surf.exact["params"])
499 self.ga = float(surf.exact["a"])
500 self.gpe = int(surf.exact["p"])
501 else:
502 self.gmode, self.gpar, self.ga, self.gpe = 0, cp.zeros(13), 0.0, 0
504 def args(self):
505 return (np.int32(self.gmode), self.coefs, np.int32(self.norder), self.gpar,
506 np.float64(self.ga), np.int32(self.gpe))
509def find_near_pairs(targ_x, cen, r2):
510 """Pairs (target, patch) with |x_t - c_j|^2 <= r2_j, grouped by target.
512 Returns (pair_t, pair_p, row_ptr) as int32 device arrays.
513 """
514 ntarg = int(targ_x.shape[0])
515 npatches = int(cen.shape[0])
516 ker = module().get_function("near_pairs")
517 counts = cp.zeros(ntarg, dtype=cp.int32)
518 nb = (ntarg + 127) // 128
519 ker((nb,), (128,), (np.int32(ntarg), targ_x, np.int32(npatches), cen, r2, counts,
520 cp.zeros(1, cp.int32), np.uintp(0)))
521 row_ptr = cp.zeros(ntarg + 1, dtype=cp.int32)
522 row_ptr[1:] = cp.cumsum(counts).astype(cp.int32)
523 npairs = int(row_ptr[-1])
524 pair_p = cp.zeros(max(npairs, 1), dtype=cp.int32)
525 ker((nb,), (128,), (np.int32(ntarg), targ_x, np.int32(npatches), cen, r2, counts, row_ptr, pair_p))
526 pair_p = pair_p[:npairs]
527 pair_t = cp.repeat(cp.arange(ntarg, dtype=cp.int32), counts.tolist()) if npairs else cp.zeros(0, cp.int32)
528 return pair_t.astype(cp.int32), pair_p, row_ptr
531def build_leaves(pair_t, pair_p, targ_x, targ_patch, geo, zk=0.0, rfac_q=1.5, nlmax=14, kR_max=3.0):
532 npairs = int(pair_t.size)
533 if npairs == 0:
534 return cp.zeros(1, cp.int32), cp.zeros((1, 3), cp.float64)
535 ker = module().get_function("build_leaves")
536 counts = cp.zeros(npairs, dtype=cp.int32)
537 nb = (npairs + 127) // 128
538 args = (np.int32(npairs), pair_t, pair_p, targ_x, targ_patch) + geo.args() + (
539 np.float64(rfac_q), np.int32(nlmax), np.float64(abs(zk)), np.float64(kR_max))
540 ker((nb,), (128,), args + (counts, cp.zeros(1, cp.int32), np.uintp(0)))
541 leaf_ptr = cp.zeros(npairs + 1, dtype=cp.int32)
542 leaf_ptr[1:] = cp.cumsum(counts).astype(cp.int32)
543 nleaves = int(leaf_ptr[-1])
544 leaves = cp.zeros((max(nleaves, 1), 3), dtype=cp.float64)
545 ker((nb,), (128,), args + (counts, leaf_ptr, leaves))
546 return leaf_ptr, leaves
549def near_weights(pair_t, pair_p, targ_x, targ_n, targ_patch, targ_uv, geo, leaf_ptr, leaves,
550 umat, zk, nq=8, nqr=10, nqx=8, sub_len=2.0, osc_len=2.0):
551 """Near weights W (npairs, 6, npols), complex128 on the device."""
552 npairs = int(pair_t.size)
553 npols = (geo.norder + 1) ** 2
554 W = cp.zeros((max(npairs, 1), 6, npols), dtype=cp.complex128)
555 if npairs == 0:
556 return W
557 gx, gw = np.polynomial.legendre.leggauss(nq)
558 rx, rw = np.polynomial.legendre.leggauss(nqr)
559 rx, rw = 0.5 * (rx + 1), 0.5 * rw # on [0, 1]
560 xx, xw = np.polynomial.legendre.leggauss(nqx)
561 ker = module().get_function("near_weights")
562 ker((npairs,), (128,), (
563 np.int32(npairs), pair_t, pair_p, targ_x, targ_n, targ_patch, targ_uv) + geo.args() + (
564 leaf_ptr, leaves, np.int32(nq), _f64(gx), _f64(gw),
565 np.int32(nqr), _f64(rx), _f64(rw), np.int32(nqx), _f64(xx), _f64(xw), np.float64(sub_len),
566 np.float64(osc_len), umat, np.complex128(zk), np.float64(abs(zk)), W))
567 return W
570class FarSources:
571 """Oversampled nodes of all patches, patch-major, in the precision of the far field."""
573 def __init__(self, x, n, w, npols_o, cen, r2, dtype):
574 self.real = np.dtype(dtype)
575 self.cplx = np.complex128 if self.real == np.float64 else np.complex64
576 self.x = cp.ascontiguousarray(cp.asarray(x).astype(self.real))
577 self.n = cp.ascontiguousarray(cp.asarray(n).astype(self.real))
578 self.w = cp.ascontiguousarray(cp.asarray(w).astype(self.real))
579 self.npols_o = npols_o
580 self.cen, self.r2 = cen, r2
581 self.npatches = int(cen.shape[0])
584def dense_apply(targ_x, targ_n, far, qa, qb, zk, sel):
585 """out (6, ntarg) complex128: far-patch sums for the kernels in sel (6 entries in {0,1,2})."""
586 ntarg = int(targ_x.shape[0])
587 name = "dense_apply<double>" if far.real == np.float64 else "dense_apply<float>"
588 ker = module().get_function(name)
589 ccast = lambda a: cp.ascontiguousarray(cp.asarray(a).astype(far.cplx, copy=False))
590 out = cp.zeros((6, ntarg), dtype=cp.complex128)
591 selmask = sum(int(s) << (4 * m) for m, s in enumerate(sel))
592 nb = (ntarg + 127) // 128
593 ker((nb,), (128,), (
594 np.int32(ntarg), targ_x, cp.ascontiguousarray(targ_n.astype(far.real)), np.int32(far.npatches),
595 np.int32(far.npols_o), far.cen, far.r2, far.x, far.n, far.w, ccast(qa), ccast(qb),
596 far.cplx(zk), far.real.type(abs(zk)), np.int32(selmask), out))
597 return out
600def sparse_apply(row_ptr, pair_p, npols, W, q, reqs, out):
601 """out (6, ntarg) += near contributions; reqs = list of (kernel m, density index d); q (ndens, npts)."""
602 ntarg = int(row_ptr.size) - 1
603 if len(reqs) == 0 or int(pair_p.size) == 0:
604 return out
605 ker = module().get_function("sparse_apply")
606 req_m = _i32([m for m, _ in reqs])
607 req_d = _i32([d for _, d in reqs])
608 q = cp.ascontiguousarray(q, dtype=cp.complex128)
609 nb = (ntarg + 127) // 128
610 ker((nb,), (128,), (np.int32(ntarg), row_ptr, pair_p, np.int32(npols), np.int32(len(reqs)),
611 req_m, req_d, W, q, np.int32(q.shape[1]), out))
612 return out