Mandelbrot fractal implementation in Go as described in The Go Programming Language, Kernighan et al.

mandelbrot.go

package main

import (
	"image"
	"image/color"
	"image/png"
	"math"
	"math/cmplx"
	"os"
)

const (
	xmin, ymin, xmax, ymax = -2, -2, +2, +2
	width, height          = 1024, 1024
	samplingFactor         = 2
)

func main() {

	// Obtain supersampled grid to implement anti-aliasing.
	var sampleGrid [][]color.RGBA = calculateSampleGrid(width, height, samplingFactor)
	downsampledImage := downsampleGrid(sampleGrid, samplingFactor)

	// Declare and define variable inferring its type from the right-hand side.
	img := image.NewRGBA(image.Rect(0, 0, width, height))
	for py := 0; py < height; py++ {
		for px := 0; px < width; px++ {
			point := downsampledImage[py][px]
			img.Set(px, py, color.RGBA{R: point.R, G: point.G, B: point.B, A: 255})
		}
	}

	png.Encode(os.Stdout, img) // NOTE -- Ignoring errors
}

func downsampleGrid(grid [][]color.RGBA, samplingFactor int) [][]color.RGBA {
	gridHeight := len(grid)
	gridWidth := len(grid[0])
	downsampledHeight := gridHeight / samplingFactor
	downsampledWidth := gridWidth / samplingFactor
	image := make([][]color.RGBA, downsampledHeight, downsampledHeight)

	var avgRed, avgGreen, avgBlue float64
	var downY, downX int
	for py := 0; py < gridHeight; py += samplingFactor {
		downY = py / samplingFactor
		image[downY] = make([]color.RGBA, downsampledWidth, downsampledWidth)
		for px := 0; px < gridWidth; px += samplingFactor {
			downX = px / samplingFactor
			// Take the average of the sub-pixels, apply basic gamma correction and convert to uint8 to put in image
			avgRed = float64(grid[py][px].R+grid[py][px+1].R+grid[py+1][px].R+grid[py+1][px+1].R) / float64(samplingFactor*2)
			avgRed = 255 * math.Pow(avgRed/255, 1.0/2.2)
			avgGreen = float64(grid[py][px].G+grid[py][px+1].G+grid[py+1][px].G+grid[py+1][px+1].G) / float64(samplingFactor*2)
			avgGreen = 255 * math.Pow(avgGreen/255, 1.0/2.2)
			avgBlue = 0.0

			image[downY][downX] = color.RGBA{R: uint8(avgRed), G: uint8(avgGreen), B: uint8(avgBlue), A: 255}
		}
	}

	return image
}

func calculateSampleGrid(width int, height int, samplingFactor int) [][]color.RGBA {
	sampleGridWidth := width * samplingFactor
	sampleGridHeight := height * samplingFactor
	sampleGrid := make([][]color.RGBA, sampleGridHeight, sampleGridHeight)

	for py := 0; py < sampleGridHeight; py++ {
		y := float64(py)/float64(sampleGridHeight)*(ymax-ymin) + ymin
		sampleGrid[py] = make([]color.RGBA, sampleGridHeight, sampleGridWidth)
		for px := 0; px < sampleGridWidth; px++ {
			x := float64(px)/float64(sampleGridWidth)*(xmax-xmin) + xmin
			z := complex(x, y)

			sampleGrid[py][px] = mandelbrot(z)
		}
	}

	return sampleGrid
}

func mandelbrot(z complex128) color.RGBA {
	const iterations = 200

	// Declare but don't define
	var v complex128
	for n := uint8(0); n < iterations; n++ {
		v = v*v + z
		if cmplx.Abs(v) > 2 {
			return getHeatmapColour(n)
		}
	}

	return color.RGBA{R: 255, G: 255, B: 255, A: 255}
}

func getHeatmapColour(numSamples uint8) color.RGBA {
	const contrast = 15
	const rgbMax = 255.0

	//intensity := numSamples * contrast
	//redComponent := math.Max(float64(rgbMax-(numSamples-100)), 0.0)
	redComponent := math.Min(float64((numSamples-100)*contrast), rgbMax)
	greenComponent := math.Max(float64(rgbMax-numSamples*contrast), 0.0)

	// Colour code:
	// - 0-99 = green
	// - 100-149 = yellow
	// - 150-200 = red
	// - >200 = black
	return color.RGBA{
		R: uint8(redComponent),
		G: uint8(greenComponent),
		B: 0,
		A: 255,
	}
}