diff --git a/scripts/xenium_analysis_D216.py b/scripts/xenium_analysis_D216.py
index 18a199125e37b918ee29a165b0d726397b5127cd..8846e939acfa4bb83409b73c82acd9171ecbb2b3 100644
--- a/scripts/xenium_analysis_D216.py
+++ b/scripts/xenium_analysis_D216.py
@@ -16,10 +16,10 @@ from spatialdata_io import xenium
 import sopa
 
 # Set parameters
-xenium_data='/archive/urology/Strand_lab/shared/xenium/output-XETG00248__0034066__D216__20240425__175357/output-XETG00248__0034066__D216__20240425__175357'
+xenium_data='/archive/urology/Strand_lab/shared/xenium/output-XETG00248__0034066__D216__20240425__175357'
 zarr_dir='data_modified/spatialdata_zarr/'
 sample_name='D216'
-outdir='data_modified/xenium_D216/'
+outdir='data_modified/xenium_'+sample_name+'/'
 os.makedirs(outdir, exist_ok=True)
 
 # Load data and create zarr store
@@ -97,22 +97,23 @@ for i in res:
     )
 
 # Create UMAP plots
-umap_dir=outdir+'umaps/'
-os.makedirs(umap_dir, exist_ok=True)
-for i in res:
-    sc.pl.umap(
-        adata,
-        color='leiden_res_'+str(i),
-        save=umap_dir+'leiden_res_'+str(i)+'.png'
-    )
+grps=[ 'leiden_res_'+str(x) for x in res ]
+sc.pl.umap(
+    adata,
+    color=grps,
+    ncols=3,
+    legend_loc='on data'
+)
+plt.savefig(fname=outdir+'umap_leiden.png')
 
 # Check gene expression
 genes=['EPCAM','KRT5','NKX3-1','PIGR','AKR1C2','DCN','ACTA2','DES','CDH19','PECAM1','PTPRC']
 sc.pl.umap(
     adata,
     color=genes,
-    save=umap_dir+'lineage_markers.png'
+    ncols=3
     )
+plt.savefig(fname=outdir+'umap_lineage_markers.png')
 
 # Save modified data
 adata.write_h5ad(outdir+sample_name+'_annotated.h5ad')