Merge branch 'main' of https://github.com/MouseLand/suite2p

Author: chriski777

docs/index.md

@@ -15,8 +15,7 @@ includes the following modules:
 For examples of how the output looks and how the GUI works, check out
 this twitter [thread](https://twitter.com/marius10p/status/1032804776633880583).
 
-This code was written by Carsen Stringer and Marius Pachitariu. For
-support, please open an [issue](https://github.com/MouseLand/suite2p/issues).
+For software support, please open an [issue](https://github.com/MouseLand/suite2p/issues). For general questions, please use [forum.image.sc](https://forum.image.sc)
 
 The reference paper is [here](https://www.biorxiv.org/content/10.64898/2026.02.04.703741v1). The deconvolution algorithm is based on
 [this paper](https://journals.plos.org/ploscompbiol/article?id=10.1371/journal.pcbi.1005423), with settings based on [this

notebooks/example_run_suite2p_2026.ipynb

@@ -8,11 +8,20 @@
       "source": [
         "# Running suite2p on example data\n",
         "\n",
+        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/MouseLand/suite2p/blob/main/notebooks/example_run_suite2p_2026.ipynb)\n",
+        "\n",
         "This notebook will guide you through the various stages and outputs of suite2p by running it on a real-life dataset. This is data collected from a wild-type mouse injected with GCaMP6s in primary visual cortex. The recording was collected at 13Hz (there were 3 planes in the recording, 1 is included here).\n",
         "\n",
         "The code cell below installs and imports the necessary packages to use suite2p in this Colab notebook."
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "To use the GPU-acceleration, connect to a hosted runtime with a GPU in colab."
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -198,6 +207,7 @@
         "settings['detection']['threshold_scaling'] = 2.0 # we are increasing the threshold for finding ROIs to limit the number of non-cell ROIs found (sometimes useful in gcamp injections)\n",
         "settings['fs'] = 13 # sampling rate of recording, determines binning for cell detection\n",
         "settings['tau'] = 1.25 # timescale of gcamp to use for deconvolution\n",
+        "settings['device'] = 'cuda' if torch.cuda.is_available() else 'cpu' # use GPU if available for faster processing\n",
         "print(settings)"
       ]
     },
@@ -234,6 +244,8 @@
         "## Run entire suite2p pipeline on data\n",
         "The suite2p.run_s2p function runs the pipeline and returns a list of output dictionaries containing the pipeline parameters used and extra data calculated along the way, one for each plane.\n",
         "\n",
+        "To see the logs during running, you now need to run `logger_setup`, optionally providing the `save_path` for the logs to be written to a text file.\n",
+        "\n",
         "The following cell might take a couple of minutes to run."
       ]
     },
@@ -249,6 +261,8 @@
       },
       "outputs": [],
       "source": [
+        "from suite2p.run_s2p import logger_setup\n",
+        "logger_setup()\n",
         "suite2p.run_s2p(settings=settings, db=db)"
       ]
     },
@@ -355,7 +369,8 @@
       },
       "outputs": [],
       "source": [
-        "device = torch.device(\"cuda\")"
+        "settings['device'] = 'cuda' if torch.cuda.is_available() else 'cpu' # use GPU if available for faster processing\n",
+        "device = torch.device(settings['device'])"
       ]
     },
     {
@@ -607,13 +622,9 @@
       "outputs": [],
       "source": [
         "# Baseline correction of fluorescence traces\n",
-        "if settings[\"dcnv_preprocess\"][\"baseline\"] == \"maximin\":\n",
-        "    dF =  extraction.baseline_maximin(dF, win_baseline=settings[\"dcnv_preprocess\"][\"win_baseline\"],\n",
-        "                                      sig_baseline=settings[\"dcnv_preprocess\"][\"sig_baseline\"],\n",
-        "                                      fs=settings[\"fs\"], device=device)\n",
-        "else:\n",
-        "    dF = extraction.preprocess(F=dF, fs=settings[\"fs\"],\n",
-        "                                **settings[\"dcnv_preprocess\"])"
+        "dF = extraction.preprocess(F=dF, fs=settings[\"fs\"], device=device,\n",
+        "                           batch_size=settings[\"extraction\"][\"batch_size\"],\n",
+        "                           **settings[\"dcnv_preprocess\"])"
       ]
     },
     {

suite2p/parameters.py

@@ -495,7 +495,15 @@
             "default": False,
             "description": "Whether to save chan2 registered tiffs.",
         },
-        
+        "upsample_meanImg": {
+            "gui_name": "Upsample meanImg",
+            "type": list,
+            "min": None,
+            "max": None,
+            "default": None,
+            "description": "Upsampling factor for super-resolution meanImg computation. If 0 or None, no upsampling is performed.",
+        },
+
     },
     "detection": {
         "algorithm": {

suite2p/registration/nonrigid.py

@@ -370,7 +370,8 @@ def phasecorr(data, blocks, maskMul, maskOffset, cfRefImg, snr_thresh,
 
     return ymax1.T.float(), xmax1.T.float(), cmax1.T, ccsm, ccb
 
-def transform_data(data, nblocks, xblock, yblock, ymax1, xmax1):
+def transform_data(data, nblocks, xblock, yblock, ymax1, xmax1, 
+                   data_ups=None, counts_ups=None):
     """
     Apply bilinear interpolation to transform image data using block-wise shifts.
     This function performs non-rigid image registration by interpolating block-wise 
@@ -399,7 +400,7 @@ def transform_data(data, nblocks, xblock, yblock, ymax1, xmax1):
         Shifted image data of shape (nimg, Ly, Lx) with dtype int16 (short).
         The input images are warped according to the interpolated displacement field.
     """
-    _, Ly, Lx = data.shape
+    n_frames, Ly, Lx = data.shape
     #device = torch.device("cuda")
     #data = torch.from_numpy(data).to(device).float()
     device = data.device
@@ -414,12 +415,33 @@ def transform_data(data, nblocks, xblock, yblock, ymax1, xmax1):
                          size=(Lyc, Lxc), mode="bilinear", align_corners=True)
     yxup = F.pad(yxup, (int(xb.min()), Lx - int(xb.max()), 
                         int(yb.min()), Ly - int(yb.max())), mode="replicate")
+    
+    if data_ups is not None and counts_ups is not None:
+        ups = torch.Tensor([data_ups.shape[0] // Ly, data_ups.shape[1] // Lx]).to(device)
+        yxup_round = -1*yxup.clone() + torch.stack((mshy, mshx), dim=0)
+        yxup_round = torch.floor(0.5 + (yxup_round * ups.unsqueeze(-1).unsqueeze(-1))).long()
+        yxup_round[:,0] = torch.clamp(yxup_round[:,0], min=0, max=Ly*ups[0] - 1)
+        yxup_round[:,1] = torch.clamp(yxup_round[:,1], min=0, max=Lx*ups[1] - 1)
+        for t in range(n_frames):
+            mat = torch.sparse_coo_tensor(indices=yxup_round[t].reshape(2, -1), values=data[t].flatten(), 
+                                            size=(int(Ly*ups[0]), int(Lx*ups[1]))).to_dense()
+            data_ups += mat 
+            mat = torch.sparse_coo_tensor(indices=yxup_round[t].reshape(2, -1), 
+                                          values=torch.ones(data[t].numel(), device=device, dtype=torch.long), 
+                                          size=(int(Ly*ups[0]), int(Lx*ups[1]))).to_dense()
+            counts_ups += mat
+            
+            # data_ups[yxup_round[t,0], yxup_round[t,1]] += data[t]
+            # counts_ups[yxup_round[t,0], yxup_round[t,1]] += 1
+
+    # rescale for grid_sample
     yxup[:,0] += mshy
     yxup[:,1] += mshx
     yxup /= torch.Tensor([Ly-1, Lx-1]).to(device).unsqueeze(-1).unsqueeze(-1)
     yxup *= 2 
     yxup -= 1
     yxup = yxup.permute(0, 2, 3, 1)
+
     if device.type == "mps":
         # Manually pad the input tensor with the border values
         data_padded = F.pad(data.float().unsqueeze(1), (1, 1, 1, 1), mode="replicate")
@@ -437,4 +459,6 @@ def transform_data(data, nblocks, xblock, yblock, ymax1, xmax1):
     else:
         fr_shift = F.grid_sample(data.float().unsqueeze(1), yxup[:,:,:,[1,0]], 
                              mode="bilinear", padding_mode="border", align_corners=True)
+        
+    
     return fr_shift.squeeze().short()#.cpu().numpy()

suite2p/registration/register.py

@@ -11,6 +11,7 @@
 import numpy as np
 import torch
 from scipy.signal import medfilt
+from scipy.ndimage import gaussian_filter
 
 import logging 
 logger = logging.getLogger(__name__)
@@ -402,7 +403,8 @@ def compute_shifts(refAndMasks, fr_reg, maxregshift=0.1, smooth_sigma_time=0,
 
     return ymax, xmax, cmax, ymax1, xmax1, cmax1, None, None
 
-def shift_frames(fr_torch, yoff, xoff, yoff1=None, xoff1=None, blocks=None, device=torch.device("cuda")):
+def shift_frames(fr_torch, yoff, xoff, yoff1=None, xoff1=None, blocks=None, 
+                 mean_img_ups=None, counts_ups=None, device=torch.device("cuda")):
     """
     Apply rigid and optionally nonrigid shifts to frames and return as numpy int16.
 
@@ -432,7 +434,7 @@ def shift_frames(fr_torch, yoff, xoff, yoff1=None, xoff1=None, blocks=None, devi
     """
     fr_torch = torch.stack([torch.roll(frame, shifts=(-dy, -dx), dims=(0, 1))
                                for frame, dy, dx in zip(fr_torch, yoff, xoff)], axis=0)
-
+    
     if yoff1 is not None:
         if isinstance(yoff1, np.ndarray):
             if fr_torch.device.type == "cuda":
@@ -441,7 +443,9 @@ def shift_frames(fr_torch, yoff, xoff, yoff1=None, xoff1=None, blocks=None, devi
             else:
                 yoff1 = torch.from_numpy(yoff1).to(device)
                 xoff1 = torch.from_numpy(xoff1).to(device)
-        fr_torch = nonrigid.transform_data(fr_torch, blocks[2], blocks[1], blocks[0], yoff1, xoff1)
+    
+        fr_torch = nonrigid.transform_data(fr_torch, blocks[2], blocks[1], blocks[0], 
+                                           yoff1, xoff1, data_ups=mean_img_ups, counts_ups=counts_ups)
 
     frames_out = np.empty(fr_torch.shape, dtype="int16")
     frames_out = fr_torch.cpu().numpy()
@@ -472,11 +476,11 @@ def normalize_reference_image(refImg):
 
 
 def register_frames(f_align_in, refImg, f_align_out=None, batch_size=100, 
-                    bidiphase=0, 
-                    norm_frames=True, smooth_sigma=1.15, spatial_taper=3.45, 
+                    bidiphase=0, norm_frames=True, smooth_sigma=1.15, spatial_taper=3.45, 
                     block_size=(128,128), nonrigid=True, maxregshift=0.1, 
                     smooth_sigma_time=0, snr_thresh=1.2, maxregshiftNR=5,
-                    device=torch.device("cuda"), tif_root=None, apply_shifts=True):
+                    device=torch.device("cuda"), tif_root=None, apply_shifts=True,
+                    upsample_meanImg=False):
     """
     Register frames to a reference image using rigid and optionally nonrigid shifts.
 
@@ -523,6 +527,12 @@ def register_frames(f_align_in, refImg, f_align_out=None, batch_size=100,
         If provided, save registered frames as tiffs in this directory.
     apply_shifts : bool
         If True, apply computed shifts to frames. If False, only compute shifts.
+    upsample_meanImg : bool, int, list, or tuple
+        Upsampling factor for super-resolution mean image computation.
+        If False or None, no upsampling is performed. If int, same factor is used
+        for both Y and X. If list/tuple of length 2, specifies [Y_factor, X_factor].
+        The mean image is computed by accumulating registered frames at subpixel
+        locations and normalizing by pixel counts.
 
     Returns
     -------
@@ -537,6 +547,15 @@ def register_frames(f_align_in, refImg, f_align_out=None, batch_size=100,
         concatenated across all batches.
     blocks : list
         Block definitions from nonrigid.make_blocks.
+    mean_img_ups : torch.Tensor or None
+        Raw upsampled mean image tensor of shape (Ly*upsample[0], Lx*upsample[1])
+        before normalization. None if upsample_meanImg is False.
+    counts_ups : torch.Tensor or None
+        Pixel counts tensor of shape (Ly*upsample[0], Lx*upsample[1]) indicating
+        how many frames contributed to each upsampled pixel. None if upsample_meanImg is False.
+    meanImg_ups : np.ndarray or None
+        Super-resolution mean image of shape (Ly*upsample[0], Lx*upsample[1])
+        after Gaussian smoothing and normalization by counts. None if upsample_meanImg is False.
     """
 
     n_frames, Ly, Lx = f_align_in.shape
@@ -558,6 +577,13 @@ def register_frames(f_align_in, refImg, f_align_out=None, batch_size=100,
     ### ------------- register frames to reference image ------------ ###
 
     mean_img = np.zeros((Ly, Lx), "float32")
+    if upsample_meanImg:
+        if not isinstance(upsample_meanImg, (np.ndarray, list, tuple)):
+            upsample_meanImg = [upsample_meanImg, upsample_meanImg]
+        mean_img_ups = torch.zeros((int(Ly*upsample_meanImg[0]), int(Lx*upsample_meanImg[1])), dtype=torch.double, device=device)
+        counts_ups = torch.zeros((int(Ly*upsample_meanImg[0]), int(Lx*upsample_meanImg[1])), dtype=torch.int, device=device)
+    else:
+        mean_img_ups, counts_ups, meanImg_ups = None, None, None
 
     n_batches = int(np.ceil(n_frames / batch_size))
     logger.info(f"Registering {n_frames} frames in {n_batches} batches")
@@ -580,8 +606,9 @@ def register_frames(f_align_in, refImg, f_align_out=None, batch_size=100,
         ymax, xmax, cmax, ymax1, xmax1, cmax1, zest, cmax_all = offsets
 
         if apply_shifts:
-            frames = shift_frames(fr_torch, ymax, xmax, ymax1, xmax1, blocks, device)
-        
+            frames = shift_frames(fr_torch, ymax, xmax, ymax1, xmax1, blocks, 
+                                  mean_img_ups=mean_img_ups, counts_ups=counts_ups, device=device)
+            
         # convert to numpy and concatenate offsets
         ymax, xmax, cmax = ymax.cpu().numpy(), xmax.cpu().numpy(), cmax.cpu().numpy()
         if ymax1 is not None:
@@ -608,7 +635,16 @@ def register_frames(f_align_in, refImg, f_align_out=None, batch_size=100,
                 fname = os.path.join(tif_root, f"file{n : 05d}.tif")
                 save_tiff(mov=frames, fname=fname)
 
-    return rmin, rmax, mean_img, offsets_all, blocks
+    if upsample_meanImg:
+        # apply Gaussian smoothing and normalize by counts
+        mimg = mean_img_ups.cpu().numpy()
+        cimg = counts_ups.cpu().numpy()
+        sig = 1 
+        mimg = gaussian_filter(mimg, sig)
+        cimg = gaussian_filter(cimg, sig)
+        meanImg_ups = mimg / cimg
+
+    return rmin, rmax, mean_img, offsets_all, blocks, mean_img_ups, counts_ups, meanImg_ups
 
 def check_offsets(yoff, xoff, yoff1, xoff1, n_frames):
     """
@@ -844,8 +880,9 @@ def registration_wrapper(f_reg, f_raw=None, f_reg_chan2=None, f_raw_chan2=None,
         Dictionary containing registration results with keys: "refImg", "rmin",
         "rmax", "meanImg", "yoff", "xoff", "corrXY", "yoff1", "xoff1",
         "corrXY1", "meanImg_chan2", "badframes", "badframes0", "yrange",
-        "xrange", "bidiphase", "meanImgE", and optionally "zpos_registration"
-        and "cmax_registration".
+        "xrange", "bidiphase", "meanImgE", and optionally "zpos_registration",
+        "cmax_registration", "meanImg_upsample", "mean_img_ups", 
+        and "counts_ups".
     """
     out = assign_reg_io(f_reg, f_raw, f_reg_chan2, f_raw_chan2, align_by_chan2,
                         save_path, settings["reg_tif"], settings["reg_tif_chan2"])
@@ -874,7 +911,7 @@ def registration_wrapper(f_reg, f_raw=None, f_reg_chan2=None, f_raw_chan2=None,
         bidiphase = settings["bidiphase"]
 
     if bidiphase != 0 and refImg is None:
-        frames = bidi.shift(frames, int(settings["bidiphase"])) 
+        frames = bidi.shift(frames, bidiphase) 
 
     if refImg is None:
         t0 = time.time()
@@ -894,15 +931,15 @@ def registration_wrapper(f_reg, f_raw=None, f_reg_chan2=None, f_raw_chan2=None,
 
         ### ----- register frames to reference image -------------- ###
         outputs = register_frames(f_align_in, f_align_out=f_align_out, bidiphase=bidiphase,
-                                refImg=refImg, tif_root=tif_root_align, 
-                                batch_size=settings["batch_size"], 
-                                norm_frames=settings["norm_frames"], smooth_sigma=settings["smooth_sigma"], 
-                                spatial_taper=settings["spatial_taper"], block_size=settings["block_size"], 
+                                refImg=refImg, tif_root=tif_root_align,
+                                batch_size=settings["batch_size"],
+                                norm_frames=settings["norm_frames"], smooth_sigma=settings["smooth_sigma"],
+                                spatial_taper=settings["spatial_taper"], block_size=settings["block_size"],
                                 nonrigid=settings["nonrigid"],
                                 maxregshift=settings["maxregshift"], smooth_sigma_time=settings["smooth_sigma_time"],
                                     snr_thresh=settings["snr_thresh"], maxregshiftNR=settings["maxregshiftNR"],
-                                    device=device)
-        rmin, rmax, mean_img, offsets_all, blocks = outputs
+                                    device=device, upsample_meanImg=settings.get("upsample_meanImg", False))
+        rmin, rmax, mean_img, offsets_all, blocks, mean_img_ups, counts_ups, meanImg_ups = outputs
         yoff, xoff, corrXY, yoff1, xoff1, corrXY1, zest, cmax_all = offsets_all
 
         # compute valid region and timepoints to exclude
@@ -932,16 +969,24 @@ def registration_wrapper(f_reg, f_raw=None, f_reg_chan2=None, f_raw_chan2=None,
     else:
         meanImg_chan2 = None
 
-    reg_outputs = registration_outputs_to_dict(refImg_orig, rmin, rmax, meanImg, 
-                                               (yoff, xoff, corrXY), 
-                                               (yoff1, xoff1, corrXY1), 
-                                               (zest, cmax_all), meanImg_chan2, 
-                                               badframes, badframes0, 
-                                               yrange, xrange, bidiphase)
-    
+    reg_outputs = registration_outputs_to_dict(refImg_orig, rmin, rmax, meanImg,
+                                               (yoff, xoff, corrXY),
+                                               (yoff1, xoff1, corrXY1),
+                                               (zest, cmax_all), meanImg_chan2,
+                                               badframes, badframes0,
+                                               yrange, xrange, bidiphase, 
+                                               )
+
     # add enhanced mean image
     meanImgE = utils.highpass_mean_image(meanImg.astype("float32"), aspect=aspect)
     reg_outputs["meanImgE"] = meanImgE
+
+    # add upsampled mean image if computed
+    if mean_img_ups is not None and counts_ups is not None:
+        reg_outputs["meanImg_upsample"] = meanImg_ups
+        reg_outputs["mean_img_ups"] = mean_img_ups.cpu().numpy()
+        reg_outputs["counts_ups"] = counts_ups.cpu().numpy()
+
     return reg_outputs
 
 def registration_outputs_to_dict(refImg, rmin, rmax, meanImg, rigid_offsets,