## Arm simulation visualization with Matplotlib

One of the downsides of switching to Python from Matlab is that it can be a pain to plot some kinds of things, and I’ve found animations to be one those things. In previous posts I’ve done the visualization of my arm simulations through Pyglet, but I recently started playing around with Matplotlib’s animation function, and the results are pretty smooth. The process is also relatively painless and quick to get up and running, so I thought I would throw up some Matplotlib code for visualizing my previously described 2 link arm MapleSim simulation.

So, let’s look at the code:

```#Written by Travis DeWolf (Sept, 2013)
#Based on code by Jake Vanderplas - http://jakevdp.github.com

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

"""
:param list u: the torque applied to each joints
"""
def __init__(self, u = [.1, 0]):
self.u = np.asarray(u, dtype='float') # control signal
self.state = np.zeros(3) # vector for current state
self.L1=0.37 # length of arm link 1 in m
self.L2=0.27 # length of arm link 2 in m
self.time_elapsed = 0

self.sim.reset(self.state)

def position(self):
"""Compute x,y position of the hand"""

x = np.cumsum([0,
self.L1 * np.cos(self.state[1]),
self.L2 * np.cos(self.state[2])])
y = np.cumsum([0,
self.L1 * np.sin(self.state[1]),
self.L2 * np.sin(self.state[2])])
return (x, y)

def step(self, dt):
"""Simulate the system and update the state"""
for i in range(1500):
self.sim.step(self.state, self.u)
self.time_elapsed = self.state[0]

#------------------------------------------------------------
# set up initial state and global variables
dt = 1./30 # 30 fps

#------------------------------------------------------------
# set up figure and animation
fig = plt.figure(figsize=(4,4))
xlim=(-1, 1), ylim=(-1, 1))
ax.grid()

line, = ax.plot([], [], 'o-', lw=4, mew=5)
time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)

def init():
"""initialize animation"""
line.set_data([], [])
time_text.set_text('')
return line, time_text

def animate(i):
"""perform animation step"""
global arm, dt
arm.step(dt)

line.set_data(*arm.position())
time_text.set_text('time = %.2f' % arm.time_elapsed)
return line, time_text

# frames=None for matplotlib 1.3
ani = animation.FuncAnimation(fig, animate, frames=None,
interval=25, blit=True,
init_func=init)

# uncomment the following line to save the video in mp4 format.
# requires either mencoder or ffmpeg to be installed
#         extra_args=['-vcodec', 'libx264'])

plt.show()
```

There’s not all too much to it, which is nice. I’ve created a class `TwoLinkArm` that wraps the actual arm simulator, stores the current state of the arm, and has a `step` function that gets called to move the system forward in time. Then, I created a `line` and stored a reference to it; this is what is going to be updated in the simulation. I then need functions that specify how the `line` will be initialized and updated. The `init` function doesn’t do anything except set the `line` and `text` data to nothing, and then the `animate` function calls the arm simulation’s `step` function and updates the `line` and `text` data.

For more details about it there’s this blog post which steps through the process a bit more. For simple arm simulations the above is all I need though, so I’ll leave it there for now!

Here’s an animation of the resulting sim visualization, I’ve removed a lot of the frames to keep the size down. It’s smoother when running the actual code, which can be found up on my github.

Hurray better visualization!

Tagged , , ,