Wednesday, September 23, 2009

Animate simulations in python

I've recently written a nice short code in python to animate the results (plots) of a simulation using chaco (part of Enthought ETS)
It is not actually very difficult, but i've just made up this one for myself, and it also has some cool features, so i thought i'd share it with others may as well benefit from it.

  • This if useful for you if you ever do coding in python and need to plot somethings which may change with time. Example lets say you are plotting the evolution of temperature over a rod with time or solving a 1D Euler equation (say shock tube problem)
  • This code provides a simple function to do it easily.
  • You can play-pause the code (simulation) at any time.
  • You can zoom-pan theplot
  • You can plot multiple subplots in a single window (Example the velocity, density and pressure in a shock tube problem)
  • You can edit axis labels, font, grids using gui
  • You can save the plots
This is the only function defined in the file  you need to know:
def animate(func, delay=0.1, total_time=0.0, time_factor=1.0, size=(800, 600), title='Plot'):
    '''function to animate the values returned by a function
    func : function which returns a tuple of x,y values to animate on each call
        x is shape (N,), y is shape (N,) or (p,N), N is number of points, p is number of plots (properties)
    delay is the time interval in which to call the func after the previous func has returned(seconds)
    total_time is the time at which to stop the animation
            animation will stop when total_time > self.time / time_factor
            total_time <= 0 will continue indefinitely
    time_factor is the time to display as title (displayed_time=time/time_factor)
    In the plot window:
        pressing 'p' key will toggle animation play-pause
        pressing Ctrl-S will open a dialog to save a rendering of the plot
        pressing 'ESC' will reset the zoom level of the plot
        double clicking on some parts of the plot allows you to edit them in a gui,
            (axis titles, grids, ticks etc)
        if the window becomes unresponsive, pause the animation for a while

The docstring explains most of the things you need to know.
Here's how you could use it:
First we see how to animate a single plot (A moving sine wave in this case)
from animate import animate
from numpy import linspace, sin
x = linspace(-10, 10, 101)
i = 0.0
def get_data():
    global i
    i += 0.1
    return x + i, sin(x + i)

The above code snippet will generate a nice moving animation window. Here's a plot from the same

Plotting multiple values is just as simple. You only need to pass on a tuple of the y values:
from animate import animate
from numpy import linspace, sin, cos
x = linspace(-10, 10, 101)
i = 0.0
def get_data2():
    global i

    i += 0.1
    return x + i, (1 / (1 + x ** 2), sin(x + i), cos(x + i), sin(x + i), cos(x + i),)

Here's a plot from the above snippet:

As you can see i have tried to make it as easy as possible without losing out on functionality.

To run this you will need the Enthought tool suite ETS (only traits with wx backend and chaco are required)
Hope this is useful for someone.
You can get the code from here that is:

No comments:

Post a Comment