Quantization RMSE

#!/usr/bin/env python

# Compare the Root Mean Square Errors between these two quantization approaches:
# a) Simple truncation
# b) Truncation combined with error diffusion

import random
from math import sqrt

# Generate a pixel vector that has at least this length
minimumNumberOfPixelsToGenerate = 500000

# Simulate a smooth gradient by appending uniform chunks of identically-valued pixels
# to the pixel vector.  Each chunk has a random width comprised between 5 and 20 pixels.
# A smooth gradient is simulated by enforcing that the value of these uniform chunks must
# randomly differ by +1 or -1 from the preceding chunk.  As an exception, if the previous
# random chunk contained zero-valued pixels, then, as we don't allow negative pixel
# values, the next chunk must containe pixels with a value of one.

# We start with an empty vector
pixelVector = []

numberOfGeneratedPixels = 0

previousPixelValue = random.randint( 0, 15 )

while numberOfGeneratedPixels < minimumNumberOfPixelsToGenerate:

    widthOfThisChunk = random.randint( 4, 20 )

    # The new chunk must contain values that differ from the previous chunk by -1 or +1
    newPixelValue = previousPixelValue + random.choice( [ -1, 1 ] )

    if newPixelValue < 0:
        # The previous pixel value must then necessarily have been zero
        newPixelValue = 1

    for i in range( widthOfThisChunk ):
        pixelVector.append( newPixelValue )

    previousPixelValue = newPixelValue
    numberOfGeneratedPixels += widthOfThisChunk

########### First case: simple truncation of the two lower bits by using the right shift (>>) operator

truncatedQuantizedVector = [ p >> 2 for p in pixelVector ]

accumulatedSquareError = 0
for i in range( numberOfGeneratedPixels ):
    error = pixelVector[ i ] - truncatedQuantizedVector[ i ] * 4
    accumulatedSquareError += error * error

print( "RMSE straightforward truncation:", sqrt( float( accumulatedSquareError ) / numberOfGeneratedPixels ) )

########### Second case: Error diffusion

errorDiffusedQuantizedVector = [ 0 for i in range( numberOfGeneratedPixels ) ]

accumulatedError = 0
for i in range( numberOfGeneratedPixels ):

    quantizedPixel = pixelVector[ i ] >> 2
    error = pixelVector[ i ] - quantizedPixel * 4
    accumulatedError += error

    if accumulatedError > 2:
        quantizedPixel += 1
        accumulatedError -= 4
    elif accumulatedError < -2:
        quantizedPixel -= 1
        accumulatedError += 4

    errorDiffusedQuantizedVector[ i ] = quantizedPixel

accumulatedSquareError = 0
for i in range( numberOfGeneratedPixels ):
    error = pixelVector[ i ] - errorDiffusedQuantizedVector[ i ] * 4
    accumulatedSquareError += error * error

print( "RMSE with error diffusion:", sqrt( float( accumulatedSquareError ) / numberOfGeneratedPixels ) )