Extend the PsychoPy ElementArrayStim object to allow non-zero pedestal value.

pedestalelementarray.py

#!/usr/bin/env python2
"""Extend ElementArrayStim to allow nonzero pedestal values of gratings.

This adapts a prior solution for adding a non-zero pedestal value to GratingStim:
https://github.com/nwilming/PedestalGrating/blob/master/pedestal_grating.py

Covered under the PsychoPy license, as it is a simple extension of the prior code:

Copyright (C) 2015 Jonathan Peirce
Distributed under the terms of the GNU General Public License (GPL).

"""
import pyglet
pyglet.options['debug_gl'] = False
import ctypes
GL = pyglet.gl

from psychopy.visual.elementarray import ElementArrayStim
from psychopy.visual.basevisual import MinimalStim, TextureMixin
from psychopy.tools.attributetools import attributeSetter
from psychopy import _shadersPyglet as _shaders

# Framgent shader for the gabor stimulus. This is needed to add the pedestal to
# the color values for each location. I'm keeping it in this file to make the
# stimulus fairly self contained and to avoid messing with anything else.
# Almost a one to one copy of the original psychopy shader.
fragSignedColorTexMask = '''
    uniform sampler2D texture, mask;
    uniform float pedestal;
    void main() {
        vec4 textureFrag = texture2D(texture,gl_TexCoord[0].st);
        vec4 maskFrag = texture2D(mask,gl_TexCoord[1].st);
        gl_FragColor.a = gl_Color.a*maskFrag.a*textureFrag.a;
        gl_FragColor.rgb =  ((pedestal+1.0)/2.0) + ((textureFrag.rgb* (gl_Color.rgb*2.0-1.0)+1.0)/2.0) -0.5;
    }
    '''


class PedestalElementArrayStim(ElementArrayStim, MinimalStim, TextureMixin):
    """
    This stimulus class defines a field of elements whose behaviour can be independently
    controlled. Suitable for creating 'global form' stimuli or more detailed random dot
    stimuli.

    This stimulus can draw thousands of elements without dropping a frame, but in order
    to achieve this performance, uses several OpenGL extensions only available on modern
    graphics cards (supporting OpenGL2.0). See the ElementArray demo.
    """
    def __init__(self,
                 win,
                 units = None,
                 fieldPos = (0.0,0.0),
                 fieldSize = (1.0,1.0),
                 fieldShape = 'circle',
                 nElements = 100,
                 sizes = 2.0,
                 xys = None,
                 rgbs = None,
                 colors=(1.0,1.0,1.0),
                 colorSpace='rgb',
                 opacities = 1.0,
                 depths = 0,
                 fieldDepth = 0,
                 oris = 0,
                 sfs=1.0,
                 contrs = 1,
                 phases=0,
                 elementTex='sin',
                 elementMask='gauss',
                 texRes=48,
                 interpolate=True,
                 name=None,
                 autoLog=None,
                 maskParams=None,
                 pedestal=0):


        super(PedestalElementArrayStim, self).__init__(
            win, units=units, fieldPos=fieldPos, fieldSize=fieldSize,
            fieldShape=fieldShape, nElements=nElements, sizes=sizes, xys=xys,
            rgbs=rgbs, colors=colors, colorSpace=colorSpace,
            opacities=opacities, depths=depths, fieldDepth=fieldDepth,
            oris=oris, sfs=sfs, contrs=contrs, phases=phases,
            elementTex=elementTex, elementMask=elementMask, texRes=texRes,
            interpolate=interpolate, name=name, autoLog=autoLog,
            maskParams=maskParams)
        self.pedestal = pedestal
        self._progSignedTexMask = _shaders.compileProgram(
            _shaders.vertSimple, fragSignedColorTexMask)

    @attributeSetter
    def pedestal(self, value):
        """
        Luminance pedestal of the grating (e.g. the background around which it
        varies)
        Should a scalar in the range -1:1.
        """
        # Recode phase to numpy array
        self.__dict__['pedestal'] = value
        self._needUpdate = True

    def draw(self, win=None):
        """
        Draw the stimulus in its relevant window. You must call
        this method after every MyWin.update() if you want the
        stimulus to appear on that frame and then update the screen
        again.
        """
        if win is None:
            win=self.win
        self._selectWindow(win)

        if self._needVertexUpdate:
            self._updateVertices()
        if self._needColorUpdate:
            self.updateElementColors()
        if self._needTexCoordUpdate:
            self.updateTextureCoords()

        #scale the drawing frame and get to centre of field
        GL.glPushMatrix()#push before drawing, pop after
        GL.glPushClientAttrib(GL.GL_CLIENT_ALL_ATTRIB_BITS)#push the data for client attributes

        #GL.glLoadIdentity()
        self.win.setScale('pix')

        GL.glColorPointer(4, GL.GL_DOUBLE, 0, self._RGBAs.ctypes.data_as(ctypes.POINTER(ctypes.c_double)))
        GL.glVertexPointer(3, GL.GL_DOUBLE, 0, self.verticesPix.ctypes.data_as(ctypes.POINTER(ctypes.c_double)))

        #setup the shaderprogram
        GL.glUseProgram(self._progSignedTexMask)
        GL.glUniform1i(GL.glGetUniformLocation(self._progSignedTexMask, "texture"), 0)
        GL.glUniform1i(GL.glGetUniformLocation(self._progSignedTexMask, "mask"), 1)
        GL.glUniform1f(GL.glGetUniformLocation(self._progSignedTexMask, "pedestal"), self.pedestal)

        #bind textures
        GL.glActiveTexture (GL.GL_TEXTURE1)
        GL.glBindTexture (GL.GL_TEXTURE_2D, self._maskID)
        GL.glEnable(GL.GL_TEXTURE_2D)
        GL.glActiveTexture (GL.GL_TEXTURE0)
        GL.glBindTexture (GL.GL_TEXTURE_2D, self._texID)
        GL.glEnable(GL.GL_TEXTURE_2D)

        #setup client texture coordinates first
        GL.glClientActiveTexture (GL.GL_TEXTURE0)
        GL.glTexCoordPointer (2, GL.GL_DOUBLE, 0, self._texCoords.ctypes)
        GL.glEnableClientState(GL.GL_TEXTURE_COORD_ARRAY)
        GL.glClientActiveTexture (GL.GL_TEXTURE1)
        GL.glTexCoordPointer (2, GL.GL_DOUBLE, 0, self._maskCoords.ctypes)
        GL.glEnableClientState(GL.GL_TEXTURE_COORD_ARRAY)

        GL.glEnableClientState(GL.GL_COLOR_ARRAY)
        GL.glEnableClientState(GL.GL_VERTEX_ARRAY)
        GL.glDrawArrays(GL.GL_QUADS, 0, self.verticesPix.shape[0]*4)

        #unbind the textures
        GL.glActiveTexture(GL.GL_TEXTURE1)
        GL.glBindTexture(GL.GL_TEXTURE_2D, 0)
        GL.glDisable(GL.GL_TEXTURE_2D)
        #main texture
        GL.glActiveTexture(GL.GL_TEXTURE0)
        GL.glBindTexture(GL.GL_TEXTURE_2D, 0)
        GL.glDisable(GL.GL_TEXTURE_2D)
        #disable states
        GL.glDisableClientState(GL.GL_COLOR_ARRAY)
        GL.glDisableClientState(GL.GL_VERTEX_ARRAY)
        GL.glDisableClientState(GL.GL_TEXTURE_COORD_ARRAY)

        GL.glUseProgram(0)
        GL.glPopClientAttrib()
        GL.glPopMatrix()