Multigrid preconditioner for stress analysis

// group strength graph
const gd = new Float64Array(ng), nbr = Array.from({ length: ng }, () => []);
const acc = new Float64Array(ng), mark = new Int32Array(ng).fill(-1), lst = [];
for (let g = 0; g < ng; g++) {
  lst.length = 0;
  for (let r = grp[g]; r < grp[g + 1]; r++) for (let k = cur.ptr[r]; k < cur.ptr[r + 1]; k++) { const h = gof[cur.col[k]]; if (mark[h] !== g) { mark[h] = g; acc[h] = 0; lst.push(h); } acc[h] += cur.val[k] * cur.val[k]; }
  for (const h of lst) { if (h === g) gd[g] = Math.sqrt(acc[h]); else nbr[g].push(h, Math.sqrt(acc[h])); }
}
const strong = nbr.map((l, g) => { const out = []; for (let i = 0; i < l.length; i += 2) if (l[i + 1] > theta * Math.sqrt(gd[g] * gd[l[i]])) out.push(l[i]); return out; });

const m = dofs.length, Q = Bc.map((b) => Float64Array.from(dofs, (d) => b[d])), R = [];
let scale = 0; for (const q of Q) for (const v of q) scale = Math.max(scale, Math.abs(v));
const kept = [];
for (let j = 0; j < nb; j++) {
  const q = Q[j], r = new Float64Array(nb);
  for (let t = 0; t < 2; t++) for (const i of kept) { let s = 0; for (let k = 0; k < m; k++) s += Q[i][k] * q[k]; r[i] += s; for (let k = 0; k < m; k++) q[k] -= s * Q[i][k]; }
  let l = 0; for (let k = 0; k < m; k++) l += q[k] * q[k]; l = Math.sqrt(l);
  if (l > 1e-10 * Math.max(scale, 1e-300) * Math.sqrt(m)) { for (let k = 0; k < m; k++) q[k] /= l; r[j] = l; kept.push(j); }
  R.push(r);
}
// coarse dofs = kept columns; B_c rows = R rows of the kept columns (R[j][i] = component of B_j on q_i)
kept.forEach((i, ci) => { const c = nc + ci; dofs.forEach((d, k) => { if (Q[i][k] !== 0) rowsP[d].push(c, Q[i][k]); }); });
for (const i of kept) { const row = new Float64Array(nb); for (let j = 0; j < nb; j++) row[j] = R[j][i]; newB.push(row); }
nc += kept.length; cgrp.push(nc);

/** λmax of D⁻¹A by power iteration (15 steps, +5 % margin) */
function powerMax(A, diag) {
  let x = new Float64Array(A.n), y = new Float64Array(A.n), lam = 1;
  let s = 7; for (let i = 0; i < A.n; i++) { s = (s * 16807) % 2147483647; x[i] = s / 2147483647 - 0.5; }
  for (let it = 0; it < 15; it++) {
    csrMul(A, x, y); for (let i = 0; i < A.n; i++) y[i] = diag[i] ? y[i] / diag[i] : 0;
    let nx = 0, ny = 0; for (let i = 0; i < A.n; i++) { nx += x[i] * x[i]; ny += y[i] * y[i]; }
    lam = Math.sqrt(ny / (nx || 1)); const l = Math.sqrt(ny) || 1; for (let i = 0; i < A.n; i++) x[i] = y[i] / l;
  }
  return 1.05 * lam;
}

// smoothing: P = (I − ω D⁻¹ A) P̂
let P = Pt;
if (smooth[levels.length - 1] ?? smooth[smooth.length - 1]) {
  const lam = powerMax(cur, lev.diag), w = 4 / (3 * lam);
  const AP = spgemm(cur, Pt, nc);
  for (let r = 0; r < cur.n; r++) { const s = lev.diag[r] > 0 ? w / lev.diag[r] : 0; for (let k = AP.ptr[r]; k < AP.ptr[r + 1]; k++) AP.val[k] *= -s; }
  P = addCsr(Pt, AP);
}
const PT = transpose(P, nc);
lev.P = P; lev.PT = PT;
cur = spgemm(PT, spgemm(cur, P, nc), nc);

/** C = A·B (CSR, B with ncols columns) */
export function spgemm(A, B, ncols) {
  const ptr = new Int32Array(A.n + 1), mark = new Int32Array(ncols).fill(-1), acc = new Float64Array(ncols), g = new Grow(A.col.length * 2);
  const list = new Int32Array(ncols);
  for (let i = 0; i < A.n; i++) {
    let m = 0;
    for (let k = A.ptr[i]; k < A.ptr[i + 1]; k++) {
      const a = A.val[k], r = A.col[k];
      for (let q = B.ptr[r]; q < B.ptr[r + 1]; q++) { const j = B.col[q]; if (mark[j] !== i) { mark[j] = i; acc[j] = 0; list[m++] = j; } acc[j] += a * B.val[q]; }
    }
    const l = list.subarray(0, m).sort();
    for (let t = 0; t < m; t++) g.push(l[t], acc[l[t]]);
    ptr[i + 1] = g.n;
  }
  return { n: A.n, ptr, col: g.c.slice(0, g.n), val: g.v.slice(0, g.n) };
}

// coarsest: dense Cholesky (with a tiny shift if singular)
const last = levels[levels.length - 1], n = last.n, Dm = new Float64Array(n * n);
for (let r = 0; r < n; r++) for (let k = last.A.ptr[r]; k < last.A.ptr[r + 1]; k++) Dm[r * n + last.A.col[k]] = last.A.val[k];
let dmax = 0; for (let i = 0; i < n; i++) dmax = Math.max(dmax, Dm[i * n + i]);
for (let i = 0; i < n; i++) {
  for (let j = 0; j <= i; j++) {
    let s = Dm[i * n + j]; for (let k = 0; k < j; k++) s -= Dm[i * n + k] * Dm[j * n + k];
    if (i === j) Dm[i * n + i] = Math.sqrt(s > 1e-12 * dmax ? s : 1e-12 * dmax + Math.abs(s)); else Dm[i * n + j] = s / Dm[j * n + j];
  }
}
const coarseSolve = (b, x) => {
  for (let i = 0; i < n; i++) { let s = b[i]; for (let k = 0; k < i; k++) s -= Dm[i * n + k] * x[k]; x[i] = s / Dm[i * n + i]; }
  for (let i = n - 1; i >= 0; i--) { let s = x[i]; for (let k = i + 1; k < n; k++) s -= Dm[k * n + i] * x[k]; x[i] = s / Dm[i * n + i]; }
  return x;
};
const work = levels.map((l) => ({ r: new Float64Array(l.n), x: new Float64Array(l.n), b: new Float64Array(l.n) }));

/** One Gauss–Seidel sweep on full CSR M (diagonal diag): forward (rows up) or backward (rows down); x updated in place.
 *  The smoother of the V-cycle (with csrMul, the only kernels the cycle runs besides the coarse Cholesky). */
export function gaussSeidel(M, diag, b, x, forward) {
  const nn = M.n;
  if (forward) for (let r = 0; r < nn; r++) { let s = b[r]; for (let k = M.ptr[r]; k < M.ptr[r + 1]; k++) { const c = M.col[k]; if (c !== r) s -= M.val[k] * x[c]; } x[r] = diag[r] ? s / diag[r] : 0; }
  else for (let r = nn - 1; r >= 0; r--) { let s = b[r]; for (let k = M.ptr[r]; k < M.ptr[r + 1]; k++) { const c = M.col[k]; if (c !== r) s -= M.val[k] * x[c]; } x[r] = diag[r] ? s / diag[r] : 0; }
  return x;
}

const cycle = (l, b, x) => {
  const lev = levels[l];
  if (l === levels.length - 1) return coarseSolve(b, x);
  x.fill(0);
  gs(lev, b, x, true);
  const r = work[l].r; csrMul(lev.A, x, r); for (let i = 0; i < lev.n; i++) r[i] = b[i] - r[i];
  const bc = work[l + 1].b, xc = work[l + 1].x;
  csrMul(lev.PT, r, bc);
  cycle(l + 1, bc, xc);
  const e = work[l].r; csrMul(lev.P, xc, e); for (let i = 0; i < lev.n; i++) x[i] += e[i];
  gs(lev, b, x, false);
  return x;
};