lbm_suite2p_python 3.2.0__py3-none-any.whl → 3.2.2__py3-none-any.whl

lbm_suite2p_python/cli.py

@@ -23,6 +23,7 @@ Examples:
 """
 
 import argparse
+import json
 import sys
 from pathlib import Path
 from typing import Any
@@ -67,6 +68,7 @@ def _get_ops_help() -> dict[str, str]:
         "cellprob_threshold": "cellpose cell probability threshold (lower = more cells)",
         "flow_threshold": "cellpose flow error threshold",
         "anatomical_only": "cellpose detection mode: 0=off, 1=max_proj, 2=mean, 3=enhanced, 4=max",
+        "algorithm": "detection algorithm: sparsery, sourcery, or cellpose",
         "pretrained_model": "cellpose model name (e.g., cpsam, cyto2, nuclei)",
         "do_registration": "whether to run motion correction",
         "nonrigid": "use nonrigid (piecewise) registration",
@@ -103,6 +105,10 @@ Examples:
   lsp data/ output/ --planes 1 2 3            # specific planes (1-indexed)
   lsp data/ output/ --num-timepoints 500      # quick test with 500 frames
   lsp data/ output/ --diameter 8              # custom cell diameter
+  lsp data/ output/ --tau 1.3 --algorithm cellpose         # decay + detection algorithm
+  lsp data/ output/ --fix-phase --phasecorr-method median  # reader phase correction
+  lsp data/ output/ --reader-kwargs '{"roi": 2}'           # arbitrary imread kwargs
+  lsp data/ output/ --ops-file my_ops.json    # load saved ops (CLI flags override)
   lsp --list-ops                              # show all suite2p parameters
         """,
     )
@@ -134,13 +140,21 @@ Examples:
         "--planes", nargs="*", type=int, dest="planes",
         help="z-planes to process (1-indexed, e.g., --planes 1 2 3)"
     )
+    pipeline.add_argument(
+        "--timepoints", nargs="*", type=int, dest="timepoints",
+        help="timepoints to process (1-indexed, e.g., --timepoints 1 50 100)"
+    )
     pipeline.add_argument(
         "--roi-mode", "--roi", type=int, dest="roi_mode",
         help="ROI mode: None=stitch, 0=split all, N=specific ROI"
     )
     pipeline.add_argument(
         "--num-timepoints", "--frames", type=int, dest="num_timepoints",
-        help="number of frames/timepoints to process (for quick testing)"
+        help="number of timepoints to process (first N, for quick testing)"
+    )
+    pipeline.add_argument(
+        "--num-zplanes", type=int, dest="num_zplanes",
+        help="number of z-planes to process (first N)"
     )
     pipeline.add_argument(
         "--overwrite", action="store_true",
@@ -237,15 +251,57 @@ Examples:
         help="maximum cell diameter in pixels"
     )
 
-    # reader options (for raw data)
+    # reader options (forwarded to mbo_utilities.imread)
     reader = parser.add_argument_group("reader options (raw scanimage data)")
     reader.add_argument(
-        "--fix-phase", action="store_true",
-        help="apply bidirectional phase correction"
+        "--fix-phase", action=argparse.BooleanOptionalAction, default=None,
+        help="bidirectional phase correction (reader default: on)"
+    )
+    reader.add_argument(
+        "--use-fft", action=argparse.BooleanOptionalAction, default=None,
+        help="FFT-based subpixel phase correction"
     )
     reader.add_argument(
-        "--use-fft", action="store_true",
-        help="use FFT for subpixel phase correction"
+        "--phasecorr-method", choices=["mean", "median", "max"], default=None,
+        help="phase-correction reduction method (default: mean)"
+    )
+    reader.add_argument(
+        "--channel", type=int, default=None,
+        help="zero-based color channel to read (multi-channel sources)"
+    )
+    reader.add_argument(
+        "--reader-kwargs", type=str, default=None, metavar="JSON",
+        help='extra imread kwargs as JSON, e.g. \'{"roi": 2}\''
+    )
+
+    # writer options (forwarded to the binary/zarr writer)
+    writer = parser.add_argument_group("writer options")
+    writer.add_argument(
+        "--writer-kwargs", type=str, default=None, metavar="JSON",
+        help='extra writer kwargs as JSON, e.g. \'{"target_chunk_mb": 200}\''
+    )
+
+    # rastermap options
+    rmap = parser.add_argument_group("rastermap options")
+    rmap.add_argument(
+        "--rastermap", action="store_true",
+        help="enable rastermap (planar + volumetric) with default settings"
+    )
+    rmap.add_argument(
+        "--rastermap-kwargs", type=str, default=None, metavar="JSON",
+        help='rastermap config JSON with "planar"/"volumetric" keys, '
+             'e.g. \'{"planar": {"n_clusters": 50}}\''
+    )
+
+    # advanced
+    advanced = parser.add_argument_group("advanced")
+    advanced.add_argument(
+        "--ops-file", type=str, default=None,
+        help="base ops from a .npy/.json file or suite2p dir; CLI flags override"
+    )
+    advanced.add_argument(
+        "--replot", action=argparse.BooleanOptionalAction, default=True,
+        help="regenerate per-plane figures (default: on)"
     )
 
     # dynamically add all ops parameters
@@ -316,8 +372,9 @@ def list_ops():
         "Main Settings": ["nplanes", "nchannels", "fs", "tau", "frames_include"],
         "Registration": ["do_registration", "nonrigid", "batch_size", "maxregshift",
                         "smooth_sigma", "nimg_init", "subpixel"],
-        "Cell Detection": ["roidetect", "sparse_mode", "spatial_scale", "threshold_scaling",
-                          "max_overlap", "connected", "nbinned", "max_iterations"],
+        "Cell Detection": ["roidetect", "algorithm", "sparse_mode", "spatial_scale",
+                          "threshold_scaling", "max_overlap", "connected", "nbinned",
+                          "max_iterations"],
         "Cellpose": ["anatomical_only", "diameter", "cellprob_threshold", "flow_threshold",
                     "pretrained_model", "spatial_hp_cp"],
         "Signal Extraction": ["neuropil_extract", "neucoeff", "spikedetect",
@@ -397,6 +454,65 @@ def build_cell_filters(args) -> list | None:
     return filters if filters else None
 
 
+def _parse_json_arg(flag: str, value: str) -> dict:
+    """parse a JSON object from a CLI flag value; exit cleanly on error."""
+    try:
+        parsed = json.loads(value)
+    except json.JSONDecodeError as e:
+        raise SystemExit(f"Error: {flag} is not valid JSON: {e}")
+    if not isinstance(parsed, dict):
+        raise SystemExit(f"Error: {flag} must be a JSON object")
+    return parsed
+
+
+def _load_ops_file(path: str) -> dict:
+    """load a base ops dict from a .json/.npy file or a suite2p directory."""
+    from lbm_suite2p_python import load_ops
+
+    p = Path(path).expanduser()
+    if p.suffix.lower() == ".json":
+        if not p.exists():
+            raise SystemExit(f"Error: --ops-file not found: {p}")
+        with open(p, "r", encoding="utf-8") as fh:
+            data = json.load(fh)
+        if not isinstance(data, dict):
+            raise SystemExit("Error: --ops-file JSON must be an object")
+        return data
+    return load_ops(p)
+
+
+def build_reader_kwargs(args) -> dict | None:
+    """build imread kwargs from CLI reader args (None if empty)."""
+    kw = {}
+    if args.fix_phase is not None:
+        kw["fix_phase"] = args.fix_phase
+    if args.use_fft is not None:
+        kw["use_fft"] = args.use_fft
+    if args.phasecorr_method is not None:
+        kw["phasecorr_method"] = args.phasecorr_method
+    if args.channel is not None:
+        kw["channel"] = args.channel
+    if args.reader_kwargs:
+        kw.update(_parse_json_arg("--reader-kwargs", args.reader_kwargs))
+    return kw or None
+
+
+def build_writer_kwargs(args) -> dict | None:
+    """build writer kwargs from CLI writer args (None if empty)."""
+    if args.writer_kwargs:
+        return _parse_json_arg("--writer-kwargs", args.writer_kwargs)
+    return None
+
+
+def build_rastermap_kwargs(args) -> dict | None:
+    """build rastermap_kwargs from CLI args (None if disabled)."""
+    if args.rastermap_kwargs:
+        return _parse_json_arg("--rastermap-kwargs", args.rastermap_kwargs)
+    if args.rastermap:
+        return {"planar": {}, "volumetric": {}}
+    return None
+
+
 def build_ops(args, base_ops: dict) -> dict:
     """build ops dict from CLI args, overriding base_ops."""
     from lbm_suite2p_python.default_ops import s2p_ops
@@ -504,6 +620,13 @@ def main():
 
     output_path.mkdir(parents=True, exist_ok=True)
 
+    # parse config/advanced args up front so invalid JSON or a missing ops
+    # file fails before any logging or processing work begins
+    base_extra_ops = _load_ops_file(args.ops_file) if args.ops_file else None
+    reader_kwargs = build_reader_kwargs(args)
+    writer_kwargs = build_writer_kwargs(args)
+    rastermap_kwargs = build_rastermap_kwargs(args)
+
     log_path = output_path / "log.txt"
     log_file = open(log_path, "w", encoding="utf-8", buffering=1)
     _orig_stdout, _orig_stderr = sys.stdout, sys.stderr
@@ -530,17 +653,10 @@ def main():
     print(f"Input:  {input_path}")
     print(f"Output: {output_path}")
 
-    # build ops
-    base_ops = lsp.default_ops()
+    # build ops (optionally starting from a user-supplied ops file)
+    base_ops = lsp.default_ops(ops=base_extra_ops) if base_extra_ops else lsp.default_ops()
     ops = build_ops(args, base_ops)
 
-    # build reader kwargs
-    reader_kwargs = {}
-    if args.fix_phase:
-        reader_kwargs["fix_phase"] = True
-    if args.use_fft:
-        reader_kwargs["use_fft"] = True
-
     # build cell filters
     cell_filters = build_cell_filters(args)
 
@@ -554,6 +670,10 @@ def main():
     print(f"  Cellpose model: {ops.get('pretrained_model', 'cpsam')}")
     if cell_filters:
         print(f"  Cell filters: {cell_filters}")
+    if reader_kwargs:
+        print(f"  Reader: {reader_kwargs}")
+    if rastermap_kwargs:
+        print(f"  Rastermap: {sorted(rastermap_kwargs)}")
 
     print(f"\n{'='*60}\n")
 
@@ -566,8 +686,10 @@ def main():
             save_path=output_path,
             ops=ops,
             planes=args.planes,
+            timepoints=args.timepoints,
             roi_mode=args.roi_mode,
             num_timepoints=args.num_timepoints,
+            num_zplanes=args.num_zplanes,
             keep_reg=args.keep_reg,
             keep_raw=args.keep_raw,
             force_reg=args.force_reg or args.overwrite,
@@ -580,7 +702,10 @@ def main():
             cell_filters=cell_filters,
             accept_all_cells=args.accept_all_cells,
             save_json=args.save_json,
-            reader_kwargs=reader_kwargs if reader_kwargs else None,
+            reader_kwargs=reader_kwargs,
+            writer_kwargs=writer_kwargs,
+            rastermap_kwargs=rastermap_kwargs,
+            replot=args.replot,
             workers=args.workers,
             skip_volumetric=args.skip_volumetric,
             threads_per_worker=args.threads_per_worker,

lbm_suite2p_python/run_lsp.py

@@ -446,14 +446,15 @@ def _is_valid_torch_checkpoint(path) -> bool:
 def _prewarm_cellpose_model(ops) -> None:
     """Download the cellpose model once, in the parent, before workers fan out.
 
-    cellpose's cache_CPSAM_model_path downloads to a temp file then renames
-    with no cross-process lock. Multiple workers hitting an empty cache at once
+    cellpose's model cache downloads to a temp file then renames with no
+    cross-process lock. Multiple workers hitting an empty cache at once
     race: one wins the rename, the rest fail (Windows WinError 32/183) or read a
     half-written file (PytorchStreamReader miniz error). Warming here serializes
     the download so workers only ever read a complete file. A corrupt leftover
     from a prior failed run is removed and re-downloaded.
     """
-    if not (ops.get("roidetect", True) and ops.get("anatomical_only", 0) > 0):
+    if not (ops.get("roidetect", True)
+            and (ops.get("anatomical_only", 0) > 0 or ops.get("algorithm") == "cellpose")):
         return
     try:
         from cellpose import models as cp_models
@@ -467,7 +468,11 @@ def _prewarm_cellpose_model(ops) -> None:
         except OSError:
             pass
     try:
-        cp_models.cache_CPSAM_model_path()
+        # cellpose 4.x renamed cache_CPSAM_model_path() -> cache_model_path(backbone)
+        if hasattr(cp_models, "cache_model_path"):
+            cp_models.cache_model_path("cpsam")
+        else:
+            cp_models.cache_CPSAM_model_path()
     except Exception as exc:
         print(
             f"Warning: could not pre-download cellpose model ({exc}); "
@@ -780,6 +785,38 @@ def _prepare_plane_ops(*, base_ops, plane_idx, num_planes, input_arr,
     return current_ops
 
 
+def _resolve_timepoints(timepoints=None, frames=None, frame_indices=None):
+    """Resolve the canonical 1-based ``timepoints`` selection.
+
+    ``frames`` (1-based) and ``frame_indices`` (0-based) are deprecated
+    aliases and emit a DeprecationWarning. Returns a 1-based list, or
+    None for all timepoints.
+    """
+    import warnings
+
+    if frames is not None:
+        warnings.warn(
+            "'frames' is deprecated, use 'timepoints' (1-based)",
+            DeprecationWarning,
+            stacklevel=3,
+        )
+        if timepoints is None:
+            timepoints = frames
+    if frame_indices is not None:
+        warnings.warn(
+            "'frame_indices' is deprecated, use 'timepoints' (1-based)",
+            DeprecationWarning,
+            stacklevel=3,
+        )
+        if timepoints is None:
+            timepoints = [int(i) + 1 for i in frame_indices]
+    if timepoints is None:
+        return None
+    if isinstance(timepoints, (int, np.integer)):
+        return [int(timepoints)]
+    return [int(t) for t in timepoints]
+
+
 def pipeline(
     input_data,
     save_path: str | Path = None,
@@ -793,7 +830,9 @@ def pipeline(
     force_reg: bool = False,
     force_detect: bool = False,
     replot: bool = True,
+    timepoints: list | int | None = None,
     num_timepoints: int = None,
+    num_zplanes: int = None,
     frame_indices: list | None = None,
     dff_window_size: int = None,
     dff_percentile: int = 20,
@@ -859,16 +898,21 @@ def pipeline(
         Regenerate per-plane figures. Set False to skip per-plane figure
         regeneration (e.g. the volumetric aggregate over already-plotted
         planes); suite2p and the volumetric plots are unaffected.
-    num_timepoints : int, optional
-        Limit processing to first N frames (truncation only). For an
-        explicit set of frames or a strided selection, use
-        ``frame_indices`` instead.
-    frame_indices : list[int], optional
-        Explicit 0-based frame indices to process. Supports stride
-        (e.g. ``list(range(0, 1574, 2))`` for every other frame).
+    timepoints : list[int] or int, optional
+        Explicit 1-based timepoints to process. Supports stride
+        (e.g. ``list(range(1, 1575, 2))`` for every other timepoint).
         When provided, the implicit stride is used by `OutputMetadata`
         to reactively scale `fs` (e.g. stride of 2 → `fs / 2` in the
         output ops.npy). Takes precedence over ``num_timepoints``.
+    num_timepoints : int, optional
+        Limit processing to first N timepoints (truncation only). For an
+        explicit set or a strided selection, use ``timepoints`` instead.
+    num_zplanes : int, optional
+        Limit processing to the first N z-planes. Shortcut for
+        ``planes=[1..N]``; ignored when ``planes`` is given.
+    frame_indices : list[int], optional
+        Deprecated alias for ``timepoints`` (0-based). Emits a
+        DeprecationWarning.
     dff_window_size : int, optional
         Window size for rolling percentile dF/F baseline (frames).
         If None, auto-calculated as ~10 * tau * fs.
@@ -981,7 +1025,15 @@ def pipeline(
             DeprecationWarning,
             stacklevel=2,
         )
-        num_timepoints = num_frames
+        if num_timepoints is None:
+            num_timepoints = num_frames
+
+    # canonical 1-based timepoint selection (frames/frame_indices deprecated).
+    timepoints = _resolve_timepoints(timepoints, kwargs.pop("frames", None), frame_indices)
+    frame_indices = None
+    # num_zplanes is a count shortcut for planes=[1..N].
+    if num_zplanes is not None and planes is None:
+        planes = list(range(1, int(num_zplanes) + 1))
 
     # flatten (db, settings) into ops so downstream run_volume / run_plane
     # don't each need to forward the pair. explicit ops keys still win.
@@ -991,14 +1043,10 @@ def pipeline(
 
     reader_kwargs = reader_kwargs or {}
     writer_kwargs = writer_kwargs or {}
+    # num_timepoints truncation reaches the writer; an explicit `timepoints`
+    # selection is forwarded as a param and rebuilt by run_plane.
     if num_timepoints is not None:
-        writer_kwargs["num_frames"] = num_timepoints
-
-    # 1-based frame numbers.
-    if frame_indices is not None:
-        writer_kwargs["frames"] = [int(i) + 1 for i in frame_indices]
-        # don't double-pass num_frames; len(frame_indices) is implicit
-        writer_kwargs.pop("num_frames", None)
+        writer_kwargs["num_timepoints"] = num_timepoints
 
     # Always load array to check dimensions and ensure downstream functions have the array shape
     # If input is already array, this is fast. If path or list of paths, it loads lazy array.
@@ -1034,7 +1082,7 @@ def pipeline(
             force_reg=force_reg,
             force_detect=force_detect,
             replot=replot,
-            frame_indices=frame_indices,
+            timepoints=timepoints,
             dff_window_size=dff_window_size,
             dff_percentile=dff_percentile,
             dff_smooth_window=dff_smooth_window,
@@ -1074,7 +1122,7 @@ def pipeline(
             force_reg=force_reg,
             force_detect=force_detect,
             replot=replot,
-            frame_indices=frame_indices,
+            timepoints=timepoints,
             dff_window_size=dff_window_size,
             dff_percentile=dff_percentile,
             dff_smooth_window=dff_smooth_window,
@@ -1246,6 +1294,9 @@ def run_volume(
     force_reg: bool = False,
     force_detect: bool = False,
     replot: bool = True,
+    timepoints: list | int | None = None,
+    num_timepoints: int = None,
+    num_zplanes: int = None,
     frame_indices: list | None = None,
     dff_window_size: int = None,
     dff_percentile: int = 20,
@@ -1337,6 +1388,17 @@ def run_volume(
     _resolve_gpu_env()
     _apply_thread_limits(threads_per_worker)
 
+    # canonical 1-based timepoints (frames/frame_indices deprecated); keep a
+    # 0-based frame_indices for this function's reactive-metadata plumbing,
+    # and forward `timepoints` to run_plane.
+    timepoints = _resolve_timepoints(timepoints, kwargs.pop("frames", None), frame_indices)
+    frame_indices = [int(t) - 1 for t in timepoints] if timepoints is not None else None
+    if num_zplanes is not None and planes is None:
+        planes = list(range(1, int(num_zplanes) + 1))
+    writer_kwargs = dict(writer_kwargs or {})
+    if num_timepoints is not None:
+        writer_kwargs.setdefault("num_timepoints", num_timepoints)
+
     # Handle input data
     input_arr = None
     input_paths = []
@@ -1430,7 +1492,7 @@ def run_volume(
         force_reg=force_reg,
         force_detect=force_detect,
         replot=replot,
-        frame_indices=frame_indices,
+        timepoints=timepoints,
         dff_window_size=dff_window_size,
         dff_percentile=dff_percentile,
         dff_smooth_window=dff_smooth_window,
@@ -2283,6 +2345,8 @@ def run_plane(
     force_reg: bool = False,
     force_detect: bool = False,
     replot: bool = True,
+    timepoints: list | int | None = None,
+    num_timepoints: int = None,
     frame_indices: list | None = None,
     dff_window_size: int = None,
     dff_percentile: int = 20,
@@ -2353,13 +2417,18 @@ def run_plane(
         Example: {"n_clusters": 50, "n_PCs": 64}.
     save_json : bool, default False
         Save ops as JSON.
-    frame_indices : list[int], optional
-        Explicit 0-based timepoint indices. Supports stride
-        (e.g. ``list(range(0, 1574, 2))`` for every other frame).
+    timepoints : list[int] or int, optional
+        Explicit 1-based timepoints. Supports stride
+        (e.g. ``list(range(1, 1575, 2))`` for every other timepoint).
         When provided, the binary on disk contains exactly these
-        frames, and `OutputMetadata` reactively scales `fs` in ops.npy
-        based on the implicit stride. Takes precedence over any
-        ``num_frames`` in ``writer_kwargs``.
+        timepoints, and `OutputMetadata` reactively scales `fs` in ops.npy
+        based on the implicit stride. Takes precedence over
+        ``num_timepoints``.
+    num_timepoints : int, optional
+        Limit processing to first N timepoints (truncation only).
+    frame_indices : list[int], optional
+        Deprecated alias for ``timepoints`` (0-based). Emits a
+        DeprecationWarning.
     plane_name : str, optional
         Custom name for the plane subdirectory.
     reader_kwargs : dict, optional
@@ -2379,6 +2448,14 @@ def run_plane(
 
     _resolve_gpu_env()
 
+    # canonical 1-based timepoints (frames/frame_indices deprecated); convert to
+    # the 0-based frame_indices this function consumes internally.
+    timepoints = _resolve_timepoints(timepoints, kwargs.pop("frames", None), frame_indices)
+    frame_indices = [int(t) - 1 for t in timepoints] if timepoints is not None else None
+    writer_kwargs = dict(writer_kwargs or {})
+    if num_timepoints is not None:
+        writer_kwargs.setdefault("num_timepoints", num_timepoints)
+
     progress_callback = kwargs.pop("progress_callback", None)
 
     if "debug" in kwargs:
@@ -2637,7 +2714,8 @@ def run_plane(
         else:
             # prefer the user-specified frame limit over raw array shape
             nframes_hint = (
-                writer_kwargs.get("num_frames")
+                writer_kwargs.get("num_timepoints")
+                or writer_kwargs.get("num_frames")
                 or ops.get("nframes")
             )
             if not nframes_hint and input_arr is not None and hasattr(input_arr, "shape"):
@@ -2666,8 +2744,8 @@ def run_plane(
         ops_file = plane_dir / "ops.npy"
 
         # extract expected dims from input for cache validation
-        # honors writer_kwargs["num_frames"] limit if user requested fewer frames
-        exp_nframes = writer_kwargs.get("num_frames")
+        # honors num_timepoints truncation if user requested fewer frames
+        exp_nframes = writer_kwargs.get("num_timepoints") or writer_kwargs.get("num_frames")
         exp_ly = exp_lx = None
         if input_arr is not None and hasattr(input_arr, "shape"):
             if exp_nframes is None:
@@ -2796,12 +2874,13 @@ def run_plane(
         write_planes = [plane] if _get_num_planes(file) > 1 else None
 
         write_kw = dict(writer_kwargs)
-        # If the caller gave us explicit frame indices, pass them as
-        # `frames=` (1-based) to imwrite. This wins over any stale
-        # `num_frames` truncation in writer_kwargs — strided semantics
-        # require an explicit index list, not a count.
+        # If the caller gave us explicit timepoints, pass them as
+        # `timepoints=` (1-based) to imwrite. This wins over any stale
+        # truncation count in writer_kwargs — strided semantics require an
+        # explicit index list, not a count.
         if frame_indices is not None:
-            write_kw["frames"] = [int(i) + 1 for i in frame_indices]
+            write_kw["timepoints"] = [int(i) + 1 for i in frame_indices]
+            write_kw.pop("num_timepoints", None)
             write_kw.pop("num_frames", None)
 
         imwrite(

lbm_suite2p_python/utils.py

@@ -111,6 +111,91 @@ def get_common_path(ops_files: list | tuple):
         return Path(os.path.commonpath(ops_files))
 
 
+def estimate_peak_memory(ops, Ly, Lx, n_frames, device="cuda", workers=1):
+    """
+    Estimate peak memory for one Suite2p plane from its parameters.
+
+    Registration and detection run sequentially in suite2p's pipeline, so
+    the peak for a plane is ``max(registration, detection)``, not the sum.
+    The two stages load different pools:
+
+    - Registration compute runs on ``device``. When ``device`` is cuda the
+      per-batch float32/FFT buffers live in VRAM, scaled by
+      ``batch_size * Ly * Lx``. The reference-image correlation
+      (``pick_initial_reference``) and the binary read stay on host.
+    - Detection's binned movie is a host numpy array, and the default
+      ``sparsery`` / ``sourcery`` detectors run on CPU. ``device`` is only
+      used by the cellpose path, which sees the 2D meanImg / max_proj, not
+      the movie. So the binned movie never enters VRAM.
+
+    Host RAM therefore peaks during detection (binned movie plus the
+    high-pass / sparsery copies, ~2.5x); VRAM peaks during registration
+    (or cellpose inference, when enabled). Neither VRAM term scales with
+    ``n_frames``; host detection plateaus once ``n_frames // bin_size``
+    exceeds ``nbins``.
+
+    Parameters
+    ----------
+    ops : dict
+        Flat Suite2p ops. Reads ``nimg_init``, ``batch_size`` (registration
+        batch), ``nbins``, ``bin_size``, ``tau``, ``fs``, ``nchannels``, and
+        ``anatomical_only``. Missing keys fall back to suite2p defaults.
+    Ly, Lx : int
+        Frame height and width in pixels. The detection crop (yrange/xrange)
+        is unknown before registration, so full ``Ly``/``Lx`` are used as an
+        upper bound.
+    n_frames : int
+        Number of frames in the plane.
+    device : str, optional (default "cuda")
+        Torch device. VRAM terms are reported only when this starts with
+        "cuda".
+    workers : int, optional (default 1)
+        Concurrent plane workers. Per-plane peaks are multiplied by this for
+        the ``*_total`` fields.
+
+    Returns
+    -------
+    dict
+        Bytes for ``host_per_plane``, ``vram_per_plane``, ``host_total``,
+        ``vram_total``. VRAM fields are 0 when ``device`` is not cuda.
+
+    Notes
+    -----
+    The 2.5x detection multiplier and the cpsam VRAM constant are rough; the
+    real values depend on data and hardware. Calibrate with
+    ``torch.cuda.max_memory_allocated()`` and an RSS sample around the
+    registration / detection calls for tight bounds.
+    """
+    cuda = str(device).startswith("cuda")
+
+    nimg_init = min(int(ops.get("nimg_init", 400)), n_frames)
+    reg_host = nimg_init * Ly * Lx * 10  # int16 frames + float64 ref corr (CPU)
+
+    nbins = int(ops.get("nbins", 5000))
+    bin_size = ops.get("bin_size") or max(
+        1, n_frames // nbins, round(ops.get("tau", 1.0) * ops.get("fs", 10.0))
+    )
+    nbinned = min(nbins, n_frames // bin_size)
+    detect_host = int(2.5 * nbinned * Ly * Lx * 4)  # binned movie + hp/sparsery copies (CPU)
+
+    host_peak = max(reg_host, detect_host)
+
+    vram_peak = 0
+    if cuda:
+        reg_batch = int(ops.get("batch_size", 100))
+        nchan = int(ops.get("nchannels", 1))
+        reg_vram = nchan * 8 * reg_batch * Ly * Lx * 4  # ~8 float32/FFT buffers per batch
+        cpsam_vram = 1_500_000_000 if ops.get("anatomical_only", 0) else 0  # cpsam weights + activations
+        vram_peak = max(reg_vram, cpsam_vram)
+
+    return {
+        "host_per_plane": host_peak,
+        "vram_per_plane": vram_peak,
+        "host_total": host_peak * max(1, workers),
+        "vram_total": vram_peak * max(1, workers),
+    }
+
+
 def bin1d(X, bin_size, axis=0):
     """
     Mean bin over `axis` of `X` with bin `bin_size`.

{lbm_suite2p_python-3.2.0.dist-info → lbm_suite2p_python-3.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lbm_suite2p_python
-Version: 3.2.0
+Version: 3.2.2
 Summary: Calcium Imaging Pipeline built with Suite2p, Cellpose and Rastermap
 License-Expression: BSD-3-Clause
 Project-URL: homepage, https://github.com/MillerBrainObservatory/LBM-Suite2p-Python
@@ -11,9 +11,11 @@ Classifier: Programming Language :: Python :: 3 :: Only
 Requires-Python: <3.14,>=3.12.7
 Description-Content-Type: text/markdown
 License-File: LICENSE.md
-Requires-Dist: mbo_utilities>=3.2.1
+Requires-Dist: mbo_utilities>=3.2.7
 Requires-Dist: suite2p>=1.0.0.1
 Requires-Dist: setuptools<81
+Requires-Dist: torch
+Requires-Dist: torchvision
 Provides-Extra: rastermap
 Requires-Dist: rastermap; extra == "rastermap"
 Provides-Extra: cellpose

{lbm_suite2p_python-3.2.0.dist-info → lbm_suite2p_python-3.2.2.dist-info}/RECORD

@@ -2,7 +2,7 @@ lbm_suite2p_python/__init__.py,sha256=vllkrB3KK_xr6INPgzfqPHNmiH1MxfHObycRJbP7Tm
 lbm_suite2p_python/__main__.py,sha256=lMnZnMkCbv8OTh5i5uZlG1sF33PEs42zITw8aG8YKwI,131
 lbm_suite2p_python/_benchmarking.py,sha256=oi7q5VDShYa-YqEiZluJEmOeq-_X2jzz-bjxIalhLZ4,172
 lbm_suite2p_python/cellpose.py,sha256=a7H1-458YYUiZhAGbqDdAosPmQlsDbkllm1ec_joCc0,63246
-lbm_suite2p_python/cli.py,sha256=Hf1F3DioABcdXSt4pRdyjHiNILx3g_OPmom3ntkdU3I,21301
+lbm_suite2p_python/cli.py,sha256=QUrW6OKNJjhk5IU9GNeNdNuwVW3uMl3W60QBsXWduHE,26629
 lbm_suite2p_python/conversion.py,sha256=PnrtlpzBXIX0QYu2WxjuJfdWwc-KIr9DhCIo7hSi0f4,31910
 lbm_suite2p_python/db_settings.py,sha256=kYJuv_-7eqPdngdEj5uYlJjavEoZyRWxCEeByABX0QE,22519
 lbm_suite2p_python/default_ops.py,sha256=muReAJIzuBTsS-pmvPeYacmUz4DZptKteQNQhDUoago,3594
@@ -10,13 +10,13 @@ lbm_suite2p_python/grid_search.py,sha256=4Sgl2Nx95jHrpm4S9yj0z4R6tR8CrslPRAM8xiD
 lbm_suite2p_python/gui.py,sha256=jA8dg1yEM0ikMNJ0UrE_6joQkuqsmDj9I_pIb-E0YpA,10594
 lbm_suite2p_python/merging.py,sha256=qcQmAnaJPPniSMDDVeu_uZTcstXMYNfvWUBU4zVKLEI,13928
 lbm_suite2p_python/postprocessing.py,sha256=ESCQCIpxyjwAJTk-pS-Y1LPJ0f1zeXkc7sPYXwbscEA,59560
-lbm_suite2p_python/run_lsp.py,sha256=iTHo7gX7j9BL-NLT04CgkP6yL2OpuudBMGQY-zR7s9M,130401
-lbm_suite2p_python/utils.py,sha256=gYoRyEdZ_YBRH62wuOX2zkXcd57NaY7pKKoxvOi42HM,3892
+lbm_suite2p_python/run_lsp.py,sha256=Z_r-cqika-YBTPQYOB2LccpsRwmWjmXeZdehEZFx0Wo,134031
+lbm_suite2p_python/utils.py,sha256=MmUe5lhOIQb-TBB3KtW_wu4l4jw8hbpkwNUOFGF9S9o,7557
 lbm_suite2p_python/volume.py,sha256=kk_YLXQBPGkzdTVhxfH_Sc4TJLN9I_WztkHZPu8x-tc,80784
 lbm_suite2p_python/zplane.py,sha256=BQ6fZidHHPl2GHoPY2R_Po6givVZjAEzH4I2tQfnI8s,177426
-lbm_suite2p_python-3.2.0.dist-info/licenses/LICENSE.md,sha256=iZa6bLG-5EfV2LXTfQQ7wdDwslaIc9SrOfyKRtzvCoQ,1981
-lbm_suite2p_python-3.2.0.dist-info/METADATA,sha256=bX9ui65k7OJj3ck3KEiglystIkrWqAndXo_oVKV0EjY,8830
-lbm_suite2p_python-3.2.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
-lbm_suite2p_python-3.2.0.dist-info/entry_points.txt,sha256=Ow6TsVyBZI0viAeQP6Ty7GJAmoXD7RmMhIxongAWAVM,100
-lbm_suite2p_python-3.2.0.dist-info/top_level.txt,sha256=WhCsVG68YLKWeTNJVwueG_4H3uuiI22I_TERsF-5mGk,19
-lbm_suite2p_python-3.2.0.dist-info/RECORD,,
+lbm_suite2p_python-3.2.2.dist-info/licenses/LICENSE.md,sha256=iZa6bLG-5EfV2LXTfQQ7wdDwslaIc9SrOfyKRtzvCoQ,1981
+lbm_suite2p_python-3.2.2.dist-info/METADATA,sha256=NDrNEkrUim62CZub3T8HiGlV6eI0xMO2F3s7AdAC83Y,8880
+lbm_suite2p_python-3.2.2.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+lbm_suite2p_python-3.2.2.dist-info/entry_points.txt,sha256=Ow6TsVyBZI0viAeQP6Ty7GJAmoXD7RmMhIxongAWAVM,100
+lbm_suite2p_python-3.2.2.dist-info/top_level.txt,sha256=WhCsVG68YLKWeTNJVwueG_4H3uuiI22I_TERsF-5mGk,19
+lbm_suite2p_python-3.2.2.dist-info/RECORD,,