Fix viewer gimbal lock and add ligand-aware chain coloring

.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

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" />

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",

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 => {

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();
+}

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() {

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');
+});