Mantle Convection (3D)

// Mantle convection — 2D Boussinesq fluid (Rayleigh-Bénard) rendered as a
// textured slab inside a wireframe 3D cell. Drag horizontally to orbit,
// drag vertically to scrub the Rayleigh number across the critical value
// Ra_c ≈ 1708. Below Ra_c the fluid sits still (pure conduction, smooth
// vertical gradient). Above it, convection rolls organize. Push Ra high
// enough and the rolls turn unsteady, shedding plumes.
//
// Numerics: colocated grid, semi-Lagrangian advection for u/v/T, explicit
// viscous + thermal diffusion, buoyancy injected into v, then Jacobi
// pressure projection (∇²p = ∇·u/dt, then u ← u − ∇p · dt) to enforce
// ∇·u = 0. Periodic in x, no-slip + fixed-T walls in y.

const NX = 64;
const NY = 40;
const PROJ_ITERS = 24;
const SUBSTEPS = 2;
const DT = 0.018;

// Domain length scales: width = 2, height = 1 (aspect 2:1).
const DX = 2.0 / NX;
const DY = 1.0 / NY;

// Dimensionless. Ra is the knob; Pr fixed.
const PR = 1.0;
let RA = 5000;     // mutable via vertical drag
const RA_MIN = 200;
const RA_MAX = 80000;
const RA_C = 1708;

// Fields
let u, v, T;        // primary
let uTmp, vTmp, tTmp;
let p, div;         // for pressure projection

// THREE objects
let slabMesh, slabTex, slabData;
let cellBox;
let raSprite, raCtx, raCanvas, raTex;

// Camera
let userYaw = 0;
let pitch = -0.35;
let isDragging = false;
let lastMouseX = 0, lastMouseY = 0;
let idleTime = 0;
let elapsed = 0;

function idx(i, j) { return j * NX + i; }
function wrapX(i) { return ((i % NX) + NX) % NX; }
function clampY(j) { return j < 0 ? 0 : (j > NY - 1 ? NY - 1 : j); }

function makeField(initVal) {
  const a = new Float32Array(NX * NY);
  if (initVal !== 0) a.fill(initVal);
  return a;
}

function resetFields() {
  u = makeField(0);
  v = makeField(0);
  T = makeField(0);
  uTmp = makeField(0);
  vTmp = makeField(0);
  tTmp = makeField(0);
  p = makeField(0);
  div = makeField(0);
  // Initial temperature: linear gradient (cold top, hot bottom) plus tiny noise to seed instability.
  for (let j = 0; j < NY; j++) {
    const base = 1 - (j + 0.5) / NY;
    for (let i = 0; i < NX; i++) {
      const k = idx(i, j);
      T[k] = base + (Math.random() - 0.5) * 0.01;
    }
  }
}

// Bilinear sample, periodic in x, clamped in y.
function sampleField(f, x, y) {
  // x, y in grid units (cell centers at i+0.5, j+0.5)
  let xi = x - 0.5;
  let yj = y - 0.5;
  let i0 = Math.floor(xi);
  let j0 = Math.floor(yj);
  const fx = xi - i0;
  const fy = yj - j0;
  const i0w = wrapX(i0);
  const i1w = wrapX(i0 + 1);
  const j0c = clampY(j0);
  const j1c = clampY(j0 + 1);
  const a = f[idx(i0w, j0c)];
  const b = f[idx(i1w, j0c)];
  const c = f[idx(i0w, j1c)];
  const d = f[idx(i1w, j1c)];
  const ab = a + (b - a) * fx;
  const cd = c + (d - c) * fx;
  return ab + (cd - ab) * fy;
}

// Sample T with Dirichlet walls (T=1 below domain, T=0 above).
function sampleT(x, y) {
  let xi = x - 0.5;
  let yj = y - 0.5;
  let i0 = Math.floor(xi);
  let j0 = Math.floor(yj);
  const fx = xi - i0;
  const fy = yj - j0;
  const i0w = wrapX(i0);
  const i1w = wrapX(i0 + 1);
  const getRow = (jr) => {
    if (jr < 0) return [1, 1];
    if (jr > NY - 1) return [0, 0];
    return [T[idx(i0w, jr)], T[idx(i1w, jr)]];
  };
  const [a, b] = getRow(j0);
  const [c, d] = getRow(j0 + 1);
  const ab = a + (b - a) * fx;
  const cd = c + (d - c) * fx;
  return ab + (cd - ab) * fy;
}

function advect(dt) {
  // Semi-Lagrangian: trace back from each cell center along velocity.
  for (let j = 0; j < NY; j++) {
    for (let i = 0; i < NX; i++) {
      const k = idx(i, j);
      // Cell center position in grid units.
      const x = i + 0.5;
      const y = j + 0.5;
      // Convert velocity (in dimensionless physical units, scaled by DX/DY)
      // to grid-unit displacement. We treat one grid cell == one unit here:
      // u is already in physical units where 1 grid cell width = DX. So
      // back-trace by velocity * dt / DX in x and DY in y.
      const px = x - u[k] * dt / DX;
      const py = y - v[k] * dt / DY;
      uTmp[k] = sampleField(u, px, py);
      vTmp[k] = sampleField(v, px, py);
      tTmp[k] = sampleT(px, py);
    }
  }
  // Swap
  let tmp = u; u = uTmp; uTmp = tmp;
  tmp = v; v = vTmp; vTmp = tmp;
  tmp = T; T = tTmp; tTmp = tmp;
}

function diffuseAndBuoyancy(dt) {
  // Pr = nu / kappa. Pick units so kappa is what Ra reflects.
  // Effective diffusion coefficients (lattice-stable explicit step):
  // nu_grid = nu * dt / DX^2 must stay small. We pick small base values.
  const nu = PR * 0.004;
  const kappa = 0.004;
  // Buoyancy coefficient: derive from Ra. Ra = g*beta*dT*L^3 / (nu*kappa);
  // with L=1, dT=1, so g*beta = Ra * nu * kappa.
  const buoy = RA * nu * kappa;

  // Velocity viscous diffusion (5-point Laplacian, periodic x, no-slip y means u=v=0 at walls).
  for (let j = 0; j < NY; j++) {
    for (let i = 0; i < NX; i++) {
      const k = idx(i, j);
      const iL = idx(wrapX(i - 1), j);
      const iR = idx(wrapX(i + 1), j);
      const uU = (j === NY - 1) ? -u[k] : u[idx(i, j + 1)]; // no-slip: ghost = -interior
      const uD = (j === 0)      ? -u[k] : u[idx(i, j - 1)];
      const vU = (j === NY - 1) ? -v[k] : v[idx(i, j + 1)];
      const vD = (j === 0)      ? -v[k] : v[idx(i, j - 1)];
      const lapU = (u[iL] + u[iR] - 2 * u[k]) / (DX * DX)
                 + (uU + uD - 2 * u[k]) / (DY * DY);
      const lapV = (v[iL] + v[iR] - 2 * v[k]) / (DX * DX)
                 + (vU + vD - 2 * v[k]) / (DY * DY);
      uTmp[k] = u[k] + dt * nu * lapU;
      vTmp[k] = v[k] + dt * nu * lapV;
    }
  }
  let tmp = u; u = uTmp; uTmp = tmp;
  tmp = v; v = vTmp; vTmp = tmp;

  // Thermal diffusion (Dirichlet T=1 at j=-1, T=0 at j=NY).
  for (let j = 0; j < NY; j++) {
    for (let i = 0; i < NX; i++) {
      const k = idx(i, j);
      const iL = idx(wrapX(i - 1), j);
      const iR = idx(wrapX(i + 1), j);
      const tU = (j === NY - 1) ? (2 * 0 - T[k]) : T[idx(i, j + 1)];
      const tD = (j === 0)      ? (2 * 1 - T[k]) : T[idx(i, j - 1)];
      const lapT = (T[iL] + T[iR] - 2 * T[k]) / (DX * DX)
                 + (tU + tD - 2 * T[k]) / (DY * DY);
      tTmp[k] = T[k] + dt * kappa * lapT;
    }
  }
  tmp = T; T = tTmp; tTmp = tmp;

  // Buoyancy: dv/dt += buoy * (T - <T_horiz_mean>). Subtract horizontal mean
  // so the constant conduction profile doesn't generate spurious flow.
  for (let j = 0; j < NY; j++) {
    // horizontal mean at this row
    let s = 0;
    for (let i = 0; i < NX; i++) s += T[idx(i, j)];
    const mean = s / NX;
    for (let i = 0; i < NX; i++) {
      const k = idx(i, j);
      v[k] += dt * buoy * (T[k] - mean);
    }
  }
}

function project(dt) {
  // ∇·u in dimensionless grid form
  for (let j = 0; j < NY; j++) {
    for (let i = 0; i < NX; i++) {
      const k = idx(i, j);
      const uR = u[idx(wrapX(i + 1), j)];
      const uL = u[idx(wrapX(i - 1), j)];
      const vU = (j === NY - 1) ? 0 : v[idx(i, j + 1)];
      const vD = (j === 0)      ? 0 : v[idx(i, j - 1)];
      div[k] = ((uR - uL) / (2 * DX) + (vU - vD) / (2 * DY)) / dt;
      p[k] = 0;
    }
  }
  // Jacobi iterations: ∇²p = div  →  p_ij = (avg neighbors - h² div)/4 (with anisotropic spacing)
  const ax = 1 / (DX * DX), ay = 1 / (DY * DY);
  const ac = 2 * (ax + ay);
  for (let it = 0; it < PROJ_ITERS; it++) {
    for (let j = 0; j < NY; j++) {
      for (let i = 0; i < NX; i++) {
        const k = idx(i, j);
        const pL = p[idx(wrapX(i - 1), j)];
        const pR = p[idx(wrapX(i + 1), j)];
        // Neumann at y walls (dp/dy = 0)
        const pU = (j === NY - 1) ? p[k] : p[idx(i, j + 1)];
        const pD = (j === 0)      ? p[k] : p[idx(i, j - 1)];
        p[k] = (ax * (pL + pR) + ay * (pU + pD) - div[k]) / ac;
      }
    }
  }
  // Subtract gradient.
  for (let j = 0; j < NY; j++) {
    for (let i = 0; i < NX; i++) {
      const k = idx(i, j);
      const pR = p[idx(wrapX(i + 1), j)];
      const pL = p[idx(wrapX(i - 1), j)];
      const pU = (j === NY - 1) ? p[k] : p[idx(i, j + 1)];
      const pD = (j === 0)      ? p[k] : p[idx(i, j - 1)];
      u[k] -= dt * (pR - pL) / (2 * DX);
      v[k] -= dt * (pU - pD) / (2 * DY);
      // Re-enforce no-slip at walls.
      if (j === 0 || j === NY - 1) { u[k] = 0; v[k] = 0; }
    }
  }
}

function step(dt) {
  advect(dt);
  diffuseAndBuoyancy(dt);
  project(dt);
}

// ---- temperature → color (a thermal palette) ----
function tempToRGB(t) {
  // Clamp
  const x = Math.max(0, Math.min(1, t));
  // 4-stop palette: deep blue → cyan → yellow → red
  const stops = [
    [0.00, 0.05, 0.10, 0.30],
    [0.35, 0.10, 0.55, 0.85],
    [0.65, 0.95, 0.85, 0.25],
    [1.00, 0.95, 0.20, 0.10],
  ];
  for (let s = 0; s < stops.length - 1; s++) {
    if (x <= stops[s + 1][0]) {
      const t0 = stops[s][0], t1 = stops[s + 1][0];
      const f = (x - t0) / (t1 - t0);
      return [
        stops[s][1] + (stops[s + 1][1] - stops[s][1]) * f,
        stops[s][2] + (stops[s + 1][2] - stops[s][2]) * f,
        stops[s][3] + (stops[s + 1][3] - stops[s][3]) * f,
      ];
    }
  }
  return [stops[stops.length - 1][1], stops[stops.length - 1][2], stops[stops.length - 1][3]];
}

function writeSlabTexture() {
  for (let j = 0; j < NY; j++) {
    for (let i = 0; i < NX; i++) {
      const t = T[idx(i, j)];
      const [r, g, b] = tempToRGB(t);
      const o = (j * NX + i) * 4;
      slabData[o]     = (r * 255) | 0;
      slabData[o + 1] = (g * 255) | 0;
      slabData[o + 2] = (b * 255) | 0;
      slabData[o + 3] = 255;
    }
  }
  slabTex.needsUpdate = true;
}

function regimeLabel(ra) {
  if (ra < RA_C * 0.95) return "conduction";
  if (ra < RA_C * 1.6) return "onset";
  if (ra < 12000) return "steady rolls";
  if (ra < 35000) return "unsteady rolls";
  return "plume regime";
}

function updateRaReadout() {
  raCtx.clearRect(0, 0, raCanvas.width, raCanvas.height);
  raCtx.fillStyle = "rgba(8,10,18,0.78)";
  raCtx.fillRect(0, 0, raCanvas.width, raCanvas.height);
  raCtx.fillStyle = "#e5ecff";
  raCtx.font = "bold 44px ui-monospace, monospace";
  raCtx.textBaseline = "top";
  raCtx.fillText("Ra = " + Math.round(RA).toString(), 16, 10);
  raCtx.font = "30px ui-monospace, monospace";
  raCtx.fillStyle = RA < RA_C ? "#7fb6ff" : "#ffb070";
  raCtx.fillText(regimeLabel(RA), 16, 64);
  raCtx.font = "22px ui-monospace, monospace";
  raCtx.fillStyle = "#8c93a8";
  raCtx.fillText("Ra_c ≈ " + RA_C, 16, 104);
  raTex.needsUpdate = true;
}

function init({ scene, camera, renderer, width, height }) {
  renderer.setClearColor(0x06080f, 1);
  camera.position.set(0, 0.6, 3.2);
  camera.lookAt(0, 0, 0);

  // Soft directional light — gives the wireframe box some depth cue.
  const amb = new THREE.AmbientLight(0xffffff, 0.6);
  scene.add(amb);
  const dir = new THREE.DirectionalLight(0xffffff, 0.5);
  dir.position.set(2, 3, 4);
  scene.add(dir);

  resetFields();

  // Slab: a thin plane in the X-Y plane (width 2.4, height 1.2) textured with T.
  slabData = new Uint8Array(NX * NY * 4);
  slabTex = new THREE.DataTexture(slabData, NX, NY, THREE.RGBAFormat);
  slabTex.magFilter = THREE.LinearFilter;
  slabTex.minFilter = THREE.LinearFilter;
  slabTex.wrapS = THREE.RepeatWrapping;
  slabTex.needsUpdate = true;
  const slabGeo = new THREE.PlaneGeometry(2.4, 1.2);
  const slabMat = new THREE.MeshBasicMaterial({ map: slabTex });
  slabMesh = new THREE.Mesh(slabGeo, slabMat);
  scene.add(slabMesh);

  // Wireframe cell box around it: thin box 2.4 × 1.2 × 0.04.
  const boxGeo = new THREE.BoxGeometry(2.42, 1.22, 0.06);
  const edges = new THREE.EdgesGeometry(boxGeo);
  cellBox = new THREE.LineSegments(
    edges,
    new THREE.LineBasicMaterial({ color: 0x9fb3ff, transparent: true, opacity: 0.55 }),
  );
  scene.add(cellBox);

  // Top/bottom thermal-boundary indicator bars (cold blue strip on top, hot orange below).
  const hotGeo = new THREE.BoxGeometry(2.42, 0.045, 0.07);
  const hotMat = new THREE.MeshBasicMaterial({ color: 0xff5a2a });
  const hotBar = new THREE.Mesh(hotGeo, hotMat);
  hotBar.position.set(0, -0.62, 0);
  scene.add(hotBar);
  const coldMat = new THREE.MeshBasicMaterial({ color: 0x2a78ff });
  const coldBar = new THREE.Mesh(hotGeo, coldMat);
  coldBar.position.set(0, 0.62, 0);
  scene.add(coldBar);

  // Ra readout: a CanvasTexture sprite that floats above the cell.
  raCanvas = document.createElement("canvas");
  raCanvas.width = 360;
  raCanvas.height = 140;
  raCtx = raCanvas.getContext("2d");
  raTex = new THREE.CanvasTexture(raCanvas);
  raTex.minFilter = THREE.LinearFilter;
  const raMat = new THREE.SpriteMaterial({ map: raTex, transparent: true });
  raSprite = new THREE.Sprite(raMat);
  raSprite.scale.set(1.1, 0.43, 1);
  raSprite.position.set(-0.6, 0.95, 0);
  scene.add(raSprite);
  updateRaReadout();

  // Burn-in: a few steps so the initial noise resolves before render.
  for (let i = 0; i < 8; i++) step(DT);
  writeSlabTexture();

  return { scene, camera };
}

function tick({ dt, time, scene, camera, renderer, width, height, input }) {
  elapsed += dt;
  const clicks = input.consumeClicks();
  void clicks;

  if (input.mouseDown) {
    if (!isDragging) {
      isDragging = true;
      lastMouseX = input.mouseX;
      lastMouseY = input.mouseY;
    } else {
      const dx = input.mouseX - lastMouseX;
      const dy = input.mouseY - lastMouseY;
      // Horizontal drag orbits the camera around Y.
      userYaw -= dx * 0.006;
      // Vertical drag scrubs Ra logarithmically — feels good across 200..80000.
      if (dy !== 0) {
        const logMin = Math.log(RA_MIN);
        const logMax = Math.log(RA_MAX);
        const cur = Math.log(RA);
        // dy in canvas pixels; map full canvas height ≈ full range.
        const span = logMax - logMin;
        const next = cur + dy / Math.max(120, height) * span;
        RA = Math.exp(Math.max(logMin, Math.min(logMax, next)));
        updateRaReadout();
      }
      lastMouseX = input.mouseX;
      lastMouseY = input.mouseY;
    }
    idleTime = 0;
  } else {
    isDragging = false;
    idleTime += dt;
    if (idleTime > 1.5) userYaw += dt * 0.16;
  }

  // Camera orbit
  const r = 3.0;
  const cy = Math.cos(userYaw), sy = Math.sin(userYaw);
  const cp = Math.cos(pitch), sp = Math.sin(pitch);
  camera.position.set(r * cp * sy, -r * sp, r * cp * cy);
  camera.lookAt(0, 0, 0);
  camera.aspect = width / height;
  camera.updateProjectionMatrix();
  renderer.setSize(width, height, false);

  // Keep the slab + Ra sprite facing the camera-ish: actually, the slab is
  // a flat plane and we want to *see* it edge-on sometimes — that's the
  // point of orbiting. Leave it fixed; only billboard the Ra sprite (sprites
  // are camera-aligned automatically).

  // Sim sub-steps. Skip when Ra is sub-critical and the field is essentially
  // frozen — but still advance diffusion so the conductive profile settles.
  for (let s = 0; s < SUBSTEPS; s++) step(DT);

  writeSlabTexture();

  // Subtle pulse on the cell box opacity tied to flow strength.
  let kin = 0;
  for (let k = 0; k < u.length; k += 8) kin += u[k] * u[k] + v[k] * v[k];
  const intensity = Math.min(1, Math.sqrt(kin / (u.length / 8)) * 4);
  cellBox.material.opacity = 0.4 + 0.4 * intensity;

  renderer.render(scene, camera);
}