PMD Compression

Simple PMD compression, the hello world of masknmf

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

  0%|          | 0/250 [00:00<?, ?it/s]
  3%|▎         | 7/250 [00:00<00:03, 63.00it/s]
  6%|▌         | 14/250 [00:00<00:03, 64.32it/s]
  8%|▊         | 21/250 [00:00<00:03, 64.92it/s]
 11%|█         | 28/250 [00:00<00:03, 65.67it/s]
 14%|█▍        | 35/250 [00:00<00:03, 65.94it/s]
 17%|█▋        | 42/250 [00:00<00:03, 66.12it/s]
 20%|█▉        | 49/250 [00:00<00:03, 66.73it/s]
 22%|██▏       | 56/250 [00:00<00:02, 65.69it/s]
 25%|██▌       | 63/250 [00:00<00:02, 64.45it/s]
 28%|██▊       | 70/250 [00:01<00:02, 63.06it/s]
 31%|███       | 77/250 [00:01<00:02, 60.39it/s]
 34%|███▎      | 84/250 [00:01<00:02, 59.52it/s]
 36%|███▌      | 90/250 [00:01<00:02, 58.64it/s]
 39%|███▉      | 97/250 [00:01<00:02, 60.19it/s]
 42%|████▏     | 104/250 [00:01<00:02, 61.75it/s]
 44%|████▍     | 111/250 [00:01<00:02, 62.41it/s]
 47%|████▋     | 118/250 [00:01<00:02, 63.70it/s]
 50%|█████     | 125/250 [00:01<00:01, 64.18it/s]
 53%|█████▎    | 132/250 [00:02<00:01, 65.02it/s]
 56%|█████▌    | 139/250 [00:02<00:01, 66.06it/s]
 58%|█████▊    | 146/250 [00:02<00:01, 66.37it/s]
 61%|██████    | 153/250 [00:02<00:01, 66.59it/s]
 64%|██████▍   | 160/250 [00:02<00:01, 67.13it/s]
 67%|██████▋   | 167/250 [00:02<00:01, 67.25it/s]
 70%|██████▉   | 174/250 [00:02<00:01, 66.95it/s]
 72%|███████▏  | 181/250 [00:02<00:01, 67.23it/s]
 75%|███████▌  | 188/250 [00:02<00:00, 66.80it/s]
 78%|███████▊  | 195/250 [00:03<00:00, 66.73it/s]
 81%|████████  | 202/250 [00:03<00:00, 66.75it/s]
 84%|████████▎ | 209/250 [00:03<00:00, 66.59it/s]
 86%|████████▋ | 216/250 [00:03<00:00, 66.88it/s]
 89%|████████▉ | 223/250 [00:03<00:00, 66.85it/s]
 92%|█████████▏| 230/250 [00:03<00:00, 66.71it/s]
 95%|█████████▍| 237/250 [00:03<00:00, 65.71it/s]
 98%|█████████▊| 244/250 [00:03<00:00, 65.91it/s]
100%|██████████| 250/250 [00:03<00:00, 65.03it/s]
[26-03-26 10:00:28]: Running Blockwise Decompositions

  0%|          | 0/3 [00:00<?, ?it/s]
 67%|██████▋   | 2/3 [00:00<00:00, 15.25it/s]
100%|██████████| 3/3 [00:00<00:00, 15.31it/s]
/home/runner/work/masknmf-toolbox/masknmf-toolbox/masknmf/compression/decomposition.py:1352: UserWarning: Sparse invariant checks are implicitly disabled. Memory errors (e.g. SEGFAULT) will occur when operating on a sparse tensor which violates the invariants, but checks incur performance overhead. To silence this warning, explicitly opt in or out. See `torch.sparse.check_sparse_tensor_invariants.__doc__` for guidance.  (Triggered internally at /pytorch/aten/src/ATen/Context.cpp:760.)
  u_spatial_fit = torch.sparse_coo_tensor(
[26-03-26 10:00:28]: Constructed U matrix. Rank of U is 199
[26-03-26 10:00:28]: PMD Objected constructed
/opt/hostedtoolcache/Python/3.12.13/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()