diff --git a/src/simularium/SelectionInterface.ts b/src/simularium/SelectionInterface.ts
index ce380d70..4adb009e 100644
--- a/src/simularium/SelectionInterface.ts
+++ b/src/simularium/SelectionInterface.ts
@@ -376,38 +376,6 @@ class SelectionInterface {
         });
         return colors;
     }
-
-    public getIdsAndColorsFromUIData = (
-        uiData: UIDisplayData
-    ): {
-        agentIds: number[];
-        color: string | number;
-    }[] => {
-        const colorAssignments: {
-            agentIds: number[];
-            color: string | number;
-        }[] = [];
-
-        uiData.forEach((agent) => {
-            colorAssignments.push({
-                agentIds: this.getAgentIdsByNamesAndTags([
-                    { name: agent.name, tags: [] },
-                ]),
-                color: agent.color,
-            });
-
-            agent.displayStates.forEach((newState) => {
-                colorAssignments.push({
-                    agentIds: this.getAgentIdsByNamesAndTags([
-                        { name: agent.name, tags: [newState.name] },
-                    ]),
-                    color: newState.color,
-                });
-            });
-        });
-
-        return colorAssignments;
-    };
 }
 
 export { SelectionInterface };
diff --git a/src/viewport/index.tsx b/src/viewport/index.tsx
index 779d0378..ea0de152 100644
--- a/src/viewport/index.tsx
+++ b/src/viewport/index.tsx
@@ -596,13 +596,22 @@ class Viewport extends React.Component<
     }
 
     public changeAgentsColor(appliedColors: UIDisplayData): void {
-        const colorSettings =
-            this.selectionInterface.getIdsAndColorsFromUIData(appliedColors);
-        colorSettings.forEach((setting) => {
+        appliedColors.forEach((agent) => {
             this.visGeometry.applyColorToAgents(
-                setting.agentIds,
-                setting.color
+                this.selectionInterface.getAgentIdsByNamesAndTags([
+                    { name: agent.name, tags: [] },
+                ]),
+                agent.color
             );
+
+            agent.displayStates.forEach((state) => {
+                this.visGeometry.applyColorToAgents(
+                    this.selectionInterface.getAgentIdsByNamesAndTags([
+                        { name: agent.name, tags: [state.name] },
+                    ]),
+                    state.color
+                );
+            });
         });
     }