All files / src/structure reconstruct.ts

99.02% Statements 405/409
92.02% Branches 150/163
100% Functions 20/20
99.02% Lines 405/409

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660                                                                                                                                                  2x 20x       20x 661x 661x 661x 661x 20x                     2x 3x 3x 42x 381x 381x 381x 381x 381x 381x 381x 381x 42x 42x 3x   2x 8x 8x 8x 275x 275x 6719x 6719x 6719x 6719x 6719x 6719x 6719x 6719x 6719x 6719x 275x 8x 8x                       2x 1x 1x 3x 3x 3x 3x 1x 1x     2x 1x 1x 1x 2x 2x 2x 1x 1x                           2x 10x 10x 10x 10x 10x 10x 10x 10x 10x                                 2x 12x 12x 12x 12x 12x 12x   12x 12x 12x 12x 12x 12x 12x 379x 8081x 8081x 8081x 5032x 5032x 5032x 8081x 1824x 1824x 1824x 8081x 379x 12x 12x 12x 317x 255x 315x 317x     12x 12x 317x 255x 255x 255x 62x 317x 623x 623x 623x 574x 623x 317x 317x 12x 12x 379x 8081x 8081x 8081x 8081x 8081x 379x 12x                               2x 5x 5x 89x 1569x 1569x 1569x 1569x 1569x 1569x 1569x 1569x 1569x 1569x 89x 5x 5x                       2x 8x 8x 8x 8x 8x                       2x 11x 11x 11x 11x               11x 11x 11x 430x 430x 430x 20100x 20100x 20100x 430x 430x 430x 11x 11x 11x 11x 11x 11x 880x 880x 880x 880x 880x       11x 33x 33x 33x 33x 33x 33x 33x 11x 11x 11x                     2673x 2673x 2673x 2673x 2673x 2673x 2673x 2673x 2673x 2673x 2673x 104490x 104490x 104490x 2673x 2673x 104490x 104490x 104490x 4884300x 4884300x 4884300x 104490x 104490x 2673x     33x 33x 33x 33x 1290x 1290x 1290x 1290x 1290x 33x 33x     891x 891x 891x 2673x 2673x 2673x 2673x 2673x 2673x 2673x 2673x 2673x 2673x 2673x         2673x 2673x 891x                                           2x 970x 970x 970x 970x                 970x 970x 970x 970x 970x 41260x 41260x 41260x 41260x 41260x 41260x 41260x 41260x 41260x 2040800x 1999540x 1999540x 2040800x 2040800x 1533724x 1533724x 1533724x 1533724x       2040800x 2040800x 2040800x 2040800x 2040800x 2040800x 41260x 41260x 41260x 41260x 41260x 970x 970x           2x 999x 999x 999x 999x 999x 999x 999x 999x 999x 42170x 1014765x 1014765x 1014765x 1014765x 776022x 776022x 776022x 776022x 776022x 776022x 776022x 42170x 999x 999x                   2x 27x 27x 27x 26x 26x 26x 27x 900x 900x 900x 900x 26x 26x 26x 26x 27x 900x 900x 900x 900x 900x 26x 26x 27x 899x 899x 899x 899x 25x     2x 36x 36x 35x 35x 35x 35x 36x 22903x 22903x 22903x 35x 35x 35x 35x 35x 36x 22903x 22903x 22903x 22903x 22903x 22903x 36x 36x   70x 70x 70x 70x 70x 70x 25028x 25028x 25028x 25028x 25028x 25028x 70x 70x                 2x 6x 6x 6x 6x 6x 6x 6x 6x 6x 6x 6x 6x 6x 750x 750x 750x 750x 750x 6x 6x 6x 750x 6x 6x 6x  
/**
 * A model from a map: contacts to distances, distances to places (§25.6).
 *
 * The pipeline is ShRec3D's (Lesne et al., *Nature Methods* 2014) with the
 * refinement every MDS toolkit ends with: a contact count `c` between two bins
 * is read as a distance `d = c^(−1/α)` — the same power law *Model versus map*
 * assumes in the other direction — the pairs the map never counted are filled
 * so the matrix is complete, classical MDS (Torgerson) turns the completed
 * matrix into three coordinates per bead, and SMACOF (de Leeuw's majorisation)
 * moves the beads until the distances the model has match the distances the
 * map implied, over the pairs the map counted and no other — each pair
 * **weighted by how often the map counted it**, so a pair seen ten thousand
 * times holds the beads where it says and a pair seen once barely moves them.
 *
 * That weighting is the point of refining at all. The count sets a pair's
 * target distance, but with every pair weighing the same the stress is
 * dominated by the far, weak, noisy pairs — a zero-count neighbour of a
 * 10,000-count pair is not a pair the model should honour equally — and the
 * near diagonal, where the map is most certain, is the part fitted worst.
 * `weightsFromContacts` turns the counts into weights; `smacofStep` takes
 * them; a pair the map never counted weighs nothing.
 *
 * Three departures from ShRec3D, all deliberate and all named in SPECS §25.6.
 * Its filling is a shortest-path closure over every pair, which is cubic in
 * the bead count and the one step that would not finish in a tab; the filling
 * here is the **mean observed distance at the same separation** for a pair on
 * one sequence and the mean observed trans distance for a pair across two,
 * which is what the closure converges towards on a map that decays with
 * distance and costs a pass over the matrix. And ShRec3D stops at the classical
 * solution; the refinement is what makes the model honour the counted pairs
 * rather than the filled ones, because a filled pair carries **no weight** in
 * it. And the refinement is weighted by the counts, where ShRec3D has no
 * refinement to weight.
 *
 * Everything here is arrays in and arrays out — no worker, no store, no DOM —
 * so the runner can breathe between iterations and the test can hold the
 * arithmetic still: a chain whose contacts were *made* from a known model comes
 * back as that model up to a rigid motion, which is the one golden a
 * reconstruction can honestly have (§18.9).
 */
 
/**
 * A dense `n × n` matrix this module reads but does not own.
 *
 * Both widths are accepted because the matrices it *makes* are `Float32Array`
 * — see {@link distancesFromContacts} — while a caller holding a `Float64Array`
 * (a test with a hand-written matrix) should not have to convert one to ask a
 * question of it.
 */
export type DenseMatrix = Float32Array | Float64Array
 
/** A configuration: one coordinate array per axis, one entry per bead. */
export interface Coordinates {
  x: Float64Array
  y: Float64Array
  z: Float64Array
}
 
/**
 * The distances a map implies, `n × n`, `NaN` where the map counted nothing.
 *
 * Symmetrised by the mean of the two triangles, because a read block is not
 * guaranteed symmetric to the last bit and a model has one distance per pair.
 * A pair with a zero count is **missing**, not infinitely far: a zero at 1 kb
 * is a pair the experiment did not sample, and `0^(−1/α)` would put the bead
 * at infinity and every other pair at nothing.
 */
/**
 * Keep only `keep`'s rows and columns, written over the matrix they came from.
 *
 * {@link subMatrix}'s answer without {@link subMatrix}'s second matrix, for the
 * caller whose matrix is large enough for a copy to matter.
 */
export function compactInPlace(values: Float32Array, n: number, keep: Int32Array): void {
  const m = keep.length
  // Forward and monotone, and safe for exactly that reason: the destination of
  // a cell never runs ahead of its source, since `keep[a] >= a` and `m <= n`,
  // so nothing is read after it has been written over.
  for (let a = 0; a < m; a++) {
    const row = keep[a] * n
    const out = a * m
    for (let b = 0; b < m; b++) values[out + b] = values[row + keep[b]]
  }
}
 
/**
 * The same law as {@link distancesFromContacts}, written over the counts.
 *
 * The counts are wanted twice — as distances and as weights — and never again,
 * and at twelve thousand beads each matrix is half a gigabyte. So the weights
 * are taken first and the counts are then *turned into* the distances where
 * they lie, which is one matrix rather than two. Reads and writes touch only
 * the two cells of the pair they are about, so the pass is safe in place.
 */
export function distancesOverContacts(values: Float32Array, n: number, alpha: number): void {
  const power = -1 / alpha
  for (let i = 0; i < n; i++) {
    for (let j = i + 1; j < n; j++) {
      const a = values[i * n + j]
      const b = values[j * n + i]
      const okA = Number.isFinite(a) && a > 0
      const okB = Number.isFinite(b) && b > 0
      const d = !okA && !okB ? Number.NaN : Math.pow(okA && okB ? (a + b) / 2 : okA ? a : b, power)
      values[i * n + j] = d
      values[j * n + i] = d
    }
    values[i * n + i] = 0
  }
}
 
export function distancesFromContacts(values: DenseMatrix, n: number, alpha: number): Float32Array {
  const out = new Float32Array(n * n).fill(Number.NaN)
  const power = -1 / alpha
  for (let i = 0; i < n; i++) {
    out[i * n + i] = 0
    for (let j = i + 1; j < n; j++) {
      const a = values[i * n + j]
      const b = values[j * n + i]
      const okA = Number.isFinite(a) && a > 0
      const okB = Number.isFinite(b) && b > 0
      if (!okA && !okB) continue
      const count = okA && okB ? (a + b) / 2 : okA ? a : b
      const d = Math.pow(count, power)
      out[i * n + j] = d
      out[j * n + i] = d
    }
  }
  return out
}
 
/**
 * The beads the map said nothing about at all.
 *
 * A bin whose row is entirely missing — a filtered bin under balancing, a gap
 * in the assembly, a centromere — has no place the data can give it, and a
 * bead with no constraint drifts to wherever the initialisation dropped it,
 * which looks exactly like a result. Such beads are dropped **before** solving
 * and the caller says how many; the alternative was a model with beads in it
 * nothing supports.
 */
export function constrainedBeads(distances: DenseMatrix, n: number): Int32Array {
  const keep: number[] = []
  for (let i = 0; i < n; i++) {
    let any = false
    for (let j = 0; j < n && !any; j++) if (j !== i && !Number.isNaN(distances[i * n + j])) any = true
    if (any) keep.push(i)
  }
  return Int32Array.from(keep)
}
 
/** The rows and columns of `distances` at `keep`, as a smaller square. */
export function subMatrix(distances: DenseMatrix, n: number, keep: Int32Array): Float32Array {
  const m = keep.length
  const out = new Float32Array(m * m)
  for (let a = 0; a < m; a++) {
    const row = keep[a] * n
    for (let b = 0; b < m; b++) out[a * m + b] = distances[row + keep[b]]
  }
  return out
}
 
/**
 * Every pair given a distance, so classical MDS has a full matrix to start from.
 *
 * `group[i]` is the sequence bead `i` lies on and `index[i]` its position along
 * it, so a missing pair is filled from the pairs that *are* there at the same
 * separation on the same sequence — the mean, walking outward to the nearest
 * separation with any observation when this one has none — and a pair across
 * two sequences from the mean of the observed trans pairs, or the largest cis
 * mean when the map counted nothing trans at all. The result is a copy;
 * `distances` keeps its `NaN`s, which the refinement needs to know which pairs
 * were counted.
 */
export function fillDistances(
  distances: DenseMatrix,
  n: number,
  group: Int32Array,
  index: Int32Array,
): Float32Array {
  const out = Float32Array.from(distances)
  fillDistancesInPlace(out, n, group, index)
  return out
}
 
/**
 * The same filling, written over the matrix it reads.
 *
 * What the copy above exists to protect is the record of *which pairs the map
 * counted*, which the `NaN`s are — and a caller that has the counts' weights
 * beside the distances has that record already, since a weight is zero exactly
 * where the count was missing. That caller is the runner, and it is the one
 * whose matrix is large enough for a second copy to matter: at ten thousand
 * beads a matrix is four hundred megabytes, and the run needs the distances,
 * the weights and the counts in memory at once.
 *
 * Safe in place because the two passes are separate: every observed distance
 * is read into the per-separation means *before* a single missing one is
 * written, so nothing this fills is read as evidence for what it fills.
 */
export function fillDistancesInPlace(
  out: Float32Array,
  n: number,
  group: Int32Array,
  index: Int32Array,
): void {
  const distances = out
  // Mean observed distance per cis separation, and the mean of every trans pair.
  let maxSeparation = 0
  for (let i = 0; i < n; i++) maxSeparation = Math.max(maxSeparation, index[i])
  const cisSum = new Float64Array(maxSeparation + 1)
  const cisCount = new Float64Array(maxSeparation + 1)
  let transSum = 0
  let transCount = 0
  for (let i = 0; i < n; i++) {
    for (let j = i + 1; j < n; j++) {
      const d = distances[i * n + j]
      if (Number.isNaN(d)) continue
      if (group[i] === group[j]) {
        const s = Math.abs(index[i] - index[j])
        cisSum[s] += d
        cisCount[s] += 1
      } else {
        transSum += d
        transCount += 1
      }
    }
  }
  const cisMean = new Float64Array(maxSeparation + 1).fill(Number.NaN)
  let largestCis = 0
  for (let s = 0; s <= maxSeparation; s++) {
    if (cisCount[s] === 0) continue
    cisMean[s] = cisSum[s] / cisCount[s]
    if (cisMean[s] > largestCis) largestCis = cisMean[s]
  }
  // Nearest observed separation, either way, so a gap in the middle of the
  // decay reads its neighbours rather than nothing.
  const nearest = new Float64Array(maxSeparation + 1)
  for (let s = 0; s <= maxSeparation; s++) {
    if (!Number.isNaN(cisMean[s])) {
      nearest[s] = cisMean[s]
      continue
    }
    let found = Number.NaN
    for (let step = 1; step <= maxSeparation && Number.isNaN(found); step++) {
      const lo = s - step
      const hi = s + step
      if (lo >= 0 && !Number.isNaN(cisMean[lo])) found = cisMean[lo]
      else if (hi <= maxSeparation && !Number.isNaN(cisMean[hi])) found = cisMean[hi]
    }
    nearest[s] = Number.isNaN(found) ? 1 : found
  }
  const trans = transCount > 0 ? transSum / transCount : largestCis > 0 ? largestCis : 1
  for (let i = 0; i < n; i++) {
    for (let j = i + 1; j < n; j++) {
      if (!Number.isNaN(distances[i * n + j])) continue
      const d = group[i] === group[j] ? nearest[Math.abs(index[i] - index[j])] : trans
      out[i * n + j] = d
      out[j * n + i] = d
    }
  }
}
 
/** How a pair's contact count becomes its weight in the fit. */
export type WeightMode = 'contacts' | 'sqrt' | 'equal'
 
/**
 * A weight per pair from the counts, `0` where the map counted nothing.
 *
 * `contacts` is the count itself — a pair counted a thousand times is a
 * thousand times as sure as a pair counted once, which is what a Poisson
 * count's precision says. `sqrt` is the count's standard deviation's worth,
 * the gentler reading that lets a weak pair still be heard on a deep map.
 * `equal` is one per counted pair: the plain SMACOF of every toolkit, kept
 * so a reader can see what the weighting changes. Symmetrised as the
 * distances are, by the mean of the two triangles.
 */
export function weightsFromContacts(values: DenseMatrix, n: number, mode: WeightMode): Float32Array {
  const out = new Float32Array(n * n)
  for (let i = 0; i < n; i++) {
    for (let j = i + 1; j < n; j++) {
      const a = values[i * n + j]
      const b = values[j * n + i]
      const okA = Number.isFinite(a) && a > 0
      const okB = Number.isFinite(b) && b > 0
      if (!okA && !okB) continue
      const count = okA && okB ? (a + b) / 2 : okA ? a : b
      const w = mode === 'contacts' ? count : mode === 'sqrt' ? Math.sqrt(count) : 1
      out[i * n + j] = w
      out[j * n + i] = w
    }
  }
  return out
}
 
/**
 * Torgerson's classical scaling: the three leading eigenvectors of the
 * double-centred squared distances, scaled by the roots of their eigenvalues.
 *
 * The eigenvectors come from subspace iteration — three random vectors,
 * multiplied through and re-orthogonalised until they settle — rather than a
 * full decomposition, because three are all that is wanted of a matrix that
 * can be two and a half thousand a side. The result is only an initialisation;
 * the refinement below is what the model is held to.
 */
export function classicalMds(full: DenseMatrix, n: number, sweeps = 80): Coordinates {
  const run = classicalMdsSweeps(full, n, sweeps)
  let step = run.next()
  while (step.done !== true) step = run.next()
  return step.value
}
 
/**
 * The same, a sweep at a time.
 *
 * At twelve thousand beads one sweep is a hundred and fifty million pair
 * visits and eighty of them are minutes, so a caller that runs this where
 * somebody is watching needs somewhere to say how far along it is and
 * somewhere to let a cancel land. Yielding after each sweep gives it both
 * without making the arithmetic asynchronous: {@link classicalMds} drains it
 * in a loop and is the same synchronous function it always was.
 */
export function* classicalMdsSweeps(
  full: DenseMatrix,
  n: number,
  sweeps = 80,
): Generator<{ sweep: number; sweeps: number }, Coordinates, void> {
  // B = −½ J D² J: the squared distances, centred by row, by column and
  // overall. **Never written down.** Only three vectors are ever wanted of it,
  // and every one of them is reached by multiplying B into a vector — which
  // the identity below does from the distances and two row sums, at the same
  // arithmetic and without the second `n × n` array. At ten thousand beads
  // that array is four hundred megabytes standing beside the matrix it was
  // copied from.
  const rowMean = new Float64Array(n)
  let grand = 0
  for (let i = 0; i < n; i++) {
    let sum = 0
    const row = i * n
    for (let j = 0; j < n; j++) {
      const d = full[row + j]
      sum += d * d
    }
    rowMean[i] = sum / n
    grand += sum
  }
  grand /= n * n
  const apply = (v: Float64Array, out: Float64Array): void => multiplyCentred(full, n, rowMean, grand, v, out)
  const q = [seeded(n, 1), seeded(n, 2), seeded(n, 3)]
  orthonormalise(q)
  const z = [new Float64Array(n), new Float64Array(n), new Float64Array(n)]
  for (let sweep = 0; sweep < sweeps; sweep++) {
    for (let k = 0; k < 3; k++) apply(q[k], z[k])
    for (let k = 0; k < 3; k++) q[k].set(z[k])
    orthonormalise(q)
    yield { sweep: sweep + 1, sweeps }
  }
  // Rayleigh quotients for the scale of each axis; a negative one — a matrix
  // that is not quite Euclidean, which a filled one never quite is — is an
  // axis the data does not support, drawn flat rather than imaginary.
  const axes = q.map((vector, k) => {
    apply(vector, z[k])
    let lambda = 0
    for (let i = 0; i < n; i++) lambda += vector[i] * z[k][i]
    const scale = Math.sqrt(Math.max(0, lambda))
    const out = new Float64Array(n)
    for (let i = 0; i < n; i++) out[i] = vector[i] * scale
    return out
  })
  return { x: axes[0], y: axes[1], z: axes[2] }
}
 
/**
 * `out = B v`, where `B` is the double-centred squared distances of `full`.
 *
 * Centring is rank-two plus a constant — `B = −½ (D² − r 1ᵀ − 1 rᵀ + g 11ᵀ)`
 * for row means `r` and grand mean `g` — so each of those three terms folds
 * into a scalar times a vector the loop already has, and B never has to exist
 * as an array. Exactly the matrix a materialised B would have given, to the
 * last bit of the arithmetic order.
 */
function multiplyCentred(
  full: DenseMatrix,
  n: number,
  rowMean: Float64Array,
  grand: number,
  v: Float64Array,
  out: Float64Array,
): void {
  let sumV = 0
  let dotR = 0
  for (let j = 0; j < n; j++) {
    sumV += v[j]
    dotR += rowMean[j] * v[j]
  }
  const tail = grand * sumV - dotR
  for (let i = 0; i < n; i++) {
    const row = i * n
    let sum = 0
    for (let j = 0; j < n; j++) {
      const d = full[row + j]
      sum += d * d * v[j]
    }
    out[i] = -0.5 * (sum - rowMean[i] * sumV + tail)
  }
}
 
/** A deterministic pseudo-random unit vector, so two runs give one answer. */
function seeded(n: number, seed: number): Float64Array {
  const out = new Float64Array(n)
  let state = 2463534242 + seed * 7919
  for (let i = 0; i < n; i++) {
    state ^= state << 13
    state ^= state >>> 17
    state ^= state << 5
    out[i] = ((state >>> 0) / 4294967296) * 2 - 1
  }
  return out
}
 
/** Gram–Schmidt in place; a vector that collapses is reseeded so three stay three. */
function orthonormalise(vectors: Float64Array[]): void {
  const n = vectors[0].length
  for (let k = 0; k < vectors.length; k++) {
    const v = vectors[k]
    for (let p = 0; p < k; p++) {
      const u = vectors[p]
      let dot = 0
      for (let i = 0; i < n; i++) dot += u[i] * v[i]
      for (let i = 0; i < n; i++) v[i] -= dot * u[i]
    }
    let norm = 0
    for (let i = 0; i < n; i++) norm += v[i] * v[i]
    norm = Math.sqrt(norm)
    if (norm < 1e-12) {
      v.set(seeded(n, 11 + k))
      k--
      continue
    }
    for (let i = 0; i < n; i++) v[i] /= norm
  }
}
 
/**
 * One SMACOF sweep: every bead moved to where the weighted majorant of the
 * stress puts it, given the others (Gansner, Koren & North's localised
 * stress majorisation, which is de Leeuw's Guttman transform taken one bead
 * at a time).
 *
 * `x_i ← Σ_j w_ij (x_j + δ_ij (x_i − x_j) / d_ij) / Σ_j w_ij`, in place, so a
 * bead moved early in the sweep is where its later neighbours see it — the
 * Gauss–Seidel reading, which converges faster than moving all at once and
 * never raises the stress: each move minimises a function that sits above
 * the stress and touches it where the bead was. Weighted, and the weights are
 * what make this a closed form: the full weighted Guttman transform is a
 * linear solve in the weight Laplacian, which is a pseudo-inverse of a
 * 2,500² matrix or a conjugate-gradient inner loop, and the bead-at-a-time
 * form costs one pass over the pairs like the unweighted one does.
 *
 * `weights` is `null` for one per counted pair. A pair with no target, or a
 * zero weight, contributes nothing. Returns the weighted normalised stress
 * **before** the sweep: `Σ w (δ − d)² / Σ w δ²`, 0 for a perfect fit.
 */
export function smacofStep(
  target: DenseMatrix,
  n: number,
  coords: Coordinates,
  weights: DenseMatrix | null = null,
  /**
   * The stress of `coords` before this step, where the caller has it already.
   *
   * A refinement loop measures the stress after every step to know when to
   * stop, and that number *is* the stress before the next one — so recomputing
   * it here costs a third of the whole refinement for an answer already in
   * hand. Left out, it is measured as before.
   */
  stressBefore?: number,
): number {
  const before = stressBefore ?? stressOf(target, n, coords, weights)
  const { x, y, z } = coords
  for (let i = 0; i < n; i++) {
    let sx = 0
    let sy = 0
    let sz = 0
    let sw = 0
    const row = i * n
    const xi = x[i]
    const yi = y[i]
    const zi = z[i]
    for (let j = 0; j < n; j++) {
      if (j === i) continue
      const delta = target[row + j]
      if (Number.isNaN(delta)) continue
      const w = weights === null ? 1 : weights[row + j]
      if (!(w > 0)) continue
      const dx = xi - x[j]
      const dy = yi - y[j]
      const dz = zi - z[j]
      const d = Math.sqrt(dx * dx + dy * dy + dz * dz)
      // Two beads on one point with a distance to keep have no direction to
      // keep it along: this term pulls onto the other bead and the next
      // sweep separates them once any other pair has moved one.
      const ratio = d > 1e-12 ? delta / d : 0
      sx += w * (x[j] + ratio * dx)
      sy += w * (y[j] + ratio * dy)
      sz += w * (z[j] + ratio * dz)
      sw += w
    }
    if (sw <= 0) continue
    x[i] = sx / sw
    y[i] = sy / sw
    z[i] = sz / sw
  }
  return before
}
 
/**
 * The weighted stress of a configuration against its targets, without
 * moving it: `Σ w (δ − d)² / Σ w δ²` over the counted pairs.
 */
export function stressOf(
  target: DenseMatrix,
  n: number,
  coords: Coordinates,
  weights: DenseMatrix | null = null,
): number {
  const { x, y, z } = coords
  let raw = 0
  let norm = 0
  for (let i = 0; i < n; i++) {
    for (let j = i + 1; j < n; j++) {
      const delta = target[i * n + j]
      if (Number.isNaN(delta)) continue
      const w = weights === null ? 1 : weights[i * n + j]
      if (!(w > 0)) continue
      const dx = x[i] - x[j]
      const dy = y[i] - y[j]
      const dz = z[i] - z[j]
      const gap = delta - Math.sqrt(dx * dx + dy * dy + dz * dz)
      raw += w * gap * gap
      norm += w * delta * delta
    }
  }
  return norm > 0 ? raw / norm : 0
}
 
/**
 * Centre the model and scale it so its root-mean-square radius is one.
 *
 * A map carries no physical unit and `c^(−1/α)` is in none, so the scale the
 * solver lands on is an accident of the counts; one that is the same for every
 * run puts two models of two maps beside each other at the same size, and
 * *model units* under a distance track mean the same thing on both.
 */
export function normaliseCoordinates(coords: Coordinates): void {
  const { x, y, z } = coords
  const n = x.length
  if (n === 0) return
  let cx = 0
  let cy = 0
  let cz = 0
  for (let i = 0; i < n; i++) {
    cx += x[i]
    cy += y[i]
    cz += z[i]
  }
  cx /= n
  cy /= n
  cz /= n
  let sum = 0
  for (let i = 0; i < n; i++) {
    x[i] -= cx
    y[i] -= cy
    z[i] -= cz
    sum += x[i] * x[i] + y[i] * y[i] + z[i] * z[i]
  }
  const rms = Math.sqrt(sum / n)
  if (rms <= 0) return
  for (let i = 0; i < n; i++) {
    x[i] /= rms
    y[i] /= rms
    z[i] /= rms
  }
}
 
/** Spearman's ρ between two lists, for the test and the runner's report. */
export function spearman(a: ArrayLike<number>, b: ArrayLike<number>): number {
  const n = a.length
  if (n < 3) return Number.NaN
  const ra = ranks(a)
  const rb = ranks(b)
  let ma = 0
  let mb = 0
  for (let i = 0; i < n; i++) {
    ma += ra[i]
    mb += rb[i]
  }
  ma /= n
  mb /= n
  let cov = 0
  let va = 0
  let vb = 0
  for (let i = 0; i < n; i++) {
    const da = ra[i] - ma
    const db = rb[i] - mb
    cov += da * db
    va += da * da
    vb += db * db
  }
  return va > 0 && vb > 0 ? cov / Math.sqrt(va * vb) : Number.NaN
}
 
function ranks(values: ArrayLike<number>): Float64Array {
  const n = values.length
  const order = Array.from({ length: n }, (_, i) => i).sort((p, q) => values[p] - values[q])
  const out = new Float64Array(n)
  let i = 0
  while (i < n) {
    let j = i
    while (j + 1 < n && values[order[j + 1]] === values[order[i]]) j++
    const rank = (i + j) / 2 + 1
    for (let k = i; k <= j; k++) out[order[k]] = rank
    i = j + 1
  }
  return out
}
 
/**
 * The whole road, for callers that do not need to breathe between steps.
 *
 * The runner takes the steps itself so a cancel lands between two of them;
 * the test and anything small takes this. Stops early once a step improves
 * the stress by less than `tolerance` of itself.
 */
export function reconstruct(
  distances: DenseMatrix,
  n: number,
  group: Int32Array,
  index: Int32Array,
  options: { iterations?: number; tolerance?: number; weights?: DenseMatrix | null } = {},
): { coords: Coordinates; stress: number; steps: number } {
  const iterations = options.iterations ?? 200
  const tolerance = options.tolerance ?? 1e-6
  const weights = options.weights ?? null
  const coords = classicalMds(fillDistances(distances, n, group, index), n)
  let stress = stressOf(distances, n, coords, weights)
  let steps = 0
  for (; steps < iterations; steps++) {
    const before = stress
    smacofStep(distances, n, coords, weights, before)
    const after = stressOf(distances, n, coords, weights)
    stress = after
    if (before - after < tolerance * before) {
      steps++
      break
    }
  }
  normaliseCoordinates(coords)
  return { coords, stress, steps }
}