Iterating fast over h5 file and perform some calculations

2. I need a super fast solution, that needs maximally 5 seconds on the 9000 datapoints I provide in the link. Reason is because the real data is actually millions of rows.
3. Link to the h5 file: https://drive.google.com/file/d/16aI3plRFa3M6nSIiT1XioUIgsPYl1Wg8/view?usp=sharing

The task at hand is as follows: Given the coordinate data of different body parts of different mice in the h5 file. Read in the h5 file (hopefully as numpy.array not as pandas what I did underneath) and then calculate the centroid based on tail1, tail2 and tail3 body parts.

My suspicion underneath is that .loc indexing is what causes the problem and generally dataframe iteration is sub-optimal.

What I have done is standard loc indexing:

filename="look at the h5 file in the link" # h5 above
new_centroid_trackings = np.array([[0,0,0,0,0,0,0,0]]) # initialize the data to concatinate after every iteration
model_name="DLC_resnet50_4mice_new_video_no_wheelFeb17shuffle1_220000" # not relevant for task
tracking_coords = pd.read_hdf(filename) # read in the data
for frame in range(tracking_coords.shape[0]):
centroid_mouse1_x=(tracking_coords.loc[frame, model_name]["mouse1"]["tail1"]["x"]+tracking_coords.loc[frame, model_name]["mouse1"]["tail2"]["x"]+tracking_coords.loc[frame, model_name]["mouse1"]["tail3"]["x"])/3
centroid_mouse1_y=(tracking_coords.loc[frame, model_name]["mouse1"]["tail1"]["y"]+tracking_coords.loc[frame, model_name]["mouse1"]["tail2"]["y"]+tracking_coords.loc[frame, model_name]["mouse1"]["tail3"]["y"])/3
centroid_mouse2_x=(tracking_coords.loc[frame, model_name]["mouse2"]["tail1"]["x"]+tracking_coords.loc[frame, model_name]["mouse2"]["tail2"]["x"]+tracking_coords.loc[frame, model_name]["mouse2"]["tail3"]["x"])/3
centroid_mouse2_y=(tracking_coords.loc[frame, model_name]["mouse2"]["tail1"]["y"]+tracking_coords.loc[frame, model_name]["mouse2"]["tail2"]["y"]+tracking_coords.loc[frame, model_name]["mouse2"]["tail3"]["y"])/3      
centroid_mouse3_x=(tracking_coords.loc[frame, model_name]["mouse3"]["tail1"]["x"]+tracking_coords.loc[frame, model_name]["mouse3"]["tail2"]["x"]+tracking_coords.loc[frame, model_name]["mouse3"]["tail3"]["x"])/3
centroid_mouse3_y=(tracking_coords.loc[frame, model_name]["mouse3"]["tail1"]["y"]+tracking_coords.loc[frame, model_name]["mouse3"]["tail2"]["y"]+tracking_coords.loc[frame, model_name]["mouse3"]["tail3"]["y"])/3
centroid_mouse4_x=(tracking_coords.loc[frame, model_name]["mouse4"]["tail1"]["x"]+tracking_coords.loc[frame, model_name]["mouse4"]["tail4"]["x"]+tracking_coords.loc[frame, model_name]["mouse4"]["tail3"]["x"])/3
centroid_mouse4_y=(tracking_coords.loc[frame, model_name]["mouse4"]["tail1"]["y"]+tracking_coords.loc[frame, model_name]["mouse4"]["tail4"]["y"]+tracking_coords.loc[frame, model_name]["mouse4"]["tail3"]["y"])/3
# now concatinate the centroids to the previous ones
new_centroid_trackings=np.concatenate((new_centroid_trackings, np.array([[centroid_mouse1_x,centroid_mouse1_y,centroid_mouse2_x, centroid_mouse2_y, centroid_mouse3_x, centroid_mouse3_y, centroid_mouse4_x, centroid_mouse4_y]])), axis=0)

And for this around 90 seconds is needed for all the rows.

Needed solution: Numpy (or not) solution that takes maximally 5 seconds for all the rows.

You have two main issues which slow down the calculation inside the loop:

• Multi-level columns dataframe: it takes longer when you access each time each level. I solved this problem converting the multilevel columns in the dataframe to a single-level dataframe by the following:

  tracking_coords.columns = ['_'.join(w) for w in tracking_coords.columns.values]
  

• Array Concatenation: the last line contains a concatenation operation which costs a lot. Therefore, it is highly recommended using pre-initailization for the final array (new_centroid_trackings) if you know the shape of the final array beforehand. Thus, I solved it by:

  new_centroid_trackings = np.zeros((len(tracking_coords),8))
  

It finishes the entire loop in less than 5 sec. I changed slightly the columns names inside the loop to adapt the new columns names.

The entire code:

import pandas as pd 
import numpy as np
filename="file.h5" # h5 above
model_name="DLC_resnet50_4mice_new_video_no_wheelFeb17shuffle1_220000" # not relevant for task
tracking_coords = pd.read_hdf(filename) # read in the data
tracking_coords.columns = ['_'.join(w) for w in tracking_coords.columns.values]
new_centroid_trackings = np.zeros((len(tracking_coords),8))
for frame in range(tracking_coords.shape[0]):
centroid_mouse1_x=(tracking_coords.loc[frame, model_name+"_mouse1_tail1_x"]+tracking_coords.loc[frame, model_name+"_mouse1_tail2_x"]+tracking_coords.loc[frame, model_name+"_mouse1_tail3_x"])/3
centroid_mouse1_y=(tracking_coords.loc[frame, model_name+"_mouse1_tail1_y"]+tracking_coords.loc[frame, model_name+"_mouse1_tail2_y"]+tracking_coords.loc[frame, model_name+"_mouse1_tail3_y"])/3
centroid_mouse2_x=(tracking_coords.loc[frame, model_name+"_mouse2_tail1_x"]+tracking_coords.loc[frame, model_name+"_mouse2_tail2_x"]+tracking_coords.loc[frame, model_name+"_mouse2_tail3_x"])/3
centroid_mouse2_y=(tracking_coords.loc[frame, model_name+"_mouse2_tail1_y"]+tracking_coords.loc[frame, model_name+"_mouse2_tail2_y"]+tracking_coords.loc[frame, model_name+"_mouse2_tail3_y"])/3
centroid_mouse3_x=(tracking_coords.loc[frame, model_name+"_mouse3_tail1_x"]+tracking_coords.loc[frame, model_name+"_mouse3_tail2_x"]+tracking_coords.loc[frame, model_name+"_mouse3_tail3_x"])/3
centroid_mouse3_y=(tracking_coords.loc[frame, model_name+"_mouse3_tail1_y"]+tracking_coords.loc[frame, model_name+"_mouse3_tail2_y"]+tracking_coords.loc[frame, model_name+"_mouse3_tail3_y"])/3
centroid_mouse4_x=(tracking_coords.loc[frame, model_name+"_mouse4_tail1_x"]+tracking_coords.loc[frame, model_name+"_mouse4_tail2_x"]+tracking_coords.loc[frame, model_name+"_mouse4_tail3_x"])/3
centroid_mouse4_y=(tracking_coords.loc[frame, model_name+"_mouse4_tail1_y"]+tracking_coords.loc[frame, model_name+"_mouse4_tail2_y"]+tracking_coords.loc[frame, model_name+"_mouse4_tail3_y"])/3
new_centroid_trackings[frame,:] = [centroid_mouse1_x,centroid_mouse1_y,centroid_mouse2_x, centroid_mouse2_y, centroid_mouse3_x, centroid_mouse3_y, centroid_mouse4_x, centroid_mouse4_y]