Better symmetry mate generation

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

@@ -44,7 +44,10 @@ export namespace Assembly {
 interface ModelSymmetry {
     readonly assemblies: ReadonlyArray<Assembly>,
     readonly spacegroup: Spacegroup,
-    readonly isNonStandardCrytalFrame: boolean
+    readonly isNonStandardCrytalFrame: boolean,
+
+    // optionally cached operators from [-3, -3, -3] to [3, 3, 3]
+    _operators_333?: SymmetryOperator[]
 }
 
 namespace ModelSymmetry {

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

@@ -7,7 +7,7 @@
 import Structure from './structure'
 import { Selection } from '../query'
 import { ModelSymmetry } from '../model'
-import { Task } from 'mol-task';
+import { Task, RuntimeContext } from 'mol-task';
 import { SortedArray } from 'mol-data/int';
 import Unit from './unit';
 import { EquivalenceClasses, hash2 } from 'mol-data/util';
@@ -44,29 +44,68 @@ namespace StructureSymmetry {
     }
 
     export function builderSymmetryMates(structure: Structure, radius: number) {
-        // TODO: do it properly
-        return buildSymmetryRange(structure, Vec3.create(-3, -3, -3), Vec3.create(3, 3, 3));
+        return Task.create('Find Symmetry Mates', ctx => findMatesRadius(ctx, structure, radius));
     }
 
     export function buildSymmetryRange(structure: Structure, ijkMin: Vec3, ijkMax: Vec3) {
-        return Task.create('Build Assembly', async ctx => {
-            const models = Structure.getModels(structure);
-            if (models.length !== 1) throw new Error('Can only build symmetries from structures based on 1 model.');
+        return Task.create('Build Symmetry', ctx => findSymmetryRange(ctx, structure, ijkMin, ijkMax));
+    }
 
-            const { spacegroup } = models[0].symmetry;
-            if (SpacegroupCell.isZero(spacegroup.cell)) return structure;
+    function hashUnit(u: Unit) {
+        return hash2(u.invariantId, SortedArray.hashCode(u.elements));
+    }
+
+    function areUnitsEquivalent(a: Unit, b: Unit) {
+        return a.invariantId === b.invariantId && a.model.id === b.model.id && SortedArray.areEqual(a.elements, b.elements);
+    }
+
+    export function UnitEquivalenceBuilder() {
+        return EquivalenceClasses<number, Unit>(hashUnit, areUnitsEquivalent);
+    }
+
+    export function getTransformGroups(s: Structure): ReadonlyArray<Unit.SymmetryGroup> {
+        const groups = UnitEquivalenceBuilder();
+        for (const u of s.units) groups.add(u.id, u);
+
+        const ret: Unit.SymmetryGroup[] = [];
+        for (const eqUnits of groups.groups) {
+            const first = s.unitMap.get(eqUnits[0]);
+            ret.push({ elements: first.elements, units: eqUnits.map(id => s.unitMap.get(id)) });
+        }
 
-            const operators: SymmetryOperator[] = [];
+        return ret;
+    }
+}
+
+function getOperators(symmetry: ModelSymmetry, ijkMin: Vec3, ijkMax: Vec3) {
+    const operators: SymmetryOperator[] = symmetry._operators_333 || [];
+    const { spacegroup } = symmetry;
+    if (operators.length === 0) {
+        operators[0] = Spacegroup.getSymmetryOperator(spacegroup, 0, 0, 0, 0)
         for (let op = 0; op < spacegroup.operators.length; op++) {
             for (let i = ijkMin[0]; i < ijkMax[0]; i++) {
                 for (let j = ijkMin[1]; j < ijkMax[1]; j++) {
                     for (let k = ijkMin[2]; k < ijkMax[2]; k++) {
+                        // we have added identity as the 1st operator.
+                        if (op === 0 && i === 0 && j === 0 && k === 0) continue;
                         operators[operators.length] = Spacegroup.getSymmetryOperator(spacegroup, op, i, j, k);
                     }
                 }
             }
         }
+        symmetry._operators_333 = operators;
+    }
+    return operators;
+}
+
+async function findSymmetryRange(ctx: RuntimeContext, structure: Structure, ijkMin: Vec3, ijkMax: Vec3) {
+    const models = Structure.getModels(structure);
+    if (models.length !== 1) throw new Error('Can only build symmetries from structures based on 1 model.');
 
+    const { spacegroup } = models[0].symmetry;
+    if (SpacegroupCell.isZero(spacegroup.cell)) return structure;
+
+    const operators = getOperators(models[0].symmetry, ijkMin, ijkMax);
     const assembler = Structure.Builder();
 
     const { units } = structure;
@@ -74,37 +113,45 @@ namespace StructureSymmetry {
         for (const unit of units) {
             assembler.addWithOperator(unit, oper);
         }
+        if (ctx.shouldUpdate) await ctx.update('Building symmetry...');
     }
 
-
     return assembler.getStructure();
-        });
 }
 
-    function hashUnit(u: Unit) {
-        return hash2(u.invariantId, SortedArray.hashCode(u.elements));
-    }
+async function findMatesRadius(ctx: RuntimeContext, structure: Structure, radius: number) {
+    const models = Structure.getModels(structure);
+    if (models.length !== 1) throw new Error('Can only build symmetries from structures based on 1 model.');
 
-    function areUnitsEquivalent(a: Unit, b: Unit) {
-        return a.invariantId === b.invariantId && a.model.id === b.model.id && SortedArray.areEqual(a.elements, b.elements);
-    }
+    const symmetry = models[0].symmetry;
+    const { spacegroup } = symmetry;
+    if (SpacegroupCell.isZero(spacegroup.cell)) return structure;
 
-    export function UnitEquivalenceBuilder() {
-        return EquivalenceClasses<number, Unit>(hashUnit, areUnitsEquivalent);
-    }
+    if (ctx.shouldUpdate) await ctx.update('Initialing...');
+    const operators = getOperators(symmetry, Vec3.create(-3, -3, -3), Vec3.create(3, 3, 3));
+    const lookup = structure.lookup3d;
 
-    export function getTransformGroups(s: Structure): ReadonlyArray<Unit.SymmetryGroup> {
-        const groups = UnitEquivalenceBuilder();
-        for (const u of s.units) groups.add(u.id, u);
+    const assembler = Structure.Builder();
 
-        const ret: Unit.SymmetryGroup[] = [];
-        for (const eqUnits of groups.groups) {
-            const first = s.unitMap.get(eqUnits[0]);
-            ret.push({ elements: first.elements, units: eqUnits.map(id => s.unitMap.get(id)) });
-        }
+    const { units } = structure;
+    const center = Vec3.zero();
+    for (const oper of operators) {
+        for (const unit of units) {
+            const boundingSphere = unit.lookup3d.boundary.sphere;
+            Vec3.transformMat4(center, boundingSphere.center, oper.matrix);
 
-        return ret;
+            const closeUnits = lookup.findUnitIndices(center[0], center[1], center[2], boundingSphere.radius + radius);
+            for (let uI = 0, _uI = closeUnits.count; uI < _uI; uI++) {
+                const closeUnit = units[closeUnits.indices[uI]];
+                if (!closeUnit.lookup3d.check(center[0], center[1], center[2], boundingSphere.radius + radius)) continue;
+                assembler.addWithOperator(unit, oper);
+            }
         }
+        if (ctx.shouldUpdate) await ctx.update('Building symmetry...');
+    }
+
+
+    return assembler.getStructure();
 }
 
 export default StructureSymmetry;
\ No newline at end of file

src/mol-task/execution/observable.ts

@@ -133,9 +133,9 @@ function abortTree(root: Progress.Node) {
     for (const c of root.children) abortTree(c);
 }
 
-function shouldNotify(info: ProgressInfo, time: number) {
-    return time - info.lastNotified > info.updateRateMs;
-}
+// function shouldNotify(info: ProgressInfo, time: number) {
+//     return time - info.lastNotified > info.updateRateMs;
+// }
 
 function notifyObserver(info: ProgressInfo, time: number) {
     info.lastNotified = time;
@@ -194,6 +194,7 @@ class ObservableRuntimeContext implements RuntimeContext {
         this.lastUpdatedTime = now();
         this.updateProgress(progress);
 
+        // TODO: do the shouldNotify check here?
         if (!!dontNotify /*|| !shouldNotify(this.info, this.lastUpdatedTime)*/) return;
 
         notifyObserver(this.info, this.lastUpdatedTime);