PMD Compression

Simple PMD compression, the hello world of masknmf

/home/runner/work/masknmf-toolbox/masknmf-toolbox/masknmf/utils/_cuda.py:7: UserWarning: You've explicitly selected to perform computations on the cpu, performance will be significantly slower
  warn(
[25-11-21 22:34:43]: Starting compression
[25-11-21 22:34:43]: sampled from the following regions: [0]
[25-11-21 22:34:43]: We are initializing on a total of 2000 frames
[25-11-21 22:34:43]: Loading data to estimate complete spatial basis
[25-11-21 22:34:43]: skipping the pruning step for frame cutoff
[25-11-21 22:34:43]: Finding spatiotemporal roughness thresholds

  0%|          | 0/250 [00:00<?, ?it/s]
  2%|▏         | 5/250 [00:00<00:05, 46.37it/s]
  4%|▍         | 10/250 [00:00<00:04, 48.11it/s]
  6%|▌         | 15/250 [00:00<00:04, 48.12it/s]
  8%|▊         | 20/250 [00:00<00:04, 48.41it/s]
 10%|█         | 25/250 [00:00<00:04, 48.55it/s]
 12%|█▏        | 30/250 [00:00<00:04, 48.28it/s]
 14%|█▍        | 35/250 [00:00<00:04, 48.56it/s]
 16%|█▌        | 40/250 [00:00<00:04, 48.38it/s]
 18%|█▊        | 45/250 [00:00<00:04, 48.22it/s]
 20%|██        | 50/250 [00:01<00:04, 48.50it/s]
 22%|██▏       | 55/250 [00:01<00:04, 48.64it/s]
 24%|██▍       | 60/250 [00:01<00:03, 48.61it/s]
 26%|██▌       | 65/250 [00:01<00:03, 48.83it/s]
 28%|██▊       | 70/250 [00:01<00:03, 49.05it/s]
 30%|███       | 75/250 [00:01<00:03, 48.62it/s]
 32%|███▏      | 80/250 [00:01<00:03, 46.19it/s]
 34%|███▍      | 85/250 [00:01<00:03, 46.86it/s]
 36%|███▌      | 90/250 [00:01<00:03, 47.24it/s]
 38%|███▊      | 95/250 [00:01<00:03, 47.83it/s]
 40%|████      | 100/250 [00:02<00:03, 48.38it/s]
 42%|████▏     | 105/250 [00:02<00:02, 48.40it/s]
 44%|████▍     | 110/250 [00:02<00:02, 48.76it/s]
 46%|████▌     | 115/250 [00:02<00:02, 48.93it/s]
 48%|████▊     | 120/250 [00:02<00:02, 48.72it/s]
 50%|█████     | 125/250 [00:02<00:02, 48.81it/s]
 52%|█████▏    | 130/250 [00:02<00:02, 48.94it/s]
 54%|█████▍    | 135/250 [00:02<00:02, 48.51it/s]
 56%|█████▌    | 140/250 [00:02<00:02, 48.52it/s]
 58%|█████▊    | 145/250 [00:02<00:02, 48.62it/s]
 60%|██████    | 150/250 [00:03<00:02, 48.44it/s]
 62%|██████▏   | 155/250 [00:03<00:01, 48.67it/s]
 64%|██████▍   | 160/250 [00:03<00:01, 48.79it/s]
 66%|██████▌   | 165/250 [00:03<00:01, 48.51it/s]
 68%|██████▊   | 170/250 [00:03<00:01, 46.97it/s]
 70%|███████   | 175/250 [00:03<00:01, 46.77it/s]
 72%|███████▏  | 180/250 [00:03<00:01, 44.52it/s]
 74%|███████▍  | 185/250 [00:03<00:01, 45.74it/s]
 76%|███████▌  | 190/250 [00:03<00:01, 46.34it/s]
 78%|███████▊  | 195/250 [00:04<00:01, 46.85it/s]
 80%|████████  | 200/250 [00:04<00:01, 47.43it/s]
 82%|████████▏ | 205/250 [00:04<00:00, 47.80it/s]
 84%|████████▍ | 210/250 [00:04<00:00, 47.90it/s]
 86%|████████▌ | 215/250 [00:04<00:00, 48.31it/s]
 88%|████████▊ | 220/250 [00:04<00:00, 48.42it/s]
 90%|█████████ | 225/250 [00:04<00:00, 48.29it/s]
 92%|█████████▏| 230/250 [00:04<00:00, 48.37it/s]
 94%|█████████▍| 235/250 [00:04<00:00, 48.60it/s]
 96%|█████████▌| 240/250 [00:04<00:00, 48.50it/s]
 98%|█████████▊| 245/250 [00:05<00:00, 48.68it/s]
100%|██████████| 250/250 [00:05<00:00, 48.70it/s]
100%|██████████| 250/250 [00:05<00:00, 48.08it/s]
[25-11-21 22:34:48]: Running Blockwise Decompositions
[25-11-21 22:34:49]: Constructed U matrix. Rank of U is 183
[25-11-21 22:34:49]: PMD Objected constructed
/opt/hostedtoolcache/Python/3.12.12/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()