Commit 497820ca authored by lucas_miranda's avatar lucas_miranda
Browse files

Smoothing 25x faster

parent c9fde930
This source diff could not be displayed because it is too large. You can view the blob instead.
......@@ -31,9 +31,9 @@ class project:
smooth_alpha=0.1,
arena_dims=[1],
center_coords=True,
distances=False,
distances='All',
ego=False,
angles=False,
angles=True,
connectivity=None,
):
......@@ -282,7 +282,7 @@ class coordinates:
else:
return "DLC analysis of {} videos".format(len(self._videos))
def get_coords(self, center=True, polar=False):
def get_coords(self, center=True, polar=False, speed=0):
tabs = deepcopy(self._tables)
if center:
......@@ -300,6 +300,17 @@ class coordinates:
for key, tab in tabs.items():
tabs[key] = tab2polar(tab)
if speed:
for order in range(speed):
for key, tab in tabs.items():
try:
cols = tab.columns.levels[0]
except AttributeError:
cols = tab.columns
vel = rolling_speed(tab, typ="coords", order=order + 1)
vel.columns = cols
tabs[key] = vel
return table_dict(
tabs,
"coords",
......@@ -309,23 +320,50 @@ class coordinates:
polar=polar,
)
def get_distances(self):
def get_distances(self, speed=0):
tabs = deepcopy(self.distances)
if self.distances is not None:
return table_dict(self.distances, typ="dists")
if speed:
for order in range(speed):
for key, tab in tabs.items():
try:
cols = tab.columns.levels[0]
except AttributeError:
cols = tab.columns
vel = rolling_speed(tab, typ="dists", order=order + 1)
vel.columns = cols
tabs[key] = vel
return table_dict(tabs, typ="dists")
raise ValueError(
"Distances not computed. Read the documentation for more details"
)
def get_angles(self, degrees=False):
def get_angles(self, degrees=False, speed=0):
tabs = deepcopy(self.angles)
if self.angles is not None:
if degrees == True:
return table_dict(
{key: np.degrees(tab) for key, tab in self.angles.items()},
typ="angles",
)
if degrees:
tabs = {key: np.degrees(tab) for key, tab in tabs.items()}
if speed:
for order in range(speed):
for key, tab in tabs.items():
try:
cols = tab.columns.levels[0]
except AttributeError:
cols = tab.columns
vel = rolling_speed(tab, typ="dists", order=order + 1)
vel.columns = cols
tabs[key] = vel
return table_dict(tabs, typ="angles")
else:
return table_dict(self.angles, typ="angles")
raise ValueError("Angles not computed. Read the documentation for more details")
def get_videos(self, play=False):
......
......@@ -244,11 +244,26 @@ def climb_wall(arena, pos_dict, fnum, tol, mouse):
return np.linalg.norm(nose - center) > arena[2] + tol
def rolling_speed(dframe, pause=10, rounds=5):
def rolling_speed(dframe, typ, pause=10, rounds=5, order=1):
"""Returns the average speed over 10 frames in pixels per frame"""
distances = np.linalg.norm(np.array(dframe) - np.array(dframe.shift()), axis=1)
distances = pd.Series(distances, index=dframe.index)
s = dframe.shape[0]
if typ == "coords":
bp = dframe.shape[1] / 2 if order == 1 else dframe.shape[1]
d = 2 if order == 1 else 1
else:
bp = dframe.shape[1]
d = 1
distances = np.linalg.norm(
np.array(dframe).reshape(s, int(bp), d)
- np.array(dframe.shift()).reshape(s, int(bp), d),
axis=2,
)
distances = pd.DataFrame(distances, index=dframe.index)
speeds = np.round(distances.rolling(pause).mean(), rounds)
speeds[np.isnan(speeds)] = 0.0
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment