Thread: [matplotlib-devel] opengl backend

SourceForge Headquarters 1320 Columbia Street Suite 310 San Diego, CA 92101 +1 (858) 422-6466

Hi,

There was a recent discussion about opengl and matplotlib in the
context of matplotlib rendering speeds.

At the scipy sprints we put together a proof of concept renderer
for quad meshes using the opengl frame buffer object, which we
then render as a matplotlib image.  Once the data is massaged
to the correct form we were seeing rendering speeds of about
10 million quads per second. See below.

Using this technique we can get the advantage of opengl without
having to write a full backend.  Please let me know if you have
time to contribute to writing a more complete backend.

   - Paul

-- meshgl.py --
# This program is public domain
import numpy
from pyglet.gl import *
import pyglet
from ctypes import POINTER,c_long,c_ubyte,c_int

def buildmesh(x,y):
    """
    x,y are corners of the mesh quadrants
    """
    # Quadlist is a list of pairs of vertices, left edge, right edge, right
    # edge, one for each row.  The resulting map will look like:
    #    [[p00 p10 p01 p11 p02 p12 ...] [p10 p20 p11 p21 ...] ...]
    # Note that the x,y shape will be one more than the number of data values.
    nr,nc = x.shape
    quadstrips = numpy.empty((nr-1,nc*4),'f')
    quadstrips[:,0::4] = x[:-1,:]
    quadstrips[:,2::4] = x[1:,:]
    quadstrips[:,1::4] = y[:-1,:]
    quadstrips[:,3::4] = y[1:,:]
    return quadstrips

def colorindex(z,clim,cmap):
    # Calc colour indices
    ncolors = cmap.shape[0]
    step = float(clim[1]-clim[0])/(ncolors-1)
    idx = numpy.asarray((z-clim[0])/step, dtype='i')
    idx = numpy.clip(idx, 0, ncolors-1)
    return idx

def buildcolor(z,clim=[0,1],cmap=None):
    """
    z are centers of the mesh quadrants.

    Note there size of z are less than size of x,y by 1 in each dim.
    """
    # Generate colours
    idx = colorindex(z, clim, cmap)

    nr,nc = z.shape
    # Quadlist is a list of pairs of vertices, left edge, right edge, right
    # edge.  The resulting map will look like:
    #    [[p00 p10 p01 p11 p02 p12 ...] [p10 p20 p11 p21 ...] ...]
    # The z values correspond to the bottom right corner of each quad.  This
    # means we do not need the first two quads of each row and every odd
    # quad after that.  The resulting color list looks like the following:
    #    [[0 0 0 z00 0 z01 0 z02 ...] [0 0 0 z10 0 z11 0 z12 ...] ...]
    # Each i,j has a four byte colour value.
    colors = numpy.zeros((nr,2*nc+2,4),'B')
    colors[:,3::2,:] = cmap[idx,:]
    return colors

def draw_quads(x,y,z,clim=[0,1],cmap=None):
    strips = buildmesh(x,y)
    colors = buildcolor(z,clim,cmap)
    nr,nc = strips.shape
    nc = nc/2 # pair of vertices per point
    starts = numpy.arange(0,nr*nc,nc,'i')
    counts = numpy.ones(nr,'i')*nc

    glEnableClientState(GL_VERTEX_ARRAY)
    glEnableClientState(GL_COLOR_ARRAY)
    glVertexPointer(2,GL_FLOAT,0,strips.ctypes.data)
    glColorPointer(4,GL_UNSIGNED_BYTE,0,colors.ctypes.data)
    glMultiDrawArrays(GL_QUAD_STRIP, #starts, counts, nr)
                      starts.ctypes.data_as(POINTER(c_int)),
                      counts.ctypes.data_as(POINTER(c_int)),
                      nr)
    glDisableClientState(GL_VERTEX_ARRAY)
    glDisableClientState(GL_COLOR_ARRAY)

def set_scene(w,h):
    """Set up the openGL transform matrices"""
    glViewport(0, 0, w, h)
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    glMatrixMode(GL_MODELVIEW)
    glLoadIdentity()

    glShadeModel(GL_FLAT)
    glEnable(GL_BLEND)
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)

    # Clear scene
    glClearColor(1, 1, 1, 0)
    glClear(GL_COLOR_BUFFER_BIT)

def renderbuffer(w,h,draw):
    """Render buffer to off screen framebuffer object"""

    # Make sure we support frame buffers
    if not pyglet.gl.gl_info.have_extension('GL_EXT_framebuffer_object'):
        raise RuntimeError("This OpenGL does not support framebuffer objects")

    # Create a frame buffer context
    framebuffers = (GLuint*1)(0)
    glGenFramebuffersEXT(1, framebuffers)
    glBindFramebufferEXT(GL_FRAMEBUFFER_EXT, framebuffers[0])

    # Create a render buffer
    renderbuffers = (GLuint*1)(0)
    glGenRenderbuffersEXT(1, renderbuffers)
    glBindRenderbufferEXT(GL_RENDERBUFFER_EXT, renderbuffers[0])
    glRenderbufferStorageEXT(GL_RENDERBUFFER_EXT, GL_RGBA8, w, h)
    glFramebufferRenderbufferEXT(GL_FRAMEBUFFER_EXT, GL_COLOR_ATTACHMENT0_EXT, 
        GL_RENDERBUFFER_EXT, renderbuffers[0])

    status = glCheckFramebufferStatusEXT(GL_FRAMEBUFFER_EXT)
    if status != GL_FRAMEBUFFER_COMPLETE_EXT:
        raise RuntimeError("Could not create GL framebuffer")

    # Render scene
    set_scene(w,h)
    draw()
    glFlush()

    # Retrieve buffer
    buffer = numpy.empty((h,w,4),'B')
    glReadPixels(0, 0, w, h, GL_RGBA, GL_UNSIGNED_BYTE, buffer.ctypes.data)
    # View the buffer in the window
    #glBindFramebufferEXT(GL_FRAMEBUFFER_EXT, 0)
    # Release buffer
    glDeleteRenderbuffersEXT(1, renderbuffers)

    return buffer

# === DEMO CODE BELOW ===

def testmesh(shape=(3,2)):
    """sample scene"""
    # Quad mesh vertices (nr x nc)
    nr,nc = shape
    x = numpy.arange(nr)[None,:] + numpy.arange(nc)[:,None]
    y = numpy.arange(nr)[None,:] + numpy.zeros((nc,1))
    # Quad mesh values (nr-1 x nc-1)
    z = x[:-1,:-1]

    # Grayscale RGBA colormap
    R = numpy.linspace(0,255,64)
    alpha = numpy.ones(64,'B')*255
    cmap = numpy.array([R,R,R,alpha],'B').T

    if True:
        # Green box in center using default -1,1 coords
        glColor3f(0., 1., 0.)
        glRectf(-0.5, 0.5, 0.5, -0.5)

    # Set data coordinates
    glPushMatrix()
    glOrtho(0,x.max()*2,0,y.max()*2,-1,1)

    if True:
        # Blue box unit size at 1,1 in data coords
        glColor4f(0., 0., 1., 1.)
        glRectf(1, 2, 2, 1)

    draw_quads(x,y,z,clim=[z.min(),z.max()],cmap=cmap)
    glPopMatrix()

class Window(pyglet.window.Window):
    """Pyglet driver"""
    def __init__(self, w, h):
        super(Window, self).__init__(width=w, height=h, resizable=True)
        glClearColor(1.,1.,1.,1.)
        self._dirty = True

    def on_resize(self, w, h):
        print "resize"
        set_scene(w,h);
        self._dirty = True

    def draw(self):
        if self._dirty:
            print "draw"
            #window.clear()
            testmesh()
            self._dirty = False
            self.flip()

    def run(self):
        counter = 0
        while not self.has_exit and counter < 1000:
            self.dispatch_events()
            self.draw()
            counter += 1

def plot():
    """Matplotlib driver"""
    import pylab
    print "start rendering"
    b = renderbuffer(400,500,lambda: testmesh((1000,1000)))
    print "done rendering"
    pylab.imshow(b,origin='lower')
    pylab.show()

if __name__ == "__main__":
    #Window(400,500).run()
    plot()