diff --git a/scimap/__init__.py b/scimap/__init__.py
index d1992213..8ff68cfa 100644
--- a/scimap/__init__.py
+++ b/scimap/__init__.py
@@ -1,8 +1,8 @@
-from pkg_resources import get_distribution, DistributionNotFound
+import importlib.metadata
 
 try:
-    __version__ = get_distribution('scimap').version
-except DistributionNotFound:
+    __version__ = importlib.metadata.version('scimap')
+except importlib.metadata.PackageNotFoundError:
     __version__ = '(local)'
 
 import scimap
diff --git a/scimap/plotting/densityPlot2D.py b/scimap/plotting/densityPlot2D.py
index 4851f20c..dbc029a7 100644
--- a/scimap/plotting/densityPlot2D.py
+++ b/scimap/plotting/densityPlot2D.py
@@ -130,9 +130,13 @@ def densityPlot2D (adata,
     
     
     # set color
-    cp = copy.copy(cm.get_cmap(cmap))
+    #cp = copy.copy(cm.get_cmap(cmap))
+    #cp.set_under(alpha=0)
+    
+    cp = copy.copy(plt.colormaps[cmap])
     cp.set_under(alpha=0)
     
+    
     # subset data if neede
     if subset is not None:
         if isinstance (subset, str):
diff --git a/scimap/plotting/heatmap.py b/scimap/plotting/heatmap.py
index bec2f0b4..d819cda3 100644
--- a/scimap/plotting/heatmap.py
+++ b/scimap/plotting/heatmap.py
@@ -270,7 +270,7 @@ def plot_category_heatmap_vectorized(data,
             figsize_height = max(8, len(unique_categories) * base_size)
             figsize=(figsize_width, figsize_height)
         
-        fig, ax = plt.subplots(figsize=figsize)
+        fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)
         
         
         # Heatmap
@@ -309,7 +309,7 @@ def plot_category_heatmap_vectorized(data,
         ax.set_xlabel('Markers')
         ax.set_ylabel('Categories')
         
-        plt.tight_layout(rect=[0, 0, 0.9, 0.9]) # Adjust the layout
+        #plt.tight_layout(rect=[0, 0, 0.9, 0.9]) # Adjust the layout
         
         # Saving the figure if saveDir and fileName are provided
         if saveDir:
diff --git a/scimap/plotting/pie.py b/scimap/plotting/pie.py
index 2b6dcb3e..571cc471 100644
--- a/scimap/plotting/pie.py
+++ b/scimap/plotting/pie.py
@@ -155,15 +155,15 @@ def pie (adata,
 
     else:
         # if group_by is provided
-        prop = pd.DataFrame(data.groupby([group_by,phenotype]).size()).reset_index(inplace=False)
+        prop = pd.DataFrame(data.groupby([group_by,phenotype], observed=False).size()).reset_index(inplace=False)
         prop.columns = ['group_by',phenotype,'value']
         labels = np.unique(prop[phenotype])
         #
         if ncols is not None:
-            g = prop.groupby('group_by')
+            g = prop.groupby('group_by', observed=False)
             rows = int(np.ceil(len(g)/ncols))
         else:
-            g = prop.groupby('group_by')
+            g = prop.groupby('group_by', observed=False)
             rows = 1
             ncols = len(g)
     
@@ -193,6 +193,7 @@ def pie (adata,
         final_axes = list(set(total_axes) ^ set(required_axes))      
         # Plot
         fig, axes = plt.subplots(ncols=ncols, nrows=rows)
+        axes = np.atleast_1d(axes)  # This ensures axes is always an array, even if it's a single subplot
         for (c, grp), ax in zip(g, axes.flat):
             ax.pie(grp.value, labels=label, colors=colors, wedgeprops =wedgeprops)
             #ax.pie(grp.value, labels=label, colors=colors, wedgeprops = wedgeprops, **kwargs)
diff --git a/scimap/plotting/voronoi.py b/scimap/plotting/voronoi.py
index 20911bc8..3f08296d 100644
--- a/scimap/plotting/voronoi.py
+++ b/scimap/plotting/voronoi.py
@@ -28,6 +28,7 @@
 """
 
 # import lib
+import os
 import numpy as np
 import matplotlib.pyplot as plt
 import seaborn as sns
@@ -129,6 +130,8 @@ def voronoi (adata,
              overlay_point_alpha= 1, 
              overlay_point_shape=".", 
              plot_legend=True, 
+             fileName='voronoi.pdf',
+             saveDir=None,
              legend_size = 6, **kwargs):
     """
 Parameters:
@@ -202,6 +205,12 @@ def voronoi (adata,
         plot_legend (bool, optional):  
             Whether to display a legend for the plot.
         
+        fileName (str, optional): 
+            Name of the file to save the plot. Relevant only if `saveDir` is not None.
+            
+        saveDir (str, optional): 
+            Directory to save the generated plot. If None, the plot is not saved.
+        
         legend_size (float, optional):  
             The font size of the legend text.
 
@@ -233,7 +242,7 @@ def voronoi (adata,
     
     
     # create the data frame needed
-    data = adata.obs
+    data = adata.obs.copy() 
         
     # Subset the image of interest
     if subset is not None:
@@ -253,13 +262,20 @@ def voronoi (adata,
             y2 = min(data[y_coordinate])
         else:
             y2 = y_lim [1]
+            
+    # do the actuall subsetting
+    #if x_lim is not None:
+    #    data = data[data[x_coordinate] >= x1]
+    #    data = data[data[x_coordinate] <= x2]
+    #if y_lim is not None:
+    #    data = data[data[y_coordinate] <= y1]
+    #    data = data[data[y_coordinate] >= y2]
+        
     # do the actuall subsetting
     if x_lim is not None:
-        data = data[data[x_coordinate] >= x1]
-        data = data[data[x_coordinate] <= x2]
+        data = data[(data[x_coordinate] >= x1) & (data[x_coordinate] <= x2)]
     if y_lim is not None:
-        data = data[data[y_coordinate] <= y1]
-        data = data[data[y_coordinate] >= y2]
+        data = data[(data[y_coordinate] >= y1) & (data[y_coordinate] <= y2)]
 
                     
     # create an extra column with index information
@@ -419,7 +435,12 @@ def voronoi (adata,
                 patchList.append(data_key)
     
         first_legend = plt.legend(handles=patchList, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0., prop={'size': legend_size})
-        plt.tight_layout()
+        
+        try:
+            plt.tight_layout()
+        except:
+            pass
+        
         # Add the legend manually to the current Axes.
         ax = plt.gca().add_artist(first_legend)
         
@@ -433,4 +454,16 @@ def voronoi (adata,
             plt.legend(handles=patchList_scatter, bbox_to_anchor=(-0.05, 1), loc=1, borderaxespad=0., prop={'size': legend_size})
             #plt.tight_layout()
     
+    # Saving the figure if saveDir and fileName are provided
+    if saveDir:
+        if not os.path.exists(saveDir):
+            os.makedirs(saveDir)
+        full_path = os.path.join(saveDir, fileName)
+        plt.savefig(full_path, dpi=300)
+        plt.close()
+        print(f"Saved plot to {full_path}")
+    else:
+        plt.show()
+    
+    
     
diff --git a/scimap/tests/test_pl.py b/scimap/tests/test_pl.py
index 88b68d74..ea730513 100644
--- a/scimap/tests/test_pl.py
+++ b/scimap/tests/test_pl.py
@@ -9,7 +9,6 @@
 import pytest
 import os
 import anndata as ad
-import unittest
 
 
 # load data
@@ -179,26 +178,26 @@ def test_stacked_barplot (adata):
     full_path = os.path.join(saveDir, fileName)
     assert os.path.exists(full_path), f"File was not created: {full_path}"
 
-
 # pie
-# =============================================================================
-# def test_pie (adata):
-#     from scimap.plotting.pie import pie
-#     saveDir = os.getcwd() + '/testFigures'
-#     fileName = 'pie.png'
-#     
-#     pie (adata)
-#     
-#     
-#     pie (adata, x_axis='imageid', y_axis='phenotype', saveDir=saveDir, fileName=fileName)
-#     # check the file exist
-#     full_path = os.path.join(saveDir, fileName)
-#     assert os.path.exists(full_path), f"File was not created: {full_path}"
-# =============================================================================
-
+def test_pie (adata):
+    from scimap.plotting.pie import pie
+    saveDir = os.getcwd() + '/testFigures'
+    fileName = 'pie.png'
+    pie (adata, x_axis='imageid', y_axis='phenotype', saveDir=saveDir, fileName=fileName)
+    # check the file exist
+    full_path = os.path.join(saveDir, fileName)
+    assert os.path.exists(full_path), f"File was not created: {full_path}"
 
 # voronoi
-
+def test_voronoi (adata):
+    from scimap.plotting.voronoi import voronoi
+    saveDir = os.getcwd() + '/testFigures'
+    fileName = 'voronoi.png'
+    voronoi(adata, imageid='ROI', subset='ROI1', color_by='phenotype', saveDir=saveDir, fileName=fileName)
+    # check the file exist
+    full_path = os.path.join(saveDir, fileName)
+    assert os.path.exists(full_path), f"File was not created: {full_path}"
+    
 
 # image_viewer
 # addROI_image
diff --git a/testFigures/_matrixplot.pdf b/testFigures/_matrixplot.pdf
deleted file mode 100644
index eeffcd19..00000000
Binary files a/testFigures/_matrixplot.pdf and /dev/null differ
diff --git a/testFigures/_ranked_markers_per_cluster.pdf b/testFigures/_ranked_markers_per_cluster.pdf
deleted file mode 100644
index e1f682b2..00000000
Binary files a/testFigures/_ranked_markers_per_cluster.pdf and /dev/null differ
diff --git a/testFigures/_umap.pdf b/testFigures/_umap.pdf
deleted file mode 100644
index 5c54f5b0..00000000
Binary files a/testFigures/_umap.pdf and /dev/null differ
diff --git a/testFigures/densityPlot2D.png b/testFigures/densityPlot2D.png
deleted file mode 100644
index 5cbf9968..00000000
Binary files a/testFigures/densityPlot2D.png and /dev/null differ
diff --git a/testFigures/distPlot.png b/testFigures/distPlot.png
deleted file mode 100644
index c264e3de..00000000
Binary files a/testFigures/distPlot.png and /dev/null differ
diff --git a/testFigures/foldchange.png b/testFigures/foldchange.png
deleted file mode 100644
index 2783c63e..00000000
Binary files a/testFigures/foldchange.png and /dev/null differ
diff --git a/testFigures/groupCorrelation.png b/testFigures/groupCorrelation.png
deleted file mode 100644
index 95ecc409..00000000
Binary files a/testFigures/groupCorrelation.png and /dev/null differ
diff --git a/testFigures/heatmap.png b/testFigures/heatmap.png
deleted file mode 100644
index 746ddf1a..00000000
Binary files a/testFigures/heatmap.png and /dev/null differ
diff --git a/testFigures/markerCorrelation.png b/testFigures/markerCorrelation.png
deleted file mode 100644
index baa44488..00000000
Binary files a/testFigures/markerCorrelation.png and /dev/null differ
diff --git a/testFigures/spatialInteractionNetwork.png b/testFigures/spatialInteractionNetwork.png
deleted file mode 100644
index ea3120ef..00000000
Binary files a/testFigures/spatialInteractionNetwork.png and /dev/null differ
diff --git a/testFigures/spatial_distance.png b/testFigures/spatial_distance.png
deleted file mode 100644
index c9af2110..00000000
Binary files a/testFigures/spatial_distance.png and /dev/null differ
diff --git a/testFigures/spatial_interaction.png b/testFigures/spatial_interaction.png
deleted file mode 100644
index b225a8ee..00000000
Binary files a/testFigures/spatial_interaction.png and /dev/null differ
diff --git a/testFigures/spatial_pscore.png b/testFigures/spatial_pscore.png
deleted file mode 100644
index 5c861a82..00000000
Binary files a/testFigures/spatial_pscore.png and /dev/null differ
diff --git a/testFigures/spatial_scatterPlot.png b/testFigures/spatial_scatterPlot.png
deleted file mode 100644
index 3cc6f2a6..00000000
Binary files a/testFigures/spatial_scatterPlot.png and /dev/null differ
diff --git a/testFigures/stacked_barplot.png b/testFigures/stacked_barplot.png
deleted file mode 100644
index b699ff12..00000000
Binary files a/testFigures/stacked_barplot.png and /dev/null differ
diff --git a/testFigures/umap.png b/testFigures/umap.png
deleted file mode 100644
index 2c9a8e29..00000000
Binary files a/testFigures/umap.png and /dev/null differ