Commit 9b68faf3 authored by David A.. Werner's avatar David A.. Werner
Browse files

Added visualization! (with terribly formatting/comments/coding/documentation

parent d8e2286d
"""
Makes animations of manducas moving
"""
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import matplotlib.patches as patches
import os
def animate(*manducas, **kwargs):
anim = _get_animation(*manducas, **kwargs)
plt.show()
plt.close()
def animate_as_gif(file_name, *manducas, **kwargs):
anim = _get_animation(*manducas, **kwargs)
_save_as_gif(file_name, anim)
plt.close()
def animate_as_mp4(file_name, *manducas, **kwargs):
anim = _get_animation(*manducas, **kwargs)
_save_as_gif(file_name, anim)
plt.close()
def _get_animation(*manducas, **kwargs):
""" Produces an animation of all \*manducas racing
NOTE: Requires all manducas to have same time_step and num_time_steps!
"""
show_winners = kwargs.pop('show_winners', False)
show_fitnesses = kwargs.pop('show_fitnesses', False)
manducas_with_labels = []
for idx,val in enumerate(manducas):
if isinstance(val, tuple):
manducas_with_labels.append(val)
else:
manducas_with_labels.append((val, str(idx)))
if show_fitnesses or show_winners:
fitnesses = [man.fitness for man, _ in manducas_with_labels]
if show_winners:
max_f = max(fitnesses)
winners = np.array(fitnesses) == max_f
else:
winners = None
if not show_fitnesses:
fitnesses = None
for idx in range(len(manducas_with_labels)):
manduca, label = manducas_with_labels[idx]
if fitnesses is not None:
new_label = "{}[fitness={}]".format(label, fitnesses[idx])
else:
new_label = str(label)
if winners is not None:
if winners[idx]:
new_label = "{}[WINNER]".format(new_label)
manducas_with_labels[idx] = (manduca, new_label)
# Make sure the manducas are compatible!
time_steps = list(set([manduca.time_step for manduca, _ in manducas_with_labels]))
assert len(time_steps) == 1, \
"All manducas must have same time_step to be animated together"
num_time_steps = list(set([manduca.num_time_steps for manduca, _ in manducas_with_labels]))
assert len(time_steps) == 1, \
"All manducas must have same number of time steps to be animated together"
dx = time_steps[0]
positions, legs, muscles = [], [], []
# Read the file(s) and build points[n_frames][15 items][n_worms]
for manduca, labels in manducas_with_labels:
sim_kwargs = {key: kwargs[key] for key in ['sub_intervals'] if key in kwargs}
x, v, t, l, m = manduca.simulate(record_level=2, record_legs_and_muscles=True, **sim_kwargs)
positions.append(x)
legs.append(l)
muscle_strength = manduca.body_properties.muscle_strength
muscles.append([force/muscle_strength for force in m])
labels = [label for manduca, label in manducas_with_labels]
return _animate(labels, positions, legs, muscles, dx, fitnesses=fitnesses, winners=winners, **kwargs)
#TODO: Also do still-frames
def get_stills(manduca, sub_intervals=None):
pass
def _save_as_gif(output_file, animation, dpi=80, **kwargs):
animation.save(output_file, dpi=dpi, writer='imagemagick', **kwargs)
def _save_as_mp4(output_file, animation, dpi=80, **kwargs):
animation.save(output_file, dpi=dpi, writer='mencoder', **kwargs)
def _animate(labels, positions, legs, muscles, dx, **kwargs):
leg_width = kwargs.pop('leg_width', 30)
show_result_percent = kwargs.pop('show_result_percent', 0.90)
fitnesses = kwargs.pop('fitnesses', None)
winners = kwargs.pop('winners', None)
# The per-frame animation function.
# Inputs: 'points' is a full array with[n_timepoints][data][n_worms] (where
# n_timepoints is the number of frames to be displayed).
# Remember that our display axes are:
# * x ranges over the min/max x values from the simulation.
# * y ranges from 0 to n_worms; i.e., each worm is allocated a vertical
# space of 1.
def per_frame (frame_number):
test_patches = []
for manduca_number in range(len(positions)):
mx = positions[manduca_number][frame_number]
ml = legs[manduca_number][frame_number]
mm = muscles[manduca_number][frame_number]
lrs = all_leg_rects[manduca_number]
brs = all_body_rects[manduca_number]
mrs = all_muscle_rects[manduca_number]
for rect, leg, pos in zip(lrs, ml, mx):
x, w = pos, leg_width
if leg == 0:
y = manduca_number + 0.2
else:
y = manduca_number + 0.1
h = 0.75 - (y - manduca_number)
rect.set_bounds(x,y,w,h)
for br, mr, musc, pos, next_pos in zip(brs, mrs, mm, mx, mx[1:]):
br.set_x(pos+leg_width)
br.set_width(next_pos-pos-leg_width)
mr.set_x(pos+leg_width)
mr.set_width(next_pos-pos-leg_width)
mr.set_alpha(musc)
return all_patches
xmin = min([min(min(val) for val in l) for l in positions])
xmax = max([max(max(val) for val in l) for l in positions])
xmax += leg_width
subplot_kwargs = kwargs.pop('subplot_kwargs',{})
fig,axes = plt.subplots(**subplot_kwargs)
axes.axis ([xmin,xmax,0,len(positions) + 0.1])
axes.set_autoscale_on(False)
axes.get_yaxis().set_visible(False)
for i,label in enumerate(labels):
axes.text (xmin, i+0.95,label, va='top')
axes.plot([xmin,xmax],[i+0.1]*2, 'brown')
all_patches = []
all_leg_rects, all_body_rects, all_muscle_rects = [], [], []
for idx, (x, l, m) in enumerate(zip(positions, legs, muscles)):
x0, l0, m0 = x[0], l[0], m[0]
leg_rects = [patches.Rectangle ((0,0),0,0, facecolor='r') for _ in l0]
body_rects = [patches.Rectangle ((0,idx+0.25),0,0.5, facecolor='g') for _ in m0]
muscle_rects = [patches.Rectangle ((0,idx+0.375),0,0.25, facecolor='k') for _ in m0]
all_leg_rects.append(leg_rects)
all_body_rects.append(body_rects)
all_muscle_rects.append(muscle_rects)
all_patches.extend(leg_rects+body_rects+muscle_rects)
[axes.add_patch(p) for p in all_patches]
msec_per_frame = dx * 1
num_frames = len(positions[0])
ani = animation.FuncAnimation(fig, per_frame, frames=num_frames,
interval=msec_per_frame,
blit=True,
repeat=False)
return ani
Markdown is supported
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