diff --git a/model-command/src/main/java/lcsb/mapviewer/commands/ColorModelCommand.java b/model-command/src/main/java/lcsb/mapviewer/commands/ColorModelCommand.java
index 1237e8f02e4af2273bc7159866d709f8624e3231..e66453a02920b2a75c04e5d401cca03838aa9bd5 100644
--- a/model-command/src/main/java/lcsb/mapviewer/commands/ColorModelCommand.java
+++ b/model-command/src/main/java/lcsb/mapviewer/commands/ColorModelCommand.java
@@ -11,6 +11,7 @@ import org.apache.log4j.Logger;
 
 import lcsb.mapviewer.common.Pair;
 import lcsb.mapviewer.common.comparator.StringComparator;
+import lcsb.mapviewer.common.exception.InvalidArgumentException;
 import lcsb.mapviewer.common.exception.NotImplementedException;
 import lcsb.mapviewer.model.graphics.ArrowTypeData;
 import lcsb.mapviewer.model.map.BioEntity;
@@ -189,6 +190,22 @@ public class ColorModelCommand extends ModelCommand {
     }
   }
 
+  protected boolean match(BioEntity element, ColorSchema schema) {
+    if (schema.getElementId() != null && !schema.getElementId().isEmpty()) {
+      if (!element.getElementId().equalsIgnoreCase(schema.getElementId())) {
+        return false;
+      }
+    }
+    if (element instanceof Element) {
+      return match((Element) element, schema);
+    } else if (element instanceof Reaction) {
+      return match((Reaction) element, schema);
+    } else {
+      throw new InvalidArgumentException("Don't know how to handle object: " + element);
+    }
+
+  }
+
   /**
    * Checks if the coloring schema should be used for the {@link Element}.
    * 
@@ -209,11 +226,6 @@ public class ColorModelCommand extends ModelCommand {
           return false;
         }
       }
-      if (schema.getElementId() != null && !schema.getElementId().isEmpty()) {
-        if (!element.getElementId().equalsIgnoreCase(schema.getElementId())) {
-          return false;
-        }
-      }
       if (schema.getTypes().size() > 0) {
         boolean found = false;
         for (Class<?> clazz : schema.getTypes()) {
@@ -277,16 +289,12 @@ public class ColorModelCommand extends ModelCommand {
     List<ColorSchema> result = new ArrayList<ColorSchema>();
     for (ColorSchema schema : schemas) {
       boolean found = false;
-      for (Element element : getModel().getElements()) {
+
+      for (BioEntity element : getModel().getBioEntities()) {
         if (match(element, schema)) {
           found = true;
         }
       }
-      for (Reaction reaction : getModel().getReactions()) {
-        if (match(reaction, schema)) {
-          found = true;
-        }
-      }
       if (!found) {
         result.add(schema);
       }
@@ -310,28 +318,18 @@ public class ColorModelCommand extends ModelCommand {
     List<Model> models = new ArrayList<>();
     models.add(getModel());
     models.addAll(getModel().getSubmodels());
-    for (Model model2 : models) {
+    for (Model model : models) {
       for (ColorSchema schema : schemas) {
-        for (Element element : model2.getElements()) {
+        for (BioEntity element : model.getBioEntities()) {
           if (match(element, schema)) {
             if (result.get(element) != null
                 && !colorExtractor.getNormalizedColor(result.get(element)).equals(Color.WHITE)) {
               throw new InvalidColorSchemaException(
-                  eu.getElementTag(element) + "Element is colored by more than one rule.");
+                  eu.getElementTag(element) + "BioEntity is colored by more than one rule.");
             }
             result.put(element, schema);
           }
         }
-        for (Reaction reaction : model2.getReactions()) {
-          if (match(reaction, schema)) {
-            if (result.get(reaction) != null
-                && !colorExtractor.getNormalizedColor(result.get(reaction)).equals(Color.WHITE)) {
-              throw new InvalidColorSchemaException(
-                  eu.getElementTag(reaction) + "Reaction is colored by more than one rule.");
-            }
-            result.put(reaction, schema);
-          }
-        }
       }
     }
     return result;
@@ -371,18 +369,12 @@ public class ColorModelCommand extends ModelCommand {
     }
 
     for (ColorSchema schema : schemas) {
-      for (Element element : result.getElements()) {
+      for (BioEntity element : result.getBioEntities()) {
         if (match(element, schema)) {
           schema.setMatches(schema.getMatches() + 1);
           applyColor(element, schema);
         }
       }
-      for (Reaction reaction : result.getReactions()) {
-        if (match(reaction, schema)) {
-          schema.setMatches(schema.getMatches() + 1);
-          applyColor(reaction, schema);
-        }
-      }
     }
 
     if (top) {
@@ -393,6 +385,16 @@ public class ColorModelCommand extends ModelCommand {
 
   }
 
+  private void applyColor(BioEntity element, ColorSchema schema) throws InvalidColorSchemaException {
+    if (element instanceof Element) {
+      applyColor((Element) element, schema);
+    } else if (element instanceof Reaction) {
+      applyColor((Reaction) element, schema);
+    } else {
+      throw new InvalidArgumentException("Don't know how to handle: " + element);
+    }
+  }
+
   @Override
   protected void executeImplementation() throws CommandExecutionException {
     try {