Added layer argument to anndata_to_image_array function.

metaspace_converter/anndata_to_array.py

@@ -1,3 +1,4 @@
+from typing import Optional
 import numpy as np
 from anndata import AnnData
 
@@ -16,13 +17,15 @@ def _check_pixel_coordinates(adata: AnnData) -> bool:
     return np.all(np.equal(pixel_list, required_pixels))
 
 
-def anndata_to_image_array(adata: AnnData) -> np.ndarray:
+def anndata_to_image_array(adata: AnnData, layer: Optional[str]=None) -> np.ndarray:
     """
     Extracts an array of ion images from an AnnData object
     (that has been generated through the ``metaspace_to_anndata`` function).
 
     Args:
         adata: An AnnData object.
+        layer: ``AnnData.layer`` that should be extracted to an image array. 
+            Default is None, which means that ``adata.X`` will be used.
 
     Returns:
         A three-dimensional Numpy array in the following shape
@@ -36,8 +39,13 @@ def anndata_to_image_array(adata: AnnData) -> np.ndarray:
             E.g. Pixel have been removed/added and the number of pixels
             and their coordinates do not match the original image dimensions.
     """
-
-    pixel_array = adata.X.transpose().copy()
+    if layer is None:
+        pixel_array = np.array(adata.X).transpose().copy()
+    elif layer in adata.layers.keys():
+        pixel_array = np.array(adata.layers[layer]).transpose().copy()
+    else:
+        raise ValueError(f"Layer `{layer}` not found in adata.layers.")
+
     img_size = adata.uns[METASPACE_KEY]["image_size"]
 
     # Check if image dimensions are okay