Drift-diffusion balance and the Stokes-Einstein relation

export function driftDiffusionFrames1d(
  p: DriftDiffusionParameters,
  budget = REFERENCE_FTCS_BUDGET,
): Computation<DriftDiffusionFrames> {
  const {
    cells,
    width,
    frames,
    stepsPerFrame,
    dt,
    kickDiffusivity,
    mobility,
    force,
    temperature,
    profile,
    set,
  } = p;
  if (!integer(cells, 3) || !integer(frames, 1) || !integer(stepsPerFrame, 1)) {
    return invalid(
      ["cells", "frames", "stepsPerFrame"],
      "Use at least three cells and positive whole-number frame and step counts.",
    );
  }
  if (
    ![width, dt, kickDiffusivity, mobility, force, temperature].every(Number.isFinite) ||
    width <= 0 ||
    dt <= 0 ||
    kickDiffusivity < 0 ||
    mobility <= 0 ||
    temperature <= 0
  ) {
    return invalid(
      ["width", "dt", "kickDiffusivity", "mobility", "force", "temperature"],
      "Use positive width, dt, mobility, and temperature, nonnegative kick diffusivity, and a finite force.",
    );
  }
  const D = kickDiffusivity;
  const u = mobility * force;
  const dx = width / cells;
  if (!Number.isFinite(u) || !Number.isFinite(1 / dx) || dx <= 0) {
    return invalid(
      ["width", "cells", "mobility", "force"],
      "The drift speed and cell scale must be representable.",
    );
  }
  const work = cells * (frames - 1) * stepsPerFrame;
  if (work > budget.workUnits || cells * frames * 8 > budget.allocationBytes) {
    return {
      kind: "outcome",
      outcome: {
        outcome: "budget-exhausted",
        message: executionOutcomeRegistry["budget-exhausted"].message,
        retry: executionOutcomeRegistry["budget-exhausted"].retry,
        requested: { workUnits: work, allocationBytes: cells * frames * 8 },
        allowed: { ...budget },
      },
    };
  }
  const speeds = transportSpeeds(D, u, dx);
  const rate = (speeds.right + speeds.left) / dx;
  const sigma = dt * rate;
  if (!Number.isFinite(rate) || !Number.isFinite(sigma)) {
    return invalid(
      ["dt", "width", "kickDiffusivity", "force"],
      "The explicit transport rate must be finite.",
    );
  }
  if (sigma > 1) {
    const dtMax = 1 / rate;
    const driftOnly = D === 0;
    return {
      kind: "refused",
      refusal: makeRefusal(
        driftOnly ? "drift-cfl-exceeded" : "drift-diffusion-unstable",
        {
          parameterIds: driftOnly ? ["dt", "force"] : ["dt"],
          capabilityId: "diffusion.driftDiffusionFrames1d",
        },
        {
          details: driftOnly
            ? { courant: sigma, limit: 1, dtMax }
            : { ratio: sigma, limit: 1, dtMax },
          rankedRepairs: [
            {
              label: "Reduce the time step to the explicit scheme's positivity limit.",
              action: { parameterId: "dt", value: dtMax },
            },
            {
              label: "Use a coarser spatial grid.",
              action: { parameterId: "cells", value: Math.max(3, Math.floor(cells / 2)) },
            },
            ...(driftOnly
              ? [
                  {
                    label: "Use a weaker force.",
                    action: { parameterId: "force", value: force / 2 },
                  },
                ]
              : []),
          ],
        },
      ),
    };
  }
  const field = initialProfile(cells, dx, width, D, u, force, temperature, profile, set);
  if (!field) {
    return invalid(
      ["profile"],
      "The initial profile must have a finite, positive, representable mass.",
    );
  }
  const values = new Float64Array(cells * frames);
  values.set(field);
  // Reuse numerical buffers. Face objects are needed only for the accepted final frame.
  const flux = new Float64Array(cells + 1);
  const next = new Float64Array(cells);
  for (let frame = 1; frame < frames; frame++) {
    for (let s = 0; s < stepsPerFrame; s++) {
      const step = advance(field, flux, next, dx, dt, speeds);
      if (step.kind === "refused") return step;
    }
    values.set(field, frame * cells);
  }
  const pe = D === 0 ? (u === 0 ? 0 : Math.sign(u) * Number.POSITIVE_INFINITY) : u / (D / dx);
  return {
    kind: "accepted",
    data: Object.freeze({
      values,
      shape: Object.freeze([frames, cells] as const),
      dx,
      peclet: pe,
      sigma,
      driftVelocity: u,
      // Diagnostics describe this frame, not the field before its last time step.
      faceFlux: Object.freeze(faceFlux(field, speeds, D, u)),
      boundary: "zero-flux",
      ownerId: "diffusion.driftDiffusionFrames1d",
      modelIdentity: D === 0 ? DRIFT_ONLY_NO_KICKS : "drift-diffusion",
    }),
  };
}

export function equilibriumBalance(
  {
    force,
    temperature,
    eta,
    a,
    kickDiffusivity,
  }: {
    force: number;
    temperature: number;
    eta: number;
    a: number;
    kickDiffusivity: number;
    profile?: string;
  },
  set: ConstantSet,
): EquilibriumBalance | RouteAEvaluation {
  if (
    ![force, temperature, eta, a, kickDiffusivity].every(Number.isFinite) ||
    temperature <= 0 ||
    eta <= 0 ||
    a <= 0 ||
    kickDiffusivity < 0
  ) {
    return outside(
      "diffusionCoefficient",
      "equilibriumBalance",
      "T > 0, eta > 0, a > 0, D_kicks >= 0",
      "Use positive temperature, viscosity, and radius.",
      "input",
      set,
    );
  }
  const muEval = stokesMobility(eta, a);
  if (muEval.result.status !== "value") return muEval;
  const mu = muEval.result.value as number;
  const kT = thermalConstant(set).value * temperature;
  const mobilityD = mu * kT;
  const relation: ScientificResult = Object.freeze({
    ...identity("diffusionCoefficient", "equilibriumBalance"),
    status: "symbolic",
    expressionRef: "D = mu k_B T",
    unspecifiedSymbols: Object.freeze(["D", "mu", "k_B", "T"]),
  });
  if (force === 0) {
    return {
      mobilityD: number("diffusionCoefficient", "equilibriumBalance", mobilityD, set),
      balanceD: wrap({
        ...identity("kickDiffusivity", "equilibriumBalance"),
        status: "not-applicable",
        reason:
          "0/0 is not an evaluation; the equilibrium construction relates two descriptions of the same state",
      }),
      relation,
      agree: null,
      cancellationFactors: null,
      kickStrength: kickDiffusivity / mobilityD,
    };
  }
  const lambdaKin = kickDiffusivity / (mu * Math.abs(force));
  const balanceD = mu * Math.abs(force) * lambdaKin;
  const m = kickDiffusivity / mobilityD;
  return {
    mobilityD: number("diffusionCoefficient", "equilibriumBalance", mobilityD, set),
    balanceD: number("kickDiffusivity", "equilibriumBalance", balanceD, set),
    relation,
    agree: Math.abs(m - 1) <= 1e-9,
    cancellationFactors: Object.freeze([`mu*|F|`, `k_B T / |F|`]),
    kickStrength: m,
  };
}

export function decayLengths(
  {
    force,
    temperature,
    kickDiffusivity,
    mobility,
  }: { force: number; temperature: number; kickDiffusivity: number; mobility: number },
  set: ConstantSet,
): DecayLengths {
  const na = (q: Quantity, reason: string) =>
    wrap({
      ...identity(q, "decayLengths"),
      status: "not-applicable",
      reason,
    });
  if (
    ![force, temperature, kickDiffusivity, mobility].every(Number.isFinite) ||
    temperature <= 0 ||
    kickDiffusivity < 0 ||
    mobility <= 0
  ) {
    const bad = outside(
      "osmoticDecayLength",
      "decayLengths",
      "T > 0, mobility > 0, D_kicks >= 0",
      "Use a positive temperature and mobility and a nonnegative kick diffusivity.",
      "input",
      set,
    );
    return { osmotic: bad, kinetic: bad, kickStrength: null };
  }
  if (force === 0) {
    const reason = "no force, no gradient: the equilibrium profile is uniform";
    return {
      osmotic: na("osmoticDecayLength", reason),
      kinetic: na("kineticDecayLength", reason),
      kickStrength: null,
    };
  }
  const kT = thermalConstant(set).value * temperature;
  const osmotic = kT / Math.abs(force);
  const kinetic = kickDiffusivity / (mobility * Math.abs(force));
  const m = kickDiffusivity / (mobility * kT);
  return {
    osmotic: number("osmoticDecayLength", "decayLengths", osmotic, set),
    kinetic: number("kineticDecayLength", "decayLengths", kinetic, set),
    kickStrength: m,
  };
}

export function osmoticEquilibriumProfile(
  x: number,
  width: number,
  force: number,
  temperature: number,
  total: number,
  set: ConstantSet,
): RouteAEvaluation {
  if (
    ![x, width, force, temperature, total].every(Number.isFinite) ||
    width <= 0 ||
    temperature <= 0 ||
    total < 0 ||
    x < 0 ||
    x > width
  ) {
    return outside(
      "probabilityDensity",
      "osmoticEquilibriumProfile",
      "box [0, W], T > 0, N_tot >= 0",
      "Evaluate inside a positive-width box at positive temperature.",
      "input",
      set,
    );
  }
  if (force === 0)
    return number("probabilityDensity", "osmoticEquilibriumProfile", total / width, set);
  const absL = Math.abs((thermalConstant(set).value * temperature) / force);
  // n = N_tot (1/λ) e^{x/λ} / (e^{W/λ} - 1) for F > 0, mirrored for F < 0.
  // Rewrite as (N_tot/λ) e^{z-w} / (1 - e^{-w}) so large W/λ never overflows an intermediate exp.
  const z = force > 0 ? x / absL : (width - x) / absL;
  const w = width / absL;
  const denom = -Math.expm1(-w);
  if (!(denom > 0)) {
    return outside(
      "probabilityDensity",
      "osmoticEquilibriumProfile",
      "W/lambda underflow",
      "The normalized profile's denominator is not a positive finite number.",
      "numerical",
      set,
    );
  }
  const logN = Math.log(total / absL) + (z - w) - Math.log(denom);
  if (!Number.isFinite(logN) || logN > 700) {
    return outside(
      "probabilityDensity",
      "osmoticEquilibriumProfile",
      "W/lambda too large",
      "The normalized profile overflowed binary64 even in log space.",
      "numerical",
      set,
    );
  }
  return number("probabilityDensity", "osmoticEquilibriumProfile", Math.exp(logN), set);
}