Updated coarse unit model

src/apps/structure-info/model.ts

@@ -15,7 +15,6 @@ import { Model, Structure, Element, ElementSet, Unit, ElementGroup, Queries } fr
 import { OrderedSet } from 'mol-data/int';
 import { Table } from 'mol-data/db';
 import { mmCIF_Database } from 'mol-io/reader/cif/schema/mmcif';
-import CoarseGrained from 'mol-model/structure/model/properties/coarse-grained';
 import { openCif, downloadCif } from './helpers';
 
 
@@ -97,7 +96,7 @@ export function printUnits(structure: Structure) {
         if (Unit.isAtomic(l.unit)) {
             console.log(`Atomic unit ${unitId}: ${size} elements`);
         } else if (Unit.isCoarse(l.unit)) {
-            console.log(`Coarse unit ${unitId} (${l.unit.elementType === CoarseGrained.ElementType.Sphere ? 'spheres' : 'gaussians'}): ${size} elements.`);
+            console.log(`Coarse unit ${unitId} (${Unit.isCoarseSpheres(l.unit) ? 'spheres' : 'gaussians'}): ${size} elements.`);
 
             const props = Queries.props.coarse_grained;
             const seq = l.unit.model.sequence;

src/mol-model/structure/model/properties/coarse-grained.ts

@@ -45,7 +45,7 @@ namespace CoarseGrained {
         entityKey: ArrayLike<number>
     }
 
-    export type SiteBases =  Common & { [P in keyof SiteBase]: Column<SiteBase[P]> }
+    export type SitesBase =  Common & { [P in keyof SiteBase]: Column<SiteBase[P]> }
     export type Spheres =  Common & { [P in keyof Sphere]: Column<Sphere[P]> }
     export type Gaussians = Common & { matrix_space: Tensor.Space } & { [P in keyof Gaussian]: Column<Gaussian[P]> }
 }

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

@@ -5,7 +5,6 @@
  */
 
 import { Element, Unit } from '../structure'
-import CoarseGrained from '../model/properties/coarse-grained';
 
 const constant = {
     true: Element.property(l => true),
@@ -61,26 +60,22 @@ const chain = {
 }
 
 const coarse_grained = {
-    modelKey: Element.property(l => !Unit.isCoarse(l.unit) ? notCoarse() : l.unit.siteBases.modelKey[l.element]),
-    entityKey: Element.property(l => !Unit.isCoarse(l.unit) ? notCoarse() : l.unit.siteBases.entityKey[l.element]),
+    modelKey: Element.property(l => !Unit.isCoarse(l.unit) ? notCoarse() : l.unit.sites.modelKey[l.element]),
+    entityKey: Element.property(l => !Unit.isCoarse(l.unit) ? notCoarse() : l.unit.sites.entityKey[l.element]),
 
     x: atom.x,
     y: atom.y,
     z: atom.z,
 
-    asym_id: Element.property(l => !Unit.isCoarse(l.unit) ? notCoarse() : l.unit.siteBases.asym_id.value(l.element)),
-    seq_id_begin: Element.property(l => !Unit.isCoarse(l.unit) ? notCoarse() : l.unit.siteBases.seq_id_begin.value(l.element)),
-    seq_id_end: Element.property(l => !Unit.isCoarse(l.unit) ? notCoarse() : l.unit.siteBases.seq_id_end.value(l.element)),
+    asym_id: Element.property(l => !Unit.isCoarse(l.unit) ? notCoarse() : l.unit.sites.asym_id.value(l.element)),
+    seq_id_begin: Element.property(l => !Unit.isCoarse(l.unit) ? notCoarse() : l.unit.sites.seq_id_begin.value(l.element)),
+    seq_id_end: Element.property(l => !Unit.isCoarse(l.unit) ? notCoarse() : l.unit.sites.seq_id_end.value(l.element)),
 
-    sphere_radius: Element.property(l => !Unit.isCoarse(l.unit) || l.unit.elementType !== CoarseGrained.ElementType.Sphere
-        ? notCoarse('spheres') : l.unit.spheres.radius.value(l.element)),
-    sphere_rmsf: Element.property(l => !Unit.isCoarse(l.unit) || l.unit.elementType !== CoarseGrained.ElementType.Sphere
-        ? notCoarse('spheres') : l.unit.spheres.rmsf.value(l.element)),
+    sphere_radius: Element.property(l => !Unit.isCoarseSpheres(l.unit) ? notCoarse('spheres') : l.unit.sites.radius.value(l.element)),
+    sphere_rmsf: Element.property(l => !Unit.isCoarseSpheres(l.unit) ? notCoarse('spheres') : l.unit.sites.rmsf.value(l.element)),
 
-    gaussian_weight: Element.property(l => !Unit.isCoarse(l.unit) || l.unit.elementType !== CoarseGrained.ElementType.Gaussian
-        ? notCoarse('gaussians') : l.unit.gaussians.weight.value(l.element)),
-    gaussian_covariance_matrix: Element.property(l => !Unit.isCoarse(l.unit) || l.unit.elementType !== CoarseGrained.ElementType.Gaussian
-        ? notCoarse('gaussians') : l.unit.gaussians.covariance_matrix.value(l.element)),
+    gaussian_weight: Element.property(l => !Unit.isCoarseGaussians(l.unit) ? notCoarse('gaussians') : l.unit.sites.weight.value(l.element)),
+    gaussian_covariance_matrix: Element.property(l => !Unit.isCoarseGaussians(l.unit) ? notCoarse('gaussians') : l.unit.sites.covariance_matrix.value(l.element))
 }
 
 function eK(l: Element.Location) { return !Unit.isAtomic(l.unit) ? notAtomic() : l.unit.hierarchy.entityKey.value(l.unit.chainIndex[l.element]); }

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

@@ -12,7 +12,6 @@ import Unit from './unit'
 import ElementSet from './element/set'
 import ElementGroup from './element/group'
 import Element from './element'
-import CoarseGrained from '../model/properties/coarse-grained';
 
 // A structure is a pair of "units" and an element set.
 // Each unit contains the data and transformation of its corresponding elements.
@@ -36,7 +35,7 @@ namespace Structure {
 
         for (let c = 0; c < chains.count; c++) {
             const group = ElementGroup.createNew(OrderedSet.ofBounds(chains.segments[c], chains.segments[c + 1]));
-            const unit = Unit.createAtomic(model, SymmetryOperator.Default, group);
+            const unit = Unit.create(Unit.Kind.Atomic, model, SymmetryOperator.Default, group);
             builder.add(unit, unit.fullGroup);
         }
 
@@ -44,12 +43,12 @@ namespace Structure {
         if (cs.isDefined) {
             if (cs.spheres.count > 0) {
                 const group = ElementGroup.createNew(OrderedSet.ofBounds(0, cs.spheres.count));
-                const unit = Unit.createCoarse(model, SymmetryOperator.Default, group, CoarseGrained.ElementType.Sphere);
+                const unit = Unit.create(Unit.Kind.CoarseSpheres, model, SymmetryOperator.Default, group);
                 builder.add(unit, unit.fullGroup);
             }
             if (cs.gaussians.count > 0) {
                 const group = ElementGroup.createNew(OrderedSet.ofBounds(0, cs.gaussians.count));
-                const unit = Unit.createCoarse(model, SymmetryOperator.Default, group, CoarseGrained.ElementType.Gaussian);
+                const unit = Unit.create(Unit.Kind.CoarseGaussians, model, SymmetryOperator.Default, group);
                 builder.add(unit, unit.fullGroup);
             }
         }

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

@@ -13,13 +13,15 @@ import CoarseGrained from '../model/properties/coarse-grained';
 
 // A building block of a structure that corresponds to an atomic or a coarse grained representation
 // 'conveniently grouped together'.
-type Unit = Unit.Atomic | Unit.Coarse
+type Unit = Unit.Atomic | Unit.CoarseSpheres | Unit.CoarseGaussians
 
 namespace Unit {
-    export const enum Kind { Atomic, Coarse }
+    export const enum Kind { Atomic, CoarseSpheres, CoarseGaussians }
 
     export function isAtomic(u: Unit): u is Atomic { return u.kind === Kind.Atomic; }
-    export function isCoarse(u: Unit): u is Coarse { return u.kind === Kind.Coarse; }
+    export function isCoarse(u: Unit): u is CoarseSpheres | CoarseGaussians { return u.kind === Kind.CoarseSpheres || u.kind === Kind.CoarseGaussians; }
+    export function isCoarseSpheres(u: Unit): u is CoarseSpheres { return u.kind === Kind.CoarseSpheres; }
+    export function isCoarseGaussians(u: Unit): u is CoarseGaussians { return u.kind === Kind.CoarseGaussians; }
 
     export interface Base extends SymmetryOperator.ArrayMapping {
         // Provides access to the underlying data.
@@ -41,7 +43,7 @@ namespace Unit {
     // An atom set can be referenced by multiple diffrent units which
     // makes construction of assemblies and spacegroups very efficient.
     export interface Atomic extends Base {
-        readonly kind: Unit.Kind.Atomic,
+        readonly kind: Kind.Atomic,
 
         // Reference some commonly accessed things for faster access.
         readonly residueIndex: ArrayLike<number>,
@@ -51,16 +53,23 @@ namespace Unit {
     }
 
     // Coarse grained representations.
-    export interface Coarse extends Base  {
-        readonly kind: Unit.Kind.Coarse,
-        readonly elementType: CoarseGrained.ElementType,
+    export interface CoarseBase<K extends Kind, S extends CoarseGrained.SitesBase> extends Base  {
+        readonly kind: K,
+        readonly sites: S
+    }
+
+    export interface CoarseSpheres extends CoarseBase<Kind.CoarseSpheres, CoarseGrained.Spheres> { }
+    export interface CoarseGaussians extends CoarseBase<Kind.CoarseGaussians, CoarseGrained.Gaussians> { }
 
-        readonly siteBases: CoarseGrained.SiteBases,
-        readonly spheres: CoarseGrained.Spheres,
-        readonly gaussians: CoarseGrained.Gaussians
+    export function create(kind: Kind, model: Model, operator: SymmetryOperator, fullGroup: ElementGroup): Unit {
+        switch (kind) {
+            case Kind.Atomic: return createAtomic(model, operator, fullGroup);
+            case Kind.CoarseSpheres: return createCoarseSpheres(model, operator, fullGroup);
+            case Kind.CoarseGaussians: return createCoarseGaussians(model, operator, fullGroup);
+        }
     }
 
-    export function createAtomic(model: Model, operator: SymmetryOperator, fullGroup: ElementGroup): Unit.Atomic {
+    function createAtomic(model: Model, operator: SymmetryOperator, fullGroup: ElementGroup): Unit.Atomic {
         const h = model.hierarchy;
         const { invariantPosition, position, x, y, z } = SymmetryOperator.createMapping(operator, model.atomSiteConformation);
 
@@ -79,30 +88,38 @@ namespace Unit {
         };
     }
 
-    export function createCoarse(model: Model, operator: SymmetryOperator, fullGroup: ElementGroup, elementType: CoarseGrained.ElementType): Unit.Coarse {
-        const siteBases = elementType === CoarseGrained.ElementType.Sphere ? model.coarseGrained.spheres : model.coarseGrained.gaussians;
-        const { invariantPosition, position, x, y, z } = SymmetryOperator.createMapping(operator, siteBases);
+    function createCoarseSpheres(model: Model, operator: SymmetryOperator, fullGroup: ElementGroup): Unit.CoarseSpheres {
+        const { invariantPosition, position, x, y, z } = SymmetryOperator.createMapping(operator, model.coarseGrained.spheres);
 
         return {
             model,
-            kind: Kind.Coarse,
-            elementType,
+            kind: Kind.CoarseSpheres,
+            sites: model.coarseGrained.spheres,
             operator,
             fullGroup,
-            siteBases,
-            spheres: model.coarseGrained.spheres,
-            gaussians: model.coarseGrained.gaussians,
             invariantPosition,
             position,
             x, y, z
         };
     }
 
-    export function withOperator(unit: Unit, operator: SymmetryOperator): Unit {
-        switch (unit.kind) {
-            case Kind.Atomic: return createAtomic(unit.model, SymmetryOperator.compose(unit.operator, operator), unit.fullGroup);
-            case Kind.Coarse: return createCoarse(unit.model, SymmetryOperator.compose(unit.operator, operator), unit.fullGroup, unit.elementType);
+    function createCoarseGaussians(model: Model, operator: SymmetryOperator, fullGroup: ElementGroup): Unit.CoarseGaussians {
+        const { invariantPosition, position, x, y, z } = SymmetryOperator.createMapping(operator, model.coarseGrained.gaussians);
+
+        return {
+            model,
+            kind: Kind.CoarseGaussians,
+            sites: model.coarseGrained.gaussians,
+            operator,
+            fullGroup,
+            invariantPosition,
+            position,
+            x, y, z
+        };
     }
+
+    export function withOperator(unit: Unit, operator: SymmetryOperator): Unit {
+        return create(unit.kind, unit.model, SymmetryOperator.compose(unit.operator, operator), unit.fullGroup);
     }
 
     export function getLookup3d(unit: Unit, group: ElementGroup) {