loci remapping

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

@@ -141,10 +141,29 @@ namespace StructureElement {
         export function remap(loci: Loci, structure: Structure): Loci {
             if (structure === loci.structure) return loci
 
-            return Loci(structure, loci.elements.map(e => ({
-                unit: structure.unitMap.get(e.unit.id)!,
-                indices: e.indices
-            })));
+            const elements: Loci['elements'][0][] = [];
+            loci.elements.forEach(e => {
+                const unit = structure.unitMap.get(e.unit.id)
+                if (!unit) return
+
+                if (SortedArray.areEqual(e.unit.elements, unit.elements)) {
+                    elements.push({ unit, indices: e.indices })
+                } else {
+                    // TODO optimize
+                    const indices: UnitIndex[] = []
+                    OrderedSet.forEach(e.indices, (v) => {
+                        const eI = e.unit.elements[v]
+                        const uI = SortedArray.indexOf(unit.elements, eI) as UnitIndex | -1
+                        if (uI !== -1) indices.push(uI)
+                    })
+                    elements.push({
+                        unit,
+                        indices: SortedArray.ofSortedArray(indices)
+                    })
+                }
+            });
+
+            return Loci(structure, elements);
         }
 
         export function union(xs: Loci, ys: Loci): Loci {

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

@@ -62,6 +62,29 @@ namespace Link {
         return true
     }
 
+    export function remapLoci(loci: Loci, structure: Structure): Loci {
+        if (structure === loci.structure) return loci
+
+        const links: Loci['links'][0][] = [];
+        loci.links.forEach(l => {
+            const unitA = structure.unitMap.get(l.aUnit.id)
+            if (!unitA) return
+            const unitB = structure.unitMap.get(l.bUnit.id)
+            if (!unitB) return
+
+            const elementA = l.aUnit.elements[l.aIndex]
+            const indexA = SortedArray.indexOf(unitA.elements, elementA) as StructureElement.UnitIndex | -1
+            if (indexA === -1) return
+            const elementB = l.bUnit.elements[l.bIndex]
+            const indexB = SortedArray.indexOf(unitB.elements, elementB) as StructureElement.UnitIndex | -1
+            if (indexB === -1) return
+
+            links.push(Location(unitA, indexA, unitB, indexB))
+        });
+
+        return Loci(structure, links);
+    }
+
     export function toStructureElementLoci(loci: Loci): StructureElement.Loci {
         const elements: StructureElement.Loci['elements'][0][] = []
         const map = new Map<number, number[]>()

src/mol-repr/structure/visual/carbohydrate-link-cylinder.ts

@@ -120,6 +120,7 @@ function eachCarbohydrateLink(loci: Loci, structure: Structure, apply: (interval
     let changed = false
     if (Link.isLoci(loci)) {
         if (!Structure.areParentsEquivalent(loci.structure, structure)) return false
+        loci = Link.remapLoci(loci, structure)
         const { getLinkIndex } = structure.carbohydrates
         for (const l of loci.links) {
             const idx = getLinkIndex(l.aUnit, l.aUnit.elements[l.aIndex], l.bUnit, l.bUnit.elements[l.bIndex])
@@ -129,6 +130,7 @@ function eachCarbohydrateLink(loci: Loci, structure: Structure, apply: (interval
         }
     } else if (StructureElement.isLoci(loci)) {
         if (!Structure.areParentsEquivalent(loci.structure, structure)) return false
+        loci = StructureElement.Loci.remap(loci, structure)
         // TODO mark link only when both of the link elements are in a StructureElement.Loci
         const { getElementIndex, getLinkIndices, elements } = structure.carbohydrates
         for (const e of loci.elements) {

src/mol-repr/structure/visual/carbohydrate-symbol-mesh.ts

@@ -197,6 +197,7 @@ function eachCarbohydrate(loci: Loci, structure: Structure, apply: (interval: In
     let changed = false
     if (!StructureElement.isLoci(loci)) return false
     if (!Structure.areParentsEquivalent(loci.structure, structure)) return false
+    loci = StructureElement.Loci.remap(loci, structure)
     for (const e of loci.elements) {
         // TODO make more efficient by handling/grouping `e.indices` by residue index
         // TODO only call apply when the full alt-residue of the unit is part of `e`

src/mol-repr/structure/visual/carbohydrate-terminal-link-cylinder.ts

@@ -134,6 +134,7 @@ function eachTerminalLink(loci: Loci, structure: Structure, apply: (interval: In
     let changed = false
     if (Link.isLoci(loci)) {
         if (!Structure.areParentsEquivalent(loci.structure, structure)) return false
+        loci = Link.remapLoci(loci, structure)
         for (const l of loci.links) {
             const idx = getTerminalLinkIndex(l.aUnit, l.aUnit.elements[l.aIndex], l.bUnit, l.bUnit.elements[l.bIndex])
             if (idx !== undefined) {
@@ -142,6 +143,7 @@ function eachTerminalLink(loci: Loci, structure: Structure, apply: (interval: In
         }
     } else if (StructureElement.isLoci(loci)) {
         if (!Structure.areParentsEquivalent(loci.structure, structure)) return false
+        loci = StructureElement.Loci.remap(loci, structure)
         // TODO mark link only when both of the link elements are in a StructureElement.Loci
         const { getElementIndex, getTerminalLinkIndices, elements } = structure.carbohydrates
         for (const e of loci.elements) {

src/mol-repr/structure/visual/cross-link-restraint-cylinder.ts

@@ -108,6 +108,7 @@ function eachCrossLink(loci: Loci, structure: Structure, apply: (interval: Inter
     let changed = false
     if (Link.isLoci(loci)) {
         if (!Structure.areParentsEquivalent(loci.structure, structure)) return false
+        loci = Link.remapLoci(loci, structure)
         for (const b of loci.links) {
             const indices = crossLinks.getPairIndices(b.aIndex, b.aUnit, b.bIndex, b.bUnit)
             if (indices) {

src/mol-repr/structure/visual/inter-unit-link-cylinder.ts

@@ -98,6 +98,7 @@ function eachLink(loci: Loci, structure: Structure, apply: (interval: Interval)
     let changed = false
     if (Link.isLoci(loci)) {
         if (!Structure.areParentsEquivalent(loci.structure, structure)) return false
+        loci = Link.remapLoci(loci, structure)
         for (const b of loci.links) {
             const idx = structure.links.getBondIndex(b.aIndex, b.aUnit, b.bIndex, b.bUnit)
             if (idx !== -1) {
@@ -106,6 +107,7 @@ function eachLink(loci: Loci, structure: Structure, apply: (interval: Interval)
         }
     } else if (StructureElement.isLoci(loci)) {
         if (!Structure.areParentsEquivalent(loci.structure, structure)) return false
+        loci = StructureElement.Loci.remap(loci, structure)
         // TODO mark link only when both of the link elements are in a StructureElement.Loci
         for (const e of loci.elements) {
             OrderedSet.forEach(e.indices, v => {

src/mol-repr/structure/visual/intra-unit-link-cylinder.ts

@@ -122,6 +122,7 @@ function eachLink(loci: Loci, structureGroup: StructureGroup, apply: (interval:
     if (Link.isLoci(loci)) {
         const { structure, group } = structureGroup
         if (!Structure.areParentsEquivalent(loci.structure, structure)) return false
+        loci = Link.remapLoci(loci, structure)
         const unit = group.units[0]
         if (!Unit.isAtomic(unit)) return false
         const groupCount = unit.links.edgeCount * 2
@@ -137,6 +138,7 @@ function eachLink(loci: Loci, structureGroup: StructureGroup, apply: (interval:
     } else if (StructureElement.isLoci(loci)) {
         const { structure, group } = structureGroup
         if (!Structure.areParentsEquivalent(loci.structure, structure)) return false
+        loci = StructureElement.Loci.remap(loci, structure)
         const unit = group.units[0]
         if (!Unit.isAtomic(unit)) return false
         const groupCount = unit.links.edgeCount * 2

src/mol-repr/structure/visual/util/element.ts

@@ -73,6 +73,7 @@ export function eachElement(loci: Loci, structureGroup: StructureGroup, apply: (
     if (!StructureElement.isLoci(loci)) return false
     const { structure, group } = structureGroup
     if (!Structure.areParentsEquivalent(loci.structure, structure)) return false
+    loci = StructureElement.Loci.remap(loci, structure)
     const elementCount = group.elements.length
     for (const e of loci.elements) {
         const unitIdx = group.unitIndexMap.get(e.unit.id)

src/mol-repr/structure/visual/util/nucleotide.ts

@@ -46,6 +46,7 @@ export function eachNucleotideElement(loci: Loci, structureGroup: StructureGroup
     if (!StructureElement.isLoci(loci)) return false
     const { structure, group } = structureGroup
     if (!Structure.areParentsEquivalent(loci.structure, structure)) return false
+    loci = StructureElement.Loci.remap(loci, structure)
     const unit = group.units[0]
     if (!Unit.isAtomic(unit)) return false
     const { nucleotideElements, model, elements } = unit

src/mol-repr/structure/visual/util/polymer.ts

@@ -96,6 +96,7 @@ export function eachPolymerElement(loci: Loci, structureGroup: StructureGroup, a
     if (!StructureElement.isLoci(loci)) return false
     const { structure, group } = structureGroup
     if (!Structure.areParentsEquivalent(loci.structure, structure)) return false
+    loci = StructureElement.Loci.remap(loci, structure)
     const { polymerElements, model, elements } = group.units[0]
     const { index, offsets } = model.atomicHierarchy.residueAtomSegments
     const { traceElementIndex } = model.atomicHierarchy.derived.residue
@@ -157,6 +158,7 @@ export function eachPolymerGapElement(loci: Loci, structureGroup: StructureGroup
     if (Link.isLoci(loci)) {
         const { structure, group } = structureGroup
         if (!Structure.areParentsEquivalent(loci.structure, structure)) return false
+        loci = Link.remapLoci(loci, structure)
         const groupCount = group.units[0].gapElements.length
         for (const b of loci.links) {
             const unitIdx = group.unitIndexMap.get(b.aUnit.id)
@@ -171,6 +173,7 @@ export function eachPolymerGapElement(loci: Loci, structureGroup: StructureGroup
     } else if (StructureElement.isLoci(loci)) {
         const { structure, group } = structureGroup
         if (!Structure.areParentsEquivalent(loci.structure, structure)) return false
+        loci = StructureElement.Loci.remap(loci, structure)
         const groupCount = group.units[0].gapElements.length
         for (const e of loci.elements) {
             const unitIdx = group.unitIndexMap.get(e.unit.id)