def prepare_data(arena, path, smoothstr, smooth):
dat = json.load(open(path))
fly_data = dat['data']
bpath = dat.get('_base',os.path.abspath(os.path.dirname(path)))
pooled_on = {k:[] for k in "axbhwq"}
pooled_off = {k:[] for k in "axbhwq"}
pooled_lon = {k:[] for k in "axbhwq"}
for exp in fly_data:
geom, dfs = madplot.load_bagfile(
madplot.get_path(path,dat,exp["bag"]),
arena,
smooth=smooth)
l_df = dfs["targeted"]
t_df = dfs["tracked"]
h_df = dfs["ttm"]
#find when the laser was on
l_on = l_df[l_df['laser_power'] > 0]
#time of first laser on
l_on0 = l_df.index[0] + datetime.timedelta(seconds=30)
#t_off = t_df.head(3000)
#t_on = t_df.tail(3000)
#fig = plt.figure()
#ax = fig.gca()
#t_df.plot(ax=ax)
#fig.savefig("%s_%s.png" % (exp["bag"],exp["type"]), bbox_inches='tight')
#the laser was off at the start and on at the end
#tracking data when the laser was on
t_on = t_df[l_on0:]
#tracking data when the laser was off
t_off = t_df[:l_on0]
pooled_on[exp["type"]].append(t_on)
pooled_off[exp["type"]].append(t_off)
pooled_lon[exp["type"]].append(l_on)
cPickle.dump(pooled_on, open(os.path.join(bpath,'pooled_on_%s_%s.pkl' % (arena.unit,smoothstr)),'wb'), -1)
cPickle.dump(pooled_off, open(os.path.join(bpath,'pooled_off_%s_%s.pkl' % (arena.unit,smoothstr)),'wb'), -1)
cPickle.dump(pooled_lon, open(os.path.join(bpath,'pooled_lon_%s_%s.pkl' % (arena.unit,smoothstr)),'wb'), -1)
return pooled_on, pooled_off, pooled_lon
def load_bagfile_get_laseron(arena, score, smooth):
FWD = 's'
BWD = 'a'
AS_MAP = {FWD:1,BWD:-1}
geom, dfs = madplot.load_bagfile(score.bag, arena, smooth=smooth)
l_df = dfs["targeted"]
t_df = dfs["tracked"]
h_df = dfs["ttm"]
scored_ix = [t_df.index[0]]
scored_v = [AS_MAP[FWD]]
for idx,row in pd.read_csv(score.csv).iterrows():
matching = t_df[t_df['t_framenumber'] == int(row['framenumber'])]
if len(matching) == 1:
scored_ix.append(matching.index[0])
scored_v.append(AS_MAP[row['as']])
s = pd.Series(scored_v,index=scored_ix)
t_df['fwd'] = s
t_df['fwd'].fillna(method='ffill', inplace=True)
t_df['Vfwd'] = t_df['v'] * t_df['fwd']
lon = l_df[l_df['laser_power'] > 0]
if len(lon):
ton = lon.index[0]
l_off_df = l_df[ton:]
loff = l_off_df[l_off_df['laser_power'] == 0]
toff = loff.index[0]
onoffset = np.where(t_df.index >= ton)[0][0]
offoffset = np.where(t_df.index >= toff)[0][0]
return t_df, onoffset
else:
return None, None
def load_bagfile_get_laseron(arena, score, smooth):
NO_PROBOSCIS = 0
PROBOSCIS = 1
NO_WING_MOVEMENTS = 0
WING_MOVEMENTS = 1
NO_JUMPING = 0
JUMPING = 1
NO_ABDOMENING = 0
ABDOMENING = 1
geom, bag_results = madplot.load_bagfile(score.bag, arena, smooth=smooth)
l_df = bag_results['targeted']
t_df = bag_results['tracked']
h_df = bag_results['ttm']
proboscis_scored_ix = [t_df.index[0]]
proboscis_scored_v = [NO_PROBOSCIS]
wing_scored_ix = [t_df.index[0]]
wing_scored_v = [NO_WING_MOVEMENTS]
jump_scored_ix = [t_df.index[0]]
jump_scored_v = [NO_JUMPING]
abdomen_scored_ix = [t_df.index[0]]
abdomen_scored_v = [NO_ABDOMENING]
score_df = pd.read_csv(score.csv)
for idx,row in score_df.iterrows():
csv_row_frame = row['framenumber']
if np.isnan(csv_row_frame):
continue
matching = t_df[t_df['t_framenumber'] == int(csv_row_frame)]
if len(matching) ==1:
# proboscis ---------------
proboscis_scored_ix.append(matching.index[0])
rowval = row['as']
if rowval=='a':
proboscis_scored_v.append( PROBOSCIS )
elif rowval=='s':
proboscis_scored_v.append( NO_PROBOSCIS )
elif np.isnan(rowval):
#proboscis_scored_v.append( NO_PROBOSCIS )
proboscis_scored_v.append(np.nan)
else:
raise ValueError('unknown row value: %r'%rowval)
# wing ---------------
wing_scored_ix.append(matching.index[0])
rowval = row['zx']
if rowval=='z':
wing_scored_v.append( WING_MOVEMENTS )
elif rowval=='x':
wing_scored_v.append( NO_WING_MOVEMENTS )
elif np.isnan(rowval):
#wing_scored_v.append( NO_WING_MOVEMENTS )
wing_scored_v.append(np.nan)
else:
raise ValueError('unknown row value: %r'%rowval)
# jump ---------------
jump_scored_ix.append(matching.index[0])
rowval = row['cv']
if rowval=='c':
jump_scored_v.append( JUMPING )
elif rowval=='v':
jump_scored_v.append( NO_JUMPING )
elif np.isnan(rowval):
#jump_scored_v.append( NO_JUMPING )
jump_scored_v.append(np.nan)
else:
raise ValueError('unknown row value: %r'%rowval)
# abdomen ---------------
abdomen_scored_ix.append(matching.index[0])
rowval = row['qw']
if rowval=='q':
abdomen_scored_v.append( ABDOMENING )
elif rowval=='w':
abdomen_scored_v.append( NO_ABDOMENING )
elif np.isnan(rowval):
#abdomen_scored_v.append( NO_ABDOMENING )
abdomen_scored_v.append(np.nan)
else:
raise ValueError('unknown row value: %r'%rowval)
elif len(matching)==0:
continue
else:
print "len(matching)",len(matching)
# why would we ever get here?
1/0
s = pd.Series(proboscis_scored_v,index=proboscis_scored_ix)
# if len(s.index) > 1000:
# raise RuntimeError('file %r seems corrupt. Delete and re-score it.'%score.csv)
t_df['proboscis'] = s
s = pd.Series(wing_scored_v,index=wing_scored_ix)
t_df['wing'] = s
s = pd.Series(jump_scored_v,index=jump_scored_ix)
t_df['jump'] = s
s = pd.Series(abdomen_scored_v,index=abdomen_scored_ix)
t_df['abdomen'] = s
t_df['proboscis'].fillna(method='ffill', inplace=True)
t_df['wing'].fillna(method='ffill', inplace=True)
t_df['jump'].fillna(method='ffill', inplace=True)
t_df['abdomen'].fillna(method='ffill', inplace=True)
# t_df['Vfwd'] = t_df['v'] * t_df['fwd']
lon = l_df[l_df['laser_power'] > 0]
if len(lon):
laser_on = lon.index[0]
# l_off_df = l_df[laser_on:]
# loff = l_off_df[l_off_df['laser_power'] == 0]
# toff = loff.index[0]
# onoffset = np.where(t_df.index >= laser_on)[0][0]
# offoffset = np.where(t_df.index >= toff)[0][0]
return t_df, laser_on
else:
return None, None
ax.add_patch(arena.get_patch(color='k', alpha=0.1))
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('path', nargs=1, help='path to bag file')
parser.add_argument('--unit', default=False)
parser.add_argument('--no-smooth', dest='smooth', action='store_false', default=True)
parser.add_argument('--tracking-stats', action='store_true', default=False)
args = parser.parse_args()
path = args.path[0]
arena = madplot.Arena(args.unit)
geom, dfs = madplot.load_bagfile(path, arena, smooth=args.smooth)
ldf = dfs["targeted"]
tdf = dfs["tracked"]
hdf = dfs["ttm"]
axt = plt.figure("Trajectory").add_subplot(1,1,1)
plot_trajectory(axt,arena,ldf,tdf,hdf,geom, debug=args.tracking_stats)
axl = plt.figure("Laser State").add_subplot(1,1,1)
plot_laser(axl,arena,ldf,tdf,hdf,geom)
for ax in (axt,axl):
ax.set_xlabel('x (%s)' % arena.unit)
ax.set_ylabel('y (%s)' % arena.unit)
if args.tracking_stats:
import sys
import json
import os.path
import matplotlib.pyplot as plt
import roslib; roslib.load_manifest('flymad')
import flymad.madplot as madplot
BAG = '/mnt/strawscience/data/FlyMAD/velocity_test/2013-11-21-18-00-00.bag'
oid = 52251
smooth = True
smoothstr = {True:'smooth',False:'nosmooth'}[smooth]
arena = madplot.Arena('mm')
geom, dfs = madplot.load_bagfile(BAG, arena, smooth=smooth)
l_df = dfs["targeted"]
t_df = dfs["tracked"]
h_df = dfs["ttm"]
fig = plt.figure('traj %s' % smoothstr)
fig.suptitle('trajectory %s' % smoothstr)
ax = fig.add_subplot(1,1,1)
madplot.plot_tracked_trajectory(ax, t_df, arena, debug_plot=False)
ax.add_patch(arena.get_patch(fill=False))
ax.set_xlabel('position (%s)' % arena.unit)
fig.savefig('dorothea_pos_%s.png' % smoothstr)
df = t_df[t_df['tobj_id'] == oid].resample('100L')
#print df.iloc[39].index[0] - df.iloc[0]
def get_df(smooth):
smoothstr = {True:'smooth',False:'nosmooth'}[smooth]
geom, dfs = madplot.load_bagfile(BAG, arena, smooth=smooth)
t_df = dfs["tracked"]
fig = plt.figure('traj %s' % smoothstr)
fig.suptitle('trajectory %s' % smoothstr)
ax = fig.add_subplot(1,1,1)
madplot.plot_tracked_trajectory(ax, t_df, arena, debug_plot=False)
ax.add_patch(arena.get_patch(fill=False))
ax.set_xlabel('position (%s)' % arena.unit)
fig.savefig('testsmooth_mwb_pos_%s.png' % smoothstr)
df = t_df[t_df['tobj_id'] == oid]
#print df.iloc[39].index[0] - df.iloc[0]
times = [(ix - df.index[0]).total_seconds() for ix in df.index]
#print times[39]
######yyyyyyyyyy
fig = plt.figure('y %s' % smoothstr,figsize=(10,8))
fig.suptitle('y %s' % smoothstr)
ax = fig.add_subplot(2,2,1)
ax.set_title('position')
ax.plot(times, df['y_px'].values)
ax.set_ylabel('y (px/s)')
ax = fig.add_subplot(2,2,2)
ax.set_title('converted position')
ax.plot(times, df['y'].values)
ax.set_ylabel('y (%s)' % arena.unit)
ax = fig.add_subplot(2,2,3)
ax.set_title('velocity')
ax.plot(times, df['vy_px'].values)
ax.set_ylabel('y-velocity (px/s)')
ax = fig.add_subplot(2,2,4)
ax.set_title('converted velocity')
ax.plot(times, df['vy'].values)
ax.set_ylabel('y-velocity (%s/s)' % arena.unit)
fig.savefig('testsmooth_mwb_y_%s.png' % smoothstr)
######xxxxxxxxxxx
fig = plt.figure('x %s' % smoothstr,figsize=(10,8))
fig.suptitle('x %s' % smoothstr)
ax = fig.add_subplot(2,2,1)
ax.set_title('position')
ax.plot(times, df['x_px'].values)
ax.set_ylabel('x (px/s)')
ax = fig.add_subplot(2,2,2)
ax.set_title('converted position')
ax.plot(times, df['x'].values)
ax.set_ylabel('x (%s)' % arena.unit)
ax = fig.add_subplot(2,2,3)
ax.set_title('velocity')
ax.plot(times, df['vx_px'].values)
ax.set_ylabel('x-velocity (px/s)')
ax = fig.add_subplot(2,2,4)
ax.set_title('converted velocity')
ax.plot(times, df['vx'].values)
ax.set_ylabel('x-velocity (%s/s)' % arena.unit)
fig.savefig('testsmooth_mwb_x_%s.png' % smoothstr)
#######
df = t_df[t_df['tobj_id'] == oid].resample('500L')
times = [(ix - df.index[0]).total_seconds() for ix in df.index]
fig = plt.figure('speed %s' % smoothstr,figsize=(8,4))
fig.suptitle('speed %s' % smoothstr)
ax = fig.add_subplot(1,2,1)
ax.plot(times, df['v_px'].values)
ax.set_ylabel('speed (px/s)')
ax = fig.add_subplot(1,2,2)
ax.plot(times, df['v'].values)
ax.set_ylabel('speed (%s/s)' % arena.unit)
fig.savefig('testsmooth_mwb_speed_%s.png' % smoothstr)
return t_df[t_df['tobj_id'] == oid]
def load_bagfile_get_laseron(arena, score, smooth):
NO_PROBOSCIS = 0
PROBOSCIS = 1
NO_WING_MOVEMENTS = 0
WING_MOVEMENTS = 1
NO_JUMPING = 0
JUMPING = 1
NO_ABDOMENING = 0
ABDOMENING = 1
geom, bag_results = madplot.load_bagfile(score.bag, arena, smooth=smooth)
l_df = bag_results['targeted']
t_df = bag_results['tracked']
h_df = bag_results['ttm']
proboscis_scored_ix = [t_df.index[0]]
proboscis_scored_v = [NO_PROBOSCIS]
wing_scored_ix = [t_df.index[0]]
wing_scored_v = [NO_WING_MOVEMENTS]
jump_scored_ix = [t_df.index[0]]
jump_scored_v = [NO_JUMPING]
abdomen_scored_ix = [t_df.index[0]]
abdomen_scored_v = [NO_ABDOMENING]
score_df = pd.read_csv(score.csv)
for idx, row in score_df.iterrows():
csv_row_frame = row['framenumber']
if np.isnan(csv_row_frame):
continue
matching = t_df[t_df['t_framenumber'] == int(csv_row_frame)]
if len(matching) == 1:
# proboscis ---------------
proboscis_scored_ix.append(matching.index[0])
rowval = row['as']
if rowval == 'a':
proboscis_scored_v.append(PROBOSCIS)
elif rowval == 's':
proboscis_scored_v.append(NO_PROBOSCIS)
elif np.isnan(rowval):
#proboscis_scored_v.append( NO_PROBOSCIS )
proboscis_scored_v.append(np.nan)
else:
raise ValueError('unknown row value: %r' % rowval)
# wing ---------------
wing_scored_ix.append(matching.index[0])
rowval = row['zx']
if rowval == 'z':
wing_scored_v.append(WING_MOVEMENTS)
elif rowval == 'x':
wing_scored_v.append(NO_WING_MOVEMENTS)
elif np.isnan(rowval):
#wing_scored_v.append( NO_WING_MOVEMENTS )
wing_scored_v.append(np.nan)
else:
raise ValueError('unknown row value: %r' % rowval)
# jump ---------------
jump_scored_ix.append(matching.index[0])
rowval = row['cv']
if rowval == 'c':
jump_scored_v.append(JUMPING)
elif rowval == 'v':
jump_scored_v.append(NO_JUMPING)
elif np.isnan(rowval):
#jump_scored_v.append( NO_JUMPING )
jump_scored_v.append(np.nan)
else:
raise ValueError('unknown row value: %r' % rowval)
# abdomen ---------------
abdomen_scored_ix.append(matching.index[0])
rowval = row['qw']
if rowval == 'q':
abdomen_scored_v.append(ABDOMENING)
elif rowval == 'w':
abdomen_scored_v.append(NO_ABDOMENING)
elif np.isnan(rowval):
#abdomen_scored_v.append( NO_ABDOMENING )
abdomen_scored_v.append(np.nan)
else:
raise ValueError('unknown row value: %r' % rowval)
elif len(matching) == 0:
continue
else:
print "len(matching)", len(matching)
# why would we ever get here?
1 / 0
s = pd.Series(proboscis_scored_v, index=proboscis_scored_ix)
# if len(s.index) > 1000:
# raise RuntimeError('file %r seems corrupt. Delete and re-score it.'%score.csv)
t_df['proboscis'] = s
s = pd.Series(wing_scored_v, index=wing_scored_ix)
t_df['wing'] = s
s = pd.Series(jump_scored_v, index=jump_scored_ix)
t_df['jump'] = s
s = pd.Series(abdomen_scored_v, index=abdomen_scored_ix)
t_df['abdomen'] = s
t_df['proboscis'].fillna(method='ffill', inplace=True)
t_df['wing'].fillna(method='ffill', inplace=True)
t_df['jump'].fillna(method='ffill', inplace=True)
t_df['abdomen'].fillna(method='ffill', inplace=True)
# t_df['Vfwd'] = t_df['v'] * t_df['fwd']
lon = l_df[l_df['laser_power'] > 0]
if len(lon):
laser_on = lon.index[0]
# l_off_df = l_df[laser_on:]
# loff = l_off_df[l_off_df['laser_power'] == 0]
# toff = loff.index[0]
# onoffset = np.where(t_df.index >= laser_on)[0][0]
# offoffset = np.where(t_df.index >= toff)[0][0]
return t_df, laser_on
else:
return None, None
args = parser.parse_args()
path = args.path[0]
if not os.path.isdir(path):
parser.error('must be a directory')
bags = glob.glob(os.path.join(path, "*.bag"))
fmffile = fmf.FlyMovie(glob.glob(os.path.join(path, "*.fmf"))[0])
ts = fmffile.get_all_timestamps().tolist()
arena = madplot.Arena(False)
cuts = []
for bag in sorted(bags):
geom, dfs = madplot.load_bagfile(bag, arena)
h_df = dfs["ttm"]
fn = h_df['h_framenumber']
start = fn.ix[0]
stop = fn.ix[-1]
dn = os.path.dirname(bag)
btime = madplot.strptime_bagfile(bag)
destfmfname = "%s%d%02d%02d_%02d%02d%02d.fmf" % (os.path.basename(dn),
btime.tm_year,btime.tm_mon,btime.tm_mday,
btime.tm_hour,btime.tm_min,btime.tm_sec)
destfmf = os.path.join(dn,'..',destfmfname)
cuts.append( Cut(start,stop,os.path.abspath(destfmf)) )
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('path', nargs=1, help='path to bag file')
parser.add_argument('--unit', default=False)
parser.add_argument('--no-smooth',
dest='smooth',
action='store_false',
default=True)
parser.add_argument('--tracking-stats', action='store_true', default=False)
args = parser.parse_args()
path = args.path[0]
arena = madplot.Arena(args.unit)
geom, dfs = madplot.load_bagfile(path, arena, smooth=args.smooth)
ldf = dfs["targeted"]
tdf = dfs["tracked"]
hdf = dfs["ttm"]
axt = plt.figure("Trajectory").add_subplot(1, 1, 1)
plot_trajectory(axt, arena, ldf, tdf, hdf, geom, debug=args.tracking_stats)
axl = plt.figure("Laser State").add_subplot(1, 1, 1)
plot_laser(axl, arena, ldf, tdf, hdf, geom)
for ax in (axt, axl):
ax.set_xlabel('x (%s)' % arena.unit)
ax.set_ylabel('y (%s)' % arena.unit)
if args.tracking_stats:
def prepare_data(arena, path, smoothstr, smooth):
dat = json.load(open(path))
fly_data = dat['data']
bpath = dat.get('_base', os.path.abspath(os.path.dirname(path)))
pooled_on = {k: [] for k in "axbhwq"}
pooled_off = {k: [] for k in "axbhwq"}
pooled_lon = {k: [] for k in "axbhwq"}
for exp in fly_data:
geom, dfs = madplot.load_bagfile(madplot.get_path(
path, dat, exp["bag"]),
arena,
smooth=smooth)
l_df = dfs["targeted"]
t_df = dfs["tracked"]
h_df = dfs["ttm"]
#find when the laser was on
l_on = l_df[l_df['laser_power'] > 0]
#time of first laser on
l_on0 = l_df.index[0] + datetime.timedelta(seconds=30)
#t_off = t_df.head(3000)
#t_on = t_df.tail(3000)
#fig = plt.figure()
#ax = fig.gca()
#t_df.plot(ax=ax)
#fig.savefig("%s_%s.png" % (exp["bag"],exp["type"]), bbox_inches='tight')
#the laser was off at the start and on at the end
#tracking data when the laser was on
t_on = t_df[l_on0:]
#tracking data when the laser was off
t_off = t_df[:l_on0]
pooled_on[exp["type"]].append(t_on)
pooled_off[exp["type"]].append(t_off)
pooled_lon[exp["type"]].append(l_on)
cPickle.dump(
pooled_on,
open(
os.path.join(bpath,
'pooled_on_%s_%s.pkl' % (arena.unit, smoothstr)),
'wb'), -1)
cPickle.dump(
pooled_off,
open(
os.path.join(bpath,
'pooled_off_%s_%s.pkl' % (arena.unit, smoothstr)),
'wb'), -1)
cPickle.dump(
pooled_lon,
open(
os.path.join(bpath,
'pooled_lon_%s_%s.pkl' % (arena.unit, smoothstr)),
'wb'), -1)
return pooled_on, pooled_off, pooled_lon
import os.path
import matplotlib.pyplot as plt
import roslib
roslib.load_manifest('flymad')
import flymad.madplot as madplot
BAG = '/mnt/strawscience/data/FlyMAD/velocity_test/2013-11-21-18-00-00.bag'
oid = 52251
smooth = True
smoothstr = {True: 'smooth', False: 'nosmooth'}[smooth]
arena = madplot.Arena('mm')
geom, dfs = madplot.load_bagfile(BAG, arena, smooth=smooth)
l_df = dfs["targeted"]
t_df = dfs["tracked"]
h_df = dfs["ttm"]
fig = plt.figure('traj %s' % smoothstr)
fig.suptitle('trajectory %s' % smoothstr)
ax = fig.add_subplot(1, 1, 1)
madplot.plot_tracked_trajectory(ax, t_df, arena, debug_plot=False)
ax.add_patch(arena.get_patch(fill=False))
ax.set_xlabel('position (%s)' % arena.unit)
fig.savefig('dorothea_pos_%s.png' % smoothstr)
df = t_df[t_df['tobj_id'] == oid].resample('100L')
#print df.iloc[39].index[0] - df.iloc[0]