Julia, Pyglet, and a GPU


A Julia fractal. Note the smooth, anti-aliased rendering. The color is picked as function of the estimated function to the Julia boundary.

A Julia fractal. Note the smooth, anti-aliased rendering. The color is picked as function of the estimated function to the Julia boundary.

Pyglet is a very good Python module, wrapping OpenGL and providing access to mouse and keyboards events. The API is really easy to learn, well-documented, with that "clean & lean" feel to it. I started to use Pyglet for work, and I positively love it.

So, Pyglet, allows you to summon the power of your graphic card using Python. And those days, even a puny low-end graphic card packs a lot of power : GPU, massively parallel computing. Although it's been around for years now, I never bothered tinkering with GPU, lack of time, other interests. Now that I feel comfortable with OpenGL, Python and Pyglet, however, harnessing this GPU goodness require very little effort. With years of delay, let's jump in that band-wagon already !

For a first try, I went for the done, done again, and overdone : fractals. More precisely, Julia fractals. You can see this fractal as function to iterate for each pixel of a picture, that takes as input the coordinate of the pixel.

Z_{n+1} = Z_{n}^{2} + C

Z_n is a complex variable, with Z_0 being the coordinate of the pixel in the complex plane, and C a complex constant. This function is iterated until |Z_N|<2 or a maximum count of iterations is reached. You then color the pixel depending on the number of iterations spent. Consider visiting the Wikipedia page about Julia fractals, there is a lot to learn about it there.

And you get a nice picture, a fractal. What happen if you play with the parameter C ? The Julia fractal is fairly intensive to compute, rendering it in real-time requires some significant engineering... or a GPU, as the computations for each pixel are completely parallel. So what I did is to render the Julia fractal as a fragment shader, and control the C parameter with the mouse. The code for this turns out to be dead simple, assuming you are familiar with the OpenGL API

import pyglet
from pyglet.gl import *
from shader import Shader

class JuliaWindow(pyglet.window.Window):
  def __init__(self):
    super(JuliaWindow, self).__init__(caption = 'julia', width = 512, height = 512)
    self.C = (-0.70176, -0.3842)
    shader_path = 'julia'
    self.shader = Shader(
      ' '.join(open('%s.v.glsl' % shader_path)),
      ' '.join(open('%s.f.glsl' % shader_path))

  def on_mouse_motion(self, x, y, dx, dy):
    self.C = (6. * ((float(x) / window.width) - .5), 6 * ((float(y) / window.height) - .5))

  def on_draw(self):
    glOrtho(-1., 1., 1., -1., 0., 1.)


    self.shader.uniformf('C', *self.C)

    glVertex2i(-1, -1)
    glTexCoord2i(-2, -2)
    glVertex2f(1, -1)
    glTexCoord2i(2, -2)
    glVertex2i(1, 1)
    glTexCoord2i(2, 2)
    glVertex2i(-1, 1)
    glTexCoord2i(-2, 2)


window = JuliaWindow()

This code create a window, and display a quad polygon that fills the window. The texture for the quad is a procedural one, computed by a shader. Here, I use a neat utility class that encapsulate OpenGL API for the shaders : shader.py, by Tristam MacDonald. The C parameter of the fractal  is controlled by the mouse position in the window.

Of course, most of the magic happens inside the shaders. The vertex shader,   julia.v.glsl, is not exactly rocket science.

#version 110

varying vec2 uv;

main() {
  gl_Position = gl_ModelViewProjectionMatrix * gl_Vertex;
  uv = vec2(gl_MultiTexCoord0);

All the actual job is done in the fragment shader, julia.f.glsl. Basically, I count how much iterations are required to have |Z_N|<2, and use this to pick a color. C can be controlled from outside the shader, as uniforn variable. For sake of simplicity, I hard-coded the color scheme. One nice improvement would be to pick the color from a 1D texture, itself procedurally generated.

#version 110

varying vec2 uv;

int max_iter_count;
uniform vec2 C;

int iter_count;
vec2 Z, dZ;
vec2 sqr_Z;
float sqr_norm_Z, sqr_norm_dZ;

main() {
  max_iter_count = 1024;

  Z = uv;
  dZ = vec2(1., 0.);

  for(iter_count = 0; iter_count < max_iter_count; ++iter_count) {
    sqr_Z = vec2(Z.x * Z.x, Z.y * Z.y);
    if (sqr_Z.x + sqr_Z.y > 1e7)

    dZ = vec2(
      2. * (Z.x * dZ.x - Z.y * dZ.y) + 1.,
      2. * (Z.x * dZ.y + Z.y * dZ.x));
      Z = vec2(
        Z.x * Z.x - Z.y * Z.y,
        2. * Z.x * Z.y) + C;

  sqr_norm_Z = Z.x * Z.x + Z.y * Z.y;
  sqr_norm_dZ = dZ.x * dZ.x + dZ.y * dZ.y;
  vec4 color0 = vec4(1.0, 0.0, 0.0, 1.0);
  vec4 color1 = vec4(1.0, 1.0, 0.0, 1.0);

  gl_FragColor = mix(
   sqrt(sqrt(sqr_norm_Z / sqr_norm_dZ) * .5 * log(sqr_norm_Z))

Note that here I use a variant on the classical Julia : I render a distance to the boundary of the fractal, the trick being from the amazing Inigo Quilez. It gives me anti-aliasing at a much cheaper cost than super-sampling. A few minutes of coding, and I got that instant enlightenment, with a feeling of achievement on the top of it. Those GPU are indeed great little pieces of silicon !

The full code is available for download julia-explorer.tar.bz2