2 * Unrestarted GMRES for (zid I + K) x = rhs with the weighted inner product
3 * <u, v> = sum conj(u) v w, following fmm3dbie's zgmres_guru. The Krylov
4 * vectors and the Hessenberg system are kept in double precision on the CPU;
5 * only the operator K is applied on the GPU.
6 */
7import type { CVec } from './problem.ts';
9export type Apply = (x: CVec) => Promise<CVec>;
11export interface GmresResult {
12 x: CVec;
13 niter: number;
14 rres: number;
15 residuals: number[];
16}
18export async function gmres(
19 applyK: Apply,
20 rhs: CVec,
21 wts: Float64Array,
22 zid: [number, number],
23 eps: number,
24 numit: number,
25 onIter?: (it: number, rres: number) => void,
26): Promise<GmresResult> {
27 const n = wts.length;
28 const wnorm = (v: CVec) => {
29 let s = 0;
30 for (let i = 0; i < n; i++) s += (v[2 * i] ** 2 + v[2 * i + 1] ** 2) * wts[i];
31 return Math.sqrt(s);
32 };
33 const rb = wnorm(rhs);
34 const V: CVec[] = [];
35 const m1 = numit + 1;
36 const Hr = new Float64Array(m1 * numit);
37 const Hi = new Float64Array(m1 * numit);
38 const csr = new Float64Array(numit);
39 const csi = new Float64Array(numit);
40 const snr = new Float64Array(numit);
41 const sni = new Float64Array(numit);
42 const sr = new Float64Array(m1);
43 const si = new Float64Array(m1);
44 sr[0] = rb;
45 const v0 = new Float64Array(2 * n);
46 for (let i = 0; i < 2 * n; i++) v0[i] = rhs[i] / rb;
47 V.push(v0);
48 const residuals: number[] = [];
49 let itDone = 0;
50 for (let it = 0; it < numit; it++) {
51 let w = await applyK(V[it]);
52 w = new Float64Array(w); // do not alias the operator's output
53 // classical Gram-Schmidt against all previous vectors, repeated once
54 for (let pass = 0; pass < 2; pass++) {
55 for (let k = 0; k <= it; k++) {
56 const vk = V[k];
57 let hr = 0;
58 let hi = 0;
59 for (let i = 0; i < n; i++) {
60 const ar = vk[2 * i];
61 const ai = vk[2 * i + 1];
62 const br = w[2 * i] * wts[i];
63 const bi = w[2 * i + 1] * wts[i];
64 hr += ar * br + ai * bi; // conj(a) b
65 hi += ar * bi - ai * br;
66 }
67 Hr[k * numit + it] += hr;
68 Hi[k * numit + it] += hi;
69 for (let i = 0; i < n; i++) {
70 const vr = vk[2 * i];
71 const vi = vk[2 * i + 1];
72 w[2 * i] -= hr * vr - hi * vi;
73 w[2 * i + 1] -= hr * vi + hi * vr;
74 }
75 }
76 }
77 Hr[it * numit + it] += zid[0];
78 Hi[it * numit + it] += zid[1];
79 const hn = wnorm(w);
80 Hr[(it + 1) * numit + it] = hn;
81 Hi[(it + 1) * numit + it] = 0;
82 const vn = new Float64Array(2 * n);
83 if (hn > 0) for (let i = 0; i < 2 * n; i++) vn[i] = w[i] / hn;
84 V.push(vn);
85 // apply the previous Givens rotations to the new column
86 for (let k = 0; k < it; k++) {
87 const ar = Hr[k * numit + it];
88 const ai = Hi[k * numit + it];
89 const br = Hr[(k + 1) * numit + it];
90 const bi = Hi[(k + 1) * numit + it];
91 // t1 = cs a + sn b ; t2 = -conj(sn) a + cs b
92 const t1r = csr[k] * ar - csi[k] * ai + snr[k] * br - sni[k] * bi;
93 const t1i = csr[k] * ai + csi[k] * ar + snr[k] * bi + sni[k] * br;
94 const t2r = -(snr[k] * ar + sni[k] * ai) + csr[k] * br - csi[k] * bi;
95 const t2i = -(snr[k] * ai - sni[k] * ar) + csr[k] * bi + csi[k] * br;
96 Hr[k * numit + it] = t1r;
97 Hi[k * numit + it] = t1i;
98 Hr[(k + 1) * numit + it] = t2r;
99 Hi[(k + 1) * numit + it] = t2i;
100 }
101 const ar = Hr[it * numit + it];
102 const ai = Hi[it * numit + it];
103 const br = Hr[(it + 1) * numit + it];
104 const bi = Hi[(it + 1) * numit + it];
105 const aa = Math.hypot(ar, ai);
106 const r = Math.hypot(aa, Math.hypot(br, bi));
107 if (r === 0) {
108 csr[it] = 1;
109 csi[it] = 0;
110 snr[it] = 0;
111 sni[it] = 0;
112 } else {
113 csr[it] = aa / r;
114 csi[it] = 0;
115 // sn = (a/|a|) conj(b) / r
116 let pr = 1;
117 let pi = 0;
118 if (aa > 0) {
119 pr = ar / aa;
120 pi = ai / aa;
121 }
122 snr[it] = (pr * br + pi * bi) / r;
123 sni[it] = (pi * br - pr * bi) / r;
124 }
125 // H[it, it] = cs a + sn b ; H[it+1, it] = 0
126 Hr[it * numit + it] = csr[it] * ar - csi[it] * ai + snr[it] * br - sni[it] * bi;
127 Hi[it * numit + it] = csr[it] * ai + csi[it] * ar + snr[it] * bi + sni[it] * br;
128 Hr[(it + 1) * numit + it] = 0;
129 Hi[(it + 1) * numit + it] = 0;
130 // s[it+1] = -conj(sn) s[it] ; s[it] = cs s[it]
131 const s0r = sr[it];
132 const s0i = si[it];
133 sr[it + 1] = -(snr[it] * s0r + sni[it] * s0i);
134 si[it + 1] = -(snr[it] * s0i - sni[it] * s0r);
135 sr[it] = csr[it] * s0r - csi[it] * s0i;
136 si[it] = csr[it] * s0i + csi[it] * s0r;
137 itDone = it + 1;
138 const rres = Math.hypot(sr[it + 1], si[it + 1]) / rb;
139 residuals.push(rres);
140 onIter?.(itDone, rres);
141 if (rres < eps) break;
142 }
143 // back substitution on the upper triangular system
144 const yr = new Float64Array(itDone);
145 const yi = new Float64Array(itDone);
146 for (let i = itDone - 1; i >= 0; i--) {
147 let tr = sr[i];
148 let ti = si[i];
149 for (let j = i + 1; j < itDone; j++) {
150 const hr = Hr[i * numit + j];
151 const hi = Hi[i * numit + j];
152 tr -= hr * yr[j] - hi * yi[j];
153 ti -= hr * yi[j] + hi * yr[j];
154 }
155 const hr = Hr[i * numit + i];
156 const hi = Hi[i * numit + i];
157 const d = hr * hr + hi * hi;
158 yr[i] = (tr * hr + ti * hi) / d;
159 yi[i] = (ti * hr - tr * hi) / d;
160 }
161 const x = new Float64Array(2 * n);
162 for (let j = 0; j < itDone; j++) {
163 const vj = V[j];
164 for (let i = 0; i < n; i++) {
165 const vr = vj[2 * i];
166 const vi = vj[2 * i + 1];
167 x[2 * i] += yr[j] * vr - yi[j] * vi;
168 x[2 * i + 1] += yr[j] * vi + yi[j] * vr;
169 }
170 }
171 const Kx = await applyK(x);
172 const res = new Float64Array(2 * n);
173 for (let i = 0; i < n; i++) {
174 res[2 * i] = zid[0] * x[2 * i] - zid[1] * x[2 * i + 1] + Kx[2 * i] - rhs[2 * i];
175 res[2 * i + 1] = zid[0] * x[2 * i + 1] + zid[1] * x[2 * i] + Kx[2 * i + 1] - rhs[2 * i + 1];
176 }
177 return { x, niter: itDone, rres: wnorm(res) / rb, residuals };
178}