Compute C-order memmap, projections, and correlation image in parallel, efficiently, avoiding OOM

.gitignore

@@ -136,6 +136,3 @@ dmypy.json
 
 # test files
 tests/test data
-
-# VSCode
-.vscode/

mesmerize_core/algorithms/cnmf.py

@@ -4,24 +4,35 @@
 import caiman as cm
 from caiman.source_extraction.cnmf import cnmf as cnmf
 from caiman.source_extraction.cnmf.params import CNMFParams
-import psutil
+from caiman.paths import decode_mmap_filename_dict
 import numpy as np
 import traceback
 from pathlib import Path, PurePosixPath
 from shutil import move as move_file
-import os
 import time
 
 # prevent circular import
 if __name__ in ["__main__", "__mp_main__"]:  # when running in subprocess
     from mesmerize_core import set_parent_raw_data_path, load_batch
     from mesmerize_core.utils import IS_WINDOWS
+    from mesmerize_core.algorithms._utils import (
+        ensure_server,
+        save_projections_parallel,
+        save_correlation_parallel,
+        save_c_order_mmap_parallel,
+    )
 else:  # when running with local backend
     from ..batch_utils import set_parent_raw_data_path, load_batch
     from ..utils import IS_WINDOWS
+    from ._utils import (
+        ensure_server,
+        save_projections_parallel,
+        save_correlation_parallel,
+        save_c_order_mmap_parallel,
+    )
 
 
-def run_algo(batch_path, uuid, data_path: str = None):
+def run_algo(batch_path, uuid, data_path: str = None, dview=None):
     algo_start = time.time()
     set_parent_raw_data_path(data_path)
 
@@ -42,103 +53,95 @@ def run_algo(batch_path, uuid, data_path: str = None):
         f"Starting CNMF item:\n{item}\nWith params:{params}"
     )
 
-    # adapted from current demo notebook
-    if "MESMERIZE_N_PROCESSES" in os.environ.keys():
-        try:
-            n_processes = int(os.environ["MESMERIZE_N_PROCESSES"])
-        except:
-            n_processes = psutil.cpu_count() - 1
-    else:
-        n_processes = psutil.cpu_count() - 1
-    # Start cluster for parallel processing
-    c, dview, n_processes = cm.cluster.setup_cluster(
-        backend="local", n_processes=n_processes, single_thread=False
-    )
-
-    # merge cnmf and eval kwargs into one dict
-    cnmf_params = CNMFParams(params_dict=params["main"])
-    # Run CNMF, denote boolean 'success' if CNMF completes w/out error
-    try:
-        fname_new = cm.save_memmap(
-            [input_movie_path], base_name=f"{uuid}_cnmf-memmap_", order="C", dview=dview
-        )
-
-        print("making memmap")
+    with ensure_server(dview) as (dview, n_processes):
 
-        Yr, dims, T = cm.load_memmap(fname_new)
-        images = np.reshape(Yr.T, [T] + list(dims), order="F")
-
-        proj_paths = dict()
-        for proj_type in ["mean", "std", "max"]:
-            p_img = getattr(np, f"nan{proj_type}")(images, axis=0)
-            proj_paths[proj_type] = output_dir.joinpath(
-                f"{uuid}_{proj_type}_projection.npy"
-            )
-            np.save(str(proj_paths[proj_type]), p_img)
-
-        # in fname new load in memmap order C
-        cm.stop_server(dview=dview)
-        c, dview, n_processes = cm.cluster.setup_cluster(
-            backend="local", n_processes=None, single_thread=False
-        )
-
-        print("performing CNMF")
-        cnm = cnmf.CNMF(n_processes, params=cnmf_params, dview=dview)
-
-        print("fitting images")
-        cnm.fit(images)
-        #
-        if "refit" in params.keys():
-            if params["refit"] is True:
-                print("refitting")
-                cnm = cnm.refit(images, dview=dview)
-
-        print("performing eval")
-        cnm.estimates.evaluate_components(images, cnm.params, dview=dview)
-
-        output_path = output_dir.joinpath(f"{uuid}.hdf5")
-
-        cnm.save(str(output_path))
-
-        Cn = cm.local_correlations(images, swap_dim=False)
-        Cn[np.isnan(Cn)] = 0
-
-        corr_img_path = output_dir.joinpath(f"{uuid}_cn.npy")
-        np.save(str(corr_img_path), Cn, allow_pickle=False)
-
-        # output dict for dataframe row (pd.Series)
-        d = dict()
-
-        cnmf_memmap_path = output_dir.joinpath(Path(fname_new).name)
-        if IS_WINDOWS:
-            Yr._mmap.close()  # accessing private attr but windows is annoying otherwise
-        move_file(fname_new, cnmf_memmap_path)
-
-        # save paths as relative path strings with forward slashes
-        cnmf_hdf5_path = str(PurePosixPath(output_path.relative_to(output_dir.parent)))
-        cnmf_memmap_path = str(
-            PurePosixPath(cnmf_memmap_path.relative_to(output_dir.parent))
-        )
-        corr_img_path = str(PurePosixPath(corr_img_path.relative_to(output_dir.parent)))
-        for proj_type in proj_paths.keys():
-            d[f"{proj_type}-projection-path"] = str(
-                PurePosixPath(proj_paths[proj_type].relative_to(output_dir.parent))
+        # merge cnmf and eval kwargs into one dict
+        cnmf_params = CNMFParams(params_dict=params["main"])
+        # Run CNMF, denote boolean 'success' if CNMF completes w/out error
+        try:
+            # only re-save memmap if necessary
+            save_new_mmap = True
+            if Path(input_movie_path).suffix == ".mmap":
+                mmap_info = decode_mmap_filename_dict(input_movie_path)
+                save_new_mmap = "order" not in mmap_info or mmap_info["order"] != "C"
+
+            if save_new_mmap:
+                print("making memmap")
+                fname_new = save_c_order_mmap_parallel(
+                    input_movie_path,
+                    base_name=f"{uuid}_cnmf-memmap_",
+                    dview=dview,
+                    var_name_hdf5=cnmf_params.data['var_name_hdf5']
+                )
+                cnmf_memmap_path = output_dir.joinpath(Path(fname_new).name)
+                move_file(fname_new, cnmf_memmap_path)
+            else:
+                cnmf_memmap_path = Path(input_movie_path)
+
+            Yr, dims, T = cm.load_memmap(str(cnmf_memmap_path))
+            images = np.reshape(Yr.T, [T] + list(dims), order="F")
+
+            print("computing projections")
+            proj_paths = save_projections_parallel(
+                uuid=uuid,
+                movie_path=cnmf_memmap_path,
+                output_dir=output_dir,
+                dview=dview,
             )
 
-        d.update(
-            {
-                "cnmf-hdf5-path": cnmf_hdf5_path,
-                "cnmf-memmap-path": cnmf_memmap_path,
-                "corr-img-path": corr_img_path,
-                "success": True,
-                "traceback": None,
-            }
-        )
+            print("computing correlation image")
+            corr_img_path = save_correlation_parallel(
+                uuid=uuid,
+                movie_path=cnmf_memmap_path,
+                output_dir=output_dir,
+                dview=dview,
+            )
 
-    except:
-        d = {"success": False, "traceback": traceback.format_exc()}
+            print("performing CNMF")
+            cnm = cnmf.CNMF(n_processes, params=cnmf_params, dview=dview)
+
+            print("fitting images")
+            cnm.fit(images)
+            #
+            if "refit" in params.keys():
+                if params["refit"] is True:
+                    print("refitting")
+                    cnm = cnm.refit(images, dview=dview)
+
+            print("performing eval")
+            cnm.estimates.evaluate_components(images, cnm.params, dview=dview)
+
+            output_path = output_dir.joinpath(f"{uuid}.hdf5")
+
+            cnm.save(str(output_path))
+
+            # output dict for dataframe row (pd.Series)
+            d = dict()
+
+            if IS_WINDOWS:
+                Yr._mmap.close()  # accessing private attr but windows is annoying otherwise
+
+            # save paths as relative path strings with forward slashes
+            cnmf_hdf5_path = str(PurePosixPath(output_path.relative_to(output_dir.parent)))
+            cnmf_memmap_path = str(PurePosixPath(df.paths.split(cnmf_memmap_path)[1]))  # still work if outside output dir
+            corr_img_path = str(PurePosixPath(corr_img_path.relative_to(output_dir.parent)))
+            for proj_type in proj_paths.keys():
+                d[f"{proj_type}-projection-path"] = str(
+                    PurePosixPath(proj_paths[proj_type].relative_to(output_dir.parent))
+                )
+
+            d.update(
+                {
+                    "cnmf-hdf5-path": cnmf_hdf5_path,
+                    "cnmf-memmap-path": cnmf_memmap_path,
+                    "corr-img-path": corr_img_path,
+                    "success": True,
+                    "traceback": None,
+                }
+            )
 
-    cm.stop_server(dview=dview)
+        except:
+            d = {"success": False, "traceback": traceback.format_exc()}
 
     runtime = round(time.time() - algo_start, 2)
     df.caiman.update_item_with_results(uuid, d, runtime)