refine wb impl/vis

src/mol-model-props/computed/interactions/water-bridges.ts

@@ -10,7 +10,7 @@ import { Vec3 } from '../../../mol-math/linear-algebra';
 import { ParamDefinition as PD } from '../../../mol-util/param-definition';
 import { DataLocation } from '../../../mol-model/location';
 import { DataLoci } from '../../../mol-model/loci';
-import { MoleculeType } from '../../../mol-model/structure/model/types';
+import { MoleculeType, NucleicBackboneAtoms, ProteinBackboneAtoms } from '../../../mol-model/structure/model/types';
 import { StructureLookup3DResultContext } from '../../../mol-model/structure/structure/util/lookup3d';
 import { Sphere3D } from '../../../mol-math/geometry';
 import { CentroidHelper } from '../../../mol-math/geometry/centroid-helper';
@@ -19,6 +19,7 @@ import { FeatureType } from './common';
 import { GeometryOptions, checkGeometry } from './hydrogen-bonds';
 import { degToRad } from '../../../mol-math/misc';
 import { elementLabel } from '../../../mol-theme/label';
+import { cantorPairing } from '../../../mol-data/util/hash-functions';
 
 export type { WaterBridgeContact, WaterBridgeContacts };
 
@@ -59,20 +60,87 @@ export type WaterBridgesProps = PD.Values<WaterBridgesParams>;
 function isWater(unit: Unit.Atomic, index: StructureElement.UnitIndex): boolean {
     return unit.model.atomicHierarchy.derived.residue.moleculeType[
         unit.residueIndex[unit.elements[index]]
-    ] === MoleculeType.Water;
+        ] === MoleculeType.Water;
+}
+
+function isBackboneAtom(unit: Unit.Atomic, index: StructureElement.UnitIndex): boolean {
+    const element = unit.elements[index];
+    const moleculeType = unit.model.atomicHierarchy.derived.residue.moleculeType[unit.residueIndex[element]];
+    if (moleculeType !== MoleculeType.Protein && moleculeType !== MoleculeType.RNA && moleculeType !== MoleculeType.DNA) {
+        return false;
+    }
+
+    const atomId = unit.model.atomicHierarchy.atoms.label_atom_id.value(element);
+    if (moleculeType === MoleculeType.Protein) {
+        return ProteinBackboneAtoms.has(atomId);
+    }
+
+    return NucleicBackboneAtoms.has(atomId);
 }
 
 const _lookupCtx = StructureLookup3DResultContext();
-const _vA = Vec3();
-const _vB = Vec3();
 
-function checkOmega(posA: Vec3, posW: Vec3, posB: Vec3, omegaMinRad: number, omegaMaxRad: number): boolean {
-    Vec3.sub(_vA, posA, posW);
-    Vec3.sub(_vB, posB, posW);
-    Vec3.normalize(_vA, _vA);
-    Vec3.normalize(_vB, _vB);
-    const omega = Math.acos(Math.max(-1, Math.min(1, Vec3.dot(_vA, _vB))));
-    return omega >= omegaMinRad && omega <= omegaMaxRad;
+type Candidate = {
+    unit: Unit.Atomic
+    featureIdx: Features.FeatureIndex
+    memberIdx: StructureElement.UnitIndex
+    x: number
+    y: number
+    z: number
+    distSq: number
+};
+
+type FeatureKey = number;
+
+function featureKey(unitId: number, featureIndex: Features.FeatureIndex): FeatureKey {
+    return cantorPairing(unitId, featureIndex);
+}
+
+type BestBridge = { contact: WaterBridgeContact; combinedDistSq: number };
+type BestBridgeMap = Map<FeatureKey, Map<FeatureKey, BestBridge>>;
+
+function getBestBridge(best: BestBridgeMap, donorKey: FeatureKey, acceptorKey: FeatureKey): BestBridge | undefined {
+    return best.get(donorKey)?.get(acceptorKey);
+}
+
+function setBestBridge(best: BestBridgeMap, donorKey: FeatureKey, acceptorKey: FeatureKey, value: BestBridge) {
+    let acceptors = best.get(donorKey);
+    if (acceptors === undefined) {
+        acceptors = new Map();
+        best.set(donorKey, acceptors);
+    }
+    acceptors.set(acceptorKey, value);
+}
+
+function bestBridgeValues(best: BestBridgeMap): BestBridge[] {
+    const values: BestBridge[] = [];
+    for (const acceptors of best.values()) {
+        for (const value of acceptors.values()) values.push(value);
+    }
+    return values;
+}
+
+function checkOmega(don: Candidate, posW: Vec3, acc: Candidate, cosOmegaMin: number, cosOmegaMax: number): boolean {
+    const ax = don.x - posW[0];
+    const ay = don.y - posW[1];
+    const az = don.z - posW[2];
+
+    const bx = acc.x - posW[0];
+    const by = acc.y - posW[1];
+    const bz = acc.z - posW[2];
+
+    const aLenSq = ax * ax + ay * ay + az * az;
+    const bLenSq = bx * bx + by * by + bz * bz;
+
+    if (aLenSq === 0 || bLenSq === 0) return false;
+
+    const cosOmega = (ax * bx + ay * by + az * bz) / Math.sqrt(aLenSq * bLenSq);
+
+    // cos decreases monotonically on [0, pi], so:
+    // omega >= omegaMin && omega <= omegaMax
+    // is equivalent to:
+    // cos(omega) <= cos(omegaMin) && cos(omega) >= cos(omegaMax)
+    return cosOmega <= cosOmegaMin && cosOmega >= cosOmegaMax;
 }
 
 export function findWaterBridgeContacts(
@@ -88,19 +156,27 @@ export function findWaterBridgeContacts(
         maxAccOutOfPlaneAngle: degToRad(props.accOutOfPlaneAngleMax),
         maxDonOutOfPlaneAngle: degToRad(props.donOutOfPlaneAngleMax),
     };
+
     const legDistMinSq = props.legDistMin * props.legDistMin;
     const legDistMaxSq = props.legDistMax * props.legDistMax;
+
     const omegaMinRad = degToRad(props.omegaMin);
     const omegaMaxRad = degToRad(props.omegaMax);
 
-    type Candidate = { unit: Unit.Atomic; featureIdx: Features.FeatureIndex; pos: Vec3; distSq: number };
+    if (omegaMinRad > omegaMaxRad) return [];
+
+    const cosOmegaMin = Math.cos(omegaMinRad);
+    const cosOmegaMax = Math.cos(omegaMaxRad);
+
+    // Best bridge per unique donor/acceptor feature pair across all water molecules.
+    const best: BestBridgeMap = new Map();
 
-    // Best bridge per unique (donor, acceptor) feature pair across all water molecules.
-    const best = new Map<string, { contact: WaterBridgeContact; combinedDistSq: number }>();
     const wPos = Vec3();
+    const candidatePos = Vec3();
 
     for (const unitW of structure.units) {
         if (!Unit.isAtomic(unitW)) continue;
+
         const featW = unitsFeatures.get(unitW.id);
         if (!featW || featW.count === 0) continue;
 
@@ -113,66 +189,104 @@ export function findWaterBridgeContacts(
         for (let fi = 0 as Features.FeatureIndex; fi < featW.count; fi++) {
             const mi = featW.members[featW.offsets[fi]] as StructureElement.UnitIndex;
             if (!isWater(unitW, mi)) continue;
+
             const t = featW.types[fi];
             if (t !== FeatureType.HydrogenAcceptor && t !== FeatureType.HydrogenDonor) continue;
+
             let e = waterMap.get(mi);
             if (!e) waterMap.set(mi, (e = { acc: undefined, don: undefined }));
+
             if (t === FeatureType.HydrogenAcceptor) e.acc = fi;
             else e.don = fi;
         }
 
+        if (waterMap.size === 0) continue;
+
+        const infoWAcc = Features.Info(structure, unitW, featW);
+        const infoWDon = Features.Info(structure, unitW, featW);
+
         for (const [waterAtomIdx, { acc: accFW, don: donFW }] of waterMap) {
             if (accFW === undefined || donFW === undefined) continue;
 
             unitW.conformation.position(unitW.elements[waterAtomIdx], wPos);
 
-            const infoWAcc = Features.Info(structure, unitW, featW);
             infoWAcc.feature = accFW;
-            const infoWDon = Features.Info(structure, unitW, featW);
             infoWDon.feature = donFW;
 
-            const { count, indices, units: hitUnits, squaredDistances } =
+            const { count, indices, units: hitUnits } =
                 structure.lookup3d.find(wPos[0], wPos[1], wPos[2], props.legDistMax, _lookupCtx);
 
             const donors: Candidate[] = [];
             const acceptors: Candidate[] = [];
 
+            const donorKeys = new Set<FeatureKey>();
+            const acceptorKeys = new Set<FeatureKey>();
+
             for (let r = 0; r < count; r++) {
                 const hitUnit = hitUnits[r];
                 if (!Unit.isAtomic(hitUnit)) continue;
-                if (hitUnit === unitW) continue;
 
+                const atomicUnit = hitUnit as Unit.Atomic;
                 const hitLocalIdx = indices[r] as StructureElement.UnitIndex;
-                if (isWater(hitUnit as Unit.Atomic, hitLocalIdx)) continue;
 
-                const dSq = squaredDistances[r];
-                if (dSq < legDistMinSq || dSq > legDistMaxSq) continue;
+                // Only skip the water atom itself. Other atoms in the same unit can still be valid.
+                if (atomicUnit === unitW && hitLocalIdx === waterAtomIdx) continue;
+                if (isWater(atomicUnit, hitLocalIdx)) continue;
 
-                const hitFeat = unitsFeatures.get(hitUnit.id);
+                const hitFeat = unitsFeatures.get(atomicUnit.id);
                 if (!hitFeat || hitFeat.count === 0) continue;
 
+                const infoHit = Features.Info(structure, atomicUnit, hitFeat);
+
                 const { indices: fIdxs, offsets: fOff } = hitFeat.elementsIndex;
                 for (let k = fOff[hitLocalIdx], kl = fOff[hitLocalIdx + 1]; k < kl; k++) {
                     const fi = fIdxs[k] as Features.FeatureIndex;
                     const fType = hitFeat.types[fi];
+
                     if (fType !== FeatureType.HydrogenDonor && fType !== FeatureType.HydrogenAcceptor) continue;
 
-                    const atomicUnit = hitUnit as Unit.Atomic;
                     const memberIdx = hitFeat.members[hitFeat.offsets[fi]] as StructureElement.UnitIndex;
-                    const candidatePos = Vec3();
+
+                    if (!props.backbone && isBackboneAtom(atomicUnit, memberIdx)) continue;
+
                     atomicUnit.conformation.position(atomicUnit.elements[memberIdx], candidatePos);
 
+                    const distSq = Vec3.squaredDistance(candidatePos, wPos);
+                    if (distSq < legDistMinSq || distSq > legDistMaxSq) continue;
+
+                    infoHit.feature = fi;
+
                     if (fType === FeatureType.HydrogenDonor) {
-                        const infoDon = Features.Info(structure, atomicUnit, hitFeat);
-                        infoDon.feature = fi;
-                        if (checkGeometry(structure, infoDon, infoWAcc, legOpts)) {
-                            donors.push({ unit: atomicUnit, featureIdx: fi, pos: candidatePos, distSq: dSq });
+                        const key = featureKey(atomicUnit.id, fi);
+                        if (donorKeys.has(key)) continue;
+
+                        if (checkGeometry(structure, infoHit, infoWAcc, legOpts)) {
+                            donorKeys.add(key);
+                            donors.push({
+                                unit: atomicUnit,
+                                featureIdx: fi,
+                                memberIdx,
+                                x: candidatePos[0],
+                                y: candidatePos[1],
+                                z: candidatePos[2],
+                                distSq,
+                            });
                         }
                     } else {
-                        const infoAcc = Features.Info(structure, atomicUnit, hitFeat);
-                        infoAcc.feature = fi;
-                        if (checkGeometry(structure, infoWDon, infoAcc, legOpts)) {
-                            acceptors.push({ unit: atomicUnit, featureIdx: fi, pos: candidatePos, distSq: dSq });
+                        const key = featureKey(atomicUnit.id, fi);
+                        if (acceptorKeys.has(key)) continue;
+
+                        if (checkGeometry(structure, infoWDon, infoHit, legOpts)) {
+                            acceptorKeys.add(key);
+                            acceptors.push({
+                                unit: atomicUnit,
+                                featureIdx: fi,
+                                memberIdx,
+                                x: candidatePos[0],
+                                y: candidatePos[1],
+                                z: candidatePos[2],
+                                distSq,
+                            });
                         }
                     }
                 }
@@ -180,19 +294,27 @@ export function findWaterBridgeContacts(
 
             for (const don of donors) {
                 for (const acc of acceptors) {
-                    if (don.unit === acc.unit && don.featureIdx === acc.featureIdx) continue;
-                    if (!checkOmega(don.pos, wPos, acc.pos, omegaMinRad, omegaMaxRad)) continue;
+                    // Reject bridges where donor and acceptor are the same physical atom
+                    // represented by different feature indices.
+                    if (don.unit === acc.unit && don.memberIdx === acc.memberIdx) continue;
+
+                    if (!checkOmega(don, wPos, acc, cosOmegaMin, cosOmegaMax)) continue;
 
                     const combinedDistSq = don.distSq + acc.distSq;
-                    const pairKey = `${don.unit.id}|${don.featureIdx}|${acc.unit.id}|${acc.featureIdx}`;
+                    const donorKey = featureKey(don.unit.id, don.featureIdx);
+                    const acceptorKey = featureKey(acc.unit.id, acc.featureIdx);
 
-                    const existing = best.get(pairKey);
+                    const existing = getBestBridge(best, donorKey, acceptorKey);
                     if (!existing || combinedDistSq < existing.combinedDistSq) {
-                        best.set(pairKey, {
+                        setBestBridge(best, donorKey, acceptorKey, {
                             contact: {
-                                unitA: don.unit.id, indexA: don.featureIdx,
-                                unitB: acc.unit.id, indexB: acc.featureIdx,
-                                unitW: unitW.id, indexWA: accFW, indexWD: donFW,
+                                unitA: don.unit.id,
+                                indexA: don.featureIdx,
+                                unitB: acc.unit.id,
+                                indexB: acc.featureIdx,
+                                unitW: unitW.id,
+                                indexWA: accFW,
+                                indexWD: donFW,
                             },
                             combinedDistSq,
                         });
@@ -202,7 +324,7 @@ export function findWaterBridgeContacts(
         }
     }
 
-    return Array.from(best.values()).map(e => e.contact);
+    return bestBridgeValues(best).map(e => e.contact);
 }
 
 // ---------------------------------------------------------------------------
@@ -241,6 +363,7 @@ namespace WaterBridges {
     function getLabel(data: Data, elements: ReadonlyArray<Element>): string {
         const e = elements[0];
         if (e === undefined) return '';
+
         const { structure, waterBridges, unitsFeatures } = data;
         const wb = waterBridges[e.bridgeIndex];
 
@@ -266,14 +389,29 @@ namespace WaterBridges {
     }
 
     function getBoundingSphere(data: Data, elements: ReadonlyArray<Element>, boundingSphere: Sphere3D) {
-        return CentroidHelper.fromPairProvider(elements.length, (i, pA, pB) => {
-            const wb = data.waterBridges[elements[i].bridgeIndex];
+        return CentroidHelper.fromPairProvider(elements.length * 2, (i, pA, pB) => {
+            const wb = data.waterBridges[elements[i >> 1].bridgeIndex];
+
             const uA = data.structure.unitMap.get(wb.unitA) as Unit.Atomic;
             const fA = data.unitsFeatures.get(wb.unitA);
-            uA.conformation.position(uA.elements[fA.members[fA.offsets[wb.indexA]]], pA);
+
+            const uW = data.structure.unitMap.get(wb.unitW) as Unit.Atomic;
+            const fW = data.unitsFeatures.get(wb.unitW);
+
             const uB = data.structure.unitMap.get(wb.unitB) as Unit.Atomic;
             const fB = data.unitsFeatures.get(wb.unitB);
-            uB.conformation.position(uB.elements[fB.members[fB.offsets[wb.indexB]]], pB);
+
+            const aIdx = fA.members[fA.offsets[wb.indexA]] as StructureElement.UnitIndex;
+            const wIdx = fW.members[fW.offsets[wb.indexWA]] as StructureElement.UnitIndex;
+            const bIdx = fB.members[fB.offsets[wb.indexB]] as StructureElement.UnitIndex;
+
+            if ((i & 1) === 0) {
+                uA.conformation.position(uA.elements[aIdx], pA);
+                uW.conformation.position(uW.elements[wIdx], pB);
+            } else {
+                uW.conformation.position(uW.elements[wIdx], pA);
+                uB.conformation.position(uB.elements[bIdx], pB);
+            }
         }, boundingSphere);
     }
-}
+}

src/mol-model-props/computed/representations/interactions-water-bridge-cylinder.ts

@@ -23,10 +23,111 @@ import { Sphere3D } from '../../../mol-math/geometry';
 import { InteractionsSharedParams } from './shared';
 import { Features } from '../interactions/features';
 
+type WaterBridgeContacts = WaterBridges.Data['waterBridges'];
+
+type CanonicalLegIndices = {
+    donor: Int32Array
+    acceptor: Int32Array
+};
+
+const CanonicalLegIndicesCache = new WeakMap<WaterBridgeContacts, CanonicalLegIndices>();
+
+function getCanonicalLegIndices(waterBridges: WaterBridgeContacts): CanonicalLegIndices {
+    const cached = CanonicalLegIndicesCache.get(waterBridges);
+    if (cached) return cached;
+
+    const n = waterBridges.length;
+    const donor = new Int32Array(n);
+    const acceptor = new Int32Array(n);
+
+    const donorLegs = new Map<string, number>();
+    const acceptorLegs = new Map<string, number>();
+
+    for (let i = 0; i < n; i++) {
+        const wb = waterBridges[i];
+
+        const dk = `${wb.unitA}|${wb.indexA}|${wb.unitW}|${wb.indexWA}`;
+        const ak = `${wb.unitW}|${wb.indexWD}|${wb.unitB}|${wb.indexB}`;
+
+        let di = donorLegs.get(dk);
+        if (di === undefined) {
+            di = i;
+            donorLegs.set(dk, i);
+        }
+        donor[i] = di;
+
+        let ai = acceptorLegs.get(ak);
+        if (ai === undefined) {
+            ai = i;
+            acceptorLegs.set(ak, i);
+        }
+        acceptor[i] = ai;
+    }
+
+    const indices = { donor, acceptor };
+    CanonicalLegIndicesCache.set(waterBridges, indices);
+    return indices;
+}
+
+function getFeatureMember(features: Features, featureIndex: Features.FeatureIndex): StructureElement.UnitIndex {
+    return features.members[features.offsets[featureIndex]] as StructureElement.UnitIndex;
+}
+
+function atomPosition(unit: Unit.Atomic, features: Features, featureIndex: Features.FeatureIndex, out: Vec3) {
+    const atomLocalIdx = getFeatureMember(features, featureIndex);
+    unit.conformation.position(unit.elements[atomLocalIdx], out);
+}
+
+function setFeatureLocation(
+    structure: Structure,
+    location: StructureElement.Location,
+    unitId: number,
+    features: Features,
+    featureIndex: Features.FeatureIndex
+) {
+    const unit = structure.unitMap.get(unitId) as Unit.Atomic;
+    const atomLocalIdx = getFeatureMember(features, featureIndex);
+
+    location.unit = unit;
+    location.element = unit.elements[atomLocalIdx];
+}
+
+function applyDonorLeg(
+    bridgeIndex: number,
+    bridgeCount: number,
+    canonical: CanonicalLegIndices,
+    apply: (interval: Interval) => boolean
+) {
+    let changed = false;
+    const i = canonical.donor[bridgeIndex];
+
+    if (apply(Interval.ofSingleton(i))) changed = true;
+    if (apply(Interval.ofSingleton(i + bridgeCount))) changed = true;
+
+    return changed;
+}
+
+function applyAcceptorLeg(
+    bridgeIndex: number,
+    bridgeCount: number,
+    canonical: CanonicalLegIndices,
+    apply: (interval: Interval) => boolean
+) {
+    let changed = false;
+    const i = canonical.acceptor[bridgeIndex];
+
+    if (apply(Interval.ofSingleton(i + 2 * bridgeCount))) changed = true;
+    if (apply(Interval.ofSingleton(i + 3 * bridgeCount))) changed = true;
+
+    return changed;
+}
+
 function createWaterBridgeCylinderMesh(ctx: VisualContext, structure: Structure, theme: Theme, props: PD.Values<WaterBridgeInterUnitParams>, mesh?: Mesh) {
     if (!structure.hasAtomic) return Mesh.createEmpty(mesh);
 
-    const interactions = InteractionsProvider.get(structure).value!;
+    const interactions = InteractionsProvider.get(structure).value;
+    if (!interactions) return Mesh.createEmpty(mesh);
+
     const { waterBridges, unitsFeatures } = interactions;
 
     const n = waterBridges.length;
@@ -34,31 +135,20 @@ function createWaterBridgeCylinderMesh(ctx: VisualContext, structure: Structure,
 
     const l = StructureElement.Location.create(structure);
     const { sizeFactor } = props;
-
-    // When multiple bridges share the same physical leg (e.g. the same donor
-    // bonded to the same water oxygen, bridging two different acceptors) the
-    // half-cylinders for that leg must be drawn only once. Drawing them twice
-    // stacks two cylinders at the same position, causing a barely-visible hover
-    // halo and non-deterministic picking.
-    const seenDonorLegs = new Set<string>();
-    const seenAcceptorLegs = new Set<string>();
-    const skipDonorLeg = new Uint8Array(n);
-    const skipAcceptorLeg = new Uint8Array(n);
-    for (let i = 0; i < n; i++) {
-        const wb = waterBridges[i];
-        const dk = `${wb.unitA}|${wb.indexA}|${wb.unitW}|${wb.indexWA}`;
-        const ak = `${wb.unitW}|${wb.indexWD}|${wb.unitB}|${wb.indexB}`;
-        if (seenDonorLegs.has(dk)) { skipDonorLeg[i] = 1; } else { seenDonorLegs.add(dk); }
-        if (seenAcceptorLegs.has(ak)) { skipAcceptorLeg[i] = 1; } else { seenAcceptorLegs.add(ak); }
-    }
+    const canonical = getCanonicalLegIndices(waterBridges);
 
     const builderProps = {
-        // Four half-cylinders per bridge (createLinkCylinderMesh draws only the A-side half per call):
+        // Four half-cylinders per bridge; createLinkCylinderMesh draws the A-side half per call:
         //   [0,   n): donor→water,    forward  (donor side)
         //   [n,  2n): donor→water,    backward (water side)
         //   [2n, 3n): water→acceptor, forward  (water side)
         //   [3n, 4n): water→acceptor, backward (acceptor side)
+        //
+        // When multiple bridges share the same physical leg, only the first
+        // occurrence is drawn. Marking later maps duplicate legs back to the
+        // canonical drawn edge index.
         linkCount: 4 * n,
+
         position: (posA: Vec3, posB: Vec3, edgeIndex: number) => {
             const wb = waterBridges[edgeIndex % n];
             const uW = structure.unitMap.get(wb.unitW) as Unit.Atomic;
@@ -72,7 +162,7 @@ function createWaterBridgeCylinderMesh(ctx: VisualContext, structure: Structure,
                 atomPosition(uA, fA, wb.indexA, posA);
                 atomPosition(uW, fW, wb.indexWA, posB);
             } else if (leg === 1) {
-                // donor→water, B-side: draw water→mid (posA/posB swapped)
+                // donor→water, B-side: draw water→mid
                 const uA = structure.unitMap.get(wb.unitA) as Unit.Atomic;
                 const fA = unitsFeatures.get(wb.unitA);
                 atomPosition(uW, fW, wb.indexWA, posA);
@@ -84,39 +174,53 @@ function createWaterBridgeCylinderMesh(ctx: VisualContext, structure: Structure,
                 atomPosition(uW, fW, wb.indexWD, posA);
                 atomPosition(uB, fB, wb.indexB, posB);
             } else {
-                // water→acceptor, B-side: draw acceptor→mid (posA/posB swapped)
+                // water→acceptor, B-side: draw acceptor→mid
                 const uB = structure.unitMap.get(wb.unitB) as Unit.Atomic;
                 const fB = unitsFeatures.get(wb.unitB);
                 atomPosition(uB, fB, wb.indexB, posA);
                 atomPosition(uW, fW, wb.indexWD, posB);
             }
         },
+
         ignore: (edgeIndex: number) => {
             const bi = edgeIndex % n;
             const leg = Math.floor(edgeIndex / n);
-            return leg <= 1 ? skipDonorLeg[bi] === 1 : skipAcceptorLeg[bi] === 1;
+
+            return leg <= 1
+                ? canonical.donor[bi] !== bi
+                : canonical.acceptor[bi] !== bi;
         },
+
         style: (_edgeIndex: number) => LinkStyle.Dashed,
+
         radius: (edgeIndex: number) => {
             const wb = waterBridges[edgeIndex % n];
             const leg = Math.floor(edgeIndex / n);
             const isDonorWaterLeg = leg <= 1;
 
-            const fA = unitsFeatures.get(isDonorWaterLeg ? wb.unitA : wb.unitW);
-            const unitIdA = isDonorWaterLeg ? wb.unitA : wb.unitW;
-            const indexA = isDonorWaterLeg ? wb.indexA : wb.indexWD;
-            l.unit = structure.unitMap.get(unitIdA);
-            l.element = l.unit.elements[fA.members[fA.offsets[indexA]]];
-            const sizeA = theme.size.size(l);
+            if (isDonorWaterLeg) {
+                const fA = unitsFeatures.get(wb.unitA);
+                const fW = unitsFeatures.get(wb.unitW);
 
-            const fB = unitsFeatures.get(isDonorWaterLeg ? wb.unitW : wb.unitB);
-            const unitIdB = isDonorWaterLeg ? wb.unitW : wb.unitB;
-            const indexB = isDonorWaterLeg ? wb.indexWA : wb.indexB;
-            l.unit = structure.unitMap.get(unitIdB);
-            l.element = l.unit.elements[fB.members[fB.offsets[indexB]]];
-            const sizeB = theme.size.size(l);
+                setFeatureLocation(structure, l, wb.unitA, fA, wb.indexA);
+                const sizeA = theme.size.size(l);
 
-            return Math.min(sizeA, sizeB) * sizeFactor;
+                setFeatureLocation(structure, l, wb.unitW, fW, wb.indexWA);
+                const sizeW = theme.size.size(l);
+
+                return Math.min(sizeA, sizeW) * sizeFactor;
+            } else {
+                const fW = unitsFeatures.get(wb.unitW);
+                const fB = unitsFeatures.get(wb.unitB);
+
+                setFeatureLocation(structure, l, wb.unitW, fW, wb.indexWD);
+                const sizeW = theme.size.size(l);
+
+                setFeatureLocation(structure, l, wb.unitB, fB, wb.indexB);
+                const sizeB = theme.size.size(l);
+
+                return Math.min(sizeW, sizeB) * sizeFactor;
+            }
         },
     };
 
@@ -125,18 +229,13 @@ function createWaterBridgeCylinderMesh(ctx: VisualContext, structure: Structure,
     if (boundingSphere) {
         m.setBoundingSphere(boundingSphere);
     } else if (m.triangleCount > 0) {
-        const sphere = Sphere3D.expand(Sphere3D(), structure.boundary.sphere, 1 * sizeFactor);
+        const sphere = Sphere3D.expand(Sphere3D(), structure.boundary.sphere, sizeFactor);
         m.setBoundingSphere(sphere);
     }
 
     return m;
 }
 
-function atomPosition(unit: Unit.Atomic, features: Features, featureIndex: Features.FeatureIndex, out: Vec3) {
-    const atomLocalIdx = features.members[features.offsets[featureIndex]];
-    unit.conformation.position(unit.elements[atomLocalIdx], out);
-}
-
 export const WaterBridgeInterUnitParams = {
     ...ComplexMeshParams,
     ...LinkCylinderParams,
@@ -151,13 +250,23 @@ export function WaterBridgeInterUnitVisual(materialId: number): ComplexVisual<Wa
         createLocationIterator: createWaterBridgeIterator,
         getLoci: getWaterBridgeLoci,
         eachLocation: eachWaterBridgeInteraction,
-        setUpdateState: (state: VisualUpdateState, newProps: PD.Values<WaterBridgeInterUnitParams>, currentProps: PD.Values<WaterBridgeInterUnitParams>, newTheme: Theme, currentTheme: Theme, newStructure: Structure, currentStructure: Structure) => {
+
+        setUpdateState: (
+            state: VisualUpdateState,
+            newProps: PD.Values<WaterBridgeInterUnitParams>,
+            currentProps: PD.Values<WaterBridgeInterUnitParams>,
+            newTheme: Theme,
+            currentTheme: Theme,
+            newStructure: Structure,
+            _currentStructure: Structure
+        ) => {
             state.createGeometry = (
                 newProps.sizeFactor !== currentProps.sizeFactor ||
                 newProps.dashCount !== currentProps.dashCount ||
                 newProps.dashScale !== currentProps.dashScale ||
                 newProps.dashCap !== currentProps.dashCap ||
-                newProps.radialSegments !== currentProps.radialSegments
+                newProps.radialSegments !== currentProps.radialSegments ||
+                newTheme.size !== currentTheme.size
             );
 
             const interactionsHash = InteractionsProvider.get(newStructure).version;
@@ -175,9 +284,15 @@ function getWaterBridgeLoci(pickingId: PickingId, structure: Structure, id: numb
     const { objectId, groupId } = pickingId;
     if (id !== objectId) return EmptyLoci;
 
-    const interactions = InteractionsProvider.get(structure).value!;
+    const interactions = InteractionsProvider.get(structure).value;
+    if (!interactions) return EmptyLoci;
+
     const { waterBridges, unitsFeatures } = interactions;
-    const bridgeIndex = groupId % waterBridges.length;
+    const n = waterBridges.length;
+
+    if (!n || groupId < 0 || groupId >= 4 * n) return EmptyLoci;
+
+    const bridgeIndex = groupId % n;
 
     return WaterBridges.Loci({ structure, waterBridges, unitsFeatures }, [{ bridgeIndex }]);
 }
@@ -189,15 +304,20 @@ function eachWaterBridgeInteraction(loci: Loci, structure: Structure, apply: (in
 
     if (WaterBridges.isLoci(loci)) {
         if (!Structure.areEquivalent(loci.data.structure, structure)) return false;
-        const interactions = InteractionsProvider.get(structure).value!;
+
+        const interactions = InteractionsProvider.get(structure).value;
+        if (!interactions) return false;
+
         const n = interactions.waterBridges.length;
+        if (!n) return false;
+
+        const canonical = getCanonicalLegIndices(interactions.waterBridges);
 
         for (const e of loci.elements) {
-            // Apply all four half-cylinders for this bridge.
-            if (apply(Interval.ofSingleton(e.bridgeIndex))) changed = true;
-            if (apply(Interval.ofSingleton(e.bridgeIndex + n))) changed = true;
-            if (apply(Interval.ofSingleton(e.bridgeIndex + 2 * n))) changed = true;
-            if (apply(Interval.ofSingleton(e.bridgeIndex + 3 * n))) changed = true;
+            if (e.bridgeIndex < 0 || e.bridgeIndex >= n) continue;
+
+            if (applyDonorLeg(e.bridgeIndex, n, canonical, apply)) changed = true;
+            if (applyAcceptorLeg(e.bridgeIndex, n, canonical, apply)) changed = true;
         }
     } else if (StructureElement.Loci.is(loci)) {
         if (!Structure.areEquivalent(loci.structure, structure)) return false;
@@ -207,36 +327,52 @@ function eachWaterBridgeInteraction(loci: Loci, structure: Structure, apply: (in
 
         const { waterBridges, unitsFeatures } = interactions;
         const n = waterBridges.length;
+        if (!n) return false;
 
-        for (const e of loci.elements) __unitMap.set(e.unit.id, e.indices);
+        const canonical = getCanonicalLegIndices(waterBridges);
+
+        __unitMap.clear();
+        for (const e of loci.elements) {
+            __unitMap.set(e.unit.id, e.indices);
+        }
 
         for (let i = 0; i < n; i++) {
             const wb = waterBridges[i];
+
             const indicesA = __unitMap.get(wb.unitA);
+            const indicesW = __unitMap.get(wb.unitW);
             const indicesB = __unitMap.get(wb.unitB);
-            if (!indicesA && !indicesB) continue;
+
+            if (!indicesA && !indicesW && !indicesB) continue;
 
             let hitA = false;
             if (indicesA) {
                 const fA = unitsFeatures.get(wb.unitA);
-                const mi = fA.members[fA.offsets[wb.indexA]] as StructureElement.UnitIndex;
+                const mi = getFeatureMember(fA, wb.indexA);
                 hitA = OrderedSet.has(indicesA, mi);
             }
 
+            let hitW = false;
+            if (indicesW) {
+                const fW = unitsFeatures.get(wb.unitW);
+                const miA = getFeatureMember(fW, wb.indexWA);
+                const miD = getFeatureMember(fW, wb.indexWD);
+                hitW = OrderedSet.has(indicesW, miA) || OrderedSet.has(indicesW, miD);
+            }
+
             let hitB = false;
             if (indicesB) {
                 const fB = unitsFeatures.get(wb.unitB);
-                const mi = fB.members[fB.offsets[wb.indexB]] as StructureElement.UnitIndex;
+                const mi = getFeatureMember(fB, wb.indexB);
                 hitB = OrderedSet.has(indicesB, mi);
             }
 
-            if (hitA) {
-                if (apply(Interval.ofSingleton(i))) changed = true;
-                if (apply(Interval.ofSingleton(i + n))) changed = true;
+            if (hitA || hitW) {
+                if (applyDonorLeg(i, n, canonical, apply)) changed = true;
             }
-            if (hitB) {
-                if (apply(Interval.ofSingleton(i + 2 * n))) changed = true;
-                if (apply(Interval.ofSingleton(i + 3 * n))) changed = true;
+
+            if (hitB || hitW) {
+                if (applyAcceptorLeg(i, n, canonical, apply)) changed = true;
             }
         }
 
@@ -249,15 +385,19 @@ function eachWaterBridgeInteraction(loci: Loci, structure: Structure, apply: (in
 function createWaterBridgeIterator(structure: Structure): LocationIterator {
     const interactions = InteractionsProvider.get(structure).value!;
     const { waterBridges, unitsFeatures } = interactions;
+
     const n = waterBridges.length;
     const groupCount = 4 * n;
     const instanceCount = 1;
+
     const data: WaterBridges.Data = { structure, waterBridges, unitsFeatures };
     const location = WaterBridges.Location(data);
     const { element } = location;
+
     const getLocation = (groupIndex: number) => {
-        element.bridgeIndex = groupIndex % n;
+        element.bridgeIndex = n === 0 ? 0 : groupIndex % n;
         return location;
     };
+
     return LocationIterator(groupCount, instanceCount, 1, getLocation, true);
-}
+}