HP-Model Protein Fold (3D)

// HP-model protein fold caricature.
//
// A 1D chain of beads, each labeled H (hydrophobic) or P (polar), is "folded"
// by a soft-bodied simulated-annealing-ish dynamics: bond springs, weak H–H
// attraction inside a cutoff (the "burial" driver), and short-range repulsion
// for excluded volume. A "temperature" (thermal noise amplitude) cools from
// hot to cold over ~15 s; the straggly chain collapses into a hydrophobic-core
// / polar-shell blob. Hold for 3 s, then reheat with a fresh sequence.
//
// NOTE: this is *not* a real protein simulation — it's a 3D, off-lattice
// caricature of Ken Dill's HP lattice model, optimized for visual punch.

const N = 30;                  // residues
const BOND_LEN = 1.6;
const K_BOND = 60;             // bond spring stiffness
const K_HH = 1.4;               // H–H attraction strength (within cutoff)
const HH_CUTOFF = 4.0;          // H–H attraction radius
const K_REP = 25;               // excluded-volume repulsion
const REP_RADIUS = 1.4;
const DAMPING = 4.0;            // velocity damping (per second)
const T_HOT = 6.0;              // starting noise stddev (units of accel)
const T_COLD = 0.05;
const COOL_TIME = 15.0;         // seconds to cool from T_HOT to T_COLD
const HOLD_TIME = 3.0;          // seconds folded before reheat
const FOLD_THRESHOLD = 0.15;    // T below this counts as "folded"
const DT_SIM = 1 / 120;         // sim substep
const BEAD_RADIUS = 0.45;

// State.
let seq;                        // Uint8Array of 0=P, 1=H
let pos;                        // Float32Array (3*N), positions
let vel;                        // Float32Array (3*N)
let acc;                        // Float32Array (3*N)
let chainGroup;                 // THREE.Group of meshes + line
let beadMeshes = [];
let bondLine;
let bondPositions;              // Float32Array (3*N) for the polyline
let temperature = T_HOT;
let phase = "fold";             // "fold" | "hold"
let phaseTimer = 0;
let coolElapsed = 0;
let userYaw = 0, userPitch = 0.2;
let isDragging = false;
let lastMouseX = 0, lastMouseY = 0;
let idleTime = 0;
let simAccum = 0;

function randSeq() {
  // Roughly 50/50 H/P, but with a guaranteed minimum H count so the core
  // is visible.
  const out = new Uint8Array(N);
  let hCount = 0;
  for (let i = 0; i < N; i++) {
    out[i] = Math.random() < 0.55 ? 1 : 0;
    if (out[i]) hCount++;
  }
  // If too few H, flip random Ps until we have ~12.
  while (hCount < 12) {
    const i = (Math.random() * N) | 0;
    if (!out[i]) { out[i] = 1; hCount++; }
  }
  return out;
}

function initChain() {
  // Lay the chain out in a wobbly stretched line so it starts "denatured".
  for (let i = 0; i < N; i++) {
    const t = i - (N - 1) / 2;
    pos[3*i]   = t * BOND_LEN * 0.9 + (Math.random() - 0.5) * 0.3;
    pos[3*i+1] = (Math.random() - 0.5) * 0.6;
    pos[3*i+2] = (Math.random() - 0.5) * 0.6;
    vel[3*i] = vel[3*i+1] = vel[3*i+2] = 0;
  }
}

function rebuildBeadColors() {
  for (let i = 0; i < N; i++) {
    const m = beadMeshes[i];
    // H = warm red, P = cool blue.
    if (seq[i]) m.material.color.setHex(0xff5a4a);
    else m.material.color.setHex(0x4aa8ff);
  }
}

function init({ scene, camera, renderer, width, height }) {
  renderer.setClearColor(0x080a12, 1);
  camera.position.set(0, 4, 22);
  camera.lookAt(0, 0, 0);

  // Lights.
  const amb = new THREE.AmbientLight(0xffffff, 0.45);
  scene.add(amb);
  const dir = new THREE.DirectionalLight(0xffffff, 0.85);
  dir.position.set(6, 10, 7);
  scene.add(dir);
  const fill = new THREE.DirectionalLight(0x88aaff, 0.25);
  fill.position.set(-8, -3, -4);
  scene.add(fill);

  seq = randSeq();
  pos = new Float32Array(3 * N);
  vel = new Float32Array(3 * N);
  acc = new Float32Array(3 * N);
  initChain();

  chainGroup = new THREE.Group();
  scene.add(chainGroup);

  const sphereGeo = new THREE.SphereGeometry(BEAD_RADIUS, 16, 12);
  for (let i = 0; i < N; i++) {
    const mat = new THREE.MeshStandardMaterial({
      color: 0xffffff,
      roughness: 0.4,
      metalness: 0.1,
    });
    const mesh = new THREE.Mesh(sphereGeo, mat);
    mesh.position.set(pos[3*i], pos[3*i+1], pos[3*i+2]);
    chainGroup.add(mesh);
    beadMeshes.push(mesh);
  }
  rebuildBeadColors();

  // Bond polyline.
  bondPositions = new Float32Array(3 * N);
  for (let i = 0; i < N; i++) {
    bondPositions[3*i]   = pos[3*i];
    bondPositions[3*i+1] = pos[3*i+1];
    bondPositions[3*i+2] = pos[3*i+2];
  }
  const bondGeo = new THREE.BufferGeometry();
  bondGeo.setAttribute(
    "position",
    new THREE.BufferAttribute(bondPositions, 3).setUsage(THREE.DynamicDrawUsage),
  );
  const bondMat = new THREE.LineBasicMaterial({
    color: 0xcfd6e6,
    transparent: true,
    opacity: 0.55,
  });
  bondLine = new THREE.Line(bondGeo, bondMat);
  chainGroup.add(bondLine);

  temperature = T_HOT;
  phase = "fold";
  phaseTimer = 0;
  coolElapsed = 0;

  return { scene, camera };
}

function gaussian() {
  // Box–Muller.
  let u = 0, v = 0;
  while (u === 0) u = Math.random();
  while (v === 0) v = Math.random();
  return Math.sqrt(-2 * Math.log(u)) * Math.cos(2 * Math.PI * v);
}

function step(h) {
  // Zero accelerations.
  acc.fill(0);

  // Bond springs (adjacent residues).
  for (let i = 0; i < N - 1; i++) {
    const ax = pos[3*i],   ay = pos[3*i+1],   az = pos[3*i+2];
    const bx = pos[3*i+3], by = pos[3*i+4],   bz = pos[3*i+5];
    const dx = bx - ax, dy = by - ay, dz = bz - az;
    const r = Math.sqrt(dx*dx + dy*dy + dz*dz) + 1e-9;
    const f = K_BOND * (r - BOND_LEN) / r;
    const fx = f * dx, fy = f * dy, fz = f * dz;
    acc[3*i]     += fx; acc[3*i+1] += fy; acc[3*i+2] += fz;
    acc[3*i+3]   -= fx; acc[3*i+4] -= fy; acc[3*i+5] -= fz;
  }

  // Pairwise: excluded volume always; H–H attraction inside cutoff.
  const cutoff2 = HH_CUTOFF * HH_CUTOFF;
  const rep2 = REP_RADIUS * REP_RADIUS;
  for (let i = 0; i < N; i++) {
    for (let j = i + 2; j < N; j++) {
      const dx = pos[3*j]   - pos[3*i];
      const dy = pos[3*j+1] - pos[3*i+1];
      const dz = pos[3*j+2] - pos[3*i+2];
      const d2 = dx*dx + dy*dy + dz*dz + 1e-9;

      // Repulsion: soft, 1/r^2-ish, only when overlapping.
      if (d2 < rep2) {
        const d = Math.sqrt(d2);
        const overlap = REP_RADIUS - d;
        const f = K_REP * overlap / d;
        const fx = f * dx, fy = f * dy, fz = f * dz;
        acc[3*i]   -= fx; acc[3*i+1] -= fy; acc[3*i+2] -= fz;
        acc[3*j]   += fx; acc[3*j+1] += fy; acc[3*j+2] += fz;
      }

      // H–H attraction: constant pull within cutoff (with a smooth falloff).
      if (seq[i] && seq[j] && d2 < cutoff2) {
        const d = Math.sqrt(d2);
        // Falloff factor 1 at d=REP_RADIUS, 0 at d=HH_CUTOFF.
        const t = Math.max(0, Math.min(1, (HH_CUTOFF - d) / (HH_CUTOFF - REP_RADIUS)));
        const f = K_HH * t / d;
        const fx = f * dx, fy = f * dy, fz = f * dz;
        acc[3*i]   += fx; acc[3*i+1] += fy; acc[3*i+2] += fz;
        acc[3*j]   -= fx; acc[3*j+1] -= fy; acc[3*j+2] -= fz;
      }
    }
  }

  // Thermal noise + damping + integrate.
  const noise = temperature;
  const damp = Math.exp(-DAMPING * h);
  for (let i = 0; i < N; i++) {
    acc[3*i]   += gaussian() * noise;
    acc[3*i+1] += gaussian() * noise;
    acc[3*i+2] += gaussian() * noise;

    vel[3*i]   = (vel[3*i]   + acc[3*i]   * h) * damp;
    vel[3*i+1] = (vel[3*i+1] + acc[3*i+1] * h) * damp;
    vel[3*i+2] = (vel[3*i+2] + acc[3*i+2] * h) * damp;

    pos[3*i]   += vel[3*i]   * h;
    pos[3*i+1] += vel[3*i+1] * h;
    pos[3*i+2] += vel[3*i+2] * h;
  }

  // Gentle pull toward origin so the molecule doesn't drift off-camera.
  let cx = 0, cy = 0, cz = 0;
  for (let i = 0; i < N; i++) {
    cx += pos[3*i]; cy += pos[3*i+1]; cz += pos[3*i+2];
  }
  cx /= N; cy /= N; cz /= N;
  for (let i = 0; i < N; i++) {
    pos[3*i]   -= cx;
    pos[3*i+1] -= cy;
    pos[3*i+2] -= cz;
  }
}

function reheat() {
  seq = randSeq();
  rebuildBeadColors();
  initChain();
  temperature = T_HOT;
  phase = "fold";
  phaseTimer = 0;
  coolElapsed = 0;
}

function tick({ dt, scene, camera, renderer, width, height, input }) {
  // Drain clicks; nothing to do with them here.
  input.consumeClicks();

  // Camera orbit.
  if (input.mouseDown) {
    if (!isDragging) {
      isDragging = true;
      lastMouseX = input.mouseX;
      lastMouseY = input.mouseY;
    } else {
      const dx = input.mouseX - lastMouseX;
      const dy = input.mouseY - lastMouseY;
      userYaw -= dx * 0.005;
      userPitch -= dy * 0.005;
      userPitch = Math.max(-1.2, Math.min(1.2, userPitch));
      lastMouseX = input.mouseX;
      lastMouseY = input.mouseY;
    }
    idleTime = 0;
  } else {
    isDragging = false;
    idleTime += dt;
    if (idleTime > 0.5) {
      userYaw += dt * 0.15;
    }
  }

  const camR = 22;
  const cy0 = Math.cos(userYaw), sy0 = Math.sin(userYaw);
  const cp = Math.cos(userPitch), sp = Math.sin(userPitch);
  camera.position.set(camR * cp * sy0, camR * sp, camR * cp * cy0);
  camera.lookAt(0, 0, 0);
  camera.aspect = width / height;
  camera.updateProjectionMatrix();
  renderer.setSize(width, height, false);

  // Annealing schedule.
  if (phase === "fold") {
    coolElapsed += dt;
    const t = Math.min(1, coolElapsed / COOL_TIME);
    // Geometric cooling between T_HOT and T_COLD (linear in log).
    const logT = Math.log(T_HOT) + (Math.log(T_COLD) - Math.log(T_HOT)) * t;
    temperature = Math.exp(logT);
    if (temperature < FOLD_THRESHOLD || coolElapsed >= COOL_TIME) {
      temperature = T_COLD;
      phase = "hold";
      phaseTimer = 0;
    }
  } else if (phase === "hold") {
    temperature = T_COLD;
    phaseTimer += dt;
    if (phaseTimer >= HOLD_TIME) {
      reheat();
    }
  }

  // Fixed-step sim, clamped to keep things stable when tab regains focus.
  simAccum += Math.min(dt, 0.1);
  let steps = 0;
  while (simAccum >= DT_SIM && steps < 30) {
    step(DT_SIM);
    simAccum -= DT_SIM;
    steps++;
  }

  // Sync bead meshes + bond line.
  for (let i = 0; i < N; i++) {
    beadMeshes[i].position.set(pos[3*i], pos[3*i+1], pos[3*i+2]);
    bondPositions[3*i]     = pos[3*i];
    bondPositions[3*i+1]   = pos[3*i+1];
    bondPositions[3*i+2]   = pos[3*i+2];
  }
  bondLine.geometry.attributes.position.needsUpdate = true;

  renderer.render(scene, camera);
}