From 2a148377f1e99cdc94260fba90d1c36f05ea6140 Mon Sep 17 00:00:00 2001
From: fdeguire03 <150461249+fdeguire03@users.noreply.github.com>
Date: Wed, 8 Jan 2025 18:10:34 -0800
Subject: [PATCH] Add npy support to load_memmap

---
 caiman/mmapping.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/caiman/mmapping.py b/caiman/mmapping.py
index 1f15f7988..601f3a4b4 100644
--- a/caiman/mmapping.py
+++ b/caiman/mmapping.py
@@ -46,9 +46,12 @@ def load_memmap(filename: str, mode: str = 'r') -> tuple[Any, tuple, int]:
 
     """
     logger = logging.getLogger("caiman")
-    if pathlib.Path(filename).suffix != '.mmap':
+    allowed_extensions = {'.mmap', '.npy'}
+    
+    extension = pathlib.Path(filename).suffix
+    if extension not in allowed_extensions:
         logger.error(f"Unknown extension for file {filename}")
-        raise ValueError(f'Unknown file extension for file {filename} (should be .mmap)')
+        raise ValueError(f'Unknown file extension for file {filename} (should be .mmap or .npy)')
     # Strip path components and use CAIMAN_DATA/example_movies
     # TODO: Eventually get the code to save these in a different dir
     #fn_without_path = os.path.split(filename)[-1]
@@ -63,7 +66,11 @@ def load_memmap(filename: str, mode: str = 'r') -> tuple[Any, tuple, int]:
     #d1, d2, d3, T, order = int(fpart[-9]), int(fpart[-7]), int(fpart[-5]), int(fpart[-1]), fpart[-3]
 
     filename = caiman.paths.fn_relocated(filename)
-    Yr = np.memmap(filename, mode=mode, shape=prepare_shape((d1 * d2 * d3, T)), dtype=np.float32, order=order)
+    if extension == '.mmap':
+        Yr = np.memmap(filename, mode=mode, shape=prepare_shape((d1 * d2 * d3, T)), dtype=np.float32, order=order)
+    elif extension == '.npy':
+        Yr = np.load(filename, mmap_mode=mode)
+
     if d3 == 1:
         return (Yr, (d1, d2), T)
     else: