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,41 +44,11 @@ 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.');
-
-            const { spacegroup } = models[0].symmetry;
-            if (SpacegroupCell.isZero(spacegroup.cell)) return structure;
-
-            const operators: SymmetryOperator[] = [];
-            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++) {
-                            operators[operators.length] = Spacegroup.getSymmetryOperator(spacegroup, op, i, j, k);
-                        }
-                    }
-                }
-            }
-
-            const assembler = Structure.Builder();
-
-            const { units } = structure;
-            for (const oper of operators) {
-                for (const unit of units) {
-                    assembler.addWithOperator(unit, oper);
-                }
-            }
-
-
-            return assembler.getStructure();
-        });
+        return Task.create('Build Symmetry', ctx => findSymmetryRange(ctx, structure, ijkMin, ijkMax));
     }
 
     function hashUnit(u: Unit) {
@@ -107,4 +77,81 @@ namespace StructureSymmetry {
     }
 }
 
+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;
+    for (const oper of operators) {
+        for (const unit of units) {
+            assembler.addWithOperator(unit, oper);
+        }
+        if (ctx.shouldUpdate) await ctx.update('Building symmetry...');
+    }
+
+    return assembler.getStructure();
+}
+
+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.');
+
+    const symmetry = models[0].symmetry;
+    const { spacegroup } = symmetry;
+    if (SpacegroupCell.isZero(spacegroup.cell)) return structure;
+
+    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;
+
+    const assembler = Structure.Builder();
+
+    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);
+
+            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);