From fb4d8025672a3e8af57e1d0239ef65c059caeee3 Mon Sep 17 00:00:00 2001
From: Ben Murrell <murrellb@gmail.com>
Date: Thu, 1 Oct 2026 10:32:20 +0200
Subject: [PATCH] Fix viewer pole locking and add ligand-aware chain coloring

Use TrackballControls for unrestricted rotation and add an element-color
preservation setting for ligands in chain-color mode. Include regression
tests and a read-only viewer CI workflow.
---
 .github/workflows/viewer-tests.yml |  44 +++++++
 index.html                         |   6 +-
 package.json                       |   2 +
 src/main.js                        |  18 +--
 src/viewer-color-controls.js       |  20 +++
 src/viewer.js                      |  90 ++++++++++++--
 tests/viewer-color-controls.mjs    |  70 +++++++++++
 tests/viewer-fixture.mjs           |  17 +++
 tests/viewer.mjs                   | 150 +++++++++++++++++++++++
 tests/viewer_browser.mjs           | 189 +++++++++++++++++++++++++++++
 10 files changed, 578 insertions(+), 28 deletions(-)
 create mode 100644 .github/workflows/viewer-tests.yml
 create mode 100644 src/viewer-color-controls.js
 create mode 100644 tests/viewer-color-controls.mjs
 create mode 100644 tests/viewer-fixture.mjs
 create mode 100644 tests/viewer.mjs
 create mode 100644 tests/viewer_browser.mjs

diff --git a/.github/workflows/viewer-tests.yml b/.github/workflows/viewer-tests.yml
new file mode 100644
index 0000000..b4f539a
--- /dev/null
+++ b/.github/workflows/viewer-tests.yml
@@ -0,0 +1,44 @@
+name: Viewer tests
+
+on:
+  pull_request:
+    paths:
+      - 'src/**'
+      - 'tests/**'
+      - 'index.html'
+      - 'package*.json'
+      - 'vendor/**'
+      - 'tools/prepare_mpnn_runtime.mjs'
+      - '.github/workflows/viewer-tests.yml'
+  push:
+    branches: [main]
+    paths:
+      - 'src/**'
+      - 'tests/**'
+      - 'index.html'
+      - 'package*.json'
+      - 'vendor/**'
+      - 'tools/prepare_mpnn_runtime.mjs'
+      - '.github/workflows/viewer-tests.yml'
+
+permissions:
+  contents: read
+
+jobs:
+  viewer:
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+      - uses: actions/setup-node@v4
+        with:
+          node-version: 22
+          cache: npm
+      - run: npm ci
+      - run: npm run test:viewer
+      - run: npx playwright install --with-deps chromium
+      - run: npm run test:viewer-browser
+      - run: npm run test:prep
+      - run: npm run build
diff --git a/index.html b/index.html
index b563ec3..4ab84dd 100644
--- a/index.html
+++ b/index.html
@@ -117,7 +117,11 @@ <h1>Interactive docking</h1>
           <p id="result-identities" class="structure-label"></p>
         </section>
         <button id="reset-camera" class="secondary" type="button" disabled>Reset view</button>
-        <label>Protein color<select id="protein-color"><option value="role">Atom role</option><option value="chain">Chain</option></select></label>
+        <label>Color by<select id="protein-color"><option value="role">Atom role</option><option value="chain">Chain</option></select></label>
+        <label class="toggle-row">
+          <input id="ligand-element-colors" type="checkbox" checked disabled />
+          <span>Keep ligand colors by element</span>
+        </label>
         <div id="chain-colors" class="chain-colors" hidden></div>
         <label class="toggle-row">
           <input id="show-reference" type="checkbox" />
diff --git a/package.json b/package.json
index d911909..f22ee91 100644
--- a/package.json
+++ b/package.json
@@ -9,6 +9,8 @@
     "dev": "vite --host 0.0.0.0",
     "build": "vite build",
     "test:prep": "node --test tests/prep.mjs tests/sequence.mjs tests/mpnn.mjs tests/pdb-export.mjs tests/backbone.mjs",
+    "test:viewer": "node --test tests/viewer.mjs tests/viewer-color-controls.mjs",
+    "test:viewer-browser": "node tests/viewer_browser.mjs",
     "test:campaign": "node tests/campaign_browser.mjs",
     "test:sequence-display": "node tests/sequence_display_browser.mjs",
     "test:mpnn": "node tests/mpnn_browser.mjs",
diff --git a/src/main.js b/src/main.js
index 406fcd3..deae23b 100644
--- a/src/main.js
+++ b/src/main.js
@@ -1,5 +1,6 @@
 import { DockingWebGpuModel } from "./model.js";
 import { MolecularViewer } from "./viewer.js";
+import { bindViewerColorControls } from "./viewer-color-controls.js";
 import { graphFromSmiles } from "./chemistry.js";
 import { loadCcdGraphs } from "./ccd.js";
 import { assetUrl, MODEL_MANIFEST_URL } from "./config.js";
@@ -21,7 +22,7 @@ const ui = Object.fromEntries([
   "example-tab", "custom-tab", "example-panel", "custom-panel", "pdb-input",
   "open-pdb", "chain-list", "ligand-list", "selection-count", "prepare-selection", "smiles-input", "replace-ligand", "structure-label",
   "pdb-id-input", "fetch-pdb", "show-reference",
-  "protein-color", "chain-colors",
+  "protein-color", "ligand-element-colors", "chain-colors",
   "active-ligand", "remove-ligand",
   "campaign-count", "run-campaign", "stop-campaign", "campaign-results", "campaign-result",
   "previous-result", "next-result", "download-result", "result-identities", "export-format",
@@ -67,15 +68,7 @@ mpnnControls = new MpnnControls({
   onSuggest: () => suggestResidues(),
 });
 createIcons({ icons: { ChevronLeft, ChevronRight, Download, Upload, X } });
-viewer.onSampleChange = () => {
-  ui['chain-colors'].hidden = viewer.proteinColorMode !== 'chain';
-  ui['chain-colors'].replaceChildren(...[...viewer.chainColors].map(([id, color]) => {
-    const entry = document.createElement('span'), swatch = document.createElement('i');
-    swatch.style.backgroundColor = `#${color.getHexString()}`;
-    entry.append(swatch, document.createTextNode(id || '(blank)')); return entry;
-  }));
-};
-ui['protein-color'].onchange = () => viewer.setProteinColorMode(ui['protein-color'].value);
+bindViewerColorControls(viewer, ui);
 const ligandLabel = g => `${g.options[0].atoms[0].rawResidue}${g.options.length > 1 ? ` +${g.options.length - 1}` : ''} / ${g.options[0].atoms[0].chain} (${g.atoms})`;
 
 function clearCampaign() {
@@ -87,10 +80,7 @@ function activeHighlight() {
   const labels = new Set(group?.options.flatMap(o => o.atoms.map(atomLocator)) ?? []);
   const active = new Set((viewer.sample?.atom_labels ?? []).flatMap((label, i) =>
     labels.has(label) || (group && label.startsWith(`replacement:${group.id}:`)) ? [i] : []));
-  for (const item of viewer.ligandMeshes ?? []) {
-    item.atoms.forEach((atom, i) => item.mesh.setColorAt(i, viewer.selectionColor(active.has(atom))));
-    if (item.mesh.instanceColor) item.mesh.instanceColor.needsUpdate = true;
-  }
+  viewer.highlightLigandAtoms(active);
 }
 
 viewer.onAtomPick = atom => {
diff --git a/src/viewer-color-controls.js b/src/viewer-color-controls.js
new file mode 100644
index 0000000..e0fa16e
--- /dev/null
+++ b/src/viewer-color-controls.js
@@ -0,0 +1,20 @@
+export function bindViewerColorControls(viewer, ui) {
+  viewer.onSampleChange = () => {
+    const byChain = viewer.proteinColorMode === 'chain';
+    ui['protein-color'].value = viewer.proteinColorMode;
+    ui['ligand-element-colors'].checked = viewer.ligandColorsByElement;
+    ui['ligand-element-colors'].disabled = !byChain;
+    ui['chain-colors'].hidden = !byChain;
+    const colors = [...viewer.chainColors].filter(([id]) =>
+      !viewer.ligandColorsByElement || viewer.proteinChainIds.has(id));
+    ui['chain-colors'].replaceChildren(...colors.map(([id, color]) => {
+      const entry = document.createElement('span'), swatch = document.createElement('i');
+      swatch.style.backgroundColor = `#${color.getHexString()}`;
+      entry.append(swatch, document.createTextNode(id || '(blank)')); return entry;
+    }));
+  };
+  ui['protein-color'].onchange = () => viewer.setProteinColorMode(ui['protein-color'].value);
+  ui['ligand-element-colors'].onchange = () =>
+    viewer.setLigandColorsByElement(ui['ligand-element-colors'].checked);
+  viewer.onSampleChange();
+}
diff --git a/src/viewer.js b/src/viewer.js
index a4c05b0..7c27da4 100644
--- a/src/viewer.js
+++ b/src/viewer.js
@@ -1,5 +1,5 @@
 import * as THREE from "three";
-import { OrbitControls } from "three/addons/controls/OrbitControls.js";
+import { TrackballControls } from "three/addons/controls/TrackballControls.js";
 
 const ROLE_BACKBONE = 1;
 const ROLE_SIDECHAIN = 2;
@@ -18,7 +18,10 @@ const residueKey = (sample, atom) => {
   const label = sample.atom_labels?.[atom]?.split('|');
   return label?.length === 5 ? JSON.stringify(label.slice(2)) : String(sample.residue_ids[atom]);
 };
-const chainKey = (sample, atom) => sample.atom_labels?.[atom]?.split('|')[2] ?? String(sample.chain_ids?.[atom] ?? '');
+const chainKey = (sample, atom) => {
+  const label = sample.atom_labels?.[atom]?.split('|');
+  return label?.length === 5 ? label[2] : String(sample.chain_ids?.[atom] ?? '');
+};
 
 const axis = new THREE.Vector3(0, 1, 0);
 const a = new THREE.Vector3();
@@ -122,13 +125,15 @@ export class MolecularViewer {
     this.renderer.setPixelRatio(Math.min(devicePixelRatio, 1.7));
     this.renderer.outputColorSpace = THREE.SRGBColorSpace;
     container.replaceChildren(this.renderer.domElement);
-    this.controls = new OrbitControls(this.camera, this.renderer.domElement);
-    this.controls.enableDamping = true;
-    this.controls.dampingFactor = 0.075;
-    this.controls.rotateSpeed = 0.55;
+    this.createControls();
     const raycaster = new THREE.Raycaster();
     let pointerStart;
-    this.renderer.domElement.addEventListener('pointerdown', event => { pointerStart = [event.clientX, event.clientY]; });
+    this.renderer.domElement.addEventListener('pointerdown', event => {
+      pointerStart = [event.clientX, event.clientY];
+      // Preserve modified left-drag panning alongside right-drag and touch pan.
+      this.controls.mouseButtons.LEFT = event.shiftKey || event.ctrlKey || event.metaKey
+        ? THREE.MOUSE.PAN : THREE.MOUSE.ROTATE;
+    }, { capture: true });
     this.renderer.domElement.addEventListener('pointerup', event => {
       if (!pointerStart || Math.hypot(event.clientX - pointerStart[0], event.clientY - pointerStart[1]) > 4) return;
       const rect = this.renderer.domElement.getBoundingClientRect();
@@ -158,12 +163,26 @@ export class MolecularViewer {
     this.backbone = null; this.sideBonds = null;
     this.referenceVisible = false;
     this.proteinColorMode = 'role'; this.highlightedResidues = new Set(); this.chainColors = new Map();
+    this.proteinChainIds = new Set();
+    this.ligandColorsByElement = true;
+    this.highlightedLigandAtoms = new Set();
     this.resizeObserver = new ResizeObserver(() => this.resize());
     this.resizeObserver.observe(container);
     this.animateFrame = this.animateFrame.bind(this);
     requestAnimationFrame(this.animateFrame);
   }
 
+  createControls() {
+    this.controls?.dispose();
+    // Unlike fixed-up orbiting, trackball rotation can pass through both poles.
+    this.controls = new TrackballControls(this.camera, this.renderer.domElement);
+    this.controls.staticMoving = false;
+    this.controls.dynamicDampingFactor = 0.15;
+    this.controls.rotateSpeed = 1;
+    // Do not let typing A/S/D in sequence or SMILES inputs change mouse modes.
+    this.controls.keys = [];
+  }
+
   clear() {
     if (!this.group) return;
     this.scene.remove(this.group);
@@ -180,9 +199,14 @@ export class MolecularViewer {
   setSample(sample, coords, preserveCamera = false) {
     this.clear();
     this.sample = sample;
+    this.highlightedLigandAtoms.clear();
     this.atomResidueKeys = sample.roles.map((_, i) => residueKey(sample, i));
     this.atomChainKeys = sample.roles.map((_, i) => chainKey(sample, i));
-    const chains = [...new Set(sample.roles.flatMap((role, i) => [1, 2].includes(role) ? [chainKey(sample, i)] : []))].sort();
+    this.proteinChainIds = new Set(sample.roles.flatMap((role, i) =>
+      [ROLE_BACKBONE, ROLE_SIDECHAIN].includes(role) ? [this.atomChainKeys[i]] : []));
+    // Append ligand-only chains without shifting the existing protein palette.
+    const chains = [...this.proteinChainIds].sort();
+    chains.push(...[...new Set(this.atomChainKeys)].filter(id => !this.proteinChainIds.has(id)).sort());
     this.chainColors = new Map(chains.map((id, i) => [id, new THREE.Color(CHAIN_COLORS[i % CHAIN_COLORS.length])]));
     this.group = new THREE.Group();
     this.scene.add(this.group);
@@ -203,12 +227,12 @@ export class MolecularViewer {
     this.sideAtoms = new THREE.InstancedMesh(sphere, material(0xffffff, 0.9), sample.roles.filter((role) => role === ROLE_SIDECHAIN).length);
     this.sideBonds = new THREE.InstancedMesh(cylinder, material(0xffffff, 0.85), sample.sidechain_bonds.length);
     this.moleculePairs = [...sample.ligand_bonds, ...sample.molecule_bonds];
-    this.ligandBonds = new THREE.InstancedMesh(cylinder, material(0xd65358), this.moleculePairs.length);
+    this.ligandBonds = new THREE.InstancedMesh(cylinder, material(0xffffff), this.moleculePairs.length);
     this.ligandByElement = groupedAtoms(sample, ROLE_LIGAND);
     this.ligandMeshes = [...this.ligandByElement].map(([element, atoms]) => ({
       element,
       atoms,
-      mesh: new THREE.InstancedMesh(sphere, material(ELEMENT_COLORS.get(element) ?? 0xcd83d2), atoms.length),
+      mesh: new THREE.InstancedMesh(sphere, material(0xffffff), atoms.length),
     }));
     this.referenceGroup = new THREE.Group();
     this.referenceGroup.visible = this.referenceVisible;
@@ -273,7 +297,39 @@ export class MolecularViewer {
 
   setProteinColorMode(mode) {
     this.proteinColorMode = mode === 'chain' ? 'chain' : 'role';
-    this.updateProteinColors(); this.onSampleChange?.();
+    this.updateProteinColors(); this.updateLigandColors(); this.onSampleChange?.();
+  }
+
+  setLigandColorsByElement(enabled) {
+    this.ligandColorsByElement = Boolean(enabled);
+    this.updateLigandColors(); this.onSampleChange?.();
+  }
+
+  highlightLigandAtoms(atoms) {
+    this.highlightedLigandAtoms = new Set(atoms);
+    this.updateLigandColors();
+  }
+
+  updateLigandColors() {
+    if (!this.sample) return;
+    const byChain = this.proteinColorMode === 'chain' && !this.ligandColorsByElement;
+    const color = new THREE.Color();
+    const tint = this.selectionColor(true);
+    for (const { mesh, atoms, element } of this.ligandMeshes) {
+      const elementColor = new THREE.Color(ELEMENT_COLORS.get(element) ?? 0xcd83d2);
+      atoms.forEach((atom, i) => {
+        color.copy(byChain ? this.chainColors.get(this.atomChainKeys[atom]) : elementColor);
+        if (this.highlightedLigandAtoms.has(atom)) color.multiply(tint);
+        mesh.setColorAt(i, color);
+      });
+      if (mesh.instanceColor) mesh.instanceColor.needsUpdate = true;
+    }
+    const bondColor = new THREE.Color(0xd65358);
+    this.drawnMoleculePairs.forEach(([first, second], i) => {
+      const atom = [ROLE_LIGAND, 0].includes(this.sample.roles[first]) ? first : second;
+      this.ligandBonds.setColorAt(i, byChain ? this.chainColors.get(this.atomChainKeys[atom]) : bondColor);
+    });
+    if (this.ligandBonds.instanceColor) this.ligandBonds.instanceColor.needsUpdate = true;
   }
 
   updateProteinColors() {
@@ -331,8 +387,11 @@ export class MolecularViewer {
     this.sideBonds.instanceMatrix.needsUpdate = true;
     this.sideBonds.boundingSphere = null;
     cursor = 0;
+    this.drawnMoleculePairs = [];
     for (const [first, second] of this.moleculePairs) {
-      if (bondMatrix(this.ligandBonds, cursor, coords, first, second, 0.068)) cursor += 1;
+      if (bondMatrix(this.ligandBonds, cursor, coords, first, second, 0.068)) {
+        this.drawnMoleculePairs.push([first, second]); cursor += 1;
+      }
     }
     this.ligandBonds.count = cursor;
     this.ligandBonds.instanceMatrix.needsUpdate = true;
@@ -342,6 +401,7 @@ export class MolecularViewer {
       mesh.boundingSphere = null;
     }
     this.updateProteinColors();
+    this.updateLigandColors();
   }
 
   setReferenceVisible(visible) {
@@ -371,11 +431,14 @@ export class MolecularViewer {
     for (let atom = 0; atom < this.sample.atoms; atom += 1) box.expandByPoint(atomPosition(coords, atom, new THREE.Vector3()));
     const center = box.getCenter(new THREE.Vector3());
     const size = box.getSize(new THREE.Vector3()).length();
-    this.controls.target.copy(center);
+    this.camera.up.set(0, 1, 0);
     this.camera.position.copy(center).add(new THREE.Vector3(size * 0.8, size * 0.6, size * 1.05));
     this.camera.near = Math.max(0.05, size / 500);
     this.camera.far = Math.max(1000, size * 10);
     this.camera.updateProjectionMatrix();
+    // Recreate controls to discard pending rotate/zoom/pan inertia on reset.
+    this.createControls();
+    this.controls.target.copy(center);
     this.controls.update();
   }
 
@@ -386,6 +449,7 @@ export class MolecularViewer {
     this.renderer.setSize(width, height, false);
     this.camera.aspect = width / height;
     this.camera.updateProjectionMatrix();
+    this.controls.handleResize();
   }
 
   animateFrame() {
diff --git a/tests/viewer-color-controls.mjs b/tests/viewer-color-controls.mjs
new file mode 100644
index 0000000..428cf09
--- /dev/null
+++ b/tests/viewer-color-controls.mjs
@@ -0,0 +1,70 @@
+import assert from 'node:assert/strict';
+import test from 'node:test';
+import { bindViewerColorControls } from '../src/viewer-color-controls.js';
+
+function setup(t) {
+  const node = () => ({ children: [], style: {},
+    append(...children) { this.children.push(...children); },
+    replaceChildren(...children) { this.children = children; },
+  });
+  const previous = globalThis.document;
+  globalThis.document = { createElement: node, createTextNode: text => ({ textContent: text }) };
+  t.after(() => {
+    if (previous === undefined) delete globalThis.document;
+    else globalThis.document = previous;
+  });
+  const ui = Object.fromEntries(['protein-color', 'ligand-element-colors', 'chain-colors'].map(id => [id, node()]));
+  const viewer = {
+    proteinColorMode: 'role', ligandColorsByElement: true,
+    proteinChainIds: new Set(['B']),
+    chainColors: new Map([['B', { getHexString: () => '63cdb5' }], ['L', { getHexString: () => 'd98bd3' }]]),
+    setProteinColorMode(mode) { this.proteinColorMode = mode; this.onSampleChange(); },
+    setLigandColorsByElement(value) { this.ligandColorsByElement = value; this.onSampleChange(); },
+  };
+  bindViewerColorControls(viewer, ui);
+  return { viewer, ui };
+}
+
+test('color controls initialize from the viewer and preserve elements by default', t => {
+  const { ui } = setup(t);
+  assert.equal(ui['protein-color'].value, 'role');
+  assert.equal(ui['ligand-element-colors'].checked, true);
+  assert.equal(ui['ligand-element-colors'].disabled, true);
+  assert.equal(ui['chain-colors'].hidden, true);
+});
+
+test('UI changes call the viewer setters and show only relevant chain swatches', t => {
+  const { viewer, ui } = setup(t);
+  ui['protein-color'].value = 'chain'; ui['protein-color'].onchange();
+  assert.equal(viewer.proteinColorMode, 'chain');
+  assert.equal(ui['ligand-element-colors'].disabled, false);
+  assert.equal(ui['chain-colors'].hidden, false);
+  assert.equal(ui['chain-colors'].children.length, 1);
+  ui['ligand-element-colors'].checked = false; ui['ligand-element-colors'].onchange();
+  assert.equal(viewer.ligandColorsByElement, false);
+  assert.equal(ui['chain-colors'].children.length, 2);
+  ui['ligand-element-colors'].checked = true; ui['ligand-element-colors'].onchange();
+  assert.equal(ui['chain-colors'].children.length, 1);
+});
+
+test('switching out of chain mode retains the ligand preference for the next switch', t => {
+  const { viewer, ui } = setup(t);
+  viewer.setProteinColorMode('chain'); viewer.setLigandColorsByElement(false);
+  viewer.setProteinColorMode('role');
+  assert.equal(ui['ligand-element-colors'].disabled, true);
+  assert.equal(ui['ligand-element-colors'].checked, false);
+  viewer.setProteinColorMode('chain');
+  assert.equal(ui['ligand-element-colors'].disabled, false);
+  assert.equal(ui['ligand-element-colors'].checked, false);
+});
+
+test('sample changes refresh legend labels as text, including blank chain identifiers', t => {
+  const { viewer, ui } = setup(t);
+  const label = '<img src=x onerror=alert(1)>';
+  viewer.proteinChainIds = new Set(['', label]);
+  viewer.chainColors = new Map([['', { getHexString: () => 'ffffff' }], [label, { getHexString: () => '000000' }]]);
+  viewer.setProteinColorMode('chain');
+  assert.equal(ui['chain-colors'].children[0].children[1].textContent, '(blank)');
+  assert.equal(ui['chain-colors'].children[1].children[1].textContent, label);
+  assert.equal(ui['chain-colors'].children[0].children[0].style.backgroundColor, '#ffffff');
+});
diff --git a/tests/viewer-fixture.mjs b/tests/viewer-fixture.mjs
new file mode 100644
index 0000000..6b44314
--- /dev/null
+++ b/tests/viewer-fixture.mjs
@@ -0,0 +1,17 @@
+export function viewerSample() {
+  const chains = ['B', 'B', 'D', 'D', 'B', 'B', 'A', 'A', 'C', 'C'];
+  return {
+    atoms: 10,
+    roles: [1, 2, 1, 2, 3, 3, 3, 3, 0, 0],
+    atomic_numbers: [6, 6, 6, 8, 6, 8, 6, 7, 6, 30],
+    residue_ids: [0, 0, 1, 1, -1, -1, -1, -1, -1, -1],
+    chain_ids: [1, 1, 3, 3, 1, 1, 0, 0, 2, 2],
+    atom_labels: chains.map((chain, i) => `RES|A${i}|${chain}|${Math.floor(i / 2)}|`),
+    target_coords: [[-6, -2, 0], [-5, -2, 1], [1, -1, 0], [2, -1, 1],
+      [-3, 3, 0], [-2, 3, 1], [5, 3, 0], [6, 3, 1], [0, 6, 0], [1, 6, 1]],
+    backbone_trace_pairs: [[0, 2]],
+    sidechain_bonds: [[0, 1], [2, 3]],
+    ligand_bonds: [[4, 5], [6, 7]],
+    molecule_bonds: [[8, 9]],
+  };
+}
diff --git a/tests/viewer.mjs b/tests/viewer.mjs
new file mode 100644
index 0000000..f1ff1c2
--- /dev/null
+++ b/tests/viewer.mjs
@@ -0,0 +1,150 @@
+import assert from 'node:assert/strict';
+import test from 'node:test';
+import * as THREE from 'three';
+import { MolecularViewer } from '../src/viewer.js';
+import { viewerSample } from './viewer-fixture.mjs';
+
+// Exercise real Three.js geometry/materials without requiring a WebGL context.
+function viewer(t, sample = viewerSample()) {
+  const v = Object.assign(Object.create(MolecularViewer.prototype), {
+    scene: new THREE.Scene(), group: null, sample: null, referenceVisible: false,
+    proteinColorMode: 'role', ligandColorsByElement: true,
+    highlightedResidues: new Set(), highlightedLigandAtoms: new Set(),
+    chainColors: new Map(), proteinChainIds: new Set(),
+  });
+  v.setSample(sample, new Float32Array(sample.target_coords.flat()), true);
+  t.after(() => v.clear());
+  return v;
+}
+
+function meshColor(mesh, index) {
+  const color = new THREE.Color();
+  mesh.getColorAt(index, color);
+  return color.multiply(mesh.material.color).getHex();
+}
+
+function ligandColor(v, atom) {
+  const { mesh, atoms } = v.ligandMeshes.find(item => item.atoms.includes(atom));
+  return meshColor(mesh, atoms.indexOf(atom));
+}
+
+const chainColor = (v, atom) => v.chainColors.get(v.atomChainKeys[atom]).getHex();
+
+test('default ligand element colors and bond colors are unchanged', t => {
+  const v = viewer(t);
+  assert.equal(ligandColor(v, 4), 0xc7cdcf);
+  assert.equal(ligandColor(v, 5), 0xf06560);
+  assert.equal(ligandColor(v, 7), 0x668be3);
+  assert.equal(ligandColor(v, 9), 0xcd83d2); // unknown-element fallback
+  assert.equal(meshColor(v.ligandBonds, 0), 0xd65358);
+});
+
+test('chain coloring applies to protein atoms and bonds, preserving ligand elements by default', t => {
+  const v = viewer(t);
+  v.setProteinColorMode('chain');
+  assert.notEqual(chainColor(v, 0), chainColor(v, 2));
+  assert.equal(meshColor(v.sideAtoms, 0), chainColor(v, 0));
+  assert.equal(meshColor(v.sideAtoms, 1), chainColor(v, 2));
+  assert.equal(meshColor(v.sideBonds, 1), chainColor(v, 2));
+  assert.equal(ligandColor(v, 4), 0xc7cdcf);
+  assert.equal(ligandColor(v, 5), 0xf06560);
+});
+
+test('ligands can follow their chain, including ligand-only chains and role-zero molecules', t => {
+  const v = viewer(t);
+  const proteinPalette = [...v.proteinChainIds].map(id => v.chainColors.get(id).getHex());
+  v.setProteinColorMode('chain');
+  v.setLigandColorsByElement(false);
+  for (let atom = 4; atom < 10; atom++) assert.equal(ligandColor(v, atom), chainColor(v, atom));
+  assert.equal(ligandColor(v, 4), meshColor(v.sideAtoms, 0));
+  assert.notEqual(ligandColor(v, 4), ligandColor(v, 6));
+  assert.equal(meshColor(v.ligandBonds, 1), chainColor(v, 6));
+  assert.equal(meshColor(v.ligandBonds, 2), chainColor(v, 8));
+  assert.deepEqual([...v.proteinChainIds].map(id => v.chainColors.get(id).getHex()), proteinPalette);
+});
+
+test('ligand highlighting composes with both color modes and survives coordinate updates', t => {
+  const v = viewer(t);
+  v.highlightLigandAtoms(new Set([5]));
+  const tintedElement = new THREE.Color(0xf06560).multiply(v.selectionColor(true)).getHex();
+  assert.equal(ligandColor(v, 5), tintedElement);
+  v.setProteinColorMode('chain');
+  v.setLigandColorsByElement(false);
+  const tintedChain = v.chainColors.get('B').clone().multiply(v.selectionColor(true)).getHex();
+  assert.equal(ligandColor(v, 5), tintedChain);
+  v.update(new Float32Array(v.sample.target_coords.flat().map(n => n + 1)));
+  assert.equal(ligandColor(v, 5), tintedChain);
+  v.setLigandColorsByElement(true);
+  assert.equal(ligandColor(v, 5), tintedElement);
+  v.highlightLigandAtoms(new Set());
+  assert.equal(ligandColor(v, 5), 0xf06560);
+});
+
+test('protein residue highlights and reference overlay survive color changes', t => {
+  const v = viewer(t);
+  v.highlightResidues(new Set([0]));
+  v.setProteinColorMode('chain');
+  v.setLigandColorsByElement(false);
+  assert.equal(meshColor(v.sideAtoms, 0), 0xffc533);
+  assert.equal(v.referenceAtomMesh.material.color.getHex(), 0xf0bd68);
+  assert.equal(v.referenceBondMesh.material.color.getHex(), 0xdca85e);
+  v.highlightResidues(new Set());
+  assert.equal(meshColor(v.sideAtoms, 0), chainColor(v, 0));
+});
+
+test('color preferences survive sample replacement without stale ligand selection indices', t => {
+  const v = viewer(t);
+  v.setProteinColorMode('chain');
+  v.setLigandColorsByElement(false);
+  v.highlightLigandAtoms(new Set([5]));
+  const next = viewerSample();
+  v.setSample(next, new Float32Array(next.target_coords.flat()), true);
+  assert.equal(v.proteinColorMode, 'chain');
+  assert.equal(v.ligandColorsByElement, false);
+  assert.equal(v.highlightedLigandAtoms.size, 0);
+  assert.equal(ligandColor(v, 5), chainColor(v, 5));
+  v.setProteinColorMode('role');
+  assert.equal(ligandColor(v, 5), 0xf06560);
+  v.setProteinColorMode('chain');
+  assert.equal(ligandColor(v, 5), chainColor(v, 5));
+});
+
+test('numeric and missing chain identifiers have safe color fallbacks', t => {
+  const sample = viewerSample();
+  delete sample.atom_labels;
+  const v = viewer(t, sample);
+  v.setProteinColorMode('chain');
+  v.setLigandColorsByElement(false);
+  assert.equal(v.atomChainKeys[6], '0');
+  assert.equal(ligandColor(v, 6), chainColor(v, 6));
+  delete sample.chain_ids;
+  v.setSample(sample, new Float32Array(sample.target_coords.flat()), true);
+  assert.equal(v.atomChainKeys[6], '');
+  assert.equal(ligandColor(v, 6), chainColor(v, 6));
+});
+
+test('replacement labels fall back to chain_ids instead of being parsed as atom locators', t => {
+  const sample = viewerSample();
+  sample.atom_labels[6] = 'replacement:group:0';
+  sample.atom_labels[7] = 'replacement:group:1';
+  const v = viewer(t, sample);
+  v.setProteinColorMode('chain');
+  v.setLigandColorsByElement(false);
+  assert.equal(v.atomChainKeys[6], '0');
+  assert.equal(ligandColor(v, 6), ligandColor(v, 7));
+});
+
+test('chain bond colors follow compacted instances when zero-length bonds are skipped', t => {
+  const sample = viewerSample();
+  sample.target_coords[5] = [...sample.target_coords[4]];
+  const v = viewer(t, sample);
+  v.setProteinColorMode('chain');
+  v.setLigandColorsByElement(false);
+  assert.equal(v.ligandBonds.count, 2);
+  assert.equal(meshColor(v.ligandBonds, 0), chainColor(v, 6));
+  assert.equal(meshColor(v.ligandBonds, 1), chainColor(v, 8));
+  sample.target_coords[5][0] += 1;
+  v.update(new Float32Array(sample.target_coords.flat()));
+  assert.equal(v.ligandBonds.count, 3);
+  assert.equal(meshColor(v.ligandBonds, 0), chainColor(v, 4));
+});
diff --git a/tests/viewer_browser.mjs b/tests/viewer_browser.mjs
new file mode 100644
index 0000000..d601a52
--- /dev/null
+++ b/tests/viewer_browser.mjs
@@ -0,0 +1,189 @@
+import assert from 'node:assert/strict';
+import { readFileSync } from 'node:fs';
+import { fileURLToPath } from 'node:url';
+import { chromium } from 'playwright';
+import { createServer } from 'vite';
+
+// Serve the real viewer and UI bindings, without model downloads or WebGPU.
+const root = fileURLToPath(new URL('../', import.meta.url));
+const index = readFileSync(new URL('../index.html', import.meta.url), 'utf8');
+const colorMarkup = index.match(/<label>Color by<select[\s\S]*?<div id="chain-colors"[^>]*><\/div>/)?.[0];
+assert.ok(colorMarkup, 'viewer color controls must exist in index.html');
+const html = `<!doctype html><html><head><meta name="viewport" content="width=device-width, initial-scale=1">
+<style>body{margin:16px}#viewport{width:min(640px,calc(100vw - 32px));height:480px}
+canvas{display:block;width:100%;height:100%}label{display:block}</style></head><body>
+<div id="viewport"></div>${colorMarkup}<script type="module">
+import { MolecularViewer } from '/src/viewer.js';
+import { bindViewerColorControls } from '/src/viewer-color-controls.js';
+import { viewerSample } from '/tests/viewer-fixture.mjs';
+const viewer = window.viewer = new MolecularViewer(document.getElementById('viewport'));
+const ui = Object.fromEntries(['protein-color', 'ligand-element-colors', 'chain-colors']
+  .map(id => [id, document.getElementById(id)]));
+bindViewerColorControls(viewer, ui);
+const sample = viewerSample();
+viewer.setSample(sample, new Float32Array(sample.target_coords.flat()));
+window.picks = []; viewer.onAtomPick = atom => window.picks.push(atom);
+window.ready = true;
+</script></body></html>`;
+const server = await createServer({ root, configFile: false,
+  server: { host: '127.0.0.1', port: 0 },
+  plugins: [{ name: 'viewer-test-fixture', configureServer(s) {
+    s.middlewares.use('/__viewer_test__', (_req, res) => {
+      res.setHeader('Content-Type', 'text/html'); res.end(html);
+    });
+  } }],
+});
+let browser;
+try {
+  await server.listen();
+  browser = await chromium.launch({ headless: true,
+    ...(process.env.PLAYWRIGHT_CHROMIUM_EXECUTABLE_PATH
+      ? { executablePath: process.env.PLAYWRIGHT_CHROMIUM_EXECUTABLE_PATH } : {}),
+    args: ['--use-angle=swiftshader', '--enable-unsafe-swiftshader'],
+  });
+  const page = await browser.newPage({ viewport: { width: 900, height: 700 }, hasTouch: true });
+  const errors = [];
+  page.on('pageerror', e => errors.push(e.message));
+  await page.goto(`http://127.0.0.1:${server.httpServer.address().port}/__viewer_test__`);
+  await page.waitForFunction(() => window.ready);
+  const frames = (n = 2) => page.evaluate(async n => {
+    for (let i = 0; i < n; i++) await new Promise(requestAnimationFrame);
+  }, n);
+  const pose = () => page.evaluate(() => {
+    const v = window.viewer;
+    return { position: v.camera.position.toArray(), up: v.camera.up.toArray(),
+      quaternion: v.camera.quaternion.toArray(), target: v.controls.target.toArray(),
+      radius: v.camera.position.distanceTo(v.controls.target) };
+  });
+  const drag = async (dx, dy, button = 'left') => {
+    const r = await page.locator('canvas').boundingBox();
+    const x = r.x + r.width / 2 - dx / 2, y = r.y + r.height / 2 - dy / 2;
+    await page.mouse.move(x, y); await page.mouse.down({ button });
+    await page.mouse.move(x + dx, y + dy, { steps: 16 });
+    await page.mouse.up({ button }); await frames();
+  };
+  await frames();
+
+  // UI: opt into chain colors; preserve elements by default, or let ligands follow chains.
+  assert.ok(await page.locator('#ligand-element-colors').isDisabled());
+  await page.selectOption('#protein-color', 'chain');
+  assert.ok(await page.locator('#ligand-element-colors').isChecked());
+  assert.ok(await page.locator('#ligand-element-colors').isEnabled());
+  assert.equal(await page.locator('#chain-colors span').count(), 2);
+  await page.uncheck('#ligand-element-colors');
+  assert.equal(await page.evaluate(() => window.viewer.ligandColorsByElement), false);
+  assert.equal(await page.locator('#chain-colors span').count(), 4);
+  await page.selectOption('#protein-color', 'role');
+  assert.ok(await page.locator('#ligand-element-colors').isDisabled());
+  await page.selectOption('#protein-color', 'chain');
+  assert.equal(await page.locator('#ligand-element-colors').isChecked(), false);
+  await page.check('#ligand-element-colors');
+  assert.equal(await page.locator('#chain-colors span').count(), 2);
+
+  // Multiple complete tumbles must cross both poles without a fixed-up clamp.
+  await page.evaluate(() => {
+    const v = window.viewer; v.controls.staticMoving = true;
+    v.camera.position.copy(v.controls.target).add(v.camera.up.clone().set(0, 0, 30));
+    v.camera.up.set(0, 1, 0); v.controls.update();
+  });
+  const ups = [];
+  for (let i = 0; i < 8; i++) {
+    await drag(0, 360);
+    const p = await pose(); ups.push(p.up[1]);
+    assert.ok([...p.position, ...p.quaternion, ...p.up].every(Number.isFinite));
+    assert.ok(Math.abs(p.radius - 30) < 1e-6);
+  }
+  assert.ok(Math.min(...ups) < -0.8, 'camera must tumble upside down through a pole');
+  assert.ok(Math.max(...ups.slice(3)) > 0.8, 'camera must continue through the other pole');
+
+  // Horizontal rotation remains usable when looking exactly down either world-Y pole.
+  for (const pole of [-1, 1]) {
+    await page.evaluate(pole => {
+      const v = window.viewer;
+      v.camera.position.copy(v.controls.target).add(v.camera.up.clone().set(0, pole * 30, 0));
+      v.camera.up.set(0, 0, -pole); v.controls.update();
+    }, pole);
+    const before = await pose(); await drag(100, 0); const after = await pose();
+    assert.notDeepEqual(after.position, before.position);
+    assert.ok([...after.position, ...after.quaternion].every(Number.isFinite));
+    assert.ok(Math.abs(after.radius - before.radius) < 1e-6);
+  }
+
+  // Reset also removes roll and all pending damping, not just camera position.
+  await page.evaluate(() => { window.viewer.controls.staticMoving = false; });
+  await drag(100, 80);
+  await page.evaluate(() => window.viewer.resetCamera());
+  const reset = await pose(); await frames(12);
+  assert.deepEqual(await pose(), reset);
+  assert.deepEqual(reset.up, [0, 1, 0]);
+  await page.evaluate(() => {
+    const v = window.viewer; const controls = v.controls;
+    v.setSample(v.sample, new Float32Array(v.sample.target_coords.flat()), true);
+    if (v.controls !== controls) throw new Error('preserveCamera replaced the controls');
+    v.controls.staticMoving = true;
+  });
+  assert.deepEqual(await pose(), reset);
+
+  // Zoom, right-drag pan and modified left-drag pan still work.
+  await page.mouse.move(300, 200);
+  await page.mouse.wheel(0, 120); await frames();
+  assert.notEqual((await pose()).radius, reset.radius);
+  let before = await pose(); await drag(80, 30, 'right'); let after = await pose();
+  assert.notDeepEqual(after.target, before.target);
+  assert.ok(Math.abs(after.radius - before.radius) < 1e-6);
+  before = after;
+  await page.keyboard.down('Shift'); await drag(60, 20); await page.keyboard.up('Shift');
+  after = await pose();
+  assert.notDeepEqual(after.target, before.target);
+  assert.ok(Math.abs(after.radius - before.radius) < 1e-6);
+
+  // Pick a ligand after a rolled view; a drag itself must not trigger selection.
+  assert.deepEqual(await page.evaluate(() => window.picks), []);
+  await page.evaluate(() => {
+    const v = window.viewer; v.resetCamera();
+    const rotation = v.camera.quaternion.clone().setFromAxisAngle(v.camera.up.clone().set(0, 0, 1), 0.7);
+    v.camera.position.sub(v.controls.target).applyQuaternion(rotation).add(v.controls.target);
+    v.camera.up.applyQuaternion(rotation); v.controls.update();
+  });
+  await frames();
+  const pick = await page.evaluate(() => {
+    const v = window.viewer, atom = 7;
+    const p = v.camera.position.clone().set(...v.sample.target_coords[atom]).project(v.camera);
+    const r = v.renderer.domElement.getBoundingClientRect();
+    return { x: r.left + (p.x + 1) * r.width / 2, y: r.top + (1 - p.y) * r.height / 2, atom };
+  });
+  await page.mouse.click(pick.x, pick.y);
+  assert.deepEqual(await page.evaluate(() => window.picks), [pick.atom]);
+
+  // Responsive resize must update TrackballControls' cached pointer bounds.
+  await page.setViewportSize({ width: 390, height: 700 }); await frames();
+  assert.equal(await page.evaluate(() => {
+    const v = window.viewer;
+    return v.controls.screen.width === v.renderer.domElement.getBoundingClientRect().width;
+  }), true);
+
+  // Touch rotation, pinch zoom and two-finger pan use the same unrestricted controls.
+  await page.evaluate(() => { window.viewer.resetCamera(); window.viewer.controls.staticMoving = true; });
+  const session = await page.context().newCDPSession(page);
+  const touch = async (type, points) => {
+    await session.send('Input.dispatchTouchEvent', { type, touchPoints: points }); await frames();
+  };
+  before = await pose();
+  await touch('touchStart', [{ id: 1, x: 140, y: 200 }]);
+  await touch('touchMove', [{ id: 1, x: 200, y: 260 }]);
+  await touch('touchEnd', []);
+  assert.notDeepEqual((await pose()).position, before.position);
+  before = await pose();
+  await touch('touchStart', [{ id: 1, x: 120, y: 200 }, { id: 2, x: 220, y: 200 }]);
+  await touch('touchMove', [{ id: 1, x: 100, y: 240 }, { id: 2, x: 250, y: 240 }]);
+  await touch('touchEnd', []);
+  after = await pose();
+  assert.notEqual(after.radius, before.radius);
+  assert.notDeepEqual(after.target, before.target);
+  await session.detach();
+  assert.deepEqual(errors, []);
+  console.log('Viewer browser checks passed: pole crossings, reset, pan/zoom, picking, resize, touch, and color settings.');
+} finally {
+  await browser?.close();
+  await server.close();
+}