Query wholeResidues, includeSurroundings

src/apps/structure-info/model.ts

@@ -119,9 +119,9 @@ export function printUnits(structure: Structure) {
         const size = OrderedSet.size(elements);
 
         if (Unit.isAtomic(l.unit)) {
-            console.log(`Atomic unit ${unit.id}: ${size} elements`);
+            console.log(`Atomic unit ${unit.id} ${unit.conformation.operator.name}: ${size} elements`);
         } else if (Unit.isCoarse(l.unit)) {
-            console.log(`Coarse unit ${unit.id} (${Unit.isSpheres(l.unit) ? 'spheres' : 'gaussians'}): ${size} elements.`);
+            console.log(`Coarse unit ${unit.id} ${unit.conformation.operator.name} (${Unit.isSpheres(l.unit) ? 'spheres' : 'gaussians'}): ${size} elements.`);
 
             const props = Queries.props.coarse;
             const seq = l.unit.model.sequence;

src/mol-data/generic/unique-array.ts

@@ -15,9 +15,10 @@ namespace UniqueArray {
     }
 
     export function add<K, T>({ keys, array }: UniqueArray<K, T>, key: K, value: T) {
-        if (keys.has(key)) return;
+        if (keys.has(key)) return false;
         keys.add(key);
         array[array.length] = value;
+        return true;
     }
 }
 

src/mol-math/geometry/spacegroup/construction.ts

@@ -87,8 +87,6 @@ namespace Spacegroup {
 
     export function getSymmetryOperator(spacegroup: Spacegroup, index: number, i: number, j: number, k: number): SymmetryOperator {
         const operator = updateOperatorMatrix(spacegroup, index, i, j, k, Mat4.zero());
-        console.log(Mat4.makeTable(operator));
-        console.log({ index, i, j, k });
         return SymmetryOperator.create(`${index + 1}_${5 + i}${5 + j}${5 + k}`, operator, Vec3.create(i, j, k));
     }
 

src/mol-model/structure/query.ts

@@ -7,11 +7,13 @@
 import Selection from './query/selection'
 import Query from './query/query'
 import * as generators from './query/generators'
+import * as modifiers from './query/modifiers'
 import props from './query/properties'
 import pred from './query/predicates'
 
 export const Queries = {
     generators,
+    modifiers,
     props,
     pred
 }

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

@@ -50,6 +50,7 @@ function atomGroupsLinear(atomTest: Element.Predicate): Query.Provider {
         const l = Element.Location();
         const builder = structure.subsetBuilder(true);
 
+        let progress = 0;
         for (const unit of units) {
             l.unit = unit;
             const elements = unit.elements;
@@ -61,7 +62,8 @@ function atomGroupsLinear(atomTest: Element.Predicate): Query.Provider {
             }
             builder.commitUnit();
 
-            if (ctx.shouldUpdate) await ctx.update({ message: 'Atom Groups', current: 0, max: units.length });
+            progress++;
+            if (ctx.shouldUpdate) await ctx.update({ message: 'Atom Groups', current: progress, max: units.length });
         }
 
         return Selection.Singletons(structure, builder.getStructure());
@@ -74,6 +76,7 @@ function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: A
         const l = Element.Location();
         const builder = structure.subsetBuilder(true);
 
+        let progress = 0;
         for (const unit of units) {
             if (unit.kind !== Unit.Kind.Atomic) continue;
 
@@ -107,7 +110,8 @@ function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: A
             }
             builder.commitUnit();
 
-            if (ctx.shouldUpdate) await ctx.update({ message: 'Atom Groups', current: 0, max: units.length });
+            progress++;
+            if (ctx.shouldUpdate) await ctx.update({ message: 'Atom Groups', current: progress, max: units.length });
         }
 
         return Selection.Singletons(structure, builder.getStructure());
@@ -120,6 +124,7 @@ function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, group
         const l = Element.Location();
         const builder = new LinearGroupingBuilder(structure);
 
+        let progress = 0;
         for (const unit of units) {
             if (unit.kind !== Unit.Kind.Atomic) continue;
 
@@ -149,7 +154,8 @@ function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, group
                 }
             }
 
-            if (ctx.shouldUpdate) await ctx.update({ message: 'Atom Groups', current: 0, max: units.length });
+            progress++;
+            if (ctx.shouldUpdate) await ctx.update({ message: 'Atom Groups', current: progress, max: units.length });
         }
 
         return builder.getSelection();

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

@@ -1,44 +0,0 @@
-// /**
-//  * Copyright (c) 2017 mol* contributors, licensed under MIT, See LICENSE file for more info.
-//  *
-//  * @author David Sehnal <david.sehnal@gmail.com>
-//  */
-
-// import Query from './query'
-// import Selection from './selection'
-// import P from './properties'
-// import { Element, Unit } from '../structure'
-// import { OrderedSet, Segmentation } from 'mol-data/int'
-// import { LinearGroupingBuilder } from './utils/builders';
-
-// export function wholeResidues(query: Query, isFlat: boolean): Query.Provider {
-//     return async (structure, ctx) => {
-//         const selection = query(structure).runAsChild(ctx);
-//         const { units } = structure;
-//         const l = Element.Location();
-//         const builder = structure.subsetBuilder(true);
-
-//         for (const unit of units) {
-//             l.unit = unit;
-//             const elements = unit.elements;
-
-//             builder.beginUnit(unit.id);
-//             for (let j = 0, _j = elements.length; j < _j; j++) {
-//                 l.element = elements[j];
-//                 if (atomTest(l)) builder.addElement(l.element);
-//             }
-//             builder.commitUnit();
-
-//             if (ctx.shouldUpdate) await ctx.update({ message: 'Atom Groups', current: 0, max: units.length });
-//         }
-
-//         return Selection.Singletons(structure, builder.getStructure());
-//     };
-// }
-
-// export interface IncludeSurroundingsParams {
-//     selection: Selection,
-//     radius: number,
-//     atomRadius?: number,
-//     wholeResidues?: boolean
-// }

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

@@ -0,0 +1,102 @@
+/**
+ * Copyright (c) 2017 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { Segmentation } from 'mol-data/int';
+import { RuntimeContext } from 'mol-task';
+import { Structure, Unit } from '../structure';
+import Query from './query';
+import Selection from './selection';
+import { UniqueStructuresBuilder } from './utils/builders';
+import { StructureUniqueSubsetBuilder } from '../structure/util/unique-subset-builder';
+
+function getWholeResidues(ctx: RuntimeContext, source: Structure, structure: Structure) {
+    const builder = source.subsetBuilder(true);
+    for (const unit of structure.units) {
+        if (unit.kind !== Unit.Kind.Atomic) {
+            // just copy non-atomic units.
+            builder.setUnit(unit.id, unit.elements);
+            continue;
+        }
+
+        const { residueSegments } = unit.model.atomicHierarchy;
+
+        const elements = unit.elements;
+        builder.beginUnit(unit.id);
+        const residuesIt = Segmentation.transientSegments(residueSegments, elements);
+        while (residuesIt.hasNext) {
+            const rI = residuesIt.move().index;
+            for (let j = residueSegments.segments[rI], _j = residueSegments.segments[rI + 1]; j < _j; j++) {
+                builder.addElement(j);
+            }
+        }
+        builder.commitUnit();
+    }
+    return builder.getStructure();
+}
+
+export function wholeResidues(query: Query.Provider, isFlat: boolean): Query.Provider {
+    return async (structure, ctx) => {
+        const inner = await query(structure, ctx);
+        if (Selection.isSingleton(inner)) {
+            return Selection.Singletons(structure, getWholeResidues(ctx, structure, inner.structure));
+        } else {
+            const builder = new UniqueStructuresBuilder(structure);
+            let progress = 0;
+            for (const s of inner.structures) {
+                builder.add(getWholeResidues(ctx, structure, s));
+                progress++;
+                if (ctx.shouldUpdate) await ctx.update({ message: 'Whole Residues', current: progress, max: inner.structures.length });
+            }
+            return builder.getSelection();
+        }
+    };
+}
+
+
+// export function groupBy()  ...
+
+export interface IncludeSurroundingsParams {
+    radius: number,
+    // atomRadius?: Element.Property<number>,
+    wholeResidues?: boolean
+}
+
+async function getIncludeSurroundings(ctx: RuntimeContext, source: Structure, structure: Structure, params: IncludeSurroundingsParams) {
+    const builder = new StructureUniqueSubsetBuilder(source);
+    const lookup = source.lookup3d;
+    const r = params.radius;
+
+    let progress = 0;
+
+    for (const unit of structure.units) {
+        const { x, y, z } = unit.conformation;
+        const elements = unit.elements;
+        for (let i = 0, _i = elements.length; i < _i; i++) {
+            const e = elements[i];
+            lookup.findIntoBuilder(x(e), y(e), z(e), r, builder);
+        }
+
+        progress++;
+        if (progress % 2500 === 0 && ctx.shouldUpdate) await ctx.update({ message: 'Include Surroudnings', isIndeterminate: true });
+    }
+    if (!!params.wholeResidues) return getWholeResidues(ctx, source, builder.getStructure())
+    return builder.getStructure();
+}
+
+export function includeSurroundings(query: Query.Provider, params: IncludeSurroundingsParams): Query.Provider {
+    return async (structure, ctx) => {
+        const inner = await query(structure, ctx);
+        if (Selection.isSingleton(inner)) {
+            return Selection.Singletons(structure, await getIncludeSurroundings(ctx, structure, inner.structure, params));
+        } else {
+            const builder = new UniqueStructuresBuilder(structure);
+            for (const s of inner.structures) {
+                builder.add(await getIncludeSurroundings(ctx, structure, s, params));
+            }
+            return builder.getSelection();
+        }
+    };
+}

src/mol-model/structure/query/utils/builders.ts

@@ -8,6 +8,7 @@ import { Element, Structure } from '../../structure';
 import Selection from '../selection';
 import { HashSet } from 'mol-data/generic';
 import { structureUnion } from './structure';
+import { StructureSubsetBuilder } from '../../structure/util/subset-builder';
 
 export class UniqueStructuresBuilder {
     private set = HashSet(Structure.hashCode, Structure.areEqual);
@@ -32,8 +33,8 @@ export class UniqueStructuresBuilder {
 }
 
 export class LinearGroupingBuilder {
-    private builders: Structure.SubsetBuilder[] = [];
-    private builderMap = new Map<string, Structure.SubsetBuilder>();
+    private builders: StructureSubsetBuilder[] = [];
+    private builderMap = new Map<string, StructureSubsetBuilder>();
 
     add(key: any, unit: number, element: number) {
         let b = this.builderMap.get(key);

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

@@ -6,6 +6,7 @@
 
 import { Structure, Unit } from '../../structure'
 import { SortedArray } from 'mol-data/int';
+import { StructureSubsetBuilder } from '../../structure/util/subset-builder';
 
 export function structureUnion(source: Structure, structures: Structure[]) {
     if (structures.length === 0) return Structure.Empty;
@@ -35,7 +36,7 @@ export function structureUnion(source: Structure, structures: Structure[]) {
     return builder.getStructure();
 }
 
-function buildUnion(this: Structure.SubsetBuilder, elements: SortedArray, id: number) {
+function buildUnion(this: StructureSubsetBuilder, elements: SortedArray, id: number) {
     this.setUnit(id, elements);
 }
 

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

@@ -8,12 +8,12 @@ import { IntMap, SortedArray, Iterator } from 'mol-data/int'
 import { UniqueArray } from 'mol-data/generic'
 import { SymmetryOperator } from 'mol-math/geometry/symmetry-operator'
 import { Model } from '../model'
-import { sortArray, sort, arraySwap, hash1 } from 'mol-data/util';
+import { sort, arraySwap, hash1 } from 'mol-data/util';
 import Element from './element'
 import Unit from './unit'
 import { StructureLookup3D } from './util/lookup3d';
-import StructureSymmetry from './symmetry';
 import { CoarseElements } from '../model/properties/coarse';
+import { StructureSubsetBuilder } from './util/subset-builder';
 
 class Structure {
     readonly unitMap: IntMap<Unit>;
@@ -23,7 +23,7 @@ class Structure {
     private _hashCode = 0;
 
     subsetBuilder(isSorted: boolean) {
-        return new Structure.SubsetBuilder(this, isSorted);
+        return new StructureSubsetBuilder(this, isSorted);
     }
 
     get hashCode() {
@@ -55,7 +55,7 @@ class Structure {
     private _lookup3d?: StructureLookup3D = void 0;
     get lookup3d() {
         if (this._lookup3d) return this._lookup3d;
-        this._lookup3d = StructureLookup3D.create(this);
+        this._lookup3d = new StructureLookup3D(this);
         return this._lookup3d;
     }
 
@@ -141,106 +141,6 @@ namespace Structure {
 
     export function Builder() { return new StructureBuilder(); }
 
-    export class SubsetBuilder {
-        private ids: number[] = [];
-        private unitMap = IntMap.Mutable<number[]>();
-        private parentId = -1;
-        private currentUnit: number[] = [];
-        elementCount = 0;
-
-        addToUnit(parentId: number, e: number) {
-            const unit = this.unitMap.get(parentId);
-            if (!!unit) { unit[unit.length] = e; }
-            else {
-                this.unitMap.set(parentId, [e]);
-                this.ids[this.ids.length] = parentId;
-            }
-            this.elementCount++;
-        }
-
-        beginUnit(parentId: number) {
-            this.parentId = parentId;
-            this.currentUnit = this.currentUnit.length > 0 ? [] : this.currentUnit;
-        }
-
-        addElement(e: number) {
-            this.currentUnit[this.currentUnit.length] = e;
-            this.elementCount++;
-        }
-
-        commitUnit() {
-            if (this.currentUnit.length === 0) return;
-            this.ids[this.ids.length] = this.parentId;
-            this.unitMap.set(this.parentId, this.currentUnit);
-            this.parentId = -1;
-        }
-
-        setUnit(parentId: number, elements: ArrayLike<number>) {
-            this.ids[this.ids.length] = parentId;
-            this.unitMap.set(parentId, elements as number[]);
-            this.elementCount += elements.length;
-        }
-
-        private _getStructure(deduplicateElements: boolean): Structure {
-            if (this.isEmpty) return Structure.Empty;
-
-            const newUnits: Unit[] = [];
-            sortArray(this.ids);
-
-            const symmGroups = StructureSymmetry.UnitEquivalenceBuilder();
-
-            for (let i = 0, _i = this.ids.length; i < _i; i++) {
-                const id = this.ids[i];
-                const parent = this.parent.unitMap.get(id);
-
-                let unit: ArrayLike<number> = this.unitMap.get(id);
-                let sorted = false;
-
-                if (deduplicateElements) {
-                    if (!this.isSorted) sortArray(unit);
-                    unit = SortedArray.deduplicate(SortedArray.ofSortedArray(this.currentUnit));
-                    sorted = true;
-                }
-
-                const l = unit.length;
-
-                // if the length is the same, just copy the old unit.
-                if (unit.length === parent.elements.length) {
-                    newUnits[newUnits.length] = parent;
-                    symmGroups.add(parent.id, parent);
-                    continue;
-                }
-
-                if (!this.isSorted && !sorted && l > 1) sortArray(unit);
-
-                let child = parent.getChild(SortedArray.ofSortedArray(unit));
-                const pivot = symmGroups.add(child.id, child);
-                if (child !== pivot) child = pivot.applyOperator(child.id, child.conformation.operator, true);
-                newUnits[newUnits.length] = child;
-            }
-
-            return create(newUnits);
-        }
-
-        getStructure(deduplicateElements = false) {
-            return this._getStructure(deduplicateElements);
-        }
-
-        setSingletonLocation(location: Element.Location) {
-            const id = this.ids[0];
-            location.unit = this.parent.unitMap.get(id);
-            location.element = this.unitMap.get(id)[0];
-        }
-
-        get isEmpty() {
-            return this.elementCount === 0;
-        }
-
-        constructor(private parent: Structure, private isSorted: boolean) {
-
-        }
-    }
-
     export function getModels(s: Structure) {
         const { units } = s;
         const arr = UniqueArray.create<Model['id'], Model>();
@@ -282,18 +182,18 @@ namespace Structure {
         private current = Element.Location();
         private unitIndex = 0;
         private elements: SortedArray;
-        private len = 0;
-        private idx = 0;
+        private maxIdx = 0;
+        private idx = -1;
 
         hasNext: boolean;
         move(): Element.Location {
-            this.current.element = this.elements[this.idx];
             this.advance();
+            this.current.element = this.elements[this.idx];
             return this.current;
         }
 
         private advance() {
-            if (this.idx < this.len - 1) {
+            if (this.idx < this.maxIdx) {
                 this.idx++;
                 return;
             }
@@ -307,14 +207,14 @@ namespace Structure {
 
             this.current.unit = this.structure.units[this.unitIndex];
             this.elements = this.current.unit.elements;
-            this.len = this.elements.length;
+            this.maxIdx = this.elements.length - 1;
         }
 
         constructor(private structure: Structure) {
             this.hasNext = structure.elementCount > 0;
             if (this.hasNext) {
                 this.elements = structure.units[0].elements;
-                this.len = this.elements.length;
+                this.maxIdx = this.elements.length - 1;
                 this.current.unit = structure.units[0];
             }
         }

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

@@ -43,6 +43,8 @@ namespace StructureSymmetry {
         });
     }
 
+    // TODO: build symmetry mates within radius
+
     export function buildSymmetryRange(structure: Structure, ijkMin: Vec3, ijkMax: Vec3) {
         return Task.create('Build Assembly', async ctx => {
             const models = Structure.getModels(structure);

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

@@ -10,87 +10,102 @@ import { Lookup3D, GridLookup3D, Result, Box3D, Sphere3D } from 'mol-math/geomet
 import { Vec3 } from 'mol-math/linear-algebra';
 import { computeStructureBoundary } from './boundary';
 import { OrderedSet } from 'mol-data/int';
-
-interface StructureLookup3D extends Lookup3D<Element> {}
-
-namespace StructureLookup3D {
-    class Impl implements StructureLookup3D {
-        private unitLookup: Lookup3D;
-        private result = Result.create<Element>();
-        private pivot = Vec3.zero();
-
-        find(x: number, y: number, z: number, radius: number): Result<Element> {
-            Result.reset(this.result);
-            const { units } = this.structure;
-            const closeUnits = this.unitLookup.find(x, y, z, radius);
-            if (closeUnits.count === 0) return this.result;
-
-            for (let t = 0, _t = closeUnits.count; t < _t; t++) {
-                const unit = units[closeUnits.indices[t]];
-                Vec3.set(this.pivot, x, y, z);
-                if (!unit.conformation.operator.isIdentity) {
-                    Vec3.transformMat4(this.pivot, this.pivot, unit.conformation.operator.inverse);
-                }
-                const unitLookup = unit.lookup3d;
-                const groupResult = unitLookup.find(this.pivot[0], this.pivot[1], this.pivot[2], radius);
-                for (let j = 0, _j = groupResult.count; j < _j; j++) {
-                    Result.add(this.result, Element.create(unit.id, groupResult.indices[j]), groupResult.squaredDistances[j]);
-                }
+import { StructureUniqueSubsetBuilder } from './unique-subset-builder';
+
+export class StructureLookup3D implements Lookup3D<Element> {
+    private unitLookup: Lookup3D;
+    private result = Result.create<Element>();
+    private pivot = Vec3.zero();
+
+    find(x: number, y: number, z: number, radius: number): Result<Element> {
+        Result.reset(this.result);
+        const { units } = this.structure;
+        const closeUnits = this.unitLookup.find(x, y, z, radius);
+        if (closeUnits.count === 0) return this.result;
+
+        for (let t = 0, _t = closeUnits.count; t < _t; t++) {
+            const unit = units[closeUnits.indices[t]];
+            Vec3.set(this.pivot, x, y, z);
+            if (!unit.conformation.operator.isIdentity) {
+                Vec3.transformMat4(this.pivot, this.pivot, unit.conformation.operator.inverse);
+            }
+            const unitLookup = unit.lookup3d;
+            const groupResult = unitLookup.find(this.pivot[0], this.pivot[1], this.pivot[2], radius);
+            for (let j = 0, _j = groupResult.count; j < _j; j++) {
+                Result.add(this.result, Element.create(unit.id, groupResult.indices[j]), groupResult.squaredDistances[j]);
             }
-
-            return this.result;
         }
 
-        check(x: number, y: number, z: number, radius: number): boolean {
-            const { units } = this.structure;
-            const closeUnits = this.unitLookup.find(x, y, z, radius);
-            if (closeUnits.count === 0) return false;
-
-            for (let t = 0, _t = closeUnits.count; t < _t; t++) {
-                const unit = units[closeUnits.indices[t]];
-                Vec3.set(this.pivot, x, y, z);
-                if (!unit.conformation.operator.isIdentity) {
-                    Vec3.transformMat4(this.pivot, this.pivot, unit.conformation.operator.inverse);
-                }
-                const groupLookup = unit.lookup3d;
-                if (groupLookup.check(this.pivot[0], this.pivot[1], this.pivot[2], radius)) return true;
+        return this.result;
+    }
+
+    findIntoBuilder(x: number, y: number, z: number, radius: number, builder: StructureUniqueSubsetBuilder) {
+        const { units } = this.structure;
+        const closeUnits = this.unitLookup.find(x, y, z, radius);
+        if (closeUnits.count === 0) return;
+
+        for (let t = 0, _t = closeUnits.count; t < _t; t++) {
+            const unit = units[closeUnits.indices[t]];
+            Vec3.set(this.pivot, x, y, z);
+            if (!unit.conformation.operator.isIdentity) {
+                Vec3.transformMat4(this.pivot, this.pivot, unit.conformation.operator.inverse);
             }
+            const unitLookup = unit.lookup3d;
+            const groupResult = unitLookup.find(this.pivot[0], this.pivot[1], this.pivot[2], radius);
+            if (groupResult.count === 0) continue;
+
+            const elements = unit.elements;
+            builder.beginUnit(unit.id);
+            for (let j = 0, _j = groupResult.count; j < _j; j++) {
+                builder.addElement(elements[groupResult.indices[j]]);
+            }
+            builder.commitUnit();
+        }
+    }
+
+    check(x: number, y: number, z: number, radius: number): boolean {
+        const { units } = this.structure;
+        const closeUnits = this.unitLookup.find(x, y, z, radius);
+        if (closeUnits.count === 0) return false;
 
-            return false;
+        for (let t = 0, _t = closeUnits.count; t < _t; t++) {
+            const unit = units[closeUnits.indices[t]];
+            Vec3.set(this.pivot, x, y, z);
+            if (!unit.conformation.operator.isIdentity) {
+                Vec3.transformMat4(this.pivot, this.pivot, unit.conformation.operator.inverse);
+            }
+            const groupLookup = unit.lookup3d;
+            if (groupLookup.check(this.pivot[0], this.pivot[1], this.pivot[2], radius)) return true;
         }
 
-        boundary: { box: Box3D; sphere: Sphere3D; };
+        return false;
+    }
 
-        constructor(private structure: Structure) {
-            const { units } = structure;
-            const unitCount = units.length;
-            const xs = new Float32Array(unitCount);
-            const ys = new Float32Array(unitCount);
-            const zs = new Float32Array(unitCount);
-            const radius = new Float32Array(unitCount);
+    boundary: { box: Box3D; sphere: Sphere3D; };
 
-            const center = Vec3.zero();
-            for (let i = 0; i < unitCount; i++) {
-                const unit = units[i];
-                const lookup = unit.lookup3d;
-                const s = lookup.boundary.sphere;
+    constructor(private structure: Structure) {
+        const { units } = structure;
+        const unitCount = units.length;
+        const xs = new Float32Array(unitCount);
+        const ys = new Float32Array(unitCount);
+        const zs = new Float32Array(unitCount);
+        const radius = new Float32Array(unitCount);
 
-                Vec3.transformMat4(center, s.center, unit.conformation.operator.matrix);
+        const center = Vec3.zero();
+        for (let i = 0; i < unitCount; i++) {
+            const unit = units[i];
+            const lookup = unit.lookup3d;
+            const s = lookup.boundary.sphere;
 
-                xs[i] = center[0];
-                ys[i] = center[1];
-                zs[i] = center[2];
-                radius[i] = s.radius;
-            }
+            Vec3.transformMat4(center, s.center, unit.conformation.operator.matrix);
 
-            this.unitLookup = GridLookup3D({ x: xs, y: ys, z: zs, radius, indices: OrderedSet.ofBounds(0, unitCount) });
-            this.boundary = computeStructureBoundary(structure);
+            xs[i] = center[0];
+            ys[i] = center[1];
+            zs[i] = center[2];
+            radius[i] = s.radius;
         }
-    }
 
-    export function create(s: Structure): StructureLookup3D {
-        return new Impl(s);
+        this.unitLookup = GridLookup3D({ x: xs, y: ys, z: zs, radius, indices: OrderedSet.ofBounds(0, unitCount) });
+        this.boundary = computeStructureBoundary(structure);
     }
-}
-
-export { StructureLookup3D }
+}

src/mol-model/structure/structure/util/subset-builder.ts

@@ -0,0 +1,116 @@
+/**
+ * Copyright (c) 2017 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { IntMap, SortedArray } from 'mol-data/int';
+import { sortArray } from 'mol-data/util';
+import Element from '../element';
+import StructureSymmetry from '../symmetry';
+import Unit from '../unit';
+import Structure from '../structure';
+
+export class StructureSubsetBuilder {
+    private ids: number[] = [];
+    private unitMap = IntMap.Mutable<number[]>();
+    private parentId = -1;
+    private currentUnit: number[] = [];
+    elementCount = 0;
+
+    addToUnit(parentId: number, e: number) {
+        const unit = this.unitMap.get(parentId);
+        if (!!unit) { unit[unit.length] = e; }
+        else {
+            this.unitMap.set(parentId, [e]);
+            this.ids[this.ids.length] = parentId;
+        }
+        this.elementCount++;
+    }
+
+    beginUnit(parentId: number) {
+        this.parentId = parentId;
+        this.currentUnit = this.currentUnit.length > 0 ? [] : this.currentUnit;
+    }
+
+    addElement(e: number) {
+        this.currentUnit[this.currentUnit.length] = e;
+        this.elementCount++;
+    }
+
+    commitUnit() {
+        if (this.currentUnit.length === 0) return;
+        this.ids[this.ids.length] = this.parentId;
+        this.unitMap.set(this.parentId, this.currentUnit);
+        this.parentId = -1;
+    }
+
+    setUnit(parentId: number, elements: ArrayLike<number>) {
+        this.ids[this.ids.length] = parentId;
+        this.unitMap.set(parentId, elements as number[]);
+        this.elementCount += elements.length;
+    }
+
+    private _getStructure(deduplicateElements: boolean): Structure {
+        if (this.isEmpty) return Structure.Empty;
+
+        const newUnits: Unit[] = [];
+        sortArray(this.ids);
+
+        const symmGroups = StructureSymmetry.UnitEquivalenceBuilder();
+
+        for (let i = 0, _i = this.ids.length; i < _i; i++) {
+            const id = this.ids[i];
+            const parent = this.parent.unitMap.get(id);
+
+            let unit: ArrayLike<number> = this.unitMap.get(id);
+            let sorted = false;
+
+            if (deduplicateElements) {
+                if (!this.isSorted) sortArray(unit);
+                unit = SortedArray.deduplicate(SortedArray.ofSortedArray(this.currentUnit));
+                sorted = true;
+            }
+
+            const l = unit.length;
+
+            // if the length is the same, just copy the old unit.
+            if (unit.length === parent.elements.length) {
+                newUnits[newUnits.length] = parent;
+                symmGroups.add(parent.id, parent);
+                continue;
+            }
+
+            if (!this.isSorted && !sorted && l > 1) sortArray(unit);
+
+            let child = parent.getChild(SortedArray.ofSortedArray(unit));
+            const pivot = symmGroups.add(child.id, child);
+            if (child !== pivot) child = pivot.applyOperator(child.id, child.conformation.operator, true);
+            newUnits[newUnits.length] = child;
+        }
+
+        return Structure.create(newUnits);
+    }
+
+    getStructure() {
+        return this._getStructure(false);
+    }
+
+    getStructureDeduplicate() {
+        return this._getStructure(true);
+    }
+
+    setSingletonLocation(location: Element.Location) {
+        const id = this.ids[0];
+        location.unit = this.parent.unitMap.get(id);
+        location.element = this.unitMap.get(id)[0];
+    }
+
+    get isEmpty() {
+        return this.elementCount === 0;
+    }
+
+    constructor(private parent: Structure, private isSorted: boolean) {
+
+    }
+}

src/mol-model/structure/structure/util/unique-subset-builder.ts

@@ -0,0 +1,98 @@
+/**
+ * Copyright (c) 2017 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { IntMap, SortedArray } from 'mol-data/int';
+import { sortArray } from 'mol-data/util';
+import StructureSymmetry from '../symmetry';
+import Unit from '../unit';
+import Structure from '../structure';
+import { UniqueArray } from 'mol-data/generic';
+
+type UArray = UniqueArray<number, number>
+
+export class StructureUniqueSubsetBuilder {
+    private ids: number[] = [];
+    private unitMap = IntMap.Mutable<UArray>();
+    private parentId = -1;
+    private currentUnit: UArray = UniqueArray.create();
+    elementCount = 0;
+
+    addToUnit(parentId: number, e: number) {
+        const unit = this.unitMap.get(parentId);
+        if (!!unit) {
+            if (UniqueArray.add(unit, e, e)) this.elementCount++;
+        }
+        else {
+            const arr: UArray = UniqueArray.create();
+            UniqueArray.add(arr, e, e);
+            this.unitMap.set(parentId, arr);
+            this.ids[this.ids.length] = parentId;
+            this.elementCount++;
+        }
+    }
+
+    beginUnit(parentId: number) {
+        this.parentId = parentId;
+        if (this.unitMap.has(parentId)) {
+            this.currentUnit = this.unitMap.get(parentId);
+        } else {
+            this.currentUnit = this.currentUnit.array.length > 0 ? UniqueArray.create() : this.currentUnit;
+        }
+    }
+
+    addElement(e: number) {
+        if (UniqueArray.add(this.currentUnit, e, e)) this.elementCount++;
+    }
+
+    commitUnit() {
+        if (this.currentUnit.array.length === 0 || this.unitMap.has(this.parentId)) return;
+        this.ids[this.ids.length] = this.parentId;
+        this.unitMap.set(this.parentId, this.currentUnit);
+        this.parentId = -1;
+    }
+
+    getStructure(): Structure {
+        if (this.isEmpty) return Structure.Empty;
+
+        const newUnits: Unit[] = [];
+        sortArray(this.ids);
+
+        const symmGroups = StructureSymmetry.UnitEquivalenceBuilder();
+
+        for (let i = 0, _i = this.ids.length; i < _i; i++) {
+            const id = this.ids[i];
+            const parent = this.parent.unitMap.get(id);
+
+            let unit: ArrayLike<number> = this.unitMap.get(id).array;
+
+            const l = unit.length;
+
+            // if the length is the same, just copy the old unit.
+            if (unit.length === parent.elements.length) {
+                newUnits[newUnits.length] = parent;
+                symmGroups.add(parent.id, parent);
+                continue;
+            }
+
+            if (l > 1) sortArray(unit);
+
+            let child = parent.getChild(SortedArray.ofSortedArray(unit));
+            const pivot = symmGroups.add(child.id, child);
+            if (child !== pivot) child = pivot.applyOperator(child.id, child.conformation.operator, true);
+            newUnits[newUnits.length] = child;
+        }
+
+        return Structure.create(newUnits);
+    }
+
+    get isEmpty() {
+        return this.elementCount === 0;
+    }
+
+    constructor(private parent: Structure) {
+
+    }
+}

src/perf-tests/structure.ts

@@ -16,6 +16,7 @@ import { Structure, Model, Queries as Q, Element, Selection, StructureSymmetry,
 
 import to_mmCIF from 'mol-model/structure/export/mmcif'
 import { Vec3 } from 'mol-math/linear-algebra';
+//import { printUnits } from 'apps/structure-info/model';
 //import { EquivalenceClasses } from 'mol-data/util';
 
 require('util.promisify').shim();
@@ -315,6 +316,42 @@ export namespace PropertyAccess {
         console.log('exported');
     }
 
+    export async function testIncludeSurroundings(id: string, s: Structure) {
+        //const a = s; 
+        console.time('symmetry')
+        const a = await StructureSymmetry.buildSymmetryRange(s, Vec3.create(-2, -2, -2), Vec3.create(2, 2, 2)).run();
+        //console.log(printUnits(a));
+
+        const auth_comp_id = Q.props.residue.auth_comp_id, op = Q.props.unit.operator_name;
+        //const q1 = Q.generators.atoms({ residueTest: l => auth_comp_id(l) === 'REA' });
+        const q1 = Q.modifiers.includeSurroundings(Q.generators.atoms({
+            chainTest: l => op(l) === '1_555',
+            residueTest: l => auth_comp_id(l) === 'REA'
+        }), {
+            radius: 5,
+            wholeResidues: true
+        });
+        const surr = Selection.unionStructure(await query(Query(q1), a));
+        console.timeEnd('symmetry')
+
+        // for (const u of surr.units) {
+        //     const { atomId } = u.model.atomicConformation;
+        //     console.log(`${u.id}, ${u.conformation.operator.name}`);
+        //     for (let i = 0; i < u.elements.length; i++) {
+        //         console.log(`  ${atomId.value(u.elements[i])}`);
+        //     }
+        // }
+
+        // const it = surr.elementLocations();
+        // while (it.hasNext) {
+        //     const e = it.move();
+        //     console.log(`${Q.props.unit.operator_name(e)} ${Q.props.atom.id(e)}`);
+        // }
+    //fs.writeFileSync(`${DATA_DIR}/${id}_surr.bcif`, to_mmCIF(id, a, true));
+        fs.writeFileSync(`${DATA_DIR}/${id}_surr.cif`, to_mmCIF(id, surr, false));
+        console.log('exported');
+    }
+
     // export async function testGrouping(structure: Structure) {
     //     const { elements, units } = await Run(Symmetry.buildAssembly(structure, '1'));
     //     console.log('grouping', units.length);
@@ -354,7 +391,8 @@ export namespace PropertyAccess {
         // console.log(mmcif.pdbx_struct_oper_list.vector.toArray());
 
         //await testAssembly('1hrv', structures[0]);
-        await testSymmetry('1cbs', structures[0]);
+        //await testSymmetry('1cbs', structures[0]);
+        await testIncludeSurroundings('1cbs', structures[0]);
         // throw '';
 
         // console.log(models[0].symmetry.assemblies);