How to render a Julia Set quickly with Pygame?

Viewed 491

Currently, I'm playing around with a script generating Julia sets and the Mandelbrot set and then using pygame to render the points.

Essentially, the screen is mapped to a smaller coordinate system where its bounded by -2.5, 2.5 on the x axis and -1, 1 on the y axis. Each of the pixels in this mapped range is then passed to a function to check whether its complex number equivalent is in the given set. This function returns the number of iterations it took to calculate whether the number is in the set or not (or the max iterations).

Then, for each pixel, I know what colour to colour it based on this iteration score and render each of the pixels one by one. This part of the process is really intensive and takes a ~30 seconds to render but can be much more depending on the complexity of the set.

Here is the code for finding out if a passed complex number and complex coordinate are in Julia set, this doesn't take long to compute at all when checking 1920 * 1080 pixels:

max_iter = 45


def julia(z, c):
    n = 0
    while abs(z) <= 2 and n < max_iter:
        z = z * z + c
        n += 1
    return n

Here is the code I use for pygame rendering, this is definitely where the problem lies:

size_ = 1920, 1080
re_ = -2.5, 2.5
im_ = -1, 1
surf = pygame.Surface(size)
colour_gradient1 = [c, c1, c2, c3, ...] # This is some list of colours generated by a gradient function

for x in range(0, size_[0]):
    for y in range(0, size_[1]):
        z = complex(re_[0] + (x / size_[0]) * (re_[1] - re_[0]),
                    im_[0] + (y / size_[1]) * (im_[1] - im_[0]))
        m = julia(z, c)
        colour = colour_gradient1[m]
        pygame.draw.rect(surf,
                         colour,
                         (x, y, 1, 1))

I think I understand why this is performance intensive in that both pygame and python aren't really optimised for rendering stuff to the screen like this. I'm currently trying to learn C++ and I understand its better for stuff like this.

I also experimented with a zoom function where I could select a box with the mouse and the script would render this selected area but implementing this was where the problem stuck out. As the zoomed in fractals got more complex, the script took too long to use this function.

So my question, is there a better way to render something like this in close to real-time using python and maybe pygame? I'm open to using a different package but if it's possible through pygame that would be ideal.

Attached below are a couple of pictures of the generated sets: This is an example of one of the fractals generated. This is an example of one of the fractals generated. This is an example of one of the fractals generated.

2 Answers

Fractal generating algorithms always slow down the further in you zoom because there needs to be ever more iterations per pixel the deeper you go (or before the bail-out is reached).

This is never going to be particularly fast in an interpreted language. Sure you can tweak it to increase the speed a little, but it will never be "real time" (say < 1 second / image) for all zoom levels.

If you want to continue in Python, you will have to just agree with yourself that it's never going to be fast.

However. However, you could split the generation of each quadrant into separate processes which would each run on their own CPU/core. That will give you a N/cores speed up.

There are some optimisations that can be performed with detecting symmetry in the image, and only calculating say half the pixels because the other side is a mirror of it (like the horizontal axis through an zoomed-out Mandelbrot set). You could probably refer to the source of the venerable Fractint Program for examples of this.

...

Aside: I wrote one of these (drawing a Mandelbrot set) in C using the nVidia CUDA library which spreads the calculation over the 1200-ish "CPU"s on the video card (using a mid-range 2018 laptop). While it worked quite fast, for sufficiently large images, or deeply "zoomed-in" fractals, it still became slow. There's just so much number crunching involved.

mandelbrot section

(This question finally made me install PyOpenGL. So thanks!)

As far as I've seen, iterating over each pixel individually,
will never give good performance (not in C++/C/Assembly/).
Vectorization (in the CPU) will help. What will really help,
is using the GPU's ability to apply one operation(/kernel),
to a whole multi-dimensional array of elements, in parallel.

Specifically: Using a fragment shader, for calculating the
color of each pixel. But that means using a graphics API
like OpenGL(/Vulkan/Direct3D/), or a GPGPU/Compute API like
OpenCL(/CUDA/).

If the resulting image is used within the graphics pipeline,
then it can stay on the GPU & be displayed directly from
there. If the resulting image needs to be used e.g. in a
GUI, saved to disk, or similar, it needs to be brought from
GPU to CPU (maybe render-to-texture, read the framebuffer,
use off-screen buffers, or other options that I don't know).

import numpy as np

from OpenGL.GL import *
from OpenGL.GL import shaders
from OpenGL.GLUT import *

# Vertex shader: Pass through (no model-view-projection).
vsSrc = '''
#version 300 es

layout (location = 0) in vec4 posIn;

void main()
{
    gl_Position = posIn;
}
'''

# Fragment shader: Compute fractal color, per-pixel.
# en.wikipedia.org/wiki/Mandelbrot_set#Computer_drawings
fsSrc = '''
#version 300 es

precision mediump float;

out vec4 colorOut;

vec2 mapLinear(
  vec2 val,
  vec2 srcMin, vec2 srcMax,
  vec2 dstMin, vec2 dstMax
) {
    vec2 valNorm = (val - srcMin) / (srcMax - srcMin);
    return valNorm * (dstMax - dstMin) + dstMin;
}

void main()
{
  // Debugging: Return fixed color; see which pixels get it.
  //colorOut = vec4(0.0, 0.5, 0.0, 1.0);
  //return;
  
  // Originally, origin is top-left. Convert to Cartesian.
  vec2 pixelMin = vec2(0.0f, 720.0f);
  vec2 pixelMax = vec2(1280.0f, 0.0f);
  
  vec2 mbMin = vec2(-2.5f, -1.0f);
  vec2 mbMax = vec2(1.0f, 1.0f);
  vec2 mbExtent = mbMax - mbMin;
  vec2 mbCenter = mbMin + (mbExtent / 2.0f);
  
  vec2 fragMapped = mapLinear(
    gl_FragCoord.xy, pixelMin, pixelMax, mbMin, mbMax
  );
  
  float real = 0.0f;
  float imag = 0.0f;
  int iter = 0;
  const int maxIter = 500;
  while (
    ((real*real + imag*imag) < 4.0f) && 
    (iter < maxIter)
  ) {
    float realTemp = real*real - imag*imag + fragMapped.x;
    imag = 2.0f*real*imag + fragMapped.y;
    real = realTemp;
    ++iter;
  }
  
  // Using generated colors, instead of indexing a palette.
  // (Don't remember anymore where this came from,
  // or if it was a heuristic.)
  vec3 chosenColor;
  float iterNorm = float(iter) / float(maxIter);
  if (iterNorm > 0.5f) {
    float iterNormInverse = 1.0f - iterNorm;
    chosenColor = vec3(
      0.0f, iterNormInverse, iterNormInverse
    );
  }
  else {
    chosenColor = vec3(0.0f, iterNorm, iterNorm);
  }
  
  colorOut = vec4(chosenColor.xyz, 1.0f);
}
'''

def compileFractalShader():
  vs = shaders.compileShader(vsSrc, GL_VERTEX_SHADER)
  fs = shaders.compileShader(fsSrc, GL_FRAGMENT_SHADER)
  return shaders.compileProgram(vs, fs)

# Geometry: Just 2 triangles, covering the display surface.
# (So that the fragment shader runs for all surface pixels.)
def drawTriangles():
  topLeftTriangle = (
    1.0, 1.0, 0.0,
    -1.0, -1.0, 0.0,
    -1.0, 1.0, 0.0
  )
  bottomRightTriangle = (
    1.0, 1.0, 0.0,
    -1.0, -1.0, 0.0,
    1.0, -1.0, 0.0
  )
  verts = np.array(
    topLeftTriangle + bottomRightTriangle,
    dtype=np.float32
  )
  
  glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, verts)
  glEnableVertexAttribArray(0)
  
  glDrawArrays(GL_TRIANGLES, 0, 6)

def printShaderException(e):
  errorMsg, shaderSrc, shaderType = e.args
  print('Shader error message:')
  for line in errorMsg.split('\\n'): print(line)
  print('--')
  #print('Shader source:')
  #for line in shaderSrc[0].split(b'\n'): print(line)
  #print('--')
  print('Shader type:', shaderType)

WIDTH = 1280
HEIGHT = 720

glutInit()
glutInitWindowSize(WIDTH, HEIGHT)
glutCreateWindow('Fractals with fragment shaders.')

# Create shaders, after creating a window / opengl-context:
try: fractalShader = compileFractalShader()
except RuntimeError as e: 
  printShaderException(e)
  exit()

glViewport(0, 0, WIDTH, HEIGHT)

glClearColor(0.5, 0.0, 0.5, 1.0)

def display():
  glClear(GL_COLOR_BUFFER_BIT)
  with fractalShader: drawTriangles()
  glutSwapBuffers()
glutDisplayFunc(display)

glutMainLoop()

enter image description here

This is entirely unoptimized.
Also, as Kingsley wrote, zooming (not shown here)
was slowing things down even in the GPU (but: unoptimized).

Related