Macaque-ganglion-cells/run_model_HPC/make_movie_BO.py

import argparse
import ipyparallel as ipp
import numpy as np
import os
import cv2 as cv2
import time
from retinatools.library import image_to_cone

#Helper function
def RF_FilterFrame(idx, luminance_frame, r_RF):
    """
    Helper function to do the parallelized filtering with RFs.
    """
    from retinatools.library import get_RF

    X = np.linspace(-128, 128, 256)
    Y = np.linspace(-128, 128, 256)
    idx_x = idx%256
    idx_y = np.int(np.floor(idx/256))
    x, y = np.meshgrid(X,Y)    
    xoff = X[idx_x]
    yoff = Y[idx_y]
    RF = get_RF(xoff, yoff, r_RF)
    product = np.sum(luminance_frame * RF)
    return product

def parallel_prediction(idx_x, RFstrip, fitparameters):    
    ''' 
    Helper function to do the parallelized processing of the filtered signal; 
    This transform center/surround signal to activity with standard parameters.
    '''
    from retinatools.library import get_BO_response

    lum_signalS = RFstrip[0, idx_x, :]
    lum_signalM = RFstrip[1, idx_x, :]
    lum_signalL = RFstrip[2, idx_x, :]
    R = get_BO_response(lum_signalS, lum_signalM, lum_signalL, fitparameters)
    return R



fitparameters = {'rcenterB': 10.0, 'rcenterY': 14.0,\
                 'tau1': 10.2, 'tau2': 86, 'delay': 29, 'beta': 0.0, 'alpha': 4.9,  \
                 'tau3': 280, 'kL': 0.01, 'kM': 0.051, 'kB': 0.12, 'offset': 0.025, 'dt': 1000/150.}
cell = "178Bon"

savepath = "/jukebox/tank/schottdorf/tmp/"

N_MAXFRAME = 4500  # 4500 to only process first 2.5 min *60sec *30fps in the mpg -> 30sec ret video

def main(profile):
    rc = ipp.Client(profile=profile)                 # ipyparallel status
    print(rc.ids)

    dview = rc[:]                                    # Set up the engines
    dview.use_cloudpickle()                          # Make sure to fix the closure over luminance_frame
    with dview.sync_imports():
        import numpy as np

    ################################
    #####   Process the movie  #####
    ################################
    vidcap = cv2.VideoCapture('/usr/people/ms81/cells/1x10_256.mpg')
    data = np.ones((3, 256, 256, N_MAXFRAME))            # for s,m,l  X 256px  X 256px  X N_frames
    counter = 0
    success = True
    t0 = time.time()
    while success and (counter < N_MAXFRAME):   
        success, image = vidcap.read()
        scone, mcone, lcone = image_to_cone(image)        
        for tt in ["S", "M", "L"]:
            if tt == "S":
                r_RF = fitparameters['rcenterB']
                luminance_frame = scone
                idx = 0
            elif tt == "M":
                r_RF = fitparameters['rcenterY']
                luminance_frame = mcone
                idx = 1
            elif tt == "L":
                r_RF = fitparameters['rcenterY']
                luminance_frame = lcone
                idx = 2

            rc[:].push(dict( luminance_frame = luminance_frame, RF_FilterFrame = RF_FilterFrame, r_RF = r_RF ))         # Seed namespace for the ipython engines. 
            tmp = rc[:].map_sync( lambda ix: RF_FilterFrame(ix, luminance_frame, r_RF), range(0,256*256) )              # Execute parallel computation on all ipython engines
            data[idx,:,:,counter] = np.reshape(tmp, (256,256))                                                          # Reorganize data back to frame
        counter += 1
        print(counter/N_MAXFRAME)
    t1 = time.time()

    # filename = "rgc-{0}-raw.npy".format(os.environ["SLURM_JOB_ID"])
    # np.save(filename, data)
    print("Part 1 DONE; Time needed [in h]:")
    print( (t1-t0)/3600 )


    ################################
    ####  Apply temp. process  #####
    ################################
    t1 = time.time()
    output_model = np.zeros((256, 256, N_MAXFRAME))
    rc[:].push(dict( parallel_prediction = parallel_prediction))
    for idx_x in range(0,256):
        RFstrip = data[:,idx_x,:,:]
        rc[:].push(dict( RFstrip = RFstrip, fitparameters = fitparameters, parallel_prediction = parallel_prediction ))
        tmp = rc[:].map_sync(lambda idx_y: parallel_prediction(idx_y, RFstrip, fitparameters), range(0,256))
        output_model[idx_x,:,:] = np.reshape(tmp, (256, N_MAXFRAME))
        print(idx_x/256)
    t2 = time.time()

    filename = "-{0}-frames.npy".format(os.environ["SLURM_JOB_ID"])
    pp = savepath + cell + filename
    np.save(pp, output_model)
    print("Part 2 DONE; Time needed [in h]:")
    print("Saved: " + pp)
    print((t2-t1)/3600)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-p", "--profile", required=True, help="Name of IPython profile to use")

    args = parser.parse_args()
    main(args.profile)
    print("Done and done.")