Mandelbrot Pygame

import pygame
from pygame.locals import *
import math
import os
import colorsys

WIDTH = 900
HEIGHT = 900
CELL_SIZE = 1
MAX_ITERATIONS = 100

def getScreenshotIndex():
    i = 0
    while os.path.isfile("screenshot_" + str(i) + ".jpg"):
        i = i + 1
    print("Screenshot index:", i)
    return i

SCREENSHOT_INDEX = getScreenshotIndex()

cells = []

def createMandlebrotSet():
    global cells

    w = 2
    h = (w * HEIGHT) / WIDTH

    xmin = -w / 2
    ymin = -h / 2

    xmax = xmin + w
    ymax = ymin + h

    dx = (xmax - xmin) / WIDTH * CELL_SIZE
    dy = (ymax - ymin) / HEIGHT * CELL_SIZE

    y = ymin

    for i in range(WIDTH // CELL_SIZE):
        cells.append([])
        x = xmin
        for j in range(HEIGHT // CELL_SIZE):
            a = x
            b = y

            ca = a  # Change these By default a
            cb = b  # Change these By default b

            n = 0
            while n < MAX_ITERATIONS:
                aa = a * a  # Change these By default a * a
                bb = b * b  # Change these By default b * b
                if aa + bb > 4:
                    break
                twoab = 2.0 * a * b
                a = aa - bb + ca
                b = twoab + cb

                n = n + 1
            if n == MAX_ITERATIONS:
                cells[i].append(0)
            else:
                n = math.sqrt(n / MAX_ITERATIONS) * 255
                cells[i].append(n)
            x = x + dx
        y = y + dy
    print("Created mandelbrot set!")

createMandlebrotSet()

pygame.init()

screen = pygame.display.set_mode((WIDTH, HEIGHT), HWSURFACE)
pygame.display.set_caption("Mandlebrot Set")

# Change here to change how the colour is calculated from the number of iterationss
def getColor(n):
    # rgb = colorsys.hsv_to_rgb(n / 255, 1, 150 / 255)
    # return (rgb[0] * 255, rgb[1] * 255, rgb[2] * 255)
    return (n, n, n)

def drawMandlebrotSet():
    global cells
    for x in range(WIDTH // CELL_SIZE):
        for y in range(HEIGHT // CELL_SIZE):
            if cells[x][y] != 0:
                pygame.draw.rect(screen, getColor(cells[x][y]), (x * CELL_SIZE, y * CELL_SIZE, CELL_SIZE + 1, CELL_SIZE + 1))

running = True
while running:
    drawMandlebrotSet()
    pygame.display.update()

    for event in pygame.event.get():
        if event.type == QUIT:
            running = False
        if event.type == KEYDOWN:
            if event.key == K_s:
                pygame.image.save(screen, "screenshot_" + str(SCREENSHOT_INDEX) + ".jpg")
                SCREENSHOT_INDEX = SCREENSHOT_INDEX + 1

pygame.quit()