diff --git a/src/mol-view/theme/color/element-index.ts b/src/mol-view/theme/color/element-index.ts
index 104d294de2b176c5d822ed48fa3f4885a64e0a41..3440f402a5dd80b4c2dbdc39d9bcd6c86e27dcf6 100644
--- a/src/mol-view/theme/color/element-index.ts
+++ b/src/mol-view/theme/color/element-index.ts
@@ -6,7 +6,7 @@
 
 import { ColorScale, Color } from 'mol-util/color';
 import { Location } from 'mol-model/location';
-import { StructureElement, Link, Unit } from 'mol-model/structure';
+import { StructureElement, Link } from 'mol-model/structure';
 import { OrderedSet } from 'mol-data/int';
 import { ColorThemeProps, ColorTheme, LocationColor } from '../color';
 
@@ -19,22 +19,24 @@ export function ElementIndexColorTheme(props: ColorThemeProps): ColorTheme {
         const { units } = props.structure
         const unitCount = units.length
         const cummulativeElementCount = new Map<number, number>()
+        const unitIdIndex = new Map<number, number>()
 
         let elementCount = 0
         for (let i = 0; i < unitCount; ++i) {
             cummulativeElementCount.set(i, elementCount)
             elementCount += units[i].elements.length
+            unitIdIndex.set(units[i].id, i)
         }
-        const scale = ColorScale.create({ domain: [ 0, elementCount ] })
+        const scale = ColorScale.create({ domain: [ 0, elementCount - 1 ] })
 
         color = (location: Location): Color => {
             if (StructureElement.isLocation(location)) {
-                const unitIndex = Unit.findUnitById(location.unit.id, units)
+                const unitIndex = unitIdIndex.get(location.unit.id)!
                 const unitElementIndex = OrderedSet.findPredecessorIndex(location.unit.elements, location.element)
-                return scale.color(cummulativeElementCount.get(unitIndex) || 0 + unitElementIndex)
+                return scale.color(cummulativeElementCount.get(unitIndex)! + unitElementIndex)
             } else if (Link.isLocation(location)) {
-                const unitId = Unit.findUnitById(location.aUnit.id, units)
-                return scale.color(cummulativeElementCount.get(unitId) || 0 + location.aIndex)
+                const unitIndex = unitIdIndex.get(location.aUnit.id)!
+                return scale.color(cummulativeElementCount.get(unitIndex)! + location.aIndex)
             }
             return DefaultColor
         }