import numpy as np def max_velocity(x) : size = x.shape res = np.zeros(size[0]) for time_step in range(size[0]) : for vertex in range(size[1]) : res[time_step] = max(res[time_step], np.linalg.norm(x[time_step, vertex, :])) return res