diff --git a/package-lock.json b/package-lock.json
index 6620401..d9bd4fa 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "@weng-lab/psychscreen-ui-components",
-  "version": "0.8.9",
+  "version": "0.9.2",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "@weng-lab/psychscreen-ui-components",
-      "version": "0.8.8",
+      "version": "0.9.2",
       "license": "MIT",
       "dependencies": {
         "@visx/tooltip": "^3.3.0",
diff --git a/package.json b/package.json
index 64402f0..5b0f6ff 100644
--- a/package.json
+++ b/package.json
@@ -2,7 +2,7 @@
   "name": "@weng-lab/psychscreen-ui-components",
   "description": "Typescript and Material UI based components used for psychSCREEN",
   "author": "SCREEN Team @ UMass Chan Medical School",
-  "version": "0.9.2",
+  "version": "0.9.3",
   "license": "MIT",
   "main": "dist/index.js",
   "typings": "dist/index.d.ts",
diff --git a/src/components/Graph/Graph.tsx b/src/components/Graph/Graph.tsx
index 79f89c7..ab44eba 100644
--- a/src/components/Graph/Graph.tsx
+++ b/src/components/Graph/Graph.tsx
@@ -44,7 +44,7 @@ const Graph: React.FC<GraphProps> = ({
   const [showLabels, setShowLabels] = useState(true);
   const [toggles, setToggles] = useState<{ [key: string]: boolean }>({});
 
-  // unique categories for legend toggles
+  // unique categories for legend toggles or the node / edge categories
   const uniqueCategories = new Set<string>();
 
   if (legendToggle !== undefined) {
@@ -69,30 +69,34 @@ const Graph: React.FC<GraphProps> = ({
     });
   }
 
-  let a = new Set<string>();
-  data.node.forEach((node) => a.add(node.category));
+  let allCategories = new Set<string>(); // holds all unique categories of nodes based on legendToggle and order
+  data.node.forEach((node) => {
+    if (legendToggle && !order) {
+      allCategories.add(legendToggle(node)); // simple legend name
+    } else {
+      allCategories.add(node.category); // node category
+    }
+  });
   data.edge.forEach((edge) => {
     if (edge.category && legendToggle) {
-      a.add(legendToggle(edge));
+      allCategories.add(legendToggle(edge));
     } else if (edge.category) {
-      a.add(edge.category);
+      allCategories.add(edge.category);
     }
   });
-  let u = Array.from(new Set(a));
+
+  let orderedCategories = Array.from(allCategories);
+
   if (order) {
-    u = u.sort((a, b) => order.indexOf(a) - order.indexOf(b));
+    orderedCategories = orderedCategories.sort(
+      (a, b) => order.indexOf(a) - order.indexOf(b)
+    );
   }
 
   const initialToggles: { [key: string]: boolean } = {};
-  if (order) {
-    u.forEach((category) => {
-      initialToggles[category] = true;
-    });
-  } else {
-    uniqueCategories.forEach((category) => {
-      initialToggles[category] = true;
-    });
-  }
+  orderedCategories.forEach((category) => {
+    initialToggles[category] = true;
+  });
 
   // DOWNLOAD SCREENSHOT
   const ref = useRef<HTMLDivElement>(null);
@@ -156,14 +160,21 @@ const Graph: React.FC<GraphProps> = ({
   }, [data]);
 
   let elem = elements.map((e) => e.category);
-
-  let unique = Array.from(new Set(elem));
+  let unique = Array.from(new Set(elem)); // holds unique node categories only
 
   if (order) {
     unique = unique.sort((a, b) => order.indexOf(a) - order.indexOf(b));
   }
 
   const simple: string[] = elements.map((e) => {
+    // used for toggles
+    if (legendToggle && !order) return legendToggle(e);
+    else {
+      return e.category;
+    }
+  });
+  const simpleCat: string[] = elements.map((e) => {
+    // used for legend calculations
     if (legendToggle) return legendToggle(e);
     else {
       return e.category;
@@ -188,7 +199,7 @@ const Graph: React.FC<GraphProps> = ({
     }
 
     // connect holds all the connections between nodes
-    // connect[0] holds the target node INDICES for the FIRST node
+    // connect[0] holds the TARGET node INDICES for the FIRST node
     edges.forEach((e) => {
       const fromIndex = allcCREs.indexOf(e.from);
       const toIndex = allcCREs.indexOf(e.to);
@@ -247,6 +258,7 @@ const Graph: React.FC<GraphProps> = ({
     for (var i = 0; i < elements.length; i++) {
       if (toggles[simple[i]] !== false) {
         if (data.centered && elements[i].id === data.centered.id) {
+          // if centered node
           cy.add({
             data: { id: createID(i) }, // create name
             position: {
@@ -291,7 +303,7 @@ const Graph: React.FC<GraphProps> = ({
       }
     }
 
-    const c = edgeTypes.every((e) => e !== 'Edge');
+    const isDirectional = edgeTypes.every((e) => e !== 'Edge');
     // ADD EDGES
     let edgeCount = 0;
     for (var j = 0; j < elements.length; j++) {
@@ -311,7 +323,7 @@ const Graph: React.FC<GraphProps> = ({
               },
               style: {
                 'line-color': getColor ? getColor(edges[j]) : 'grey',
-                'target-arrow-shape': c ? 'triangle' : null,
+                'target-arrow-shape': isDirectional ? 'triangle' : null,
                 'target-arrow-color': getColor ? getColor(edges[j]) : 'grey',
                 width: scale(scales[j]),
               },
@@ -371,6 +383,7 @@ const Graph: React.FC<GraphProps> = ({
     };
   }, [elements, scales, edgeTypes, edges, toggles, showTooltip, hideTooltip]);
 
+  // toggle labels
   useEffect(() => {
     if (!cyRef.current) return;
     let ind = 0;
@@ -412,7 +425,7 @@ const Graph: React.FC<GraphProps> = ({
     }
   };
 
-  // TOGGLE
+  // TOGGLE CATEGORIES
   const handleToggle = (category: string) => {
     setToggles((prevToggles) => ({
       ...prevToggles,
@@ -452,7 +465,7 @@ const Graph: React.FC<GraphProps> = ({
           <ControlPanel
             toggles={toggles}
             onToggle={handleToggle}
-            simpleCategories={simple}
+            simpleCategories={simpleCat}
             edgeType={data.edge.every((e) => e.category)}
             elements={elements}
             edges={edges}
diff --git a/src/components/Graph/Legend.tsx b/src/components/Graph/Legend.tsx
index 5f1f3c0..f153282 100644
--- a/src/components/Graph/Legend.tsx
+++ b/src/components/Graph/Legend.tsx
@@ -29,6 +29,7 @@ const Legend: React.FC<LegendProps> = ({
   legendToggle,
   legendNodeLabel,
   legendEdgeLabel,
+  order,
   uniqueCat,
 }) => {
   const edgeTypes = Array.from(
@@ -44,8 +45,11 @@ const Legend: React.FC<LegendProps> = ({
     )
   );
 
-  const uniqueCategories = Array.from(new Set(simpleCategories));
+  const uniqueCategories = order
+    ? order
+    : Array.from(new Set(simpleCategories)); // simple names
 
+  // loop on uniqueCat if there is order, else loop on uniqueCategories
   return (
     <div
       style={{
@@ -63,7 +67,7 @@ const Legend: React.FC<LegendProps> = ({
       {uniqueCat
         ? uniqueCat.map((category) => {
             let color = 'grey';
-            let c = '';
+            let simpleDisplay = '';
 
             elements.forEach((node) => {
               uniqueCategories.forEach((cat) => {
@@ -74,12 +78,11 @@ const Legend: React.FC<LegendProps> = ({
                   node.category === category
                 ) {
                   color = colorFunc(node);
-                  c = cat;
+                  simpleDisplay = cat; // simple display category name if legend toggle
                   return;
                 }
               });
             });
-
             return (
               <div
                 key={category}
@@ -106,7 +109,7 @@ const Legend: React.FC<LegendProps> = ({
                   }}
                   onClick={() => onToggle(category)}
                 >
-                  {c} {legendToggle ? '(' : null}
+                  {simpleDisplay} {legendToggle ? '(' : null}
                   {legendToggle ? category : null}
                   {legendToggle ? ')' : null}
                 </Typography>
@@ -115,7 +118,7 @@ const Legend: React.FC<LegendProps> = ({
           })
         : uniqueCategories.map((category) => {
             let color = 'grey';
-            let cat = '';
+            let nodeCat = '';
 
             elements.forEach((node) => {
               if (
@@ -124,7 +127,7 @@ const Legend: React.FC<LegendProps> = ({
                 legendToggle(node) === category
               ) {
                 color = colorFunc(node);
-                cat = node.category;
+                nodeCat = node.category; // category of node
               }
             });
 
@@ -155,7 +158,7 @@ const Legend: React.FC<LegendProps> = ({
                   onClick={() => onToggle(category)}
                 >
                   {category} {legendToggle ? '(' : null}
-                  {legendToggle ? cat : null}
+                  {legendToggle ? nodeCat : null}
                   {legendToggle ? ')' : null}
                 </Typography>
               </div>
diff --git a/stories/Graph.stories.tsx b/stories/Graph.stories.tsx
index ed410ef..8cb4b96 100644
--- a/stories/Graph.stories.tsx
+++ b/stories/Graph.stories.tsx
@@ -140,6 +140,7 @@ SampleGraph.args = {
   scale: (n: number) => 10 * n,
   getColor: setColor3,
   legendToggle: convertToSimple2,
+  order: ['P', 'R', 'B'],
 };
 export const PilotDataWithCentered = Template.bind({});
 PilotDataWithCentered.args = {