Home > Misc > WebGL

Shaders in three.js: plotting the Mandelbrot set

This is a tutorial on using custom shaders in three.js. There are other tutorials on this subject, but this one is mine. Much more than for other programming topics, I found myself completely bewildered by the tutorials and had to study examples and try my own before my brain latched on to the ideas. So I don't know how useful this post will be, except to future-me, who'll have similar thinking patterns to present- and recent-past-me.

The examples require WebGL, which most visitors here should have, but the Mandelbrot set should fail in IE. (At least, a very similar case failed in IE when I tested it.)

Let's look at the Mandelbrot set: the set of complex numbers c for which the sequence defined by z0 = 0 and zn+1 = zn2 + c remains bounded; click to zoom. You won't be able to zoom too far and maintain good resolution – you can increase the number of iterations but not past 216. You might crash your graphics card if you set the iterations too high and zoom on a dark area, though for me it recovered after a second or two. On the one hand, this is disappointing. On the other, it's useful to know that you can't do calculations that are too intense in the shaders. Consider yourselves warned, anyway.

Iterations:
Re(z): ,
Im(z): ,
Length:
Zoom out, Reset.

It's pretty quick! (At least on my laptop.) This is the wonder of using the GPU to do grunt work – since each pixel in the Mandelbrot set is calculated independently of the others, plotting it is a problem that benefits enormously from parallelisation, which is what the GPU provides.

So, how does it work? I'm using three.js (r81) as my interface to WebGL, and some of what I write will be specific to it, and some of it will be more general. I'll start by only making minor nods to three.js processes, before going into the code later.

The basic idea of three.js is that you define objects in a three-dimensional space, along with a camera. A fractal isn't the natural use case for three.js, since I'm just plotting on a flat plane, but hey, it works. You just make a square whose vertices have z-coordinates of zero, put an orthographic camera above it, then colour the square with the Mandelbrot set.

So let's think about the square (everything I write will also apply to quadrilaterals more generally). Any mesh in the scene is made up of triangles. You can tell three.js to make a beautifully textured sphere, but it won't be a perfect sphere – it will be made of hundreds of triangles, hopefully all small enough that the curves look smooth on the screen. A square is made up of two triangles, the two triangles sharing two vertices in common.

So to make a square, we'll need to define six vertices, and put them into a geometry; three.js will automatically interpret a set of six vertices as two triangles. The vertex locations can be chosen freely – we'll be defining the complex plane for the Mandelbrot set later – so I will make the easy choice of {[-1, -1, 0], [1, 1, 0], [-1, 1, 0]} for the first triangle and {[-1, -1, 0], [1, 1, 0], [1, -1, 0]} for the second triangle.

When working with shaders, we're able to assign attributes to each vertex independently, and access these values in the vertex shader, which is a piece of code that runs on all the vertices. But all the action takes place inside the triangles – inside each fragment, as they're called. When writing the fragment shader, we'll need to know what coordinate (in the complex plane) we're testing.

The key is as follows. The vertex shader always runs first, and you can pass variables from the vertex shader onto the fragment shader. A fragment has three vertices defining it, so what happens when you "pass" three variables to it and only have one variable name to hold these values in the fragment shader? The answer is that as the fragment shader works its way through all the relevant pixels, it automatically linearly interpolates between the variables passed from its vertices, based on how far the current pixel is from each vertex.

So, for example, if each vertex gets assigned a different colour, then the fragment will be coloured by smooth interpolation:

This means that we can make a Mandelbrot set! We just have to assign two attributes to each vertex: the real and imaginary components of the complex number that we want the vertex to represent. They will then define a square in the complex plane, and when the fragment shader runs, it will have access to the complex number that it's testing for inclusion in the Mandelbrot set, and it can do its calculations on it.

Passing variables from the vertex shader to the fragment shader is both unintuitive and easy: you declare a varying in both shaders, each of the same type and with the same name.

Let's assume that we know how to pass attributes from three.js into the vertex shader, and see how the shaders work. They're written in GLSL, which looks a lot like C, and you'll need to remember to type the decimal point when putting numbers into floats. The vertex shader runs first, and in this case it is quite short. The attributes (defined in the JavaScript) get declared, the varyings, are also declared, and then a short void main() runs. Inside main(), the values of the varyings get set to the values of the attributes, and then the position of the vertex is calculated (by a copy-paste job from other examples).

The variable position doesn't appear to be declared, because three.js automatically declares it (along with a whole host of other variables); this makes the code easier to write and more mysterious to understand. The vec4 indicates that it's creating a vector of 4 floats; position is defined as a vec3, and I don't know what the fourth entry in the vec4 represents.

/* The vertex shader for the Mandelbrot set plotter */

// The real and imaginary components of the
// complex number this vertex represents; must
// be set in the JavaScript.
attribute float vertex_z_r;
attribute float vertex_z_i;

// varyings get passed to the fragment shader; these
// will be the same as the attributes but the names
// needs to be different:
varying float c_r;
varying float c_i;

void main() {
  // Say what the varying values are.
  c_r = vertex_z_r;
  c_i = vertex_z_i;
  
  // Have to say what the position is:
  gl_Position = projectionMatrix * modelViewMatrix * vec4(position, 1.0);
}

Next comes the fragment shader. The two varyings are declared, giving us the complex number to be tested, as is a uniform – like an attribute but holds the same value across all vertices (like attributes, they are set in the JavaScript). Half of the rest of the code is then preparing a blue-yellow colour scale. An unexpected aspect of GLSL is that in a for loop, you can't compare the looping variable to a non-constant. Since I want the user to set the maximum number of iterations, that means breaking the loop in an if block.

The end of the shader routine is defining the colour in a vec4 – red, green, blue, and alpha components.

/* The fragment shader for the Mandelbrot set plotter */

// Varyings get declared in both vertex and fragment shaders:
varying float c_r;
varying float c_i;

// uniforms are defined in the JavaScript:
uniform int max_iterations;

// Defining the lower and upper ends of the colour scale:
#define r_min 0.0
#define r_max 1.0

#define g_min 0.0
#define g_max 1.0

#define b_min 0.35
#define b_max 0.0

void main() {
  float r;
  float g;
  float b;
  float t;
  float w_r;
  float w_i;
  float u;
  float v;
  
  // Initial value in the w --> w^2 + c:
  w_r = 0.0;
  w_i = 0.0;
  
  // Colour will be black if it stays converged:
  r = 0.0;
  g = 0.0;
  b = 0.0;
  
  // The looping variable can't be compared to a non-constant,
  // so to allow the user to change the maximum number of
  // iterations, I break manually.
  for (int i = 0; i < 65536; i++) {
    // The fractal calculation:
    u = w_r;
    v = w_i;
    
    w_r = u*u - v*v + c_r;
    w_i = 2.0*u*v + c_i;
    
    if (w_r*w_r + w_i*w_i > 4.0) {
      // Diverged: make a pretty colour.
      t = log(float(i + 1)) / log(float(max_iterations + 1));
      
      r = t*r_max + (1.0 - t)*r_min;
      g = t*g_max + (1.0 - t)*g_min;
      b = t*b_max + (1.0 - t)*b_min;
      
      break;
    }
    
    if (i >= max_iterations) {
      break;
    }
  }
  
  // Set the colour:
  gl_FragColor = vec4(r, g, b, 1.0);
}

Now that the shaders are written, how do we use them? In three.js, every mesh comprises a geometry (which holds the vertices) and a material; the shaders and any uniforms get attached to the material. Each shader function must be in a single string, and for this page (following others' examples) I've typed them into <script> elements with weird types so that the browser ignores them, and then in the JavaScript I grab the .textContent. It looks a bit like this:

var material = new THREE.ShaderMaterial({
  "uniforms": {"max_iterations": {"type": "i", "value": 100}},
  "vertexShader":   document.getElementById("shader_vertex").textContent,
  "fragmentShader": document.getElementById("shader_fragment").textContent,
  "side": THREE.DoubleSide
});

The argument of THREE.ShaderMaterial() is an object, with properties as shown above (the side property is set to THREE.DoubleSide so that I don't have to worry about getting the order of the vertices correct). The uniforms are defined in their own object, one object per uniform, and each uniform is defined in an object, containing properties type (the "i" means that the uniform is an integer) and value.

Since attributes are defined for each vertex, they get attached to the geometry of the mesh. The geometry must be a THREE.BufferGeometry(), not a regular THREE.Geometry(). One-dimensional arrays of the appropriate size are defined – 3 times the number of vertices for positions – and then get put in a THREE.BufferAttribute which itself is added to the geometry by calling .addAttribute(). I found it easy to make mistakes getting all this straight.

var geometry = new THREE.BufferGeometry();

var z_r = new Float32Array(6);
var z_i = new Float32Array(6);
var vertices = new Float32Array(18);

// vertices hold x-, y-, z-coordinates in world space.
// z_r and z_i are the real and imaginary components of
// the points the vertices represent in the complex
// plane that will hold the fractal.

// First triangle:
vertices[0]  = -1.0;
vertices[1]  = -1.0;
vertices[2]  =  0.0;

vertices[3]  =  1.0;
vertices[4]  =  1.0;
vertices[5]  =  0.0;

vertices[6]  = -1.0;
vertices[7]  =  1.0;
vertices[8]  =  0.0;

// Second triangle:
vertices[9]  = -1.0;
vertices[10] = -1.0;
vertices[11] =  0.0;

vertices[12] =  1.0;
vertices[13] =  1.0;
vertices[14] =  0.0;

vertices[15] =  1.0;
vertices[16] = -1.0;
vertices[17] =  0.0;

// The second argument in THREE.BufferAttribute() is how many entries in the
// array belong to each vertex.
geometry.addAttribute("position", new THREE.BufferAttribute(vertices, 3));
geometry.addAttribute("vertex_z_r", new THREE.BufferAttribute(z_r, 1));
geometry.addAttribute("vertex_z_i", new THREE.BufferAttribute(z_i, 1));

// Make the mesh!
var the_rectangle = new THREE.Mesh(geometry, material);

That code omitted setting the values of the vertex_z_r and vertex_z_i attributes. I farmed them out to another function, and writing it separately here is instructive because it also shows how to update an attribute: you make a shallow copy of the array, edit the shallow copy, then set the needsUpdate property to true.

function set_plot_bounds(mid_z_r, mid_z_i, range) {
  // mid_z_r, mid_z_i: coords at the middle of the desired
  // square in the complex plane.
  //
  // range: the half-width of the square
  
  // Make a shallow copy of the attribute arrays:
  var z_r = the_rectangle.geometry.attributes.vertex_z_r.array;
  var z_i = the_rectangle.geometry.attributes.vertex_z_i.array;
  
  // First triangle:
  z_r[0] = mid_z_r - range;
  z_i[0] = mid_z_i - range;
  
  z_r[1] = mid_z_r + range;
  z_i[1] = mid_z_i + range;
  
  z_r[2] = mid_z_r - range;
  z_i[2] = mid_z_i + range;
  
  // Second triangle:
  z_r[3] = mid_z_r - range;
  z_i[3] = mid_z_i - range;
  
  z_r[4] = mid_z_r + range;
  z_i[4] = mid_z_i + range;
  
  z_r[5] = mid_z_r + range;
  z_i[5] = mid_z_i - range;
  
  // Have to tell the shader to re-read the attributes:
  the_rectangle.geometry.attributes.vertex_z_r.needsUpdate = true;
  the_rectangle.geometry.attributes.vertex_z_i.needsUpdate = true;
  
  // I haven't defined any of this in the exceprts, but it's all
  // standard three.js:
  renderer.render(scene, camera);
}

Full details in the source to this page, where variable names are a little different.

Posted 2016-10-05.


Home > Misc > WebGL