polarchecker.py

# vim: filetype=pyopencl.python
import numpy as np
from PIL import Image
import pyopencl as cl
from pyopencl import array as clarray

SRC = '''//CL//
__constant const int CHECKER_SIZE = 10;

float2 polar(const int2);
int disc(const int2);


float2 polar(const int2 pos)
{
    const float x = pos.x;
    const float y = pos.y;
    return (float2)(sqrt(hypot(x, y)), atan2(x, y));
}


int disc(const int2 pos)
{
    const float2 ppos = polar(pos);
    return ((int)(ppos.x) + (int)(ppos.y * CHECKER_SIZE / M_PI_F)) % 2;
}


__kernel void gen(write_only image2d_t dest)
{
    const int2 pos = (int2)(get_global_id(0), get_global_id(1));
    uint4 pix;
    if (disc(pos)) {
        pix = (uint4)(0, 0, 255, 255);
    } else {
        pix = (uint4)(255, 255, 255, 255);
    }
    write_imageui(dest, pos, pix);
}
'''

ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)
prg = cl.Program(ctx, SRC).build()

width = 640
height = 480
fmt = cl.ImageFormat(cl.channel_order.RGBA, cl.channel_type.UNSIGNED_INT8)
buf = cl.Image(ctx, cl.mem_flags.WRITE_ONLY, fmt, shape=(width, height))

prg.gen(queue, (width, height), None, buf)

result = np.empty((height, width, 4), dtype=np.uint8)
cl.enqueue_copy(queue, result, buf, origin=(0, 0), region=(width, height))

image = Image.fromarray(result)
image.save('checker.png')