Macaque-ganglion-cells/run_model_HPC/make_movie_PC.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, celltype):    
    ''' 
    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_PC_response

    lum_signalM = RFstrip[0, idx_x, :]
    lum_signalL = RFstrip[1, idx_x, :]
    R = get_PC_response(lum_signalM, lum_signalL, fitparameters, celltype)
    return R



# fitparameters = {'rcenterM': 20, 'rcenterL': 5.0, \
#                     'lum_factor': 1200, 'tau1': 6.9, 'tau2': 60, 'k1': 0.017, 'tau3': 400,\
#                     'alpha': 0.68, 'k2': 4, 'o1': 0.07, 'dt': 1000/150.}  #110Ron
# celltype = "Ron"

fitparameters = {'rcenterM': 4.9, 'rcenterL': 32, \
                 'lum_factor': 290, 'tau1': 7.8, 'tau2': 100, 'k1': 0.018, 'tau3': 2000,\
                 'alpha': 0.71, 'k2': 3.0, 'o1': 0.06, 'dt': 1000/150.}  #130Gon
celltype = "Gon"



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((2, 256, 256, N_MAXFRAME))            # for center & surround   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 ["M", "L"]:
            if tt == "M":
                r_RF = fitparameters['rcenterM']
                luminance_frame = mcone
                idx = 0
            else:
                r_RF = fitparameters['rcenterL']
                luminance_frame = lcone
                idx = 1
            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, celltype = celltype, fitparameters = fitparameters, parallel_prediction = parallel_prediction ))
        tmp = rc[:].map_sync(lambda idx_y: parallel_prediction(idx_y, RFstrip, fitparameters, celltype), 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 + celltype + 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.")