wxgl sample

from __future__ import with_statement

import math
import os
import sys
import numpy as np
from contextlib import contextmanager

import wxversion
wxversion.select('2.8')
import wx
import wx.glcanvas as wxgl

import OpenGL
# OpenGL.ERROR_CHECKING = False
from OpenGL.GL import *
from OpenGL.GLU import *

fragment_src = r'''
#version 120

const float bailout = 2.0;
const int maxiter = 64;

uniform bool mandel;
uniform vec2 pos;

vec2 csqr(vec2 z)
{
    float re = z[0] * z[0] - z[1] * z[1];
    float im = 2 * z[0] * z[1];
    return vec2(re, im);
}

void main()
{
    vec2 c, z;
    if (mandel) {
        c = gl_TexCoord[0].st;
        z = vec2(0, 0);
    }
    else {
        c = pos;
        z = gl_TexCoord[0].st;
    }

    int i;
    for (i = 0; i < maxiter; i++) {
        if (length(z) >= bailout)
            break;
        z = csqr(z) + c;
    }

    float k = 5 * i / maxiter;
    switch (int(k)) {
    case 0:  gl_FragColor = mix(vec4(0,0,1,1),vec4(0,1,1,1),k-0); break;
    case 1:  gl_FragColor = mix(vec4(0,1,1,1),vec4(0,1,0,1),k-1); break;
    case 2:  gl_FragColor = mix(vec4(0,1,0,1),vec4(1,1,0,1),k-2); break;
    case 3:  gl_FragColor = mix(vec4(1,1,0,1),vec4(1,0,0,1),k-3); break;
    case 4:  gl_FragColor = mix(vec4(1,0,0,1),vec4(1,0,1,1),k-4); break;
    case 5:  gl_FragColor = vec4(0,0,0,1); break;
    }
}
'''

@contextmanager
def Attrib(attrib):
    glPushAttrib(attrib)
    yield
    glPopAttrib()

@contextmanager
def Matrix():
    glPushMatrix()
    yield
    glPopMatrix()

@contextmanager
def Primitive(type):
    glBegin(type)
    yield
    try:
        glEnd()
    except Exception, e:
        print 'exception: %s' % e

class GLCanvas(wxgl.GLCanvas):
    attrib = [wxgl.WX_GL_RGBA, wxgl.WX_GL_DOUBLEBUFFER, 0]

    def __init__(self, parent, id = -1, *args, **kwargs):
        wxgl.GLCanvas.__init__(self, parent, id, attribList = GLCanvas.attrib,
                               *args, **kwargs)
        self.realized = False
        self.ctx = None
        self.program = None
        self.SetInitialSize((640, 640))
        self.Bind(wx.EVT_SIZE, self.resize)
        self.Bind(wx.EVT_ERASE_BACKGROUND, self.erase)

    def realize(self):
        if self.realized:
            self.SetCurrent(self.ctx)
            return

        self.ctx = wxgl.GLContext(self)
        self.SetCurrent(self.ctx)
        self.init_shader()
        self.realized = True

        if not self.program:
            sys.exit(1)

    def erase(self, event):
        pass

    def resize(self, event):
        self.realize()
        w, h = event.GetSize()
        glViewport(0, 0, w, h)
        with Attrib(GL_TRANSFORM_BIT):
            glMatrixMode(GL_PROJECTION)
            glLoadIdentity()
            gluOrtho2D(-2, 2, -2, 2)

    def init_shader(self):
        if not bool(glCreateShader):
            raise NotImplementedError('shaders not supported')

        shader = glCreateShader(GL_FRAGMENT_SHADER)
        glShaderSource(shader, fragment_src)
        glCompileShader(shader)
        if not glGetShaderiv(shader, GL_COMPILE_STATUS):
            raise RuntimeError(glGetShaderInfoLog(shader))

        program = glCreateProgram()
        glAttachShader(program, shader)
        glLinkProgram(program)
        if not glGetProgramiv(program, GL_LINK_STATUS):
            raise RuntimeError(glGetProgramInfoLog(program))

        self.program = program

    @staticmethod
    def quad():
        with Primitive(GL_QUADS):
            for v in [(-2,-2), (2,-2), (2,2), (-2,2)]:
                glTexCoord2f(*v)
                glVertex2f(*v)

class Mandelbrot(GLCanvas):
    def __init__(self, *args, **kwargs):
        GLCanvas.__init__(self, *args, **kwargs)
        self.pos = (0, 0)
        self.Bind(wx.EVT_PAINT, self.expose)
        self.Bind(wx.EVT_MOTION, self.motion)

    def motion(self, event):
        self.realize()
        x, y = event.GetPosition()
        w, h = self.GetSize()
        x =  4.0 * x / w - 2
        y = -4.0 * y / h + 2
        self.pos = x, y
        self.GetParent().julia.Refresh()

    def expose(self, event):
        dc = wx.PaintDC(self)
        self.realize()
        glClear(GL_COLOR_BUFFER_BIT)
        glUseProgram(self.program)
        with Matrix():
            glLoadIdentity()
            glUniform1i(glGetUniformLocation(self.program, "mandel"), 1)
            GLCanvas.quad()
        self.SwapBuffers()

class Julia(GLCanvas):
    def __init__(self, *args, **kwargs):
        GLCanvas.__init__(self, *args, **kwargs)
        self.Bind(wx.EVT_PAINT, self.expose)

    def expose(self, event):
        dc = wx.PaintDC(self)
        self.realize()
        glClear(GL_COLOR_BUFFER_BIT)
        glUseProgram(self.program)
        with Matrix():
            glLoadIdentity()
            pos = self.GetParent().mandel.pos
            glUniform1i(glGetUniformLocation(self.program, "mandel"), 0)
            glUniform2f(glGetUniformLocation(self.program, "pos"), pos[0], pos[1])
            GLCanvas.quad()
        self.SwapBuffers()

class MyFrame(wx.Frame):
    def __init__(self):
        wx.Frame.__init__(self, None, -1, "Fractal")
        self.Create()
        self.Bind(wx.EVT_CLOSE, self.Close)

    def Close(self, event):
        self.Destroy()

    def Create(self):
        self.mandel = Mandelbrot(self)
        self.julia = Julia(self)

        sizer = wx.BoxSizer(wx.HORIZONTAL)
        sizer.Add(self.mandel, 1, wx.EXPAND)
        sizer.Add(self.julia, 1, wx.EXPAND)
        self.SetSizerAndFit(sizer)

class MyApp(wx.App):
    def __init__(self, *args, **kwargs):
        wx.App.__init__(self, redirect = False, *args, **kwargs)

    def OnInit(self):
        self.mainwin = MyFrame()
        self.mainwin.Show()
        self.SetTopWindow(self.mainwin)
        return True

if __name__ == "__main__":
    app = MyApp()
    app.MainLoop()