The Brownian tracer ensemble

/** SI: temperature K, viscosity Pa s, radius m, output m2/s. No ambient constants. */
export function stokesEinsteinD(
  {
    T,
    eta,
    a,
    medium = "liquid",
  }: { T: number; eta: number; a: number; medium?: "liquid" | "gas" },
  set: ConstantSet,
): Evaluation {
  if (![T, eta, a].every((v) => Number.isFinite(v) && v > 0))
    return outside(
      "diffusionCoefficient",
      "stokesEinsteinD",
      "T > 0, eta > 0, a > 0",
      "Enter positive finite temperature, viscosity, and particle radius.",
      "input",
      set,
    );
  if (medium !== "liquid")
    return outside(
      "diffusionCoefficient",
      "stokesEinsteinD",
      "stokes-gas-medium",
      "The liquid Stokes-drag model does not include the slip correction needed in a gas.",
      "model",
      set,
    );
  const k = thermalConstant(set);
  return number(
    "diffusionCoefficient",
    "stokesEinsteinD",
    (k.value * T) / (6 * Math.PI * eta * a),
    set,
    undefined,
    true,
  );
}

export function rmsDisplacement(D: number, t: number): Evaluation {
  if (!validDt(D, t))
    return outside(
      "rmsDisplacement1d",
      "rmsDisplacement",
      "D >= 0 and t >= 0",
      "Diffusivity and elapsed time must be finite and nonnegative.",
    );
  return number(
    "rmsDisplacement1d",
    "rmsDisplacement",
    scale(D, t),
    undefined,
    undefined,
    D > 0 && t > 0,
  );
}

export function apparentSpeed(D: number, tau: number): Evaluation {
  if (!validDt(D, tau) || tau === 0)
    return outside(
      "apparentSpeed",
      "apparentSpeed",
      "D >= 0 and tau > 0",
      "An apparent speed needs a positive observation interval.",
    );
  return number(
    "apparentSpeed",
    "apparentSpeed",
    (Math.SQRT2 * Math.sqrt(D)) / Math.sqrt(tau),
    undefined,
    undefined,
    D > 0,
  );
}

/** Compensated sums over ALL members, never a viewport-selected subset. Row-major M by d. */
export function ensembleMoments({
  displacements,
  d,
  dt,
  viewport,
}: {
  displacements: Float64Array;
  d: number;
  dt?: number;
  viewport?: ViewportBounds;
}): Computation<EnsembleMoments> {
  if (
    !(displacements instanceof Float64Array) ||
    ![1, 2, 3].includes(d) ||
    displacements.length === 0 ||
    displacements.length % d !== 0 ||
    displacements.length > 30000 ||
    !displacements.every(Number.isFinite)
  )
    return invalid(
      ["displacements", "d"],
      "Provide finite, row-major displacements with one, two or three coordinates per tracer.",
    );
  const M = displacements.length / d,
    sums = new Float64Array(d * 3 + 1),
    corrections = new Float64Array(d * 3 + 1);
  function add(i: number, v: number) {
    const corr = corrections[i] ?? 0,
      sum = sums[i] ?? 0,
      y = v - corr,
      t = sum + y;
    corrections[i] = t - sum - y;
    sums[i] = t;
  }
  for (let i = 0; i < M; i++) {
    let norm = 0;
    for (let j = 0; j < d; j++) {
      const v = displacements[i * d + j];
      if (v === undefined) return failure("A required displacement coordinate was undefined.");
      if (v !== 0 && v * v === 0)
        return failure("A nonzero squared displacement is below the representable range.");
      add(j * 3, v);
      add(j * 3 + 1, Math.abs(v));
      add(j * 3 + 2, v * v);
      norm = Math.hypot(norm, v);
    }
    add(d * 3, norm);
  }
  const axes = Array.from({ length: d }, (_, j) => {
    const mean = (sums[j * 3] ?? 0) / M;
    const meanAbsolute = (sums[j * 3 + 1] ?? 0) / M;
    const meanSquare = (sums[j * 3 + 2] ?? 0) / M;
    return {
      mean,
      meanAbsolute,
      meanSquare,
      rms: Math.sqrt(meanSquare),
    };
  });
  const meanSquareNorm = axes.reduce((s, a) => s + a.meanSquare, 0);
  if (![...sums, meanSquareNorm].every(Number.isFinite))
    return failure("The moment reduction exceeded the numerical range.");
  const rmsNorm = Math.sqrt(meanSquareNorm);
  let insideCount: number | undefined;
  let outsideCount: number | undefined;
  if (viewport) {
    let inc = 0;
    let outc = 0;
    for (let i = 0; i < M; i++) {
      let isInside = true;
      for (let j = 0; j < d; j++) {
        const v = displacements[i * d + j] ?? 0;
        const minVal = viewport.min[j] ?? -Infinity;
        const maxVal = viewport.max[j] ?? Infinity;
        if (v < minVal || v > maxVal) {
          isInside = false;
          break;
        }
      }
      if (isInside) inc++;
      else outc++;
    }
    insideCount = inc;
    outsideCount = outc;
  }
  const apparentSpeed =
    dt !== undefined && Number.isFinite(dt) && dt > 0
      ? (d === 1 ? (axes[0]?.rms ?? 0) : rmsNorm) / dt
      : undefined;
  return {
    kind: "accepted",
    data: {
      M,
      axes,
      meanNorm: (sums[d * 3] ?? 0) / M,
      meanSquareNorm,
      rmsNorm,
      ...(apparentSpeed !== undefined ? { apparentSpeed } : {}),
      ...(insideCount !== undefined && outsideCount !== undefined ? { insideCount, outsideCount } : {}),
    },
  };
}

export function displacementHistogram(
  values: Float64Array,
  edges: Float64Array,
): Computation<{ counts: Float64Array; underflow: number; overflow: number; total: number }> {
  if (
    values.length === 0 ||
    values.length > 10000 ||
    edges.length < 2 ||
    edges.length > 1001 ||
    !values.every(Number.isFinite) ||
    !edges.every((v, i) => {
      if (!Number.isFinite(v)) return false;
      if (i === 0) return true;
      const prev = edges[i - 1];
      return prev !== undefined && v > prev;
    })
  )
    return invalid(
      ["values", "edges"],
      "Use finite samples and strictly increasing histogram edges.",
    );
  const counts = new Float64Array(edges.length - 1);
  let underflow = 0,
    overflow = 0;
  const firstEdge = edges[0];
  const lastEdge = edges[edges.length - 1];
  if (firstEdge === undefined || lastEdge === undefined)
    return invalid(["edges"], "Histogram edges array must contain at least two finite bounds.");
  for (const value of values) {
    if (value < firstEdge) {
      underflow++;
      continue;
    }
    if (value > lastEdge) {
      overflow++;
      continue;
    }
    let lo = 0,
      hi = edges.length - 1;
    while (hi - lo > 1) {
      const mid = (lo + hi) >>> 1;
      const edgeMid = edges[mid];
      if (edgeMid !== undefined && value < edgeMid) hi = mid;
      else lo = mid;
    }
    const bin = Math.min(lo, counts.length - 1);
    const prevCount = counts[bin];
    if (prevCount !== undefined) {
      counts[bin] = prevCount + 1;
    }
  }
  return { kind: "accepted", data: { counts, underflow, overflow, total: values.length } };
}