ASA calc now uses lookup3d and ties results to atoms and residues

src/mol-model/structure/model/model.ts

@@ -49,8 +49,6 @@ export interface Model extends Readonly<{
         readonly chemicalComponentMap: ChemicalComponentMap
         /** maps residue name to `SaccharideComponent` data */
         readonly saccharideComponentMap: SaccharideComponentMap
-        /** brute-force way to add ASA data, prob should be a custom property */
-        readonly asa: number[]
     },
 
     customProperties: CustomProperties,

src/mol-model/structure/structure/unit.ts

@@ -18,6 +18,8 @@ import { IntMap, SortedArray } from 'mol-data/int';
 import { hash2, hashFnv32a } from 'mol-data/util';
 import { getAtomicPolymerElements, getCoarsePolymerElements, getAtomicGapElements, getCoarseGapElements } from './util/polymer';
 import { getNucleotideElements } from './util/nucleotide';
+import { computeAccessibleSurfaceArea } from './unit/accessible-surface-area/compute';
+import { AccessibleSurfaceArea } from './unit/accessible-surface-area/data';
 
 /**
  * A building block of a structure that corresponds to an atomic or
@@ -191,6 +193,12 @@ namespace Unit {
             return this.props.nucleotideElements.ref;
         }
 
+        get accessibleSurfaceArea() {
+            if (this.props.accessibleSurfaceArea.ref) return this.props.accessibleSurfaceArea.ref;
+            this.props.accessibleSurfaceArea.ref = computeAccessibleSurfaceArea(this);
+            return this.props.accessibleSurfaceArea.ref;
+        }
+
         getResidueIndex(elementIndex: StructureElement.UnitIndex) {
             return this.model.atomicHierarchy.residueAtomSegments.index[this.elements[elementIndex]];
         }
@@ -215,6 +223,7 @@ namespace Unit {
         polymerElements: ValueRef<SortedArray<ElementIndex> | undefined>
         gapElements: ValueRef<SortedArray<ElementIndex> | undefined>
         nucleotideElements: ValueRef<SortedArray<ElementIndex> | undefined>
+        accessibleSurfaceArea: ValueRef<AccessibleSurfaceArea | undefined>
     }
 
     function AtomicProperties(): AtomicProperties {
@@ -225,6 +234,7 @@ namespace Unit {
             polymerElements: ValueRef.create(void 0),
             gapElements: ValueRef.create(void 0),
             nucleotideElements: ValueRef.create(void 0),
+            accessibleSurfaceArea: ValueRef.create(void 0)
         };
     }
 

src/mol-model/structure/model/properties/utils/accessible-surface-area.ts → src/mol-model/structure/structure/unit/accessible-surface-area/compute.ts

+import Unit from '../../unit'
 import { Vec3 } from 'mol-math/linear-algebra';
-import { AtomicHierarchy, AtomicConformation } from '../atomic';
-import { ElementSymbol, MaxAsa, DefaultMaxAsa, MoleculeType } from '../../types';
-import { VdwRadius } from '../atomic/measures';
+import { AccessibleSurfaceAreaComputationParameters, AccessibleSurfaceArea } from './data';
+import { isHydrogen, getElementIdx } from '../links/common'; // TODO move these functions somewhere more common
+import { MoleculeType, ElementSymbol, MaxAsa, DefaultMaxAsa } from 'mol-model/structure/model/types';
+import { VdwRadius } from 'mol-model/structure/model/properties/atomic/measures';
 
-const defaultNumberOfPoints = 960;
-const defaultProbeSize = 1.4;
 const trigonalCarbonVdw = 1.76;
 const tetrahedralCarbonVdw = 1.87;
 const trigonalNitrogenVdw = 1.65;
@@ -12,134 +12,137 @@ const tetrahedralNitrogenVdw = 1.50;
 /** deviating radii from the definition in types.ts */
 const oxygenVdw = 1.4;
 const sulfurVdw = 1.85;
+const missingAccessibleSurfaceAreaValue = -1.0;
+
+function _computeAccessibleSurfaceArea(unit: Unit.Atomic, params: AccessibleSurfaceAreaComputationParameters): AccessibleSurfaceArea {
+    const ctx = initialize(unit, params);
+    assignRadiusForHeavyAtoms(ctx);
+    computePerResidue(ctx);
+    normalizeAccessibleSurfaceArea(ctx);
+
+    return {
+        atomRadius: ctx.atomRadius,
+        accessibleSurfaceArea: ctx.accessibleSurfaceArea,
+        relativeAccessibleSurfaceArea: ctx.relativeAccessibleSurfaceArea,
+        buried: void 0 // TODO impl - rasa < 0.16 - find Rost reference
+    };
+}
 
-export function computeModelASA(hierarchy: AtomicHierarchy, conformation: AtomicConformation) {
-    const numberOfSpherePoints = defaultNumberOfPoints;
+function normalizeAccessibleSurfaceArea(ctx: AccessibleSurfaceAreaContext) {
+    const { residues, derived } = ctx.unit.model.atomicHierarchy;
+    const { accessibleSurfaceArea, relativeAccessibleSurfaceArea } = ctx;
 
-    const ctx: ASAContext = {
-        probeSize: defaultProbeSize,
-        spherePoints: generateSpherePoints(numberOfSpherePoints),
-        cons: 4.0 * Math.PI / numberOfSpherePoints,
+    for (let i = 0; i < residues.label_comp_id.rowCount; ++i) {
+        // skip entities not part of a peptide chain
+        if (derived.residue.moleculeType[i] !== MoleculeType.protein) continue;
 
-        hierarchy,
-        conformation
+        const maxAsa = (MaxAsa as any)[residues.label_comp_id.value(i)];
+        relativeAccessibleSurfaceArea[i] = accessibleSurfaceArea[i] / (maxAsa === undefined ? DefaultMaxAsa : maxAsa);
     }
 
-    const { probeSize, spherePoints, cons } = ctx;
-    const { atoms, residues, residueAtomSegments, derived } = ctx.hierarchy;
-    const { type_symbol } = atoms;
-    const { x, y, z } = ctx.conformation;
+}
 
-    const radii: number[] = [];
-    const residueAsa: number[] = [];
+// TODO compare performance of lookup and naive approach
+function computePerResidue(ctx: AccessibleSurfaceAreaContext) {
+    const { atomRadius, accessibleSurfaceArea, maxLookupRadius, spherePoints, cons } = ctx;
+    const { probeSize } = ctx.params;
+    const { elements: atoms } = ctx.unit;
+    const { residueAtomSegments } = ctx.unit.model.atomicHierarchy;
+    const { x, y, z } = ctx.unit.model.atomicConformation;
+    const atomCount = ctx.unit.elements.length;
+    const { lookup3d } = ctx.unit;
 
     const position = (i: number, v: Vec3) => Vec3.set(v, x[i], y[i], z[i]);
     const a1Pos = Vec3.zero();
     const a2Pos = Vec3.zero();
 
-    // extract all heavy atoms
-    // TODO can be more elegant by obtaining residue info once and then using offset to navigate to next residue
-    for (let i = 0; i < type_symbol.rowCount; ++i) {
-        // skip hydrogen atoms
-        if (type_symbol.value(i) === 'H' || type_symbol.value(i) === 'D' || type_symbol.value(i) === 'T') {
-            radii[radii.length] = valueForIgnoredAtom; // -1 is used to signal atoms not to be considered downstream
-            continue;
-        }
-
-        // determine group this atom belongs to
-        const groupIndex = residueAtomSegments.index[i];
-
-        // skip entities not part of a peptide chain
-        if (derived.residue.moleculeType[groupIndex] !== MoleculeType.protein) {
-            radii[radii.length] = valueForIgnoredAtom;
-            continue;
-        }
-
-        const atomId = atoms.label_atom_id.value(i);
-        const compId = residues.label_comp_id.value(groupIndex);
-        const element = type_symbol.value(i);
-        // assign radius to all heavy atoms - depends on element and bonding patterns
-        radii[radii.length] = determineRadius(atomId, element, compId);
-        // set ASA of corresponding group to 0
-        residueAsa[groupIndex] = 0.0;
-    }
-
-    // calculate the individual ASA of each atom
-    // TODO again might be more elegant to use offsets
-    // TODO distance is symmetric, omit redudant calcuations
-    for (let i = 0; i < radii.length; ++i) {
-        const radius = radii[i];
-
-        // skip invalid entries
-        if (radius === valueForIgnoredAtom) {
-            // atomAsa[atomAsa.length] = valueForIgnoredAtom;
-            continue;
-        }
-
-        position(i, a1Pos);
-
-        // collect all neighboring atoms
-        const cutoff = probeSize + probeSize + radius;
-        const neighborIndices: number[] = [];
-        for (let k = 0; k < radii.length; ++k) {
-            const radius2 = radii[k];
-            if (i === k || radius2 === valueForIgnoredAtom)
-                continue;
-
-            position(k, a2Pos);
-
-            if (Vec3.distance(a1Pos, a2Pos) < cutoff + radius2) {
-                neighborIndices[neighborIndices.length] = k;
+    for (let _aI = 0; _aI < atomCount; ++_aI) {
+        const aI = atoms[_aI];
+        const radii1 = atomRadius[aI];
+        if (radii1 === missingAccessibleSurfaceAreaValue) continue;
+
+        // find suitable neighbors by lookup
+        const { indices, count } = lookup3d.find(x[aI], y[aI], z[aI], maxLookupRadius);
+        position(aI, a1Pos);
+
+        // refine list by actual criterion
+        const cutoff = probeSize + probeSize + radii1;
+        const filteredIndicies = [];
+        for (let ni = 0; ni < count; ni++) {
+            const _bI = indices[ni];
+            const bI = atoms[_bI];
+            const radii2 = atomRadius[bI];
+            if (bI === aI || radii2 === missingAccessibleSurfaceAreaValue) continue;
+
+            const cutoff2 = (cutoff + radii2) * (cutoff + radii2);
+            // accurately check for neighborhood
+            position(bI, a2Pos);
+            if (Vec3.squaredDistance(a1Pos, a2Pos) < cutoff2) {
+                filteredIndicies[filteredIndicies.length] = bI;
             }
         }
 
-        const r = probeSize + radius;
-        let accessiblePoints = 0;
-
-        for (let k = 0; k < spherePoints.length; ++k) {
-            const point = spherePoints[k];
+        // test sphere points
+        const r = probeSize + radii1;
+        let accessiblePointCount = 0;
+        for (let si = 0; si < spherePoints.length; ++si) {
+            const spherePoint = spherePoints[si];
+            const testPoint = [spherePoint[0] * r + a1Pos[0], spherePoint[1] * r + a1Pos[1], spherePoint[2] * r + a1Pos[2]] as Vec3;
             let accessible = true;
-            const testPoint = [point[0] * r + a1Pos[0], point[1] * r + a1Pos[1], point[2] * r + a1Pos[2]] as Vec3;
 
-            for (let j = 0; j < neighborIndices.length; ++j) {
-                const naI = neighborIndices[j];
-                // store position of neighboring atom in a2Pos
+            for (let ni = 0; ni < filteredIndicies.length; ++ni) {
+                const naI = filteredIndicies[ni];
                 position(naI, a2Pos);
-                const neighboringAtomRadius = radii[naI] + probeSize;
-                if (Vec3.squaredDistance(testPoint, a2Pos) < neighboringAtomRadius * neighboringAtomRadius) {
+                const cutoff3 = (atomRadius[naI] + probeSize) * (atomRadius[naI] + probeSize);
+                if (Vec3.squaredDistance(testPoint, a2Pos) < cutoff3) {
                     accessible = false;
                     break;
                 }
             }
 
-            if (accessible) {
-                ++accessiblePoints;
-            }
+            if (accessible) ++accessiblePointCount;
         }
 
-        const value = cons * accessiblePoints * r * r;
-        // atomAsa[atomAsa.length] = value;
-        // sum up values for each residue
-        residueAsa[residueAtomSegments.index[i]] += value;
+        const value = cons * accessiblePointCount * r * r;
+        accessibleSurfaceArea[residueAtomSegments.index[aI]] += value;
     }
+}
 
-    // console.log(residueAsa);
+function assignRadiusForHeavyAtoms(ctx: AccessibleSurfaceAreaContext) {
+    const atomCount = ctx.unit.elements.length;
+    const { elements: atoms, residueIndex } = ctx.unit;
+    const { moleculeType } = ctx.unit.model.atomicHierarchy.derived.residue;
+    const { type_symbol, label_atom_id } = ctx.unit.model.atomicHierarchy.atoms;
+    const { label_comp_id } = ctx.unit.model.atomicHierarchy.residues;
 
-    // normalize by maximum value expected for a particular amino acid - needs lookup of max values
-    for (let i = 0; i < residues.label_comp_id.rowCount; ++i) {
-        // skip entities not part of a peptide chain
-        if (derived.residue.moleculeType[i] !== 4)
+    for (let _aI = 0; _aI < atomCount; ++_aI) {
+        const aI =  atoms[_aI];
+        const raI = residueIndex[aI];
+        const aeI = getElementIdx(type_symbol.value(aI)!);
+
+        // skip hydrogen atoms
+        if (isHydrogen(aeI)) {
+            ctx.atomRadius[ctx.atomRadius.length] = missingAccessibleSurfaceAreaValue;
             continue;
+        }
 
-        const maxAsa = (MaxAsa as any)[residues.label_comp_id.value(i)];
-        residueAsa[i] /= (maxAsa === undefined ? DefaultMaxAsa : maxAsa);
-    }
+        // skip non-peptide groups
+        if (moleculeType[raI] !== MoleculeType.protein) {
+            ctx.atomRadius[ctx.atomRadius.length] = missingAccessibleSurfaceAreaValue;
+            continue;
+        }
 
-    console.log(residueAsa);
-    return residueAsa;
-}
+        const atomId = label_atom_id.value(aI);
+        const element = type_symbol.value(aI);
+        const resn = label_comp_id.value(raI)!;
 
-const valueForIgnoredAtom = -1.0;
+        ctx.atomRadius[ctx.atomRadius.length] = determineRadius(atomId, element, resn);
+
+        // while having atom->parent mapping at hand, initialize residue-level of information
+        ctx.accessibleSurfaceArea[ctx.accessibleSurfaceArea.length] = 0.0;
+        ctx.relativeAccessibleSurfaceArea[ctx.relativeAccessibleSurfaceArea.length] = 0.0;
+    }
+}
 
 /**
  * Gets the van der Waals radius of the given atom following the values defined by Chothia (1976)
@@ -174,6 +177,37 @@ function determineRadius(atomId: string, element: ElementSymbol, compId: string)
     return VdwRadius(element);
 }
 
+interface AccessibleSurfaceAreaContext {
+    unit: Unit.Atomic,
+    params: AccessibleSurfaceAreaComputationParameters,
+    spherePoints: Vec3[],
+    cons: number,
+    atomRadius: number[],
+    accessibleSurfaceArea: number[],
+    relativeAccessibleSurfaceArea: number[],
+    maxLookupRadius: number
+}
+
+function initialize(unit: Unit.Atomic, params: AccessibleSurfaceAreaComputationParameters): AccessibleSurfaceAreaContext {
+    return {
+        unit: unit,
+        params: params,
+        spherePoints: generateSpherePoints(params.numberOfSpherePoints),
+        cons: 4.0 * Math.PI / params.numberOfSpherePoints,
+        atomRadius: [],
+        accessibleSurfaceArea: [],
+        relativeAccessibleSurfaceArea: [],
+        maxLookupRadius: 1.4 + 1.4 + 1.87 + 1.87
+    }
+}
+
+function computeAccessibleSurfaceArea(unit: Unit.Atomic, params?: Partial<AccessibleSurfaceAreaComputationParameters>): AccessibleSurfaceArea {
+    return _computeAccessibleSurfaceArea(unit, {
+        numberOfSpherePoints: (params && params.numberOfSpherePoints) || 960,
+        probeSize: (params && params.probeSize) || 1.4
+    });
+}
+
 /** Creates a collection of points on a sphere by the Golden Section Spiral algorithm. */
 function generateSpherePoints(numberOfSpherePoints: number): Vec3[] {
     const points: Vec3[] = [];
@@ -188,11 +222,4 @@ function generateSpherePoints(numberOfSpherePoints: number): Vec3[] {
     return points;
 }
 
-interface ASAContext {
-    probeSize: number,
-    spherePoints: Vec3[],
-    cons: number,
-
-    hierarchy: AtomicHierarchy,
-    conformation: AtomicConformation
-}
\ No newline at end of file
+export { computeAccessibleSurfaceArea, missingAccessibleSurfaceAreaValue }
\ No newline at end of file

src/mol-model/structure/structure/unit/accessible-surface-area/data.ts

+interface AccessibleSurfaceArea {
+    readonly atomRadius: ArrayLike<number>,
+    readonly accessibleSurfaceArea: ArrayLike<number>,
+    readonly relativeAccessibleSurfaceArea: ArrayLike<number>,
+    readonly buried: any
+}
+
+interface AccessibleSurfaceAreaComputationParameters {
+    numberOfSpherePoints: number
+    probeSize: number
+}
+
+export { AccessibleSurfaceArea, AccessibleSurfaceAreaComputationParameters }
\ No newline at end of file

src/mol-theme/color/accessible-surface-area.ts

@@ -6,13 +6,14 @@
 
 import { Color } from 'mol-util/color';
 import { Location } from 'mol-model/location';
-import { ColorTheme } from '../color';
+import { ColorTheme, LocationColor } from '../color';
 import { ParamDefinition as PD } from 'mol-util/param-definition'
 import { ThemeDataContext } from '../theme';
 import { ColorListOptions, ColorListName, ColorScale } from 'mol-util/color/scale';
-import { StructureElement, Link, ElementIndex, Unit } from 'mol-model/structure';
+import { StructureElement, Unit } from 'mol-model/structure';
+import { missingAccessibleSurfaceAreaValue } from 'mol-model/structure/structure/unit/accessible-surface-area/compute';
 
-const DefaultAccessibleSurfaceAreaColor = Color(0xCCCCCC)
+const DefaultColor = Color(0xCCCCCC)
 const Description = 'Assigns a color based on the relative accessible surface area of a residue.'
 
 export const AccessibleSurfaceAreaColorThemeParams = {
@@ -24,27 +25,29 @@ export function getAccessibleSurfaceAreaColorThemeParams(ctx: ThemeDataContext)
 }
 
 export function AccessibleSurfaceAreaColorTheme(ctx: ThemeDataContext, props: PD.Values<AccessibleSurfaceAreaColorThemeParams>): ColorTheme<AccessibleSurfaceAreaColorThemeParams> {
-    const scale = ColorScale.create({
-        listOrName: props.list,
-        minLabel: 'Start',
-        maxLabel: 'End',
-    })
-    const scaleColor = scale.color
+    let color: LocationColor
+    let scale: ColorScale | undefined = undefined
 
-    function asaColor(unit: Unit, element: ElementIndex): Color {
-        if (Unit.isAtomic(unit)) {
-            return scaleColor(unit.model.properties.asa[unit.residueIndex[element]]);
-        }
-        return DefaultAccessibleSurfaceAreaColor;
-    }
+    if (ctx.structure) {
+        scale = ColorScale.create({
+            listOrName: props.list,
+            minLabel: 'Start',
+            maxLabel: 'End'
+        })
+        const scaleColor = scale.color
+
+        color = (location: Location): Color => {
+            if (StructureElement.isLocation(location)) {
+                if (Unit.isAtomic(location.unit)) {
+                    const value = location.unit.accessibleSurfaceArea.relativeAccessibleSurfaceArea[location.unit.residueIndex[location.element]];
+                    return value !== missingAccessibleSurfaceAreaValue ? scaleColor(value) : DefaultColor;
+                }
+            }
 
-    const color = (location: Location): Color => {
-        if (StructureElement.isLocation(location)) {
-            return asaColor(location.unit, location.element);
-        } else if (Link.isLocation(location)) {
-            return asaColor(location.aUnit, location.aUnit.elements[location.aIndex]);
+            return DefaultColor
         }
-        return DefaultAccessibleSurfaceAreaColor;
+    } else {
+        color = () => DefaultColor
     }
 
     return {

src/tests/browser/render-structure.ts

@@ -12,7 +12,6 @@ import { ColorTheme } from 'mol-theme/color';
 import { SizeTheme } from 'mol-theme/size';
 import { CartoonRepresentationProvider } from 'mol-repr/structure/representation/cartoon';
 import { trajectoryFromMmCIF } from 'mol-model-formats/structure/mmcif';
-import { computeModelASA } from 'mol-model/structure/model/properties/utils/accessible-surface-area';
 
 const parent = document.getElementById('app')!
 parent.style.width = '100%'
@@ -64,10 +63,10 @@ function getCartoonRepr() {
 async function init() {
     const cif = await downloadFromPdb('1acj')
     const models = await getModels(cif)
-    console.time('computeModelASA')
-    const asa = computeModelASA(models[0].atomicHierarchy, models[0].atomicConformation)
-    console.timeEnd('computeModelASA');
-    (models[0].properties as any).asa = asa
+    // console.time('computeModelASA')
+    // const asa = computeModelASA(models[0].atomicHierarchy, models[0].atomicConformation)
+    // console.timeEnd('computeModelASA');
+    // (models[0].properties as any).asa = asa
     const structure = await getStructure(models[0])
     const cartoonRepr = getCartoonRepr()