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

manduca/Visualization.py

+"""
+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