Merge pull request #47 from JonStargaryen/anvil

ANVIL / Membrane placement prediction implementation

src/apps/viewer/index.ts

@@ -24,6 +24,7 @@ import { PluginState } from '../../mol-plugin/state';
 import { DownloadDensity } from '../../mol-plugin-state/actions/volume';
 import { PluginLayoutControlsDisplay } from '../../mol-plugin/layout';
 import { BuiltInTrajectoryFormat } from '../../mol-plugin-state/formats/trajectory';
+import { MembraneOrientationData } from '../../extensions/membrane-orientation/behavior';
 import { DnatcoConfalPyramids } from '../../extensions/dnatco';
 
 require('mol-plugin-ui/skin/light.scss');
@@ -36,7 +37,8 @@ const Extensions = {
     'dnatco-confal-pyramids': PluginSpec.Behavior(DnatcoConfalPyramids),
     'pdbe-structure-quality-report': PluginSpec.Behavior(PDBeStructureQualityReport),
     'rcsb-assembly-symmetry': PluginSpec.Behavior(RCSBAssemblySymmetry),
-    'rcsb-validation-report': PluginSpec.Behavior(RCSBValidationReport)
+    'rcsb-validation-report': PluginSpec.Behavior(RCSBValidationReport),
+    'membrane-orientation': PluginSpec.Behavior(MembraneOrientationData)
 };
 
 const DefaultViewerOptions = {

src/extensions/membrane-orientation/ANVIL.ts

@@ -0,0 +1,355 @@
+/**
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Sebastian Bittrich <sebastian.bittrich@rcsb.org>
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { ANVILContext } from './anvil/common';
+import { Structure, StructureElement, StructureProperties } from '../../mol-model/structure';
+import { Task, RuntimeContext } from '../../mol-task';
+import { CentroidHelper } from '../../mol-math/geometry/centroid-helper';
+import { AccessibleSurfaceAreaProvider } from '../../mol-model-props/computed/accessible-surface-area';
+import { Vec3 } from '../../mol-math/linear-algebra';
+import { getElementMoleculeType } from '../../mol-model/structure/util';
+import { MoleculeType } from '../../mol-model/structure/model/types';
+import { AccessibleSurfaceArea } from '../../mol-model-props/computed/accessible-surface-area/shrake-rupley';
+import { ParamDefinition as PD } from '../../mol-util/param-definition';
+import { MembraneOrientation } from './membrane-orientation';
+
+export const ANVILParams = {
+    numberOfSpherePoints: PD.Numeric(120, { min: 35, max: 700, step: 1 }, { description: 'Number of spheres/directions to test for membrane placement. Original value is 350.' }),
+    stepSize: PD.Numeric(1, { min: 0.25, max: 4, step: 0.25 }, { description: 'Thickness of membrane slices that will be tested' }),
+    minThickness: PD.Numeric(20, { min: 10, max: 30, step: 1}, { description: 'Minimum membrane thickness used during refinement' }),
+    maxThickness: PD.Numeric(40, { min: 30, max: 50, step: 1}, { description: 'Maximum membrane thickness used during refinement' }),
+    asaCutoff: PD.Numeric(40, { min: 10, max: 100, step: 1 }, { description: 'Absolute ASA cutoff above which residues will be considered' })
+};
+export type ANVILParams = typeof ANVILParams
+export type ANVILProps = PD.Values<ANVILParams>
+
+/**
+ * Implements:
+ * Membrane positioning for high- and low-resolution protein structures through a binary classification approach
+ * Guillaume Postic, Yassine Ghouzam, Vincent Guiraud, and Jean-Christophe Gelly
+ * Protein Engineering, Design & Selection, 2015, 1–5
+ * doi: 10.1093/protein/gzv063
+ */
+export function computeANVIL(structure: Structure, props: ANVILProps) {
+    return Task.create('Compute Membrane Orientation', async runtime => {
+        return await calculate(runtime, structure, props);
+    });
+}
+
+const l = StructureElement.Location.create(void 0);
+const centroidHelper = new CentroidHelper();
+function initialize(structure: Structure, props: ANVILProps): ANVILContext {
+    const { label_atom_id, x, y, z } = StructureProperties.atom;
+    const elementCount = structure.polymerResidueCount;
+    centroidHelper.reset();
+    l.structure = structure;
+
+    let offsets = new Int32Array(elementCount);
+    let exposed = new Array<boolean>(elementCount);
+
+    const accessibleSurfaceArea = structure && AccessibleSurfaceAreaProvider.get(structure);
+    const asa = accessibleSurfaceArea.value!;
+
+    const vec = Vec3();
+    let m = 0;
+    for (let i = 0, il = structure.units.length; i < il; ++i) {
+        const unit = structure.units[i];
+        const { elements } = unit;
+        l.unit = unit;
+
+        for (let j = 0, jl = elements.length; j < jl; ++j) {
+            const eI = elements[j];
+            l.element = eI;
+
+            // consider only amino acids
+            if (getElementMoleculeType(unit, eI) !== MoleculeType.Protein) {
+                continue;
+            }
+
+            // only CA is considered for downstream operations
+            if (label_atom_id(l) !== 'CA') {
+                continue;
+            }
+
+            // while iterating use first pass to compute centroid
+            Vec3.set(vec, x(l), y(l), z(l));
+            centroidHelper.includeStep(vec);
+
+            // keep track of offsets and exposed state to reuse
+            offsets[m] = structure.serialMapping.getSerialIndex(l.unit, l.element);
+            exposed[m] = AccessibleSurfaceArea.getValue(l, asa) > props.asaCutoff;
+
+            m++;
+        }
+    }
+
+    // omit potentially empty tail1
+    offsets = offsets.slice(0, m);
+    exposed = exposed.slice(0, m);
+
+    // calculate centroid and extent
+    centroidHelper.finishedIncludeStep();
+    const centroid = centroidHelper.center;
+    for (let k = 0, kl = offsets.length; k < kl; k++) {
+        setLocation(l, structure, offsets[k]);
+        Vec3.set(vec, x(l), y(l), z(l));
+        centroidHelper.radiusStep(vec);
+    }
+    const extent = 1.2 * Math.sqrt(centroidHelper.radiusSq);
+
+    return {
+        ...props,
+        structure: structure,
+
+        offsets: offsets,
+        exposed: exposed,
+        centroid: centroid,
+        extent: extent
+    };
+}
+
+export async function calculate(runtime: RuntimeContext, structure: Structure, params: ANVILProps): Promise<MembraneOrientation> {
+    const { label_comp_id } = StructureProperties.atom;
+
+    const ctx = initialize(structure, params);
+    const initialHphobHphil = HphobHphil.filtered(ctx, label_comp_id);
+
+    const initialMembrane = findMembrane(ctx, generateSpherePoints(ctx, ctx.numberOfSpherePoints), initialHphobHphil, label_comp_id);
+    const alternativeMembrane = findMembrane(ctx, findProximateAxes(ctx, initialMembrane), initialHphobHphil, label_comp_id);
+
+    const membrane = initialMembrane.qmax! > alternativeMembrane.qmax! ? initialMembrane : alternativeMembrane;
+
+    return {
+        planePoint1: membrane.planePoint1,
+        planePoint2: membrane.planePoint2,
+        normalVector: membrane.normalVector!,
+        radius: ctx.extent,
+        centroid: ctx.centroid
+    };
+}
+
+interface MembraneCandidate {
+    planePoint1: Vec3,
+    planePoint2: Vec3,
+    stats: HphobHphil,
+    normalVector?: Vec3,
+    spherePoint?: Vec3,
+    qmax?: number
+}
+
+namespace MembraneCandidate {
+    export function initial(c1: Vec3, c2: Vec3, stats: HphobHphil): MembraneCandidate {
+        return {
+            planePoint1: c1,
+            planePoint2: c2,
+            stats: stats
+        };
+    }
+
+    export function scored(spherePoint: Vec3, c1: Vec3, c2: Vec3, stats: HphobHphil, qmax: number, centroid: Vec3): MembraneCandidate {
+        const diam_vect = Vec3();
+        Vec3.sub(diam_vect, centroid, spherePoint);
+        return {
+            planePoint1: c1,
+            planePoint2: c2,
+            stats: stats,
+            normalVector: diam_vect,
+            spherePoint: spherePoint,
+            qmax: qmax
+        };
+    }
+}
+
+function findMembrane(ctx: ANVILContext, spherePoints: Vec3[], initialStats: HphobHphil, label_comp_id: StructureElement.Property<string>): MembraneCandidate {
+    const { centroid, stepSize, minThickness, maxThickness } = ctx;
+    // best performing membrane
+    let membrane: MembraneCandidate;
+    // score of the best performing membrane
+    let qmax = 0;
+
+    // construct slices of thickness 1.0 along the axis connecting the centroid and the spherePoint
+    const diam = Vec3();
+    for (let i = 0, il = spherePoints.length; i < il; i++) {
+        const spherePoint = spherePoints[i];
+        Vec3.sub(diam, centroid, spherePoint);
+        Vec3.scale(diam, diam, 2);
+        const diamNorm = Vec3.magnitude(diam);
+        const qvartemp = [];
+
+        for (let i = 0, il = diamNorm - stepSize; i < il; i += stepSize) {
+            const c1 = Vec3();
+            const c2 = Vec3();
+            Vec3.scaleAndAdd(c1, spherePoint, diam, i / diamNorm);
+            Vec3.scaleAndAdd(c2, spherePoint, diam, (i + stepSize) / diamNorm);
+
+            // evaluate how well this membrane slice embeddeds the peculiar residues
+            const stats = HphobHphil.filtered(ctx, label_comp_id, (testPoint: Vec3) => isInMembranePlane(testPoint, diam, c1, c2));
+            qvartemp.push(MembraneCandidate.initial(c1, c2, stats));
+        }
+
+        let jmax = (minThickness / stepSize) - 1;
+
+        for (let width = 0, widthl = maxThickness; width < widthl;) {
+            const imax = qvartemp.length - 1 - jmax;
+
+            for (let i = 0, il = imax; i < il; i++) {
+                const c1 = qvartemp[i].planePoint1;
+                const c2 = qvartemp[i + jmax].planePoint2;
+
+                let hphob = 0;
+                let hphil = 0;
+                let total = 0;
+                for (let j = 0; j < jmax; j++) {
+                    const ij = qvartemp[i + j];
+                    if (j === 0 || j === jmax - 1) {
+                        hphob += 0.5 * ij.stats.hphob;
+                        hphil += 0.5 * ij.stats.hphil;
+                    } else {
+                        hphob += ij.stats.hphob;
+                        hphil += ij.stats.hphil;
+                    }
+                    total += ij.stats.total;
+                }
+
+                const stats = HphobHphil.of(hphob, hphil, total);
+
+                if (hphob !== 0) {
+                    const qvaltest = qValue(stats, initialStats);
+                    if (qvaltest > qmax) {
+                        qmax = qvaltest;
+                        membrane = MembraneCandidate.scored(spherePoint, c1, c2, HphobHphil.of(hphob, hphil, total), qmax, centroid);
+                    }
+                }
+            }
+            jmax++;
+            width = (jmax + 1) * stepSize;
+        }
+    }
+
+    return membrane!;
+}
+
+function qValue(currentStats: HphobHphil, initialStats: HphobHphil): number {
+    if(initialStats.hphob < 1) {
+        initialStats.hphob = 0.1;
+    }
+
+    if(initialStats.hphil < 1) {
+        initialStats.hphil += 1;
+    }
+
+    const part_tot = currentStats.hphob + currentStats.hphil;
+    return (currentStats.hphob * (initialStats.hphil - currentStats.hphil) - currentStats.hphil * (initialStats.hphob - currentStats.hphob)) /
+            Math.sqrt(part_tot * initialStats.hphob * initialStats.hphil * (initialStats.hphob + initialStats.hphil - part_tot));
+}
+
+export function isInMembranePlane(testPoint: Vec3, normalVector: Vec3, planePoint1: Vec3, planePoint2: Vec3): boolean {
+    const d1 = -Vec3.dot(normalVector, planePoint1);
+    const d2 = -Vec3.dot(normalVector, planePoint2);
+    const d = -Vec3.dot(normalVector, testPoint);
+    return d > Math.min(d1, d2) && d < Math.max(d1, d2);
+}
+
+// generates a defined number of points on a sphere with radius = extent around the specified centroid
+function generateSpherePoints(ctx: ANVILContext, numberOfSpherePoints: number): Vec3[] {
+    const { centroid, extent } = ctx;
+    const points = [];
+    let oldPhi = 0, h, theta, phi;
+    for(let k = 1, kl = numberOfSpherePoints + 1; k < kl; k++) {
+        h = -1 + 2 * (k - 1) / (numberOfSpherePoints - 1);
+        theta = Math.acos(h);
+        phi = (k === 1 || k === numberOfSpherePoints) ? 0 : (oldPhi + 3.6 / Math.sqrt(numberOfSpherePoints * (1 - h * h))) % (2 * Math.PI);
+
+        const point = Vec3.create(
+            extent * Math.sin(phi) * Math.sin(theta) + centroid[0],
+            extent * Math.cos(theta) + centroid[1],
+            extent * Math.cos(phi) * Math.sin(theta) + centroid[2]
+        );
+        points[k - 1] = point;
+        oldPhi = phi;
+    }
+
+    return points;
+}
+
+// generates sphere points close to that of the initial membrane
+function findProximateAxes(ctx: ANVILContext, membrane: MembraneCandidate): Vec3[] {
+    const { numberOfSpherePoints, extent } = ctx;
+    const points = generateSpherePoints(ctx, 30000);
+    let j = 4;
+    let sphere_pts2: Vec3[] = [];
+    while (sphere_pts2.length < numberOfSpherePoints) {
+        const d = 2 * extent / numberOfSpherePoints + j;
+        const dsq = d * d;
+        sphere_pts2 = [];
+        for (let i = 0, il = points.length; i < il; i++) {
+            if (Vec3.squaredDistance(points[i], membrane.spherePoint!) < dsq) {
+                sphere_pts2.push(points[i]);
+            }
+        }
+        j += 0.2;
+    }
+    return sphere_pts2;
+}
+
+interface HphobHphil {
+    hphob: number,
+    hphil: number,
+    total: number
+}
+
+namespace HphobHphil {
+    export function of(hphob: number, hphil: number, total?: number) {
+        return {
+            hphob: hphob,
+            hphil: hphil,
+            total: !!total ? total : hphob + hphil
+        };
+    }
+
+    const testPoint = Vec3();
+    export function filtered(ctx: ANVILContext, label_comp_id: StructureElement.Property<string>, filter?: (test: Vec3) => boolean): HphobHphil {
+        const { offsets, exposed, structure } = ctx;
+        const { x, y, z } = StructureProperties.atom;
+        let hphob = 0;
+        let hphil = 0;
+        for (let k = 0, kl = offsets.length; k < kl; k++) {
+            // ignore buried residues
+            if (!exposed[k]) {
+                continue;
+            }
+
+            setLocation(l, structure, offsets[k]);
+            Vec3.set(testPoint, x(l), y(l), z(l));
+
+            // testPoints have to be in putative membrane layer
+            if (filter && !filter(testPoint)) {
+                continue;
+            }
+
+            if (isHydrophobic(label_comp_id(l))) {
+                hphob++;
+            } else {
+                hphil++;
+            }
+        }
+        return of(hphob, hphil);
+    }
+}
+
+// ANVIL-specific (not general) definition of membrane-favoring amino acids
+const HYDROPHOBIC_AMINO_ACIDS = new Set(['ALA', 'CYS', 'GLY', 'HIS', 'ILE', 'LEU', 'MET', 'PHE', 'SER', 'THR', 'VAL']);
+export function isHydrophobic(label_comp_id: string): boolean {
+    return HYDROPHOBIC_AMINO_ACIDS.has(label_comp_id);
+}
+
+function setLocation(l: StructureElement.Location, structure: Structure, serialIndex: number) {
+    l.structure = structure;
+    l.unit = structure.units[structure.serialMapping.unitIndices[serialIndex]];
+    l.element = structure.serialMapping.elementIndices[serialIndex];
+    return l;
+}

src/extensions/membrane-orientation/anvil/common.ts

@@ -0,0 +1,23 @@
+/**
+ * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Sebastian Bittrich <sebastian.bittrich@rcsb.org>
+ */
+
+import { Structure } from '../../../mol-model/structure';
+import { Vec3 } from '../../../mol-math/linear-algebra';
+
+export interface ANVILContext {
+    structure: Structure,
+
+    numberOfSpherePoints: number,
+    stepSize: number,
+    minThickness: number,
+    maxThickness: number,
+    asaCutoff: number,
+
+    offsets: ArrayLike<number>,
+    exposed: ArrayLike<boolean>,
+    centroid: Vec3,
+    extent: number
+};

src/extensions/membrane-orientation/behavior.ts

@@ -0,0 +1,161 @@
+/**
+ * Copyright (c) 2018-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ * @author Sebastian Bittrich <sebastian.bittrich@rcsb.org>
+ */
+
+import { ParamDefinition as PD } from '../../mol-util/param-definition';
+import { StructureRepresentationPresetProvider, PresetStructureRepresentations } from '../../mol-plugin-state/builder/structure/representation-preset';
+import { MembraneOrientationProvider, isTransmembrane } from './membrane-orientation';
+import { StateObjectRef, StateAction, StateTransformer, StateTransform } from '../../mol-state';
+import { Task } from '../../mol-task';
+import { PluginBehavior } from '../../mol-plugin/behavior';
+import { MembraneOrientationRepresentationProvider, MembraneOrientationParams, MembraneOrientationRepresentation } from './representation';
+import { HydrophobicityColorThemeProvider } from '../../mol-theme/color/hydrophobicity';
+import { PluginStateObject, PluginStateTransform } from '../../mol-plugin-state/objects';
+import { PluginContext } from '../../mol-plugin/context';
+import { DefaultQueryRuntimeTable } from '../../mol-script/runtime/query/compiler';
+
+export const MembraneOrientationData = PluginBehavior.create<{ autoAttach: boolean }>({
+    name: 'membrane-orientation-prop',
+    category: 'custom-props',
+    display: {
+        name: 'Membrane Orientation',
+        description: 'Initialize orientation of membrane layers. Data calculated with ANVIL algorithm.' // TODO add ' or obtained via RCSB PDB'
+    },
+    ctor: class extends PluginBehavior.Handler<{ autoAttach: boolean }> {
+        private provider = MembraneOrientationProvider
+
+        register(): void {
+            DefaultQueryRuntimeTable.addCustomProp(this.provider.descriptor);
+
+            this.ctx.state.data.actions.add(InitMembraneOrientation3D);
+            this.ctx.customStructureProperties.register(this.provider, this.params.autoAttach);
+
+            this.ctx.representation.structure.registry.add(MembraneOrientationRepresentationProvider);
+            this.ctx.query.structure.registry.add(isTransmembrane);
+
+            this.ctx.builders.structure.representation.registerPreset(MembraneOrientationPreset);
+        }
+
+        update(p: { autoAttach: boolean }) {
+            let updated = this.params.autoAttach !== p.autoAttach;
+            this.params.autoAttach = p.autoAttach;
+            this.ctx.customStructureProperties.setDefaultAutoAttach(this.provider.descriptor.name, this.params.autoAttach);
+            return updated;
+        }
+
+        unregister() {
+            DefaultQueryRuntimeTable.removeCustomProp(this.provider.descriptor);
+
+            this.ctx.state.data.actions.remove(InitMembraneOrientation3D);
+            this.ctx.customStructureProperties.unregister(this.provider.descriptor.name);
+
+            this.ctx.representation.structure.registry.remove(MembraneOrientationRepresentationProvider);
+            this.ctx.query.structure.registry.remove(isTransmembrane);
+
+            this.ctx.builders.structure.representation.unregisterPreset(MembraneOrientationPreset);
+        }
+    },
+    params: () => ({
+        autoAttach: PD.Boolean(false)
+    })
+});
+
+export const InitMembraneOrientation3D = StateAction.build({
+    display: {
+        name: 'Membrane Orientation',
+        description: 'Initialize Membrane Orientation planes and rims. Data calculated with ANVIL algorithm.'
+    },
+    from: PluginStateObject.Molecule.Structure,
+    isApplicable: (a) => MembraneOrientationProvider.isApplicable(a.data)
+})(({ a, ref, state }, plugin: PluginContext) => Task.create('Init Membrane Orientation', async ctx => {
+    try {
+        const propCtx = { runtime: ctx, assetManager: plugin.managers.asset };
+        await MembraneOrientationProvider.attach(propCtx, a.data);
+    } catch(e) {
+        plugin.log.error(`Membrane Orientation: ${e}`);
+        return;
+    }
+    const tree = state.build().to(ref)
+        .applyOrUpdateTagged('membrane-orientation-3d', MembraneOrientation3D);
+    await state.updateTree(tree).runInContext(ctx);
+}));
+
+export { MembraneOrientation3D };
+
+type MembraneOrientation3D = typeof MembraneOrientation3D
+const MembraneOrientation3D = PluginStateTransform.BuiltIn({
+    name: 'membrane-orientation-3d',
+    display: {
+        name: 'Membrane Orientation',
+        description: 'Membrane Orientation planes and rims. Data calculated with ANVIL algorithm.'
+    },
+    from: PluginStateObject.Molecule.Structure,
+    to: PluginStateObject.Shape.Representation3D,
+    params: (a) => {
+        return {
+            ...MembraneOrientationParams,
+        };
+    }
+})({
+    canAutoUpdate({ oldParams, newParams }) {
+        return true;
+    },
+    apply({ a, params }, plugin: PluginContext) {
+        return Task.create('Membrane Orientation', async ctx => {
+            await MembraneOrientationProvider.attach({ runtime: ctx, assetManager: plugin.managers.asset }, a.data);
+            const repr = MembraneOrientationRepresentation({ webgl: plugin.canvas3d?.webgl, ...plugin.representation.structure.themes }, () => MembraneOrientationParams);
+            await repr.createOrUpdate(params, a.data).runInContext(ctx);
+            return new PluginStateObject.Shape.Representation3D({ repr, source: a }, { label: 'Membrane Orientation' });
+        });
+    },
+    update({ a, b, newParams }, plugin: PluginContext) {
+        return Task.create('Membrane Orientation', async ctx => {
+            await MembraneOrientationProvider.attach({ runtime: ctx, assetManager: plugin.managers.asset }, a.data);
+            const props = { ...b.data.repr.props, ...newParams };
+            await b.data.repr.createOrUpdate(props, a.data).runInContext(ctx);
+            return StateTransformer.UpdateResult.Updated;
+        });
+    },
+    isApplicable(a) {
+        return MembraneOrientationProvider.isApplicable(a.data);
+    }
+});
+
+export const MembraneOrientationPreset = StructureRepresentationPresetProvider({
+    id: 'preset-membrane-orientation',
+    display: {
+        name: 'Membrane Orientation', group: 'Annotation',
+        description: 'Shows orientation of membrane layers. Data calculated with ANVIL algorithm.' // TODO add ' or obtained via RCSB PDB'
+    },
+    isApplicable(a) {
+        return MembraneOrientationProvider.isApplicable(a.data);
+    },
+    params: () => StructureRepresentationPresetProvider.CommonParams,
+    async apply(ref, params, plugin) {
+        const structureCell = StateObjectRef.resolveAndCheck(plugin.state.data, ref);
+        const structure  = structureCell?.obj?.data;
+        if (!structureCell || !structure) return {};
+
+        if (!MembraneOrientationProvider.get(structure).value) {
+            await plugin.runTask(Task.create('Membrane Orientation', async runtime => {
+                await MembraneOrientationProvider.attach({ runtime, assetManager: plugin.managers.asset }, structure);
+            }));
+        }
+
+        const membraneOrientation = await tryCreateMembraneOrientation(plugin, structureCell);
+        const colorTheme = HydrophobicityColorThemeProvider.name as any;
+        const preset = await PresetStructureRepresentations.auto.apply(ref, { ...params, theme: { globalName: colorTheme, focus: { name: colorTheme } } }, plugin);
+
+        return { components: preset.components, representations: { ...preset.representations, membraneOrientation } };
+    }
+});
+
+export function tryCreateMembraneOrientation(plugin: PluginContext, structure: StateObjectRef<PluginStateObject.Molecule.Structure>, params?: StateTransformer.Params<MembraneOrientation3D>, initialState?: Partial<StateTransform.State>) {
+    const state = plugin.state.data;
+    const membraneOrientation = state.build().to(structure)
+        .applyOrUpdateTagged('membrane-orientation-3d', MembraneOrientation3D, params, { state: initialState });
+    return membraneOrientation.commit({ revertOnError: true });
+}

src/extensions/membrane-orientation/membrane-orientation.ts

@@ -0,0 +1,95 @@
+/**
+ * Copyright (c) 2019-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Sebastian Bittrich <sebastian.bittrich@rcsb.org>
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { ParamDefinition as PD } from '../../mol-util/param-definition';
+import { Structure, StructureProperties, Unit } from '../../mol-model/structure';
+import { CustomPropertyDescriptor } from '../../mol-model/custom-property';
+import { ANVILParams, ANVILProps, computeANVIL, isInMembranePlane } from './ANVIL';
+import { CustomStructureProperty } from '../../mol-model-props/common/custom-structure-property';
+import { CustomProperty } from '../../mol-model-props/common/custom-property';
+import { AccessibleSurfaceAreaProvider } from '../../mol-model-props/computed/accessible-surface-area';
+import { Vec3 } from '../../mol-math/linear-algebra';
+import { QuerySymbolRuntime } from '../../mol-script/runtime/query/base';
+import { CustomPropSymbol } from '../../mol-script/language/symbol';
+import Type from '../../mol-script/language/type';
+import { StructureSelectionQuery, StructureSelectionCategory } from '../../mol-plugin-state/helpers/structure-selection-query';
+import { MolScriptBuilder as MS } from '../../mol-script/language/builder';
+
+export const MembraneOrientationParams = {
+    ...ANVILParams
+};
+export type MembraneOrientationParams = typeof MembraneOrientationParams
+export type MembraneOrientationProps = PD.Values<MembraneOrientationParams>
+
+interface MembraneOrientation {
+    // point in membrane boundary
+    readonly planePoint1: Vec3,
+    // point in opposite side of membrane boundary
+    readonly planePoint2: Vec3,
+    // normal vector of membrane layer
+    readonly normalVector: Vec3,
+    // the radius of the membrane layer
+    readonly radius: number,
+    readonly centroid: Vec3
+}
+
+const pos = Vec3();
+export const MembraneOrientationSymbols = {
+    isTransmembrane: QuerySymbolRuntime.Dynamic(CustomPropSymbol('computed', 'membrane-orientation.is-transmembrane', Type.Bool),
+        ctx => {
+            const { unit, structure } = ctx.element;
+            const { x, y, z } = StructureProperties.atom;
+            if (!Unit.isAtomic(unit)) return 0;
+            const membraneOrientation = MembraneOrientationProvider.get(structure).value;
+            if (!membraneOrientation) return 0;
+            Vec3.set(pos, x(ctx.element), y(ctx.element), z(ctx.element));
+            const { normalVector, planePoint1, planePoint2 } = membraneOrientation!;
+            return isInMembranePlane(pos, normalVector, planePoint1, planePoint2);
+        })
+}
+
+export const isTransmembrane = StructureSelectionQuery('Residues Embedded in Membrane', MS.struct.modifier.union([
+    MS.struct.modifier.wholeResidues([
+        MS.struct.modifier.union([
+            MS.struct.generator.atomGroups({
+                'chain-test': MS.core.rel.eq([MS.ammp('objectPrimitive'), 'atomistic']),
+                'atom-test': MembraneOrientationSymbols.isTransmembrane.symbol(),
+            })
+        ])
+    ])
+]), {
+    description: 'Select residues that are embedded between the membrane layers.',
+    category: StructureSelectionCategory.Residue,
+    ensureCustomProperties: (ctx, structure) => {
+        return MembraneOrientationProvider.attach(ctx, structure);
+    }
+});
+
+export { MembraneOrientation };
+
+export const MembraneOrientationProvider: CustomStructureProperty.Provider<MembraneOrientationParams, MembraneOrientation> = CustomStructureProperty.createProvider({
+    label: 'Membrane Orientation',
+    descriptor: CustomPropertyDescriptor({
+        name: 'molstar_computed_membrane_orientation',
+        symbols: MembraneOrientationSymbols,
+        // TODO `cifExport`
+    }),
+    type: 'root',
+    defaultParams: MembraneOrientationParams,
+    getParams: (data: Structure) => MembraneOrientationParams,
+    isApplicable: (data: Structure) => true,
+    obtain: async (ctx: CustomProperty.Context, data: Structure, props: Partial<MembraneOrientationProps>) => {
+        const p = { ...PD.getDefaultValues(MembraneOrientationParams), ...props };
+        return { value: await computeAnvil(ctx, data, p) };
+    }
+});
+
+async function computeAnvil(ctx: CustomProperty.Context, data: Structure, props: Partial<ANVILProps>): Promise<MembraneOrientation> {
+    await AccessibleSurfaceAreaProvider.attach(ctx, data);
+    const p = { ...PD.getDefaultValues(ANVILParams), ...props };
+    return await computeANVIL(data, p).runInContext(ctx.runtime);
+}

src/extensions/membrane-orientation/representation.ts

@@ -0,0 +1,172 @@
+/**
+ * Copyright (c) 2018-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ * @author Sebastian Bittrich <sebastian.bittrich@rcsb.org>
+ */
+
+import { ParamDefinition as PD } from '../../mol-util/param-definition';
+import { Vec3, Mat4 } from '../../mol-math/linear-algebra';
+import { Representation, RepresentationContext, RepresentationParamsGetter } from '../../mol-repr/representation';
+import { Structure } from '../../mol-model/structure';
+import { Spheres } from '../../mol-geo/geometry/spheres/spheres';
+import { SpheresBuilder } from '../../mol-geo/geometry/spheres/spheres-builder';
+import { StructureRepresentationProvider, StructureRepresentation, StructureRepresentationStateBuilder } from '../../mol-repr/structure/representation';
+import { MembraneOrientation } from './membrane-orientation';
+import { ThemeRegistryContext } from '../../mol-theme/theme';
+import { ShapeRepresentation } from '../../mol-repr/shape/representation';
+import { Shape } from '../../mol-model/shape';
+import { ColorNames } from '../../mol-util/color/names';
+import { RuntimeContext } from '../../mol-task';
+import { Lines } from '../../mol-geo/geometry/lines/lines';
+import { Mesh } from '../../mol-geo/geometry/mesh/mesh';
+import { LinesBuilder } from '../../mol-geo/geometry/lines/lines-builder';
+import { Circle } from '../../mol-geo/primitive/circle';
+import { transformPrimitive } from '../../mol-geo/primitive/primitive';
+import { MeshBuilder } from '../../mol-geo/geometry/mesh/mesh-builder';
+import { MembraneOrientationProvider } from './membrane-orientation';
+
+const SharedParams = {
+    color: PD.Color(ColorNames.lightgrey),
+    radiusFactor: PD.Numeric(0.8333, { min: 0.1, max: 3.0, step: 0.01 }, { description: 'Scale the radius of the membrane layer' })
+};
+
+const BilayerSpheresParams = {
+    ...SharedParams,
+    ...Spheres.Params,
+    sphereSize: PD.Numeric(1, { min: 0.1, max: 10, step: 0.1 }, { description: 'Size of spheres that represent membrane planes' }),
+    density: PD.Numeric(1, { min: 0.25, max: 10, step: 0.25 }, { description: 'Distance between spheres'})
+};
+export type BilayerSpheresParams = typeof BilayerSpheresParams
+export type BilayerSpheresProps = PD.Values<BilayerSpheresParams>
+
+const BilayerPlanesParams = {
+    ...SharedParams,
+    ...Mesh.Params,
+    sectorOpacity: PD.Numeric(0.5, { min: 0, max: 1, step: 0.01 })
+};
+export type BilayerPlanesParams = typeof BilayerPlanesParams
+export type BilayerPlanesProps = PD.Values<BilayerPlanesParams>
+
+const BilayerRimsParams = {
+    ...SharedParams,
+    ...Lines.Params,
+    lineSizeAttenuation: PD.Boolean(true),
+    linesSize: PD.Numeric(0.5, { min: 0.01, max: 50, step: 0.01 }),
+    dashedLines: PD.Boolean(true)
+};
+export type BilayerRimsParams = typeof BilayerRimsParams
+export type BilayerRimsProps = PD.Values<BilayerRimsParams>
+
+const MembraneOrientationVisuals = {
+    'bilayer-spheres': (ctx: RepresentationContext, getParams: RepresentationParamsGetter<MembraneOrientation, BilayerSpheresParams>) => ShapeRepresentation(getBilayerSpheres, Spheres.Utils, { modifyState: s => ({ ...s, pickable: false }) }),
+    'bilayer-planes': (ctx: RepresentationContext, getParams: RepresentationParamsGetter<MembraneOrientation, BilayerPlanesParams>) => ShapeRepresentation(getBilayerPlanes, Mesh.Utils, { modifyProps: p => ({ ...p, alpha: p.sectorOpacity }), modifyState: s => ({ ...s, pickable: false }) }),
+    'bilayer-rims': (ctx: RepresentationContext, getParams: RepresentationParamsGetter<MembraneOrientation, BilayerRimsParams>) => ShapeRepresentation(getBilayerRims, Lines.Utils, { modifyState: s => ({ ...s, pickable: false }) })
+};
+
+export const MembraneOrientationParams = {
+    ...BilayerSpheresParams,
+    ...BilayerPlanesParams,
+    ...BilayerRimsParams,
+    visuals: PD.MultiSelect(['bilayer-planes', 'bilayer-rims'], PD.objectToOptions(MembraneOrientationVisuals)),
+};
+export type MembraneOrientationParams = typeof MembraneOrientationParams
+export type MembraneOrientationProps = PD.Values<MembraneOrientationParams>
+
+export function getMembraneOrientationParams(ctx: ThemeRegistryContext, structure: Structure) {
+    return PD.clone(MembraneOrientationParams);
+}
+
+export type MembraneOrientationRepresentation = StructureRepresentation<MembraneOrientationParams>
+export function MembraneOrientationRepresentation(ctx: RepresentationContext, getParams: RepresentationParamsGetter<Structure, MembraneOrientationParams>): MembraneOrientationRepresentation {
+    return Representation.createMulti('Membrane Orientation', ctx, getParams, StructureRepresentationStateBuilder, MembraneOrientationVisuals as unknown as Representation.Def<Structure, MembraneOrientationParams>);
+}
+
+export const MembraneOrientationRepresentationProvider = StructureRepresentationProvider({
+    name: 'membrane-orientation',
+    label: 'Membrane Orientation',
+    description: 'Displays a grid of points representing membrane layers.',
+    factory: MembraneOrientationRepresentation,
+    getParams: getMembraneOrientationParams,
+    defaultValues: PD.getDefaultValues(MembraneOrientationParams),
+    defaultColorTheme: { name: 'uniform' },
+    defaultSizeTheme: { name: 'uniform' },
+    isApplicable: (structure: Structure) => structure.elementCount > 0
+});
+
+function getBilayerRims(ctx: RuntimeContext, data: Structure, props: BilayerRimsProps): Shape<Lines> {
+    const { planePoint1: p1, planePoint2: p2, centroid, normalVector: normal, radius } = MembraneOrientationProvider.get(data).value!;
+    const scaledRadius = props.radiusFactor * radius;
+    const builder = LinesBuilder.create(128, 64);
+    getLayerCircle(builder, p1, centroid, normal, scaledRadius, props);
+    getLayerCircle(builder, p2, centroid, normal, scaledRadius, props);
+    return Shape.create(name, data, builder.getLines(), () => props.color, () => props.linesSize, () => '');
+}
+
+function getLayerCircle(builder: LinesBuilder, p: Vec3, centroid: Vec3, normal: Vec3, radius: number, props: BilayerRimsProps) {
+    const circle = getCircle(p, centroid, normal, radius);
+    const { indices, vertices } = circle;
+    for (let j = 0, jl = indices.length; j < jl; j += 3) {
+        if (props.dashedLines && j % 2 === 1) continue; // draw every other segment to get dashes
+        const start = indices[j] * 3;
+        const end = indices[j + 1] * 3;
+        const startX = vertices[start];
+        const startY = vertices[start + 1];
+        const startZ = vertices[start + 2];
+        const endX = vertices[end];
+        const endY = vertices[end + 1];
+        const endZ = vertices[end + 2];
+        builder.add(startX, startY, startZ, endX, endY, endZ, 0);
+    }
+}
+
+const tmpMat = Mat4();
+function getCircle(p: Vec3, centroid: Vec3, normal: Vec3, radius: number) {
+    Mat4.targetTo(tmpMat, p, centroid, normal);
+    Mat4.setTranslation(tmpMat, p);
+    Mat4.mul(tmpMat, tmpMat, Mat4.rotX90);
+
+    const circle = Circle({ radius });
+    return transformPrimitive(circle, tmpMat);
+}
+
+function getBilayerPlanes(ctx: RuntimeContext, data: Structure, props: BilayerPlanesProps, shape?: Shape<Mesh>): Shape<Mesh> {
+    const { planePoint1: p1, planePoint2: p2, centroid, normalVector: normal, radius } = MembraneOrientationProvider.get(data).value!;
+    const state = MeshBuilder.createState(128, 64, shape && shape.geometry);
+    const scaledRadius = props.radiusFactor * radius;
+    getLayerPlane(state, p1, centroid, normal, scaledRadius);
+    getLayerPlane(state, p2, centroid, normal, scaledRadius);
+    return Shape.create(name, data, MeshBuilder.getMesh(state), () => props.color, () => 1, () => '');
+}
+
+function getLayerPlane(state: MeshBuilder.State, p: Vec3, centroid: Vec3, normal: Vec3, radius: number) {
+    const circle = getCircle(p, centroid, normal, radius);
+    state.currentGroup = 0;
+    MeshBuilder.addPrimitive(state, Mat4.id, circle);
+    MeshBuilder.addPrimitiveFlipped(state, Mat4.id, circle);
+}
+
+function getBilayerSpheres(ctx: RuntimeContext, data: Structure, props: BilayerSpheresProps): Shape<Spheres> {
+    const { density } = props;
+    const { radius, planePoint1, planePoint2, normalVector } = MembraneOrientationProvider.get(data).value!;
+    const scaledRadius = (props.radiusFactor * radius) * (props.radiusFactor * radius);
+
+    const spheresBuilder = SpheresBuilder.create();
+    getLayerSpheres(spheresBuilder, planePoint1, normalVector, density, scaledRadius);
+    getLayerSpheres(spheresBuilder, planePoint2, normalVector, density, scaledRadius);
+    return Shape.create(name, data, spheresBuilder.getSpheres(), () => props.color, () => props.sphereSize, () => '');
+}
+
+function getLayerSpheres(spheresBuilder: SpheresBuilder, point: Vec3, normalVector: Vec3, density: number, sqRadius: number) {
+    Vec3.normalize(normalVector, normalVector);
+    const d = -Vec3.dot(normalVector, point);
+    const rep = Vec3();
+    for (let i = -1000, il = 1000; i < il; i += density) {
+        for (let j = -1000, jl = 1000; j < jl; j += density) {
+            Vec3.set(rep, i, j, normalVector[2] === 0 ? 0 : -(d + i * normalVector[0] + j * normalVector[1]) / normalVector[2]);
+            if (Vec3.squaredDistance(rep, point) < sqRadius) {
+                spheresBuilder.add(rep[0], rep[1], rep[2], 0);
+            }
+        }
+    }
+}

src/mol-theme/color/hydrophobicity.ts

@@ -31,7 +31,7 @@ export function hydrophobicity(compId: string, scaleIndex: number): number {
 }
 
 function getAtomicCompId(unit: Unit.Atomic, element: ElementIndex) {
-    return unit.model.atomicHierarchy.atoms.label_comp_id.value(unit.residueIndex[element]);
+    return unit.model.atomicHierarchy.atoms.label_comp_id.value(element);
 }
 
 function getCoarseCompId(unit: Unit.Spheres | Unit.Gaussians, element: ElementIndex) {

src/mol-theme/color/residue-name.ts

@@ -74,7 +74,7 @@ export function getResidueNameColorThemeParams(ctx: ThemeDataContext) {
 }
 
 function getAtomicCompId(unit: Unit.Atomic, element: ElementIndex) {
-    return unit.model.atomicHierarchy.atoms.label_comp_id.value(unit.residueIndex[element]);
+    return unit.model.atomicHierarchy.atoms.label_comp_id.value(element);
 }
 
 function getCoarseCompId(unit: Unit.Spheres | Unit.Gaussians, element: ElementIndex) {

src/tests/browser/render-structure.ts

@@ -26,6 +26,8 @@ import { InteractionsProvider } from '../../mol-model-props/computed/interaction
 import { SecondaryStructureProvider } from '../../mol-model-props/computed/secondary-structure';
 import { SyncRuntimeContext } from '../../mol-task/execution/synchronous';
 import { AssetManager } from '../../mol-util/assets';
+import { MembraneOrientationProvider } from '../../extensions/membrane-orientation/membrane-orientation';
+import { MembraneOrientationRepresentationProvider } from '../../extensions/membrane-orientation/representation';
 
 const parent = document.getElementById('app')!;
 parent.style.width = '100%';
@@ -116,6 +118,10 @@ function getGaussianSurfaceRepr() {
     return GaussianSurfaceRepresentationProvider.factory(reprCtx, GaussianSurfaceRepresentationProvider.getParams);
 }
 
+function getMembraneOrientationRepr() {
+    return MembraneOrientationRepresentationProvider.factory(reprCtx, MembraneOrientationRepresentationProvider.getParams);
+}
+
 async function init() {
     const ctx = { runtime: SyncRuntimeContext, assetManager: new AssetManager() };
 
@@ -127,6 +133,10 @@ async function init() {
     await SecondaryStructureProvider.attach(ctx, structure);
     console.timeEnd('compute SecondaryStructure');
 
+    console.time('compute Membrane Orientation');
+    await MembraneOrientationProvider.attach(ctx, structure);
+    console.timeEnd('compute Membrane Orientation');
+
     console.time('compute Interactions');
     await InteractionsProvider.attach(ctx, structure);
     console.timeEnd('compute Interactions');
@@ -138,6 +148,7 @@ async function init() {
         ballAndStick: true,
         molecularSurface: false,
         gaussianSurface: false,
+        membrane: true
     };
 
     const cartoonRepr = getCartoonRepr();
@@ -145,6 +156,7 @@ async function init() {
     const ballAndStickRepr = getBallAndStickRepr();
     const molecularSurfaceRepr = getMolecularSurfaceRepr();
     const gaussianSurfaceRepr = getGaussianSurfaceRepr();
+    const membraneOrientationRepr = getMembraneOrientationRepr();
 
     if (show.cartoon) {
         cartoonRepr.setTheme({
@@ -190,11 +202,16 @@ async function init() {
         console.timeEnd('gaussian surface');
     }
 
+    if (show.membrane) {
+        await membraneOrientationRepr.createOrUpdate({ ...MembraneOrientationRepresentationProvider.defaultValues, quality: 'auto' }, structure).run();
+    }
+
     if (show.cartoon) canvas3d.add(cartoonRepr);
     if (show.interaction) canvas3d.add(interactionRepr);
     if (show.ballAndStick) canvas3d.add(ballAndStickRepr);
     if (show.molecularSurface) canvas3d.add(molecularSurfaceRepr);
     if (show.gaussianSurface) canvas3d.add(gaussianSurfaceRepr);
+    if (show.membrane) canvas3d.add(membraneOrientationRepr);
     canvas3d.requestCameraReset();
     // canvas3d.setProps({ trackball: { ...canvas3d.props.trackball, spin: true } })
 }