contour.py

"""
Contouring using the Marching Squares algorithm.
George Lesica
"""

from cairo import SVGSurface, Context
 
WIDTH = 8 * 72.0
HEIGHT = 8 * 72.0

WHITE = (1.0, 1.0, 1.0)
BLACK = (0.0, 0.0, 0.0)
RED = (1.0, 0.0, 0.0)
GREEN = (0.0, 1.0, 0.0)
BLUE = (0.0, 0.0, 1.0)

def contour_plot(A, v, ctx, color=BLACK):
    """
    Create contour map of data contained in A. Contour the value v, using the
    drawing context ctx. Draw the line using the given color.
    TODO: Handle ambiguous case (opposite corners) by taking the average of 4
    """
    ctx.set_source_rgb(*color)

    nrows = len(A)
    ncols = len(A[0])

    dx = WIDTH / (ncols - 1)
    dy = HEIGHT / (nrows - 1)

    hdx = 0.5 * dx
    hdy = 0.5 * dy

    B = []
    for row in A:
        B.append([0 if i > v else 1 for i in row])

    for r in range(nrows - 1):
        for c in range(ncols - 1):
            # Determine cell type, t
            t = 0
            t = (t | B[r][c]) << 1
            t = (t | B[r][c+1]) << 1
            t = (t | B[r+1][c+1]) << 1
            t = (t | B[r+1][c])

            # Compute upper left corner of cell on image
            x = c * dx
            y = r * dy

            # Draw lines on image, depending on type
            if t == 0:
                pass
            if t == 1:
                rxb = (A[r+1][c+1] - v) / (A[r+1][c+1] - A[r+1][c])
                ryl = (A[r][c] - v) / (A[r][c] - A[r+1][c])
                ctx.move_to(x, y + ryl * dy)
                ctx.line_to(x + (1 - rxb) * dx, y + dy)
            if t == 2:
                rx = (A[r+1][c] - v) / (A[r+1][c] - A[r+1][c+1])
                ry = (A[r][c+1] - v) / (A[r][c+1] - A[r+1][c+1])
                ctx.move_to(x + rx * dx, y + dy)
                ctx.line_to(x + dx, y + ry * dy)
            if t == 3:
                ryl = (A[r][c] - v) / (A[r][c] - A[r+1][c])
                ryr = (A[r][c+1] - v) / (A[r][c+1] - A[r+1][c+1])
                ctx.move_to(x, y + ryl * dy)
                ctx.line_to(x + dx, y + ryr * dy)
            if t == 4:
                rx = (A[r][c] - v) / (A[r][c] - A[r][c+1])
                ry = (A[r+1][c+1] - v) / (A[r+1][c+1] - A[r][c+1])
                ctx.move_to(x + rx * dx, y)
                ctx.line_to(x + dx, y + (1 - ry) * dy)
            if t == 5:
                rxt = (A[r][c] - v) / (A[r][c] - A[r][c+1])
                rxb = (A[r+1][c+1] - v) / (A[r+1][c+1] - A[r+1][c])
                ryl = (A[r][c] - v) / (A[r][c] - A[r+1][c])
                ryr = (A[r+1][c+1] - v) / (A[r+1][c+1] - A[r][c+1])
                ctx.move_to(x, y + ryl * dy)
                ctx.line_to(x + rxt * dx, y)
                ctx.move_to(x + (1 - rxb) * dx, y + dy)
                ctx.line_to(x + dx, y + (1 - ryr) * dy)
            if t == 6:
                rxt = (A[r][c] - v) / (A[r][c] - A[r][c+1])
                rxb = (A[r+1][c] - v) / (A[r+1][c] - A[r+1][c+1])
                ctx.move_to(x + rxt * dx, y)
                ctx.line_to(x + rxb * dx, y + dy)
            if t == 7:
                rxt = (A[r][c] - v) / (A[r][c] - A[r][c+1])
                ryl = (A[r][c] - v) / (A[r][c] - A[r+1][c])
                ctx.move_to(x, y + ryl * dy)
                ctx.line_to(x + rxt * dx, y)
            if t == 8:
                rxt = (A[r][c+1] - v) / (A[r][c+1] - A[r][c])
                ryl = (A[r+1][c] - v) / (A[r+1][c] - A[r][c])
                ctx.move_to(x, y + (1 - ryl) * dy)
                ctx.line_to(x + (1 - rxt) * dx, y)
            if t == 9:
                rxt = (A[r][c+1] - v) / (A[r][c+1] - A[r][c])
                rxb = (A[r+1][c+1] - v) / (A[r+1][c+1] - A[r+1][c])
                ctx.move_to(x + (1 - rxt) * dx, y)
                ctx.line_to(x + (1 - rxb) * dx, y + dy)
            if t == 10:
                rxt = (A[r][c+1] - v) / (A[r][c+1] - A[r][c])
                rxb = (A[r+1][c] - v) / (A[r+1][c] - A[r+1][c+1])
                ryl = (A[r+1][c] - v) / (A[r+1][c] - A[r][c])
                ryr = (A[r][c+1] - v) / (A[r][c+1] - A[r+1][c+1])
                ctx.move_to(x + (1 - rxt) * dx, y)
                ctx.line_to(x + dx, y + ryr * dy)
                ctx.move_to(x, y + (1 - ryl) * dy)
                ctx.line_to(x + rxb * dx, y + dy)
            if t == 11:
                rxt = (A[r][c+1] - v) / (A[r][c+1] - A[r][c])
                ryr = (A[r][c+1] - v) / (A[r][c+1] - A[r+1][c+1])
                ctx.move_to(x + (1 - rxt) * dx, y)
                ctx.line_to(x + dx, y + ryr * dy)
            if t == 12:
                ryl = (A[r+1][c] - v) / (A[r+1][c] - A[r][c])
                ryr = (A[r+1][c+1] - v) / (A[r+1][c+1] - A[r][c+1])
                ctx.move_to(x, y + (1 - ryl) * dy)
                ctx.line_to(x + dx, y + (1 - ryr) * dy)
            if t == 13:
                rxb = (A[r+1][c+1] - v) / (A[r+1][c+1] - A[r+1][c])
                ryr = (A[r+1][c+1] - v) / (A[r+1][c+1] - A[r][c+1])
                ctx.move_to(x + (1 - rxb) * dx, y + dy)
                ctx.line_to(x + dx, y + (1 - ryr) * dy)
            if t == 14:
                rxb = (A[r+1][c] - v) / (A[r+1][c] - A[r+1][c+1])
                ryl = (A[r+1][c] - v) / (A[r+1][c] - A[r][c])
                ctx.move_to(x, y + (1 - ryl) * dy)
                ctx.line_to(x + rxb * dx, y + dy)
            if t == 15:
                pass

            ctx.stroke()


def colors(base, values):
    """
    Create a color for each value in values using base as the base color.
    """
    minval = min(values)
    maxval = max(values)
    scaled_values = [(v - minval) / (maxval - minval) for v in values]
    return [(0.0, 0.0, v) for v in scaled_values]


def contour(data, values, filename):
    """
    Creates a map of data contours of the values given and stores it as
    filename.
    """
    sur = SVGSurface(filename + '.svg', WIDTH, HEIGHT)
    ctx = Context(sur)
    ctx.set_line_width(2.0)
    ctx.set_source_rgb(*WHITE)
    ctx.paint()

    for value, color in zip(values, colors(None, values)):
        cells = contour_plot(data, value, ctx, color)

    sur.write_to_png(filename + '.png')
    sur.finish()


if __name__ == '__main__':
    testvals = [[1, 1, 1, 1, 1, 1, 1],
                [1, 1, 2, 2, 2, 1, 1],
                [1, 2, 2, 3, 2, 2, 1],
                [1, 2, 3, 3, 3, 2, 1],
                [1, 2, 2, 3, 2, 2, 1],
                [1, 1, 2, 2, 2, 1, 1],
                [1, 1, 1, 1, 1, 1, 1]]
    contour(testvals, [1.5, 2.5], 'test')