Generating Julia set

Question

So, I just wanted to post something on Rosetta Code, and I found this task of generating and plotting a Julia set: http://www.rosettacode.org/wiki/Julia_set. There was already one solution but it was quite inefficient and not Pythonic. Here is my attempt on this:

"""
This solution is an improved version of an efficient Julia set solver
from:
'Bauckhage C. NumPy/SciPy Recipes for Image Processing:
 Creating Fractal Images. researchgate. net, Feb. 2015.'
"""
import itertools
from functools import partial
from numbers import Complex
from typing import Callable

import matplotlib.pyplot as plt
import numpy as np


def douady_hubbard_polynomial(z: Complex,
                              *,
                              c: Complex):
    """
    Monic and centered quadratic complex polynomial
    https://en.wikipedia.org/wiki/Complex_quadratic_polynomial#Map
    """
    return z ** 2 + c


def julia_set(*,
              mapping: Callable[[Complex], Complex],
              min_coordinate: Complex,
              max_coordinate: Complex,
              width: int,
              height: int,
              iterations_count: int = 256,
              threshold: float = 2.) -> np.ndarray:
    """
    As described in https://en.wikipedia.org/wiki/Julia_set
    :param mapping: function defining Julia set
    :param min_coordinate: bottom-left complex plane coordinate
    :param max_coordinate: upper-right complex plane coordinate
    :param height: pixels in vertical axis
    :param width: pixels in horizontal axis
    :param iterations_count: number of iterations
    :param threshold: if the magnitude of z becomes greater
    than the threshold we assume that it will diverge to infinity
    :return: 2D pixels array of intensities
    """
    imaginary_axis, real_axis = np.ogrid[
                        min_coordinate.imag: max_coordinate.imag: height * 1j,
                        min_coordinate.real: max_coordinate.real: width * 1j]
    complex_plane = real_axis + 1j * imaginary_axis

    result = np.ones(complex_plane.shape)

    for _ in itertools.repeat(None, iterations_count):
        mask = np.abs(complex_plane) <= threshold
        if not mask.any():
            break
        complex_plane[mask] = mapping(complex_plane[mask])
        result[~mask] += 1

    return result


if __name__ == '__main__':
    mapping = partial(douady_hubbard_polynomial,
                      c=-0.7 + 0.27015j)  # type: Callable[[Complex], Complex]

    image = julia_set(mapping=mapping,
                      min_coordinate=-1.5 - 1j,
                      max_coordinate=1.5 + 1j,
                      width=800,
                      height=600)
    plt.axis('off')
    plt.imshow(image,
               cmap='nipy_spectral',
               origin='lower')
    plt.show()

I think it looks good, and it is definitely more efficient. There was just one thing that I was not sure about. I was thinking to take out creating a complex_plane to a separate function and pass it as a parameter to julia_set. But in this case the julia_set wouldn't be a pure function as it would mutate the complex_plane. And I prefer my functions not to have any side effects. So I decided to leave it as is.

Any comments on this matter or anything else are welcome.

Here are some examples of output:

Answer 1

1. Review

1. Some of the variable names could be improved:

   complex_plane is an array of \$z\$ values for each pixel in the image, so naming it z would help the reader relate it to the z in douady_hubbard_polynomial.

   imaginary_axis and real_axis are only used once in the very next line, so there is no need for them to have long and memorable names. I would use something short like im and re.

   result is an array of iteration counts, so it could be named something like iterations.

   mask is a Boolean array selecting pixels that have not yet diverged to infinity, so something like not_diverged or live would convey this better.

2. On each iteration, the iteration counts of the escaped pixels are incremented. This means that some pixels get incremented many times, for example a pixel that escapes on the first iteration gets its count incremented 256 times. It would be more efficient to set the iteration count for each pixel just once. A convenient time to do this is when it escapes.

3. As the number of iterations goes up, the number of pixels that have not escaped to infinity gets smaller and smaller. But the masking operations are always on the whole array. It would be more efficient to keep track of the indexes of the pixels that have not escaped, so that subsequent operations are on smaller and smaller arrays.

2. Revised code

im, re = np.ogrid[min_coordinate.imag: max_coordinate.imag: height * 1j,
                  min_coordinate.real: max_coordinate.real: width * 1j]
z = (re + 1j * im).flatten()
live, = np.indices(z.shape) # indexes of pixels that have not escaped
iterations = np.empty_like(z, dtype=int)
for i in range(iterations_count):
    z_live = z[live] = mapping(z[live])
    escaped = abs(z_live) > threshold
    iterations[live[escaped]] = i
    live = live[~escaped]
iterations[live] = iterations_count - 1
return (iterations_count - iterations).reshape((height, width))

Notes

1. This is about three times as fast as the code in the post.

2. Because we are maintaining an array of indexes, it is convenient to flatten the z array and then reshape iterations to two dimensions before returning it. If we left the array two-dimensional, there would need to be two arrays of indexes, live_i and live_j.

3. Pixels that don't escape are given the value iterations_count - 1 in order to match the code in the post. It would make more sense to use iterations_count or a larger value here.

4. The subtraction iterations_count - iterations is only there so that the returned values match the code in the post. The subtraction could be omitted if you reverse the colour map.

Answer 2

For douady_hubbard_polynomial, you're missing a return type.

This:

for _ in itertools.repeat(None, iterations_count):

can just be

for _ in range(iterations_count):

I don't see any other obvious issues.

Answer 3

This is a minor change, with a 50% improvement to @gareth's answer. Changing

escaped = abs(z_live) > threshold to

escaped = z_live.real**2 + z_live.imag**2>threshold**2

while less nice looking is about 50% faster because it saves a square-root of all the elements.

To time it, I used

t1 = time()
image = julia_set(mapping=mapping,
                  min_coordinate=-1.5 - 1j,
                  max_coordinate=1.5 + 1j,
                  iterations_count = 255,
                  width=1920,
                  height=1080)
print(time() - t1)

Before: 3.16s After 2.21s