Runing with large scale data (Stereo-seq)

import packages and load data

import numpy as np
import pandas as pd
import scanpy as sc
from scipy import sparse, linalg, stats
import sys
import tqdm
import seaborn as sns
import matplotlib.pyplot as plt
sc.settings.set_figure_params(dpi=150)
import matplotlib as mpl
mpl.rcParams['pdf.fonttype']  = 42

# subsample if necessary
scdata_raw = sc.read_h5ad('Path/to/scdata')

stdata_raw = sc.read_h5ad('Path/to/stdata')

evaluate peak mem, GPU and time

from memory_profiler import memory_usage
from time import time

import torch
import pynvml

# Initialize NVML
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)

import Polyomino as plmo

start_time = time()
def plmo_mem():
    stdata_spot = plmo.generate_grid(stdata_raw,width=5)
    plmo_object = plmo.Polyomino(scdata_raw,stdata_spot,device='cuda')
    plmo_object.allocate()
    cell_alocated_data = plmo.sc2sc(scdata_raw, stdata_raw, zm.spot_matrix, method='max')
    pass
mem_usage, retval = memory_usage((zmap_mem,), retval=True, max_usage=True)
print(f"Peak memory usage: {mem_usage} MiB")
end_time = time() - start_time
print(end_time)

Start to run the spot mapping.
total_loss: -1.623
total_loss: -1.878
total_loss: -1.887
total_loss: -1.889
total_loss: -1.890
Finish spot mapping.
Finish spot mapping.
Mapping Matrix save in zm.spot_matrix

100%|██████████| 6437/6437 [10:23<00:00, 10.33it/s]

Peak memory usage: 60007.22265625 MiB
737.9632623195648

# Get peak GPU memory usage
peak_gpu_mem = torch.cuda.max_memory_allocated() / (1024 * 1024)  # Convert to MiB
print(f"Peak GPU memory usage: {peak_gpu_mem:.2f} MiB")

Peak GPU memory usage: 5634.22 MiB