CraftScript sparse matrices and direct factorization

/** Triplets (0-based rows/cols) → CSR, duplicates summed, columns sorted; two counting sorts (column, then stable row). */
export function fromTriplets(I, J, V, n, m) {
  const nz = I.length;
  // pass 1: by column
  const cc = new Int32Array(m + 1);
  for (let k = 0; k < nz; k++) cc[J[k] + 1]++;
  for (let c = 0; c < m; c++) cc[c + 1] += cc[c];
  const o1 = new Int32Array(nz), fc = cc.slice(0, m);
  for (let k = 0; k < nz; k++) o1[fc[J[k]]++] = k;
  // pass 2: stable by row
  const rc = new Int32Array(n + 1);
  for (let k = 0; k < nz; k++) rc[I[k] + 1]++;
  for (let r = 0; r < n; r++) rc[r + 1] += rc[r];
  const o2 = new Int32Array(nz), fr = rc.slice(0, n);
  for (let q = 0; q < nz; q++) { const k = o1[q]; o2[fr[I[k]]++] = k; }
  // merge duplicates
  const rowPtr = new Int32Array(n + 1), col = new Int32Array(nz), val = new Float64Array(nz);
  let w = 0;
  for (let r = 0; r < n; r++) {
    const start = w;
    for (let q = rc[r]; q < rc[r + 1]; q++) {
      const k = o2[q], c = J[k];
      if (w > start && col[w - 1] === c) val[w - 1] += V[k];
      else { col[w] = c; val[w] = V[k]; w++; }
    }
    rowPtr[r + 1] = w;
  }
  return sp(n, m, rowPtr, col.slice(0, w), val.slice(0, w));
}

/** a·A + b·B (same size). */
export function spAdd(A, B, a = 1, b = 1) {
  if (A.n !== B.n || A.m !== B.m) throw err(t('script.linalg.sparse_sizes_differ', { a: `${A.n}×${A.m}`, b: `${B.n}×${B.m}` }));
  const rowPtr = new Int32Array(A.n + 1), cap = A.col.length + B.col.length, col = new Int32Array(cap), val = new Float64Array(cap);
  let w = 0;
  for (let r = 0; r < A.n; r++) {
    let i = A.rowPtr[r], j = B.rowPtr[r];
    const ie = A.rowPtr[r + 1], je = B.rowPtr[r + 1];
    while (i < ie || j < je) {
      if (j >= je || (i < ie && A.col[i] < B.col[j])) { col[w] = A.col[i]; val[w++] = a * A.val[i++]; }
      else if (i >= ie || B.col[j] < A.col[i]) { col[w] = B.col[j]; val[w++] = b * B.val[j++]; }
      else { col[w] = A.col[i]; val[w++] = a * A.val[i++] + b * B.val[j++]; }
    }
    rowPtr[r + 1] = w;
  }
  return sp(A.n, A.m, rowPtr, col.slice(0, w), val.slice(0, w));
}

/** y = S x (x Float64Array) */
export function spMulVec(S, x, y = new Float64Array(S.n)) {
  const { rowPtr, col, val } = S;
  for (let r = 0; r < S.n; r++) { let s = 0; for (let k = rowPtr[r]; k < rowPtr[r + 1]; k++) s += val[k] * x[col[k]]; y[r] = s; }
  return y;
}

/** Sparse × sparse (Gustavson). */
export function spMulSp(A, B) {
  if (A.m !== B.n) throw err(t('script.linalg.inner_sizes', { a: `${A.n}×${A.m}`, b: `${B.n}×${B.m}` }));
  const mark = new Int32Array(B.m).fill(-1), acc = new Float64Array(B.m), list = new Int32Array(B.m);
  const rowPtr = new Int32Array(A.n + 1); let col = new Int32Array(Math.max(16, A.col.length + B.col.length)), val = new Float64Array(col.length), w = 0;
  for (let i = 0; i < A.n; i++) {
    let m = 0;
    for (let k = A.rowPtr[i]; k < A.rowPtr[i + 1]; k++) {
      const a = A.val[k], r = A.col[k];
      for (let q = B.rowPtr[r]; q < B.rowPtr[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();
    if (w + m > col.length) { const nc = new Int32Array(2 * (w + m)), nv = new Float64Array(2 * (w + m)); nc.set(col); nv.set(val); col = nc; val = nv; }
    for (let t = 0; t < m; t++) { col[w] = l[t]; val[w++] = acc[l[t]]; }
    rowPtr[i + 1] = w;
  }
  return sp(A.n, B.m, rowPtr, col.slice(0, w), val.slice(0, w));
}

/** Is S symmetric (pattern and values within rel·max|v|)? */
export function spIsSymmetric(S, rel = 1e-10) {
  if (S.n !== S.m) return false;
  const T = spTranspose(S);
  let mx = 0; for (const v of S.val) mx = Math.max(mx, Math.abs(v));
  const tol = rel * (mx || 1);
  const D = spAdd(S, T, 1, -1);
  for (const v of D.val) if (Math.abs(v) > tol) return false;
  return true;
}

/** Envelope size after a permutation (new index of each old node). */
function envelopeOf(S, perm) {
  const n = S.n, first = new Int32Array(n);
  for (let i = 0; i < n; i++) first[i] = i;
  for (let r = 0; r < n; r++) for (let k = S.rowPtr[r]; k < S.rowPtr[r + 1]; k++) {
    const a = perm[r], b = perm[S.col[k]], hi = a > b ? a : b, lo = a > b ? b : a;
    if (lo < first[hi]) first[hi] = lo;
  }
  let tot = 0, work = 0; for (let i = 0; i < n; i++) { const w = i - first[i] + 1; tot += w; work += w * w / 2; }
  return { first, tot, work };
}

/** Factor an SPD sparse matrix (envelope Cholesky, RCM order) → { solve(b, x), n, envelope, ms } or null when it does not fit (cap entries, maxWork multiply-adds). Throws 'not positive definite'. */
export function envelopeCholesky(S, { cap = ENVELOPE_CAP, maxWork = ENVELOPE_WORK, check = () => {} } = {}) {
  const t0 = Date.now(), n = S.n;
  const g = symGraph(S), perm = n > 1 ? kRcm(n, g.ptr, g.adj) : new Int32Array(n);
  const { first, tot, work } = envelopeOf(S, perm);
  if (tot > cap || work > maxWork) return null;
  const start = new Float64Array(n + 1);   // offset of row i (Float64 index arithmetic safe beyond 2^31 is not needed below the cap)
  for (let i = 0; i < n; i++) start[i + 1] = start[i] + (i - first[i] + 1);
  const L = new Float64Array(tot);
  // scatter A (lower part, permuted) into the envelope; L[start[i] + (j − first[i])] = A_ij
  for (let r = 0; r < n; r++) for (let k = S.rowPtr[r]; k < S.rowPtr[r + 1]; k++) {
    const a = perm[r], b = perm[S.col[k]];
    if (b > a) continue;
    L[start[a] + (b - first[a])] += S.val[k];
  }
  // in-place row-oriented Cholesky
  for (let i = 0; i < n; i++) {
    if ((i & 1023) === 0) check();
    const fi = first[i], si = start[i];
    for (let j = fi; j < i; j++) {
      const fj = first[j], sj = start[j], k0 = fi > fj ? fi : fj;
      let s = L[si + (j - fi)];
      let p = si + (k0 - fi), q = sj + (k0 - fj);
      for (let k = k0; k < j; k++) s -= L[p++] * L[q++];
      L[si + (j - fi)] = s / L[sj + (j - fj)];
    }
    let d = L[si + (i - fi)];
    for (let p = si, e = si + (i - fi); p < e; p++) d -= L[p] * L[p];
    if (!(d > 0)) { const e = err(t('script.linalg.not_pd')); e.notPD = true; throw e; }
    L[si + (i - fi)] = Math.sqrt(d);
  }
  const y = new Float64Array(n);
  return {
    n, envelope: tot, ms: Date.now() - t0, kind: 'envelope Cholesky (RCM)',   // i18n-ignore: a method name scripts read from info.method
    solve(b, x = new Float64Array(n)) {
      for (let i = 0; i < n; i++) y[perm[i]] = b[i];
      for (let i = 0; i < n; i++) { const fi = first[i], si = start[i]; let s = y[i]; for (let k = fi, p = si; k < i; k++) s -= L[p++] * y[k]; y[i] = s / L[si + (i - fi)]; }
      for (let i = n - 1; i >= 0; i--) { const fi = first[i], si = start[i]; const v = (y[i] /= L[si + (i - fi)]); for (let k = fi, p = si; k < i; k++) y[k] -= L[p++] * v; }
      for (let i = 0; i < n; i++) x[i] = y[perm[i]];
      return x;
    },
  };
}

/**
 * K = Σ_e w_e · P_eᵀ Ke(e) P_e — edof (ne × nd, 1-based dof numbers), Ke (nd × nd, or nd × nd × ne), w (ne, optional),
 * n (size, default the largest dof). → sparse n × n. Pattern built once from the connectivity (no triplet arrays).
 */
export function assembleElements(edof, Ke, w = null, n = 0, check = () => {}) {
  const [ne, nd] = size2(edof, 'edof');
  const per = Ke.shape.length === 3 && Ke.shape[2] > 1;
  if (Ke.shape[0] !== nd || Ke.shape[1] !== nd || (per && Ke.shape[2] !== ne)) throw err(t('script.linalg.assemble_ke', { size: `${nd}×${nd}${per ? `×${ne}` : ''}`, nd }));
  if (w && w.length !== ne) throw err(t('script.linalg.assemble_w', { w: w.length, ne }));
  const E = new Int32Array(ne * nd);   // row-major element dofs, 0-based
  let mx = 0;
  for (let e = 0; e < ne; e++) for (let a = 0; a < nd; a++) {
    const v = edof.data[e + a * ne];
    if (!Number.isInteger(v) || v < 1) throw err(t('script.linalg.assemble_dof', { dof: v, e: e + 1 }));
    E[e * nd + a] = v - 1; if (v > mx) mx = v;
  }
  n = n || mx;
  if (mx > n) throw err(t('script.linalg.assemble_past', { dof: mx, n }));
  // node-free pattern: for each dof, the set of dofs sharing an element
  const cnt = new Int32Array(n + 1);
  for (let i = 0; i < E.length; i++) cnt[E[i] + 1]++;
  for (let i = 0; i < n; i++) cnt[i + 1] += cnt[i];
  const elOf = new Int32Array(E.length), fill = cnt.slice(0, n);
  for (let e = 0; e < ne; e++) for (let a = 0; a < nd; a++) elOf[fill[E[e * nd + a]]++] = e;
  const mark = new Int32Array(n).fill(-1), rowPtr = new Int32Array(n + 1), rows = new Array(n);
  let nnz = 0;
  for (let r = 0; r < n; r++) {
    const l = [];
    for (let q = cnt[r]; q < cnt[r + 1]; q++) { const e = elOf[q]; for (let b = 0; b < nd; b++) { const c = E[e * nd + b]; if (mark[c] !== r) { mark[c] = r; l.push(c); } } }
    rows[r] = Int32Array.from(l).sort(); nnz += l.length; rowPtr[r + 1] = nnz;
    if ((r & 4095) === 0) check();
  }
  const col = new Int32Array(nnz); for (let r = 0; r < n; r++) col.set(rows[r], rowPtr[r]);
  const val = new Float64Array(nnz), loc = new Int32Array(nd * nd), K = Ke.data;
  const find = (r, c) => { let lo = rowPtr[r], hi = rowPtr[r + 1] - 1; while (lo <= hi) { const m = (lo + hi) >> 1, v = col[m]; if (v === c) return m; if (v < c) lo = m + 1; else hi = m - 1; } return -1; };
  for (let e = 0; e < ne; e++) {
    if ((e & 4095) === 0) check();
    const s = w ? w[e] : 1; if (s === 0) continue;
    const off = per ? e * nd * nd : 0;
    for (let a = 0; a < nd; a++) for (let b = 0; b < nd; b++) loc[a * nd + b] = find(E[e * nd + a], E[e * nd + b]);
    for (let b = 0; b < nd; b++) for (let a = 0; a < nd; a++) val[loc[a * nd + b]] += s * K[off + a + b * nd];
  }
  return sp(n, n, rowPtr, col, val);
}