use unit.chainGroupId in sequence widget

src/mol-plugin/ui/sequence.tsx

@@ -38,17 +38,25 @@ function splitModelEntityId(modelEntityId: string) {
 }
 
 function getSequenceWrapper(state: SequenceViewState, structureSelection: StructureElementSelectionManager): SequenceWrapper.Any | string {
-    const { structure, modelEntityId, invariantUnitId, operatorKey } = state
+    const { structure, modelEntityId, chainGroupId, operatorKey } = state
     const l = StructureElement.Location.create()
     const [ modelIdx, entityId ] = splitModelEntityId(modelEntityId)
+
+    const units: Unit[] = []
+
     for (const unit of structure.units) {
         StructureElement.Location.set(l, unit, unit.elements[0])
         if (structure.getModelIndex(unit.model) !== modelIdx) continue
         if (SP.entity.id(l) !== entityId) continue
-        if (unit.invariantId !== invariantUnitId) continue
+        if (unit.chainGroupId !== chainGroupId) continue
         if (opKey(l) !== operatorKey) continue
 
-        const data = { structure, unit }
+        units.push(unit)
+    }
+
+    if (units.length > 0) {
+        const data = { structure, units }
+        const unit = units[0]
 
         let sw: SequenceWrapper<any>
         if (unit.polymerElements.length) {
@@ -77,8 +85,9 @@ function getSequenceWrapper(state: SequenceViewState, structureSelection: Struct
 
         sw.markResidue(structureSelection.get(structure), MarkerAction.Select)
         return sw
+    } else {
+        return 'No sequence available'
     }
-    return 'No sequence available'
 }
 
 function getModelEntityOptions(structure: Structure) {
@@ -106,11 +115,10 @@ function getModelEntityOptions(structure: Structure) {
     return options
 }
 
-function getUnitOptions(structure: Structure, modelEntityId: string) {
+function getChainOptions(structure: Structure, modelEntityId: string) {
     const options: [number, string][] = []
     const l = StructureElement.Location.create()
     const seen = new Set<number>()
-    const water = new Map<string, number>()
     const [ modelIdx, entityId ] = splitModelEntityId(modelEntityId)
 
     for (const unit of structure.units) {
@@ -118,18 +126,12 @@ function getUnitOptions(structure: Structure, modelEntityId: string) {
         if (structure.getModelIndex(unit.model) !== modelIdx) continue
         if (SP.entity.id(l) !== entityId) continue
 
-        const id = unit.invariantId
+        const id = unit.chainGroupId
         if (seen.has(id)) continue
 
-        // TODO handle special cases
+        // TODO handle special case
         // - more than one chain in a unit
-        // - chain spread over multiple units
         let label = elementLabel(l, { granularity: 'chain', hidePrefix: true, htmlStyling: false })
-        if (SP.entity.type(l) === 'water') {
-            const count = water.get(label) || 1
-            water.set(label, count + 1)
-            label += ` #${count}`
-        }
 
         options.push([ id, label ])
         seen.add(id)
@@ -139,7 +141,7 @@ function getUnitOptions(structure: Structure, modelEntityId: string) {
     return options
 }
 
-function getOperatorOptions(structure: Structure, modelEntityId: string, invariantUnitId: number) {
+function getOperatorOptions(structure: Structure, modelEntityId: string, chainGroupId: number) {
     const options: [string, string][] = []
     const l = StructureElement.Location.create()
     const seen = new Set<string>()
@@ -149,7 +151,7 @@ function getOperatorOptions(structure: Structure, modelEntityId: string, invaria
         StructureElement.Location.set(l, unit, unit.elements[0])
         if (structure.getModelIndex(unit.model) !== modelIdx) continue
         if (SP.entity.id(l) !== entityId) continue
-        if (unit.invariantId !== invariantUnitId) continue
+        if (unit.chainGroupId !== chainGroupId) continue
 
         const id = opKey(l)
         if (seen.has(id)) continue
@@ -179,12 +181,12 @@ type SequenceViewState = {
     structure: Structure,
     structureRef: string,
     modelEntityId: string,
-    invariantUnitId: number,
+    chainGroupId: number,
     operatorKey: string
 }
 
 export class SequenceView extends PluginUIComponent<{ }, SequenceViewState> {
-    state = { structure: Structure.Empty, structureRef: '', modelEntityId: '', invariantUnitId: -1, operatorKey: '' }
+    state = { structure: Structure.Empty, structureRef: '', modelEntityId: '', chainGroupId: -1, operatorKey: '' }
 
     componentDidMount() {
         if (this.plugin.state.dataState.select(StateSelection.Generators.rootsOfType(PSO.Molecule.Structure)).length > 0) this.setState(this.getInitialState())
@@ -223,26 +225,26 @@ export class SequenceView extends PluginUIComponent<{ }, SequenceViewState> {
         const structureRef = getStructureOptions(this.plugin.state.dataState)[0][0]
         const structure = this.getStructure(structureRef)
         let modelEntityId = getModelEntityOptions(structure)[0][0]
-        let invariantUnitId = getUnitOptions(structure, modelEntityId)[0][0]
-        let operatorKey = getOperatorOptions(structure, modelEntityId, invariantUnitId)[0][0]
+        let chainGroupId = getChainOptions(structure, modelEntityId)[0][0]
+        let operatorKey = getOperatorOptions(structure, modelEntityId, chainGroupId)[0][0]
         if (this.state.structure && this.state.structure === structure) {
             modelEntityId = this.state.modelEntityId
-            invariantUnitId = this.state.invariantUnitId
+            chainGroupId = this.state.chainGroupId
             operatorKey = this.state.operatorKey
         }
-        return { structure, structureRef, modelEntityId, invariantUnitId, operatorKey }
+        return { structure, structureRef, modelEntityId, chainGroupId, operatorKey }
     }
 
     private get params() {
-        const { structure, modelEntityId, invariantUnitId } = this.state
+        const { structure, modelEntityId, chainGroupId } = this.state
         const structureOptions = getStructureOptions(this.plugin.state.dataState)
         const entityOptions = getModelEntityOptions(structure)
-        const unitOptions = getUnitOptions(structure, modelEntityId)
-        const operatorOptions = getOperatorOptions(structure, modelEntityId, invariantUnitId)
+        const chainOptions = getChainOptions(structure, modelEntityId)
+        const operatorOptions = getOperatorOptions(structure, modelEntityId, chainGroupId)
         return {
             structure: PD.Select(structureOptions[0][0], structureOptions, { shortLabel: true }),
             entity: PD.Select(entityOptions[0][0], entityOptions, { shortLabel: true }),
-            unit: PD.Select(unitOptions[0][0], unitOptions, { shortLabel: true, twoColumns: true, label: 'Chain' }),
+            chain: PD.Select(chainOptions[0][0], chainOptions, { shortLabel: true, twoColumns: true, label: 'Chain' }),
             operator: PD.Select(operatorOptions[0][0], operatorOptions, { shortLabel: true, twoColumns: true })
         }
     }
@@ -251,7 +253,7 @@ export class SequenceView extends PluginUIComponent<{ }, SequenceViewState> {
         return {
             structure: this.state.structureRef,
             entity: this.state.modelEntityId,
-            unit: this.state.invariantUnitId,
+            chain: this.state.chainGroupId,
             operator: this.state.operatorKey
         }
     }
@@ -263,17 +265,17 @@ export class SequenceView extends PluginUIComponent<{ }, SequenceViewState> {
                 state.structureRef = p.value
                 state.structure = this.getStructure(p.value)
                 state.modelEntityId = getModelEntityOptions(state.structure)[0][0]
-                state.invariantUnitId = getUnitOptions(state.structure, state.modelEntityId)[0][0]
-                state.operatorKey = getOperatorOptions(state.structure, state.modelEntityId, state.invariantUnitId)[0][0]
+                state.chainGroupId = getChainOptions(state.structure, state.modelEntityId)[0][0]
+                state.operatorKey = getOperatorOptions(state.structure, state.modelEntityId, state.chainGroupId)[0][0]
                 break
             case 'entity':
                 state.modelEntityId = p.value
-                state.invariantUnitId = getUnitOptions(state.structure, state.modelEntityId)[0][0]
-                state.operatorKey = getOperatorOptions(state.structure, state.modelEntityId, state.invariantUnitId)[0][0]
+                state.chainGroupId = getChainOptions(state.structure, state.modelEntityId)[0][0]
+                state.operatorKey = getOperatorOptions(state.structure, state.modelEntityId, state.chainGroupId)[0][0]
                 break
-            case 'unit':
-                state.invariantUnitId = p.value
-                state.operatorKey = getOperatorOptions(state.structure, state.modelEntityId, state.invariantUnitId)[0][0]
+            case 'chain':
+                state.chainGroupId = p.value
+                state.operatorKey = getOperatorOptions(state.structure, state.modelEntityId, state.chainGroupId)[0][0]
                 break
             case 'operator':
                 state.operatorKey = p.value

src/mol-plugin/ui/sequence/chain.ts

@@ -12,7 +12,7 @@ import { ColorNames } from '../../../mol-util/color/names';
 
 export class ChainSequenceWrapper extends SequenceWrapper<StructureUnit> {
     private label: string
-    private indices: Interval<StructureElement.UnitIndex>
+    private unitIndices: Map<number, Interval<StructureElement.UnitIndex>>
     private loci: StructureElement.Loci
 
     residueLabel(seqIdx: number) {
@@ -24,14 +24,15 @@ export class ChainSequenceWrapper extends SequenceWrapper<StructureUnit> {
 
     eachResidue(loci: Loci, apply: (set: OrderedSet) => boolean) {
         let changed = false
-        const { structure, unit } = this.data
+        const { structure } = this.data
         if (StructureElement.Loci.is(loci)) {
             if (!Structure.areRootsEquivalent(loci.structure, structure)) return false
             loci = StructureElement.Loci.remap(loci, structure)
 
             for (const e of loci.elements) {
-                if (e.unit.id === unit.id) {
-                    if (OrderedSet.isSubset(this.indices, e.indices)) {
+                const indices = this.unitIndices.get(e.unit.id)
+                if (indices) {
+                    if (OrderedSet.isSubset(indices, e.indices)) {
                         if (apply(Interval.ofSingleton(0))) changed = true
                     }
                 }
@@ -49,21 +50,36 @@ export class ChainSequenceWrapper extends SequenceWrapper<StructureUnit> {
     }
 
     constructor(data: StructureUnit) {
+        let residueCount = 0
+        let elementCount = 0
         const counts: string[] = []
-        const l = StructureElement.Location.create(data.unit, data.unit.elements[0])
-        const entitySeq = data.unit.model.sequence.byEntityKey[StructureProperties.entity.key(l)]
-        if (entitySeq) counts.push(`${entitySeq.sequence.length} residues`)
-        counts.push(`${data.unit.elements.length} elements`)
+        const l = StructureElement.Location.create()
+
+        const unitIndices = new Map<number, Interval<StructureElement.UnitIndex>>()
+        const lociElements: StructureElement.Loci['elements'][0][] = []
+
+        for (let i = 0, il = data.units.length; i < il; ++i) {
+            const unit = data.units[i]
+            StructureElement.Location.set(l, unit, unit.elements[0])
+            const entitySeq = unit.model.sequence.byEntityKey[StructureProperties.entity.key(l)]
+            if (entitySeq) residueCount += entitySeq.sequence.length
+            elementCount += unit.elements.length
+
+            const indices = Interval.ofBounds(0, unit.elements.length)
+            unitIndices.set(unit.id, indices)
+            lociElements.push({ unit, indices })
+        }
+
+        if (residueCount > 0) counts.push(`${residueCount} residues`)
+        counts.push(`${elementCount} elements`)
 
         const length = 1
         const markerArray = new Uint8Array(length)
 
         super(data, markerArray, length)
 
-        this.label = `Whole Unit (${counts.join(', ')})`
-        this.indices = Interval.ofBounds(0, data.unit.elements.length)
-        this.loci = StructureElement.Loci(this.data.structure, [{
-            unit: this.data.unit, indices: this.indices
-        }])
+        this.label = `Whole Chain (${counts.join(', ')})`
+        this.unitIndices = unitIndices
+        this.loci = StructureElement.Loci(this.data.structure, lociElements)
     }
 }

src/mol-plugin/ui/sequence/element.ts

@@ -11,7 +11,7 @@ import { Loci } from '../../../mol-model/loci';
 import { ColorNames } from '../../../mol-util/color/names';
 
 export class ElementSequenceWrapper extends SequenceWrapper<StructureUnit> {
-    private indices: Interval<StructureElement.UnitIndex>
+    private unitIndices: Map<number, Interval<StructureElement.UnitIndex>>
 
     residueLabel(seqIdx: number) {
         return 'X'
@@ -22,14 +22,15 @@ export class ElementSequenceWrapper extends SequenceWrapper<StructureUnit> {
 
     eachResidue(loci: Loci, apply: (set: OrderedSet) => boolean) {
         let changed = false
-        const { structure, unit } = this.data
+        const { structure, units } = this.data
         if (StructureElement.Loci.is(loci)) {
             if (!Structure.areRootsEquivalent(loci.structure, structure)) return false
             loci = StructureElement.Loci.remap(loci, structure)
 
             for (const e of loci.elements) {
-                if (e.unit.id === unit.id) {
-                    if (OrderedSet.isSubset(this.indices, e.indices)) {
+                const indices = this.unitIndices.get(e.unit.id)
+                if (indices) {
+                    if (OrderedSet.isSubset(indices, e.indices)) {
                         if (apply(e.indices)) changed = true
                     }
                 }
@@ -37,25 +38,47 @@ export class ElementSequenceWrapper extends SequenceWrapper<StructureUnit> {
         } else if (Structure.isLoci(loci)) {
             if (!Structure.areRootsEquivalent(loci.structure, structure)) return false
 
-            if (apply(this.indices)) changed = true
+            for (let i = 0, il = units.length; i < il; ++i) {
+                const indices = this.unitIndices.get(units[i].id)!
+                if (apply(indices)) changed = true
+            }
         }
         return changed
     }
 
     getLoci(seqIdx: number) {
-        const elements: StructureElement.Loci['elements'][0][] = [{
-            unit: this.data.unit,
-            indices: Interval.ofSingleton(seqIdx)
-        }]
-        return StructureElement.Loci(this.data.structure, elements)
+        const { units } = this.data
+        const lociElements: StructureElement.Loci['elements'][0][] = []
+        let offset = 0
+        for (let i = 0, il = units.length; i < il; ++i) {
+            const unit = units[i]
+            if (seqIdx < offset + unit.elements.length) {
+                lociElements.push({ unit, indices: Interval.ofSingleton(seqIdx - offset) })
+                break
+            }
+            offset += unit.elements.length
+        }
+        return StructureElement.Loci(this.data.structure, lociElements)
     }
 
     constructor(data: StructureUnit) {
-        const length = data.unit.elements.length
+        let length = 0
+
+        const unitIndices = new Map<number, Interval<StructureElement.UnitIndex>>()
+        const lociElements: StructureElement.Loci['elements'][0][] = []
+
+        for (let i = 0, il = data.units.length; i < il; ++i) {
+            const unit = data.units[i]
+            length += unit.elements.length
+
+            const indices = Interval.ofBounds(0, unit.elements.length)
+            unitIndices.set(unit.id, indices)
+            lociElements.push({ unit, indices })
+        }
         const markerArray = new Uint8Array(length)
 
         super(data, markerArray, length)
 
-        this.indices = Interval.ofBounds(0, length)
+        this.unitIndices = unitIndices
     }
 }

src/mol-plugin/ui/sequence/hetero.ts

@@ -4,16 +4,18 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { Structure, StructureElement, ResidueIndex } from '../../../mol-model/structure';
+import { Structure, StructureElement, ResidueIndex, Unit } from '../../../mol-model/structure';
 import { SequenceWrapper, StructureUnit } from './wrapper';
 import { OrderedSet, Segmentation, Interval, SortedArray } from '../../../mol-data/int';
 import { Loci } from '../../../mol-model/loci';
 import { ColorNames } from '../../../mol-util/color/names';
 
 export class HeteroSequenceWrapper extends SequenceWrapper<StructureUnit> {
+    private readonly unitMap: Map<number, Unit>
     private readonly sequence: string[]
     private readonly sequenceIndices: Map<ResidueIndex, number>
     private readonly residueIndices: Map<number, ResidueIndex>
+    private readonly seqToUnit: Map<number, Unit>
 
     residueLabel(seqIdx: number) {
         return this.sequence[seqIdx]
@@ -24,13 +26,14 @@ export class HeteroSequenceWrapper extends SequenceWrapper<StructureUnit> {
 
     eachResidue(loci: Loci, apply: (set: OrderedSet) => boolean) {
         let changed = false
-        const { structure, unit } = this.data
+        const { structure } = this.data
         if (StructureElement.Loci.is(loci)) {
             if (!Structure.areRootsEquivalent(loci.structure, structure)) return false
             loci = StructureElement.Loci.remap(loci, structure)
 
             for (const e of loci.elements) {
-                if (e.unit.id === unit.id) {
+                const unit = this.unitMap.get(e.unit.id)
+                if (unit) {
                     const { index: residueIndex } = e.unit.model.atomicHierarchy.residueAtomSegments
                     OrderedSet.forEach(e.indices, v => {
                         const seqIdx = this.sequenceIndices.get(residueIndex[unit.elements[v]])
@@ -50,7 +53,7 @@ export class HeteroSequenceWrapper extends SequenceWrapper<StructureUnit> {
         const elements: StructureElement.Loci['elements'][0][] = []
         const rI = this.residueIndices.get(seqIdx)
         if (rI !== undefined) {
-            const { unit } = this.data
+            const unit = this.seqToUnit.get(seqIdx)!
             const { offsets } = unit.model.atomicHierarchy.residueAtomSegments
             const start = SortedArray.findPredecessorIndex(unit.elements, offsets[rI])
             const end = SortedArray.findPredecessorIndex(unit.elements, offsets[rI + 1])
@@ -63,13 +66,19 @@ export class HeteroSequenceWrapper extends SequenceWrapper<StructureUnit> {
         const sequence: string[] = []
         const sequenceIndices = new Map<ResidueIndex, number>()
         const residueIndices = new Map<number, ResidueIndex>()
+        const seqToUnit = new Map<number, Unit>()
 
-        const residueIt = Segmentation.transientSegments(data.unit.model.atomicHierarchy.residueAtomSegments, data.unit.elements)
-        while (residueIt.hasNext) {
-            const { index } = residueIt.move()
-            sequenceIndices.set(index, sequence.length)
-            residueIndices.set(sequence.length, index)
-            sequence.push(data.unit.model.atomicHierarchy.residues.label_comp_id.value(index))
+        for (let i = 0, il = data.units.length; i < il; ++i) {
+            const unit = data.units[i]
+            const { residueAtomSegments, residues } = unit.model.atomicHierarchy
+            const residueIt = Segmentation.transientSegments(residueAtomSegments, unit.elements)
+            while (residueIt.hasNext) {
+                const { index } = residueIt.move()
+                sequenceIndices.set(index, sequence.length)
+                residueIndices.set(sequence.length, index)
+                seqToUnit.set(sequence.length, unit)
+                sequence.push(residues.label_comp_id.value(index))
+            }
         }
 
         const length = sequence.length
@@ -77,8 +86,12 @@ export class HeteroSequenceWrapper extends SequenceWrapper<StructureUnit> {
 
         super(data, markerArray, length)
 
+        this.unitMap = new Map()
+        for (const unit of data.units) this.unitMap.set(unit.id, unit)
+
         this.sequence = sequence
         this.sequenceIndices = sequenceIndices
         this.residueIndices = residueIndices
+        this.seqToUnit = seqToUnit
     }
 }

src/mol-plugin/ui/sequence/polymer.ts

@@ -13,6 +13,7 @@ import { MissingResidues } from '../../../mol-model/structure/model/properties/c
 import { ColorNames } from '../../../mol-util/color/names';
 
 export class PolymerSequenceWrapper extends SequenceWrapper<StructureUnit> {
+    private readonly unitMap: Map<number, Unit>
     private readonly sequence: Sequence
     private readonly missing: MissingResidues
     private readonly observed: OrderedSet // sequences indices
@@ -35,14 +36,14 @@ export class PolymerSequenceWrapper extends SequenceWrapper<StructureUnit> {
 
     eachResidue(loci: Loci, apply: (set: OrderedSet) => boolean) {
         let changed = false
-        const { structure, unit } = this.data
+        const { structure } = this.data
         if (StructureElement.Loci.is(loci)) {
             if (!Structure.areRootsEquivalent(loci.structure, structure)) return false
             loci = StructureElement.Loci.remap(loci, structure)
 
             const { offset } = this.sequence
             for (const e of loci.elements) {
-                if (e.unit.id === unit.id) {
+                if (this.unitMap.has(e.unit.id)) {
                     OrderedSet.forEach(e.indices, v => {
                         if (apply(getSeqIndices(e.unit, e.unit.elements[v], offset))) changed = true
                     })
@@ -57,24 +58,27 @@ export class PolymerSequenceWrapper extends SequenceWrapper<StructureUnit> {
     }
 
     getLoci(seqIdx: number) {
-        const query = createResidueQuery(this.data.unit.id, this.seqId(seqIdx));
+        const query = createResidueQuery(this.data.units[0].chainGroupId, this.seqId(seqIdx));
         return StructureSelection.toLociWithSourceUnits(StructureQuery.run(query, this.data.structure));
     }
 
     constructor(data: StructureUnit) {
-        const l = StructureElement.Location.create(data.unit, data.unit.elements[0])
-        const entitySeq = data.unit.model.sequence.byEntityKey[SP.entity.key(l)]
+        const l = StructureElement.Location.create(data.units[0], data.units[0].elements[0])
+        const entitySeq = data.units[0].model.sequence.byEntityKey[SP.entity.key(l)]
 
         const length = entitySeq.sequence.length
         const markerArray = new Uint8Array(length)
 
         super(data, markerArray, length)
 
+        this.unitMap = new Map()
+        for (const unit of data.units) this.unitMap.set(unit.id, unit)
+
         this.sequence = entitySeq.sequence
-        this.missing = data.unit.model.properties.missingResidues
+        this.missing = data.units[0].model.properties.missingResidues
 
-        this.modelNum = data.unit.model.modelNum
-        this.asymId = Unit.isAtomic(data.unit) ? SP.chain.label_asym_id(l) : SP.coarse.asym_id(l)
+        this.modelNum = data.units[0].model.modelNum
+        this.asymId = Unit.isAtomic(data.units[0]) ? SP.chain.label_asym_id(l) : SP.coarse.asym_id(l)
 
         const missing: number[] = []
         for (let i = 0; i < length; ++i) {
@@ -84,10 +88,10 @@ export class PolymerSequenceWrapper extends SequenceWrapper<StructureUnit> {
     }
 }
 
-function createResidueQuery(unitId: number, label_seq_id: number) {
+function createResidueQuery(chainGroupId: number, label_seq_id: number) {
     return Queries.generators.atoms({
         unitTest: ctx => {
-            return SP.unit.id(ctx.element) === unitId
+            return SP.unit.chainGroupId(ctx.element) === chainGroupId
         },
         residueTest: ctx => {
             if (ctx.element.unit.kind === Unit.Kind.Atomic) {

src/mol-plugin/ui/sequence/wrapper.ts

@@ -10,7 +10,7 @@ import { MarkerAction, applyMarkerAction } from '../../../mol-util/marker-action
 import { StructureElement, Structure, Unit } from '../../../mol-model/structure';
 import { Color } from '../../../mol-util/color';
 
-export type StructureUnit = { structure: Structure, unit: Unit }
+export type StructureUnit = { structure: Structure, units: Unit[] }
 
 export { SequenceWrapper }