Avoiding large memory load and processing time when running apply_shifts_movie on 3d datasets?

[noahpettit]

Hello,

I'm working with fairly noisy volumetric 2p imaging data and it is neccessary for me to perform some smoothing (across space and time) for motion correction to work properly. I therefore have to compute the shifts on smoothed data and then apply them to the raw data.

I've found that MotionCorrect.apply_shifts_movie works very quickly for small example datasets (like the generated data in demo_caimaon_cnmf_3D.ipynb. However, for more realistic datasets (for example: 5k frames, 128 x 256 resolution, 10 planes) I find that running this function results in a large memory load (>40GB sometimes for a movie that is ~5GB) and very slow runtime (several times slower than the original motion correction). This is even the case when working with mem mapped files as the input. I am a bit confused because the runtime and memory load are way higher than when I run the motion correction itself (which takes about a minute on my machine).

Am I missing someting about how to apply shifts? I have no issue running the motion correction itself. I've tried adjusting n_processes down to 1 or 2 and I still have the same issue.

Thanks in advance for any advice!

[noahpettit]

Just did a test where I tried processing movies of increasing size and I think the issue is definitely a memory one. You can see motion correction runtime is pretty linear, and applying shifts has a huge jump which corresponds to when I see that I am running out of RAM and I suppose I start swapping which of course slows down the performance greatly and imposes an upper limit on the size of movie I can process.

Is there a workaround for this? Can you point me to how I might modify the apply_shifts_movie function to be as memory-efficient as motion_correct ?

[noahpettit]

For reference i have 32GB of RAM which i believe should be enough for motion correction and applying shifts to this size movie

[kushalkolar]

I'm confused about exactly what you're doing. What do you mean by smoothing? Are you applying a Gaussian filter across space? Are you doing a rolling average over time? Gaussian temporal filter? Can you provide code to reproduce this.

[noahpettit]

Hi, thanks for your reply!

I mentioned the smoothing only to motivate why I need to compute shifts on one version of the data and apply it to another. But yes, the smoothing I've tried is a 4d gaussian filter across space and time. I haven't really narrowed in on exactly what parameters to use for the smoothing yet, but the slowdown I'm describing is entirely independent of the smoothing - i.e. the test I ran does not incorporating any smoothing, it is simply doing motion correction and then applying those shifts to the same file.

Smoothing aside, I am guessing it is common for people to calculate shifts using one channel (like a brighter structural channel) and then apply them to the other channel in caiman. So I am assuming this is probably error on my part in how I am calling that function...

Here is some minimal code I am using to test.

from IPython import get_ipython
import logging
import matplotlib.pyplot as plt
import numpy as np
import os
from scipy.ndimage import gaussian_filter
from tifffile.tifffile import imwrite

import caiman as cm
from caiman.utils.visualization import nb_view_patches3d
import caiman.source_extraction.cnmf as cnmf
from caiman.paths import caiman_datadir

try:
    if __IPYTHON__:
        get_ipython().run_line_magic('load_ext', 'autoreload')
        get_ipython().run_line_magic('autoreload', '2')
except NameError:
    pass

import bokeh.plotting as bpl
bpl.output_notebook()

logging.basicConfig(format=
                          "%(relativeCreated)12d [%(filename)s:%(funcName)20s():%(lineno)s] [%(process)d] %(message)s",
                    # filename="/tmp/caiman.log",
                    level=logging.WARNING)

import time

#%% start a cluster for parallel prcessing
if 'dview' in locals():
    cm.stop_server(dview=dview)
c, dview, n_processes = cm.cluster.setup_cluster(
    backend='multiprocessing', n_processes=2, single_thread=False)

fname = r'/Users/noahpettit/data/flyg/local_fast/caiman_tmp/raw_ch0.tif'
Y = cm.load(fname)
print(Y.shape)

n_frames_to_process = [100,1000,4000]
fname_list = []

# here I save out movies with varying numbers of frames
for n_frames in n_frames_to_process:

    # create a new movie with the first n_frames
    fname = os.path.join(caiman_datadir(), 'example_movies', f'demoMovie3D_{n_frames}.tif')
    # if it exists, delete it
    if os.path.exists(fname):
        os.remove(fname)
    imwrite(fname, Y[0:n_frames,:,:,:], bigtiff=True)
    fname_list.append(fname)

print('done saving out movies')
del Y


time_to_motion_correct = []
time_to_apply_shifts = []

for fname in fname_list:

    print(f'Processing {fname}')

    # motion correction parameters
    opts_dict = {'fnames': fname,
                'strides': (24, 24, 6),    # start a new patch for pw-rigid motion correction every x pixels
                'overlaps': (12, 12, 2),   # overlap between patches (size of patch strides+overlaps)
                'max_shifts': (4, 4, 2),   # maximum allowed rigid shifts (in pixels)
                'max_deviation_rigid': 5,  # maximum shifts deviation allowed for patch with respect to rigid shifts
                'pw_rigid': False,         # flag for performing non-rigid motion correction
                'is3D': True}

    opts = cnmf.params.CNMFParams(params_dict=opts_dict)

    # trying saving to a memap file here just in case that helps
    mmap_fname = cm.save_memmap([fname], is_3D=True, dview=dview)

    # first we create a motion correction object with the parameters specified
    mc = cm.motion_correction.MotionCorrect(mmap_fname, dview=dview, **opts.get_group('motion'))
    # note that the file is not loaded in memory

    # time the motion correction
    time_start = time.time()
    mc.motion_correct(save_movie=False)
    time_end = time.time()
    time_to_motion_correct.append(time_end - time_start)
    print(f'Motion correction took {time_end - time_start} seconds')

    # apply the shifts to the movie
    time_start = time.time()
    reg = mc.apply_shifts_movie(mmap_fname, save_memmap=False) 
    time_end = time.time()
    time_to_apply_shifts.append(time_end - time_start)
    print(f'Applying shifts took {time_end - time_start} seconds')

    del reg
    del mc

[noahpettit]

@kushalkolar Just following up on this - Do you have any thoughts on the memory efficiency of apply_shifts_movie? I know my example uses a tiff, but I imagine you could reproduce this by just saving out a much longer example movie and looking at your python process memory usage when running motion correction and when running apply shifts. Is what i'm experiencing perhaps realted to this recent issue #1482 ?