wip, strongly typed indexing

src/mol-geo/representation/structure/visual/util/polymer.ts

@@ -4,7 +4,7 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { Unit, Element, StructureProperties } from 'mol-model/structure';
+import { Unit, Element, StructureProperties, ElementIndex } from 'mol-model/structure';
 import { Segmentation } from 'mol-data/int';
 import { MoleculeType } from 'mol-model/structure/model/types';
 import Iterator from 'mol-data/iterator';
@@ -50,7 +50,7 @@ function getTraceName(l: Element.Location) {
     return traceName
 }
 
-function setTraceElement(l: Element.Location, residueSegment: Segmentation.Segment<Element>) {
+function setTraceElement(l: Element.Location, residueSegment: Segmentation.Segment<ElementIndex>) {
     const elements = l.unit.elements
     l.element = elements[residueSegment.start]
     const traceName = getTraceName(l)
@@ -97,9 +97,9 @@ const enum AtomicPolymerBackboneIteratorState { nextPolymer, firstResidue, nextR
 export class AtomicPolymerBackboneIterator<T extends number = number> implements Iterator<PolymerBackbonePair> {
     private value: PolymerBackbonePair
 
-    private polymerIt: SegmentIterator<Element>
-    private residueIt: SegmentIterator<Element>
-    private polymerSegment: Segmentation.Segment<Element>
+    private polymerIt: SegmentIterator<ElementIndex>
+    private residueIt: SegmentIterator<ElementIndex>
+    private polymerSegment: Segmentation.Segment<ElementIndex>
     private state: AtomicPolymerBackboneIteratorState = AtomicPolymerBackboneIteratorState.nextPolymer
     private pos: SymmetryOperator.CoordinateMapper
 
@@ -161,8 +161,8 @@ const enum CoarsePolymerBackboneIteratorState { nextPolymer, firstElement, nextE
 export class CoarsePolymerBackboneIterator<T extends number = number> implements Iterator<PolymerBackbonePair> {
     private value: PolymerBackbonePair
 
-    private polymerIt: SegmentIterator<Element>
-    private polymerSegment: Segmentation.Segment<Element>
+    private polymerIt: SegmentIterator<ElementIndex>
+    private polymerSegment: Segmentation.Segment<ElementIndex>
     private state: CoarsePolymerBackboneIteratorState = CoarsePolymerBackboneIteratorState.nextPolymer
     private pos: SymmetryOperator.CoordinateMapper
     private elementIndex: number
@@ -268,8 +268,8 @@ function createPolymerTraceElement (unit: Unit) {
 export class AtomicPolymerTraceIterator<T extends number = number> implements Iterator<PolymerTraceElement> {
     private value: PolymerTraceElement
 
-    private polymerIt: SegmentIterator<Element>
-    private residueIt: SegmentIterator<Element>
+    private polymerIt: SegmentIterator<ElementIndex>
+    private residueIt: SegmentIterator<ElementIndex>
     // private polymerSegment: Segmentation.Segment<Element>
     // private state: AtomicPolymerTraceIteratorState = AtomicPolymerTraceIteratorState.nextPolymer
     // private pos: SymmetryOperator.CoordinateMapper

src/mol-model/structure/export/categories/secondary-structure.ts

@@ -13,6 +13,7 @@ import { CifExportContext } from '../mmcif';
 import CifField = CifWriter.Field
 import CifCategory = CifWriter.Category
 import { Column } from 'mol-data/db';
+import { ElementIndex } from '../../model';
 
 export function _struct_conf(ctx: CifExportContext): CifCategory {
     const elements = findElements(ctx, 'helix');
@@ -91,7 +92,7 @@ function findElements<T extends SecondaryStructure.Element>(ctx: CifExportContex
         const segs = unit.model.atomicHierarchy.residueAtomSegments;
         const residues = Segmentation.transientSegments(segs, unit.elements);
 
-        let current: Segmentation.Segment<Element>, move = true;
+        let current: Segmentation.Segment<ElementIndex>, move = true;
         while (residues.hasNext) {
             if (move) current = residues.move();
 

src/mol-model/structure/model/formats/mmcif/atomic.ts

@@ -8,7 +8,7 @@ import { Column, Table } from 'mol-data/db';
 import { Interval, Segmentation } from 'mol-data/int';
 import { mmCIF_Database } from 'mol-io/reader/cif/schema/mmcif';
 import UUID from 'mol-util/uuid';
-import { Element } from '../../../../structure';
+import { ElementIndex } from '../../../../structure';
 import Format from '../../format';
 import { Model } from '../../model';
 import { AtomicConformation, AtomicData, AtomicHierarchy, AtomicSegments, AtomsSchema, ChainsSchema, ResiduesSchema } from '../../properties/atomic';
@@ -23,11 +23,11 @@ function findHierarchyOffsets(atom_site: AtomSite) {
     if (atom_site._rowCount === 0) return { residues: [], polymers: [], chains: [] };
 
     const start = 0, end = atom_site._rowCount;
-    const residues = [start as Element], chains = [start as Element], polymers = [start as Element];
+    const residues = [start as ElementIndex], chains = [start as ElementIndex], polymers = [start as ElementIndex];
 
     const { label_entity_id, label_asym_id, label_seq_id, auth_seq_id, pdbx_PDB_ins_code, label_comp_id } = atom_site;
 
-    for (let i = start + 1; i < end; i++) {
+    for (let i = start + 1 as ElementIndex; i < end; i++) {
         const newChain = !label_entity_id.areValuesEqual(i - 1, i) || !label_asym_id.areValuesEqual(i - 1, i);
         const newResidue = newChain
             || !label_seq_id.areValuesEqual(i - 1, i)
@@ -36,9 +36,9 @@ function findHierarchyOffsets(atom_site: AtomSite) {
             || !label_comp_id.areValuesEqual(i - 1, i);
         const newPolymer = newResidue && label_seq_id.value(i - 1) !== label_seq_id.value(i) - 1;
 
-        if (newResidue) residues[residues.length] = i as Element;
-        if (newPolymer) polymers[polymers.length] = i as Element;
-        if (newChain) chains[chains.length] = i as Element;
+        if (newResidue) residues[residues.length] = i;
+        if (newPolymer) polymers[polymers.length] = i;
+        if (newChain) chains[chains.length] = i;
     }
     return { residues, polymers, chains };
 }

src/mol-model/structure/model/formats/mmcif/bonds/struct_conn.ts

@@ -6,7 +6,7 @@
  */
 
 import { Model } from '../../../model'
-import { Element, Structure } from '../../../../structure'
+import { Structure } from '../../../../structure'
 import { LinkType } from '../../../types'
 import { findEntityIdByAsymId, findAtomIndexByLabelName } from '../util'
 import { Column } from 'mol-data/db'
@@ -14,6 +14,7 @@ import { ModelPropertyDescriptor } from '../../../properties/custom';
 import { mmCIF_Database } from 'mol-io/reader/cif/schema/mmcif';
 import { SortedArray } from 'mol-data/int';
 import { CifWriter } from 'mol-io/writer/cif'
+import { ElementIndex } from '../../../indexing';
 
 export interface StructConn {
     getResidueEntries(residueAIndex: number, residueBIndex: number): ReadonlyArray<StructConn.Entry>,
@@ -34,7 +35,7 @@ export namespace StructConn {
                 const strConn = get(ctx.model);
                 if (!strConn || strConn.entries.length === 0) return [];
 
-                const foundAtoms = new Set<Element>();
+                const foundAtoms = new Set<ElementIndex>();
                 const indices: number[] = [];
                 for (const entry of strConn.entries) {
                     const { partners } = entry;
@@ -61,7 +62,7 @@ export namespace StructConn {
         }
     }
 
-    function hasAtom({ units }: Structure, element: Element) {
+    function hasAtom({ units }: Structure, element: ElementIndex) {
         for (let i = 0, _i = units.length; i < _i; i++) {
             if (SortedArray.indexOf(units[i].elements, element) >= 0) return true;
         }
@@ -132,7 +133,7 @@ export namespace StructConn {
         distance: number,
         order: number,
         flags: number,
-        partners: { residueIndex: number, atomIndex: Element, symmetry: string }[]
+        partners: { residueIndex: number, atomIndex: ElementIndex, symmetry: string }[]
     }
 
     type StructConnType =

src/mol-model/structure/model/formats/mmcif/ihm.ts

@@ -12,7 +12,7 @@ import { getCoarseKeys } from '../../properties/utils/coarse-keys';
 import { UUID } from 'mol-util';
 import { Segmentation, Interval } from 'mol-data/int';
 import { Mat3, Tensor } from 'mol-math/linear-algebra';
-import { Element } from '../../../structure'
+import { ElementIndex } from '../../indexing';
 
 export interface IHMData {
     model_id: number,
@@ -81,14 +81,14 @@ function getGaussianConformation(data: mmCIF['ihm_gaussian_obj_site']): CoarseGa
 }
 
 function getSegments(asym_id: Column<string>, seq_id_begin: Column<number>, seq_id_end: Column<number>) {
-    const polymerOffsets = [0 as Element], chainOffsets = [0 as Element];
+    const polymerOffsets = [0 as ElementIndex], chainOffsets = [0 as ElementIndex];
     for (let i = 1, _i = asym_id.rowCount; i < _i; i++) {
         const newChain = !asym_id.areValuesEqual(i - 1, i);
         const newPolymer = newChain
             || seq_id_end.value(i - 1) !== seq_id_begin.value(i) - 1;
 
-        if (newPolymer) polymerOffsets[polymerOffsets.length] = i as Element;
-        if (newChain) chainOffsets[chainOffsets.length] = i as Element;
+        if (newPolymer) polymerOffsets[polymerOffsets.length] = i as ElementIndex;
+        if (newChain) chainOffsets[chainOffsets.length] = i as ElementIndex;
     }
 
     return {

src/mol-model/structure/model/formats/mmcif/pair-restraints/cross-links.ts

@@ -8,17 +8,18 @@ import { Model } from '../../../model'
 import { Table } from 'mol-data/db'
 import { mmCIF_Schema } from 'mol-io/reader/cif/schema/mmcif';
 import { findAtomIndexByLabelName } from '../util';
-import { Element, Unit } from '../../../../structure';
+import { Unit } from '../../../../structure';
+import { ElementIndex } from '../../../indexing';
 
 function findAtomIndex(model: Model, entityId: string, asymId: string, compId: string, seqId: number, atomId: string) {
     if (!model.atomicHierarchy.atoms.auth_atom_id.isDefined) return -1
     const residueIndex = model.atomicHierarchy.findResidueKey(entityId, compId, asymId, seqId, '')
     if (residueIndex < 0) return -1
-    return findAtomIndexByLabelName(model, residueIndex, atomId, '') as Element
+    return findAtomIndexByLabelName(model, residueIndex, atomId, '') as ElementIndex
 }
 
 export interface IHMCrossLinkRestraint {
-    getIndicesByElement: (element: Element, kind: Unit.Kind) => number[]
+    getIndicesByElement: (element: ElementIndex, kind: Unit.Kind) => number[]
     data: Table<mmCIF_Schema['ihm_cross_link_restraint']>
 }
 
@@ -47,7 +48,7 @@ export namespace IHMCrossLinkRestraint {
             atom_id: ihm_cross_link_restraint.atom_id_2,
         }
 
-        function _add(map: Map<Element, number[]>, element: Element, row: number) {
+        function _add(map: Map<ElementIndex, number[]>, element: ElementIndex, row: number) {
             const indices = map.get(element)
             if (indices) indices.push(row)
             else map.set(element, [ row ])
@@ -60,14 +61,14 @@ export namespace IHMCrossLinkRestraint {
 
             if (ihm_cross_link_restraint.model_granularity.value(row) === 'by-atom') {
                 const atomicElement = findAtomIndex(model, entityId, asymId, ps.comp_id.value(row), seqId, ps.atom_id.value(row))
-                if (atomicElement >= 0) _add(atomicElementMap, atomicElement as Element, row)
+                if (atomicElement >= 0) _add(atomicElementMap, atomicElement as ElementIndex, row)
             } else if (model.coarseHierarchy.isDefined) {
                 const sphereElement = model.coarseHierarchy.spheres.findSequenceKey(entityId, asymId, seqId)
                 if (sphereElement >= 0) {
-                    _add(sphereElementMap, sphereElement as Element, row)
+                    _add(sphereElementMap, sphereElement, row)
                 } else {
                     const gaussianElement = model.coarseHierarchy.gaussians.findSequenceKey(entityId, asymId, seqId)
-                    if (gaussianElement >= 0) _add(gaussianElementMap, gaussianElement as Element, row)
+                    if (gaussianElement >= 0) _add(gaussianElementMap, gaussianElement, row)
                 }
             }
         }
@@ -81,11 +82,11 @@ export namespace IHMCrossLinkRestraint {
         }
 
         /** map from atomic element to cross link indices */
-        const atomicElementMap: Map<Element, number[]> = new Map()
+        const atomicElementMap: Map<ElementIndex, number[]> = new Map()
         /** map from sphere element to cross link indices */
-        const sphereElementMap: Map<Element, number[]> = new Map()
+        const sphereElementMap: Map<ElementIndex, number[]> = new Map()
         /** map from gaussian element to cross link indices */
-        const gaussianElementMap: Map<Element, number[]> = new Map()
+        const gaussianElementMap: Map<ElementIndex, number[]> = new Map()
 
         const emptyIndexArray: number[] = [];
 
@@ -95,7 +96,7 @@ export namespace IHMCrossLinkRestraint {
         }
 
         const crossLinkRestraint = {
-            getIndicesByElement: (element: Element, kind: Unit.Kind) => {
+            getIndicesByElement: (element: ElementIndex, kind: Unit.Kind) => {
                 const map = getMapByKind(kind)
                 const idx = map.get(element)
                 return idx !== undefined ? idx : emptyIndexArray

src/mol-model/structure/model/formats/mmcif/util.ts

@@ -5,7 +5,7 @@
  */
 
 import { Model } from '../../model'
-import { Element } from '../../../structure'
+import { ElementIndex } from '../../indexing';
 
 export function findEntityIdByAsymId(model: Model, asymId: string) {
     if (model.sourceData.kind !== 'mmCIF') return ''
@@ -16,11 +16,11 @@ export function findEntityIdByAsymId(model: Model, asymId: string) {
     return ''
 }
 
-export function findAtomIndexByLabelName(model: Model, residueIndex: number, atomName: string, altLoc: string | null): Element {
+export function findAtomIndexByLabelName(model: Model, residueIndex: number, atomName: string, altLoc: string | null): ElementIndex {
     const { offsets } = model.atomicHierarchy.residueAtomSegments;
     const { label_atom_id, label_alt_id } = model.atomicHierarchy.atoms;
     for (let i = offsets[residueIndex], n = offsets[residueIndex + 1]; i < n; ++i) {
-        if (label_atom_id.value(i) === atomName && (!altLoc || label_alt_id.value(i) === altLoc)) return i as Element;
+        if (label_atom_id.value(i) === atomName && (!altLoc || label_alt_id.value(i) === altLoc)) return i as ElementIndex;
     }
-    return -1 as Element;
+    return -1 as ElementIndex;
 }

src/mol-model/structure/model/properties/atomic/hierarchy.ts

@@ -8,8 +8,8 @@ import { Column, Table } from 'mol-data/db'
 import { Segmentation } from 'mol-data/int'
 import { mmCIF_Schema as mmCIF } from 'mol-io/reader/cif/schema/mmcif'
 import { ElementSymbol } from '../../types'
-import { Element } from '../../../structure'
-import { ChainIndex, EntityIndex, ResidueIndex } from '../../indexing';
+// import { Element } from '../../../structure'
+import { ChainIndex, EntityIndex, ResidueIndex, ElementIndex } from '../../indexing';
 
 export const AtomsSchema = {
     type_symbol: Column.Schema.Aliased<ElementSymbol>(mmCIF.atom_site.type_symbol),
@@ -50,7 +50,7 @@ export interface AtomicData {
 
 export interface AtomicSegments {
     /** Maps residueIndex to a range of atoms [segments[rI], segments[rI + 1]) */
-    residueAtomSegments: Segmentation<Element, ResidueIndex>,
+    residueAtomSegments: Segmentation<ElementIndex, ResidueIndex>,
     /**
      * Maps chainIndex to a range of atoms [segments[cI], segments[cI + 1]),
      *
@@ -59,12 +59,12 @@ export interface AtomicSegments {
      * const start = rI[offsets[i]], const end = rI[offsets[i + 1] - 1] + 1;
      * for (let j = start; j < end; i++) { }
      */
-    chainAtomSegments: Segmentation<Element, ChainIndex>,
+    chainAtomSegments: Segmentation<ElementIndex, ChainIndex>,
     /**
      * bonded/connected stretches of polymer chains, i.e. a chain will be
      * broken into multiple polymer segments if there are missing residues
      */
-    polymerAtomSegments: Segmentation<Element>
+    polymerAtomSegments: Segmentation<ElementIndex>
 
     // TODO: include entity segments?
 }

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

@@ -7,7 +7,7 @@
 
 import { Column } from 'mol-data/db'
 import { Segmentation } from 'mol-data/int';
-import { Element } from '../../../structure'
+import { ElementIndex, ChainIndex } from '../../indexing';
 
 export interface CoarsedElementKeys {
     // assign a key to each element
@@ -16,8 +16,8 @@ export interface CoarsedElementKeys {
     entityKey: ArrayLike<number>,
 
     /** find index of the residue/feature element where seq_id is included */
-    findSequenceKey(entityId: string, asym_id: string, seq_id: number): number
-    findChainKey(entityId: string, asym_id: string): number
+    findSequenceKey(entityId: string, asym_id: string, seq_id: number): ElementIndex
+    findChainKey(entityId: string, asym_id: string): ChainIndex
 }
 
 export interface CoarseElementData {
@@ -27,12 +27,12 @@ export interface CoarseElementData {
     seq_id_begin: Column<number>,
     seq_id_end: Column<number>,
 
-    chainSegments: Segmentation<Element>,
+    chainSegments: Segmentation<ElementIndex>,
     /**
      * bonded/connected stretches of polymer chains, i.e. a chain will be
      * broken into multiple polymer segments if there are missing residues
      */
-    polymerSegments: Segmentation<Element>
+    polymerSegments: Segmentation<ElementIndex>
 }
 
 export type CoarseElements = CoarsedElementKeys & CoarseElementData

src/mol-model/structure/model/properties/utils/coarse-keys.ts

@@ -7,6 +7,7 @@
 
 import { Entities } from '../common';
 import { CoarseElementData, CoarsedElementKeys } from '../coarse';
+import { ChainIndex, ElementIndex } from '../../indexing';
 
 function getElementKey(map: Map<string, number>, key: string, counter: { index: number }) {
     if (map.has(key)) return map.get(key)!;
@@ -26,22 +27,22 @@ function createLookUp(entities: Entities, chain: Map<number, Map<string, number>
     const getEntKey = entities.getEntityIndex;
     const findChainKey: CoarsedElementKeys['findChainKey'] = (e, c) => {
         const eKey = getEntKey(e);
-        if (eKey < 0) return -1;
+        if (eKey < 0) return -1 as ChainIndex;
         const cm = chain.get(eKey)!;
-        if (!cm.has(c)) return -1;
-        return cm.get(c)!;
+        if (!cm.has(c)) return -1 as ChainIndex;
+        return cm.get(c)! as ChainIndex;
     }
     // TODO consider implementing as binary search
     const findSequenceKey: CoarsedElementKeys['findSequenceKey'] = (e, c, s) => {
         const eKey = getEntKey(e);
-        if (eKey < 0) return -1;
+        if (eKey < 0) return -1 as ElementIndex;
         const cm = chain.get(eKey);
-        if (cm === undefined) return -1
+        if (cm === undefined) return -1 as ElementIndex
         const cKey = cm.get(c)
-        if (cKey === undefined) return -1
+        if (cKey === undefined) return -1 as ElementIndex
         const sm = seq.get(cKey)!
-        if (!sm.has(s)) return -1;
-        return sm.get(s)!
+        if (!sm.has(s)) return -1 as ElementIndex;
+        return sm.get(s)! as ElementIndex
     }
     return { findChainKey, findSequenceKey };
 }

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

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

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

@@ -6,12 +6,13 @@
 
 import { OrderedSet, SortedArray } from 'mol-data/int'
 import Unit from './unit'
+import { ElementIndex } from '../model';
 
 /** Element index in Model */
-type Element = { readonly '@type': 'element' } & number
+// type Element = { readonly '@type': 'element' } & number
 
 namespace Element {
-    export type Set = SortedArray<Element>
+    export type Set = SortedArray<ElementIndex>
 
     /** Index into Unit.elements */
     export type Index = { readonly '@type': 'element-index' } & number
@@ -20,9 +21,9 @@ namespace Element {
     export interface Location<U = Unit> {
         unit: U,
         /** Index into element (atomic/coarse) properties of unit.model */
-        element: Element
+        element: ElementIndex
     }
-    export function Location(unit?: Unit, element?: Element): Location { return { unit: unit as any, element: element || (0 as Element) }; }
+    export function Location(unit?: Unit, element?: ElementIndex): Location { return { unit: unit as any, element: element || (0 as ElementIndex) }; }
     export interface Property<T> { (location: Location): T }
     export interface Predicate extends Property<boolean> { }
 

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

@@ -7,7 +7,7 @@
 import { IntMap, SortedArray, Iterator, Segmentation } from 'mol-data/int'
 import { UniqueArray } from 'mol-data/generic'
 import { SymmetryOperator } from 'mol-math/geometry/symmetry-operator'
-import { Model } from '../model'
+import { Model, ElementIndex } from '../model'
 import { sort, arraySwap, hash1, sortArray } from 'mol-data/util';
 import Element from './element'
 import Unit from './unit'
@@ -131,7 +131,7 @@ namespace Structure {
                 c++;
             }
 
-            const elements = SortedArray.ofBounds(start as Element, chains.offsets[c + 1] as Element);
+            const elements = SortedArray.ofBounds(start as ElementIndex, chains.offsets[c + 1] as ElementIndex);
             builder.addUnit(Unit.Kind.Atomic, model, SymmetryOperator.Default, elements);
         }
 
@@ -151,7 +151,7 @@ namespace Structure {
     function addCoarseUnits(builder: StructureBuilder, model: Model, elements: CoarseElements, kind: Unit.Kind) {
         const { chainSegments } = elements;
         for (let cI = 0; cI < chainSegments.count; cI++) {
-            const elements = SortedArray.ofBounds(chainSegments.offsets[cI] as Element, chainSegments.offsets[cI + 1] as Element);
+            const elements = SortedArray.ofBounds<ElementIndex>(chainSegments.offsets[cI], chainSegments.offsets[cI + 1]);
             builder.addUnit(kind, model, SymmetryOperator.Default, elements);
         }
     }

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

@@ -10,15 +10,16 @@ import Element from '../element';
 import StructureSymmetry from '../symmetry';
 import Unit from '../unit';
 import Structure from '../structure';
+import { ElementIndex } from '../../model';
 
 export class StructureSubsetBuilder {
     private ids: number[] = [];
-    private unitMap = IntMap.Mutable<Element[]>();
+    private unitMap = IntMap.Mutable<ElementIndex[]>();
     private parentId = -1;
-    private currentUnit: Element[] = [];
+    private currentUnit: ElementIndex[] = [];
     elementCount = 0;
 
-    addToUnit(parentId: number, e: Element) {
+    addToUnit(parentId: number, e: ElementIndex) {
         const unit = this.unitMap.get(parentId);
         if (!!unit) { unit[unit.length] = e; }
         else {
@@ -33,7 +34,7 @@ export class StructureSubsetBuilder {
         this.currentUnit = this.currentUnit.length > 0 ? [] : this.currentUnit;
     }
 
-    addElement(e: Element) {
+    addElement(e: ElementIndex) {
         this.currentUnit[this.currentUnit.length] = e;
         this.elementCount++;
     }
@@ -45,9 +46,9 @@ export class StructureSubsetBuilder {
         this.parentId = -1;
     }
 
-    setUnit(parentId: number, elements: ArrayLike<Element>) {
+    setUnit(parentId: number, elements: ArrayLike<ElementIndex>) {
         this.ids[this.ids.length] = parentId;
-        this.unitMap.set(parentId, elements as Element[]);
+        this.unitMap.set(parentId, elements as ElementIndex[]);
         this.elementCount += elements.length;
     }