diff --git a/package.json b/package.json
index 57671a7..3f78857 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
     "name": "cluster_viz",
-    "version": "0.2.1",
+    "version": "0.2.2",
     "description": "Show a random NASA Astronomy Picture of the Day in a JupyterLab panel",
     "keywords": [
         "jupyter",
diff --git a/src/VizComponent.tsx b/src/VizComponent.tsx
index b9835e9..1c99933 100644
--- a/src/VizComponent.tsx
+++ b/src/VizComponent.tsx
@@ -47,14 +47,14 @@ const GroupedCells: React.FC<GroupedCellsProps> = ({ className, cells, onSelectN
     if (!acc[cell.cluster]) {
       acc[cell.cluster] = [];
     }
-    acc[cell.cluster].push(cell);
+    acc[cell.cluster].push(cell); 
     return acc;
   }, {} as { [key: string]: NotebookCellWithID[] });
 
   const totalCells = cells.length; // Total number of cells within the class
 
   const selectedCells = (clusterNames: string[]) => {
-    return clusterNames.flatMap((clusterName) =>
+    return clusterNames.filter(c => c in clusters).flatMap((clusterName) =>
       clusters[clusterName].map(cell => ({ clusterName, cell }))
     ).sort((a, b) => {
       if (a.cell.notebook_id === b.cell.notebook_id) {
@@ -63,6 +63,7 @@ const GroupedCells: React.FC<GroupedCellsProps> = ({ className, cells, onSelectN
       return a.cell.notebook_id - b.cell.notebook_id;
     });
   };
+
   const handleClusterClick = (clusterName: string) => {
     setOpenClusters((prev) =>
       prev.includes(clusterName) ? prev.filter((name) => name !== clusterName) : [...prev, clusterName]