PMD Compression

Simple PMD compression, the hello world of masknmf

[25-10-25 22:10:20]: Starting compression
[25-10-25 22:10:20]: sampled from the following regions: [0]
[25-10-25 22:10:20]: We are initializing on a total of 2000 frames
[25-10-25 22:10:20]: Loading data to estimate complete spatial basis
[25-10-25 22:10:20]: skipping the pruning step for frame cutoff
[25-10-25 22:10:20]: Finding spatiotemporal roughness thresholds

  0%|          | 0/250 [00:00<?, ?it/s]
  2%|▏         | 5/250 [00:00<00:05, 43.27it/s]
  4%|▍         | 10/250 [00:00<00:06, 39.23it/s]
  6%|▌         | 14/250 [00:00<00:06, 34.18it/s]
  7%|▋         | 18/250 [00:00<00:06, 33.33it/s]
  9%|▉         | 23/250 [00:00<00:06, 37.15it/s]
 11%|█         | 28/250 [00:00<00:05, 39.58it/s]
 13%|█▎        | 33/250 [00:00<00:05, 41.39it/s]
 15%|█▌        | 38/250 [00:00<00:04, 42.53it/s]
 17%|█▋        | 43/250 [00:01<00:04, 43.28it/s]
 19%|█▉        | 48/250 [00:01<00:04, 43.64it/s]
 21%|██        | 53/250 [00:01<00:04, 44.15it/s]
 23%|██▎       | 58/250 [00:01<00:04, 44.44it/s]
 25%|██▌       | 63/250 [00:01<00:04, 44.23it/s]
 27%|██▋       | 68/250 [00:01<00:04, 44.25it/s]
 29%|██▉       | 73/250 [00:01<00:03, 44.54it/s]
 31%|███       | 78/250 [00:01<00:03, 44.57it/s]
 33%|███▎      | 83/250 [00:01<00:03, 44.42it/s]
 35%|███▌      | 88/250 [00:02<00:03, 44.37it/s]
 37%|███▋      | 93/250 [00:02<00:03, 44.26it/s]
 39%|███▉      | 98/250 [00:02<00:03, 43.99it/s]
 41%|████      | 103/250 [00:02<00:03, 44.19it/s]
 43%|████▎     | 108/250 [00:02<00:03, 43.85it/s]
 45%|████▌     | 113/250 [00:02<00:03, 44.20it/s]
 47%|████▋     | 118/250 [00:02<00:03, 43.37it/s]
 49%|████▉     | 123/250 [00:02<00:02, 43.64it/s]
 51%|█████     | 128/250 [00:02<00:02, 44.08it/s]
 53%|█████▎    | 133/250 [00:03<00:02, 44.51it/s]
 55%|█████▌    | 138/250 [00:03<00:02, 44.70it/s]
 57%|█████▋    | 143/250 [00:03<00:02, 45.00it/s]
 59%|█████▉    | 148/250 [00:03<00:02, 45.08it/s]
 61%|██████    | 153/250 [00:03<00:02, 45.15it/s]
 63%|██████▎   | 158/250 [00:03<00:02, 45.17it/s]
 65%|██████▌   | 163/250 [00:03<00:01, 45.32it/s]
 67%|██████▋   | 168/250 [00:03<00:01, 45.28it/s]
 69%|██████▉   | 173/250 [00:03<00:01, 45.43it/s]
 71%|███████   | 178/250 [00:04<00:01, 45.47it/s]
 73%|███████▎  | 183/250 [00:04<00:01, 45.57it/s]
 75%|███████▌  | 188/250 [00:04<00:01, 45.64it/s]
 77%|███████▋  | 193/250 [00:04<00:01, 45.69it/s]
 79%|███████▉  | 198/250 [00:04<00:01, 45.42it/s]
 81%|████████  | 203/250 [00:04<00:01, 44.94it/s]
 83%|████████▎ | 208/250 [00:04<00:00, 44.82it/s]
 85%|████████▌ | 213/250 [00:04<00:00, 44.75it/s]
 87%|████████▋ | 218/250 [00:04<00:00, 44.56it/s]
 89%|████████▉ | 223/250 [00:05<00:00, 44.66it/s]
 91%|█████████ | 228/250 [00:05<00:00, 44.62it/s]
 93%|█████████▎| 233/250 [00:05<00:00, 44.65it/s]
 95%|█████████▌| 238/250 [00:05<00:00, 44.79it/s]
 97%|█████████▋| 243/250 [00:05<00:00, 44.85it/s]
 99%|█████████▉| 248/250 [00:05<00:00, 44.69it/s]
100%|██████████| 250/250 [00:05<00:00, 43.87it/s]
[25-10-25 22:10:26]: Running Blockwise Decompositions
[25-10-25 22:10:26]: Constructed U matrix. Rank of U is 182
[25-10-25 22:10:26]: PMD Objected constructed
/opt/hostedtoolcache/Python/3.12.11/x64/lib/python3.12/site-packages/fastplotlib/graphics/features/_base.py:18: UserWarning: casting float64 array to float32
  warn(f"casting {array.dtype} array to float32")

# test_example = true

import masknmf
import torch
import fastplotlib as fpl
from urllib.request import urlretrieve
import tifffile


urlretrieve(
    "https://github.com/flatironinstitute/CaImAn/raw/refs/heads/main/example_movies/demoMovie.tif",
    "./demo.tif"
)

# always lazy load raw data by memmaping or other methods
data = tifffile.imread("./demo.tif")

block_sizes = [32, 32]
max_components = 20

# it's recommended to use masknmf on a machine with a GPU
# device = "cuda" if torch.cuda.is_available() else "cpu"
device = "cpu"
# number of frames used to estimate the spatial basis in PMD
num_frames_for_spatial_fit = data.shape[0]

# perform PMD
pmd_result = masknmf.compression.pmd_decomposition(
    data,
    block_sizes,
    num_frames_for_spatial_fit,
    max_components=max_components,
    device=device,
    frame_batch_size=1024
)

# get the residual
pmd_residual = masknmf.PMDResidualArray(data, pmd_result)

# view the movies, note that all these array are LAZY evaluated, allowing you to view extremely large datasets!
iw = fpl.ImageWidget(
    data=[data, pmd_result, pmd_residual],
    names=["raw", "pmd", "residual"],
    figure_kwargs={"size": (1000, 340), "shape": (1, 3)},
    cmap="gnuplot2",
)

iw.show()

# use the time slider or set the frame index programmatically
iw.current_index = {"t": 1610}

# manually set vmin-vmax to emphasize noise in raw video
# you can also adjust the vmin-vmax using the histogram tool
# reset the vmin-vmax by clicking the buttons under "ImageWidget Controls"
for image in iw.managed_graphics:
    image.vmax = 3_200

# remove toolbar to reduce clutter
for subplot in iw.figure:
    subplot.toolbar = False


# ignore the remaining lines these are just for docs generation
figure = iw.figure
if __name__ == "__main__":
    print(__doc__)
    fpl.loop.run()