Julia, Pyglet, and a GPU

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 \) 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):
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    glOrtho(-1., 1., 1., -1., 0., 1.)

    glMatrixMode(GL_MODELVIEW)
    glLoadIdentity()

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

    glBegin(GL_QUADS)
    glVertex2i(-1, -1)
    glTexCoord2i(-2, -2)
    glVertex2f(1, -1)
    glTexCoord2i(2, -2)
    glVertex2i(1, 1)
    glTexCoord2i(2, 2)
    glVertex2i(-1, 1)
    glTexCoord2i(-2, 2)
    glEnd()

    self.shader.unbind()

window = JuliaWindow()
pyglet.app.run()

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;

void
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 uniform 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;

void
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)
     break;

    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(
    color0,
   color1,
   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 GPUs are indeed great little pieces of silicon !