vispy 0.15.0__cp313-cp313-macosx_11_0_arm64.whl

vispy/visuals/shaders/function.py

@@ -0,0 +1,788 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Vispy Development Team. All Rights Reserved.
+# Distributed under the (new) BSD License. See LICENSE.txt for more info.
+"""
+Classses representing GLSL objects (functions, variables, etc) that may be
+composed together to create complete shaders.
+See the docstring of Function for details.
+
+Details
+-------
+
+A complete GLSL program is composed of ShaderObjects, each of which may be used
+inline as an expression, and some of which include a definition that must be
+included on the final code. ShaderObjects keep track of a hierarchy of
+dependencies so that all necessary code is included at compile time, and
+changes made to any object may be propagated to the root of the hierarchy to
+trigger a recompile.
+"""
+
+from collections import OrderedDict
+import logging
+import re
+
+import numpy as np
+
+from ...util.eq import eq
+from ...util import logger
+from . import parsing
+from .shader_object import ShaderObject
+from .variable import Variable, Varying
+from .expression import TextExpression, FunctionCall
+
+
+class Function(ShaderObject):
+    """Representation of a GLSL function
+
+    Objects of this class can be used for re-using and composing GLSL
+    snippets. Each Function consists of a GLSL snippet in the form of
+    a function. The code may have template variables that start with
+    the dollar sign. These stubs can be replaced with expressions using
+    the index operation. Expressions can be:
+
+    * plain text that is inserted verbatim in the code
+    * a Function object or a call to a funcion
+    * a Variable (or Varying) object
+    * float, int, tuple are automatically turned into a uniform Variable
+    * a VertexBuffer is automatically turned into an attribute Variable
+
+    All functions have implicit "$pre" and "$post" placeholders that may be
+    used to insert code at the beginning and end of the function.
+
+    Examples
+    --------
+    This example shows the basic usage of the Function class::
+
+        vert_code_template = Function('''
+            void main() {
+            gl_Position = $pos;
+            gl_Position.x += $xoffset;
+            gl_Position.y += $yoffset;
+        }''')
+
+        scale_transform = Function('''
+        vec4 transform_scale(vec4 pos){
+            return pos * $scale;
+        }''')
+
+        # If you get the function from a snippet collection, always
+        # create new Function objects to ensure they are 'fresh'.
+        vert_code = Function(vert_code_template)
+        trans1 = Function(scale_transform)
+        trans2 = Function(scale_transform)  # trans2 != trans1
+
+        # Three ways to assign to template variables:
+        #
+        # 1) Assign verbatim code
+        vert_code['xoffset'] = '(3.0 / 3.1415)'
+
+        # 2) Assign a value (this creates a new uniform or attribute)
+        vert_code['yoffset'] = 5.0
+
+        # 3) Assign a function call expression
+        pos_var = Variable('attribute vec4 a_position')
+        vert_code['pos'] = trans1(trans2(pos_var))
+
+        # Transforms also need their variables set
+        trans1['scale'] = 0.5
+        trans2['scale'] = (1.0, 0.5, 1.0, 1.0)
+
+        # You can actually change any code you want, but use this with care!
+        vert_code.replace('gl_Position.y', 'gl_Position.z')
+
+        # Finally, you can set special variables explicitly. This generates
+        # a new statement at the end of the vert_code function.
+        vert_code['gl_PointSize'] = '10.'
+
+
+    If we use ``vert_code.compile()`` we get::
+
+        attribute vec4 a_position;
+        uniform float u_yoffset;
+        uniform float u_scale_1;
+        uniform vec4 u_scale_2;
+        uniform float u_pointsize;
+
+        vec4 transform_scale_1(vec4 pos){
+            return pos * u_scale_1;
+        }
+
+        vec4 transform_scale_2(vec4 pos){
+            return pos * u_scale_2;
+        }
+
+        void main() {
+            gl_Position = transform_scale_1(transform_scale_2(a_position));
+            gl_Position.x += (3.0 / 3.1415);
+            gl_Position.z += u_yoffset;
+
+            gl_PointSize = u_pointsize;
+        }
+
+    Note how the two scale function resulted in two different functions
+    and two uniforms for the scale factors.
+
+    Notes
+    -----
+
+    Function calls:
+
+    As can be seen above, the arguments with which a function is to be
+    called must be specified by calling the Function object. The
+    arguments can be any of the expressions mentioned earlier. If the
+    signature is already specified in the template code, that function
+    itself must be given.
+
+    ::
+
+        code = Function('''
+            void main() {
+                vec4 position = $pos;
+                gl_Position = $scale(position)
+            }
+        ''')
+
+        # Example of a function call with all possible three expressions
+        vert_code['pos'] = func1('3.0', 'uniform float u_param', func2())
+
+        # For scale, the sigfnature is already specified
+        code['scale'] = scale_func  # Must not specify args
+
+    Data for uniform and attribute variables:
+
+    To each variable a value can be associated. In fact, in most cases
+    the Function class is smart enough to be able to create a Variable
+    object if only the data is given.
+
+    ::
+
+        code['offset'] = Variable('uniform float offset')  # No data
+        code['offset'] = Variable('uniform float offset', 3.0)  # With data
+        code['offset'] = 3.0  # -> Uniform Variable
+        position['position'] = VertexBuffer()  # -> attribute Variable
+
+        # Updating variables
+        code['offset'].value = 4.0
+        position['position'].value.set_data(...)
+
+    """
+
+    def __init__(self, code, dependencies=None):
+        super(Function, self).__init__()
+
+        # Add depencencies is given. This is to allow people to
+        # manually define deps for a function that they use.
+        if dependencies is not None:
+            for dep in dependencies:
+                self._add_dep(dep)
+
+        self.code = code
+
+        # Expressions replace template variables (also our dependencies)
+        self._expressions = OrderedDict()
+
+        # Verbatim string replacements
+        self._replacements = OrderedDict()
+
+        # Stuff to do at the end
+        self._assignments = OrderedDict()
+
+    def __setitem__(self, key, val):
+        """Setting of replacements through a dict-like syntax.
+
+        Each replacement can be:
+        * verbatim code: ``fun1['foo'] = '3.14159'``
+        * a FunctionCall: ``fun1['foo'] = fun2()``
+        * a Variable: ``fun1['foo'] = Variable(...)`` (can be auto-generated)
+        """
+        # Check the key. Must be Varying, 'gl_X' or a known template variable
+        if isinstance(key, Variable):
+            if key.vtype == 'varying':
+                if self.name != 'main':
+                    raise Exception("Varying assignment only alowed in 'main' "
+                                    "function.")
+                storage = self._assignments
+            else:
+                raise TypeError("Variable assignment only allowed for "
+                                "varyings, not %s (in %s)"
+                                % (key.vtype, self.name))
+        elif isinstance(key, str):
+            if any(map(key.startswith,
+                       ('gl_PointSize', 'gl_Position', 'gl_FragColor'))):
+                storage = self._assignments
+            elif key in self.template_vars or key in ('pre', 'post'):
+                storage = self._expressions
+            else:
+                raise KeyError('Invalid template variable %r' % key)
+        else:
+            raise TypeError('In `function[key]` key must be a string or '
+                            'varying.')
+
+        # If values already match, bail out now
+        if eq(storage.get(key), val):
+            return
+
+        # If we are only changing the value (and not the dtype) of a uniform,
+        # we can set that value and return immediately to avoid triggering a
+        # recompile.
+        if val is not None and not isinstance(val, Variable):
+            # We are setting a value. If there is already a variable set here,
+            # try just updating its value.
+            variable = storage.get(key, None)
+            if isinstance(variable, Variable):
+                if np.any(variable.value != val):
+                    variable.value = val
+                    self.changed(value_changed=True)
+                return
+
+            # Could not set variable.value directly; instead we will need
+            # to create a new ShaderObject
+            val = ShaderObject.create(val, ref=key)
+            if variable is val:
+                # This can happen if ShaderObject.create returns the same
+                # object (such as when setting a Transform).
+                return
+
+        # Remove old references, if any
+        oldval = storage.pop(key, None)
+        if oldval is not None:
+            for obj in (key, oldval):
+                if isinstance(obj, ShaderObject):
+                    self._remove_dep(obj)
+
+        # Add new references
+        if val is not None:
+            if isinstance(key, Varying):
+                # tell this varying to inherit properties from
+                # its source attribute / expression.
+                key.link(val)
+
+            # Store value and dependencies
+            storage[key] = val
+            for obj in (key, val):
+                if isinstance(obj, ShaderObject):
+                    self._add_dep(obj)
+
+        # In case of verbatim text, we might have added new template vars
+        if isinstance(val, TextExpression):
+            for var in parsing.find_template_variables(val.expression()):
+                if var not in self.template_vars:
+                    self.template_vars.add(var.lstrip('$'))
+
+        self.changed(code_changed=True, value_changed=True)
+        if logger.level <= logging.DEBUG:
+            import traceback
+            last = traceback.format_list(traceback.extract_stack()[-2:-1])
+            logger.debug("Assignment would trigger shader recompile:\n"
+                         "Original: %r\nReplacement: %r\nSource: %s",
+                         oldval, val, ''.join(last))
+
+    def __getitem__(self, key):
+        """Return a reference to a program variable from this function.
+
+        This allows variables between functions to be linked together::
+
+            func1['var_name'] = func2['other_var_name']
+
+        In the example above, the two local variables would be assigned to the
+        same program variable whenever func1 and func2 are attached to the same
+        program.
+        """
+        try:
+            return self._expressions[key]
+        except KeyError:
+            pass
+
+        try:
+            return self._assignments[key]
+        except KeyError:
+            pass
+
+        if key not in self.template_vars:
+            raise KeyError('Invalid template variable %r' % key)
+        else:
+            raise KeyError('No value known for key %r' % key)
+
+    def __call__(self, *args):
+        """Set the signature for this function and return an FunctionCall
+        object. Each argument can be verbatim code or a FunctionCall object.
+        """
+        return FunctionCall(self, args)
+
+    def __contains__(self, key):
+        return key in self.template_vars
+
+    # Public API methods
+
+    @property
+    def signature(self):
+        if self._signature is None:
+            try:
+                self._signature = parsing.parse_function_signature(self._code)
+            except Exception as err:
+                raise ValueError('Invalid code: ' + str(err))
+        return self._signature
+
+    @property
+    def name(self):
+        """The function name. The name may be mangled in the final code
+        to avoid name clashes.
+        """
+        return self.signature[0]
+
+    @property
+    def args(self):
+        """
+        List of input arguments in the function signature::
+            [(arg_name, arg_type), ...]
+        """
+        return self.signature[1]
+
+    @property
+    def rtype(self):
+        """The return type of this function."""
+        return self.signature[2]
+
+    @property
+    def code(self):
+        """The template code used to generate the definition for this function."""
+        return self._code
+
+    @code.setter
+    def code(self, code):
+        # Get and strip code
+        if isinstance(code, Function):
+            code = code._code
+        elif not isinstance(code, str):
+            raise ValueError('Function needs a string or Function; got %s.' %
+                             type(code))
+        self._code = self._clean_code(code)
+
+        # (name, args, rval)
+        self._signature = None
+
+        # $placeholders parsed from the code
+        self._template_vars = None
+
+        # Create static Variable instances for any global variables declared
+        # in the code
+        self._static_vars = None
+
+    @property
+    def template_vars(self):
+        if self._template_vars is None:
+            self._template_vars = self._parse_template_vars()
+        return self._template_vars
+
+    def static_names(self):
+        if self._static_vars is None:
+            self._static_vars = parsing.find_program_variables(self._code)
+        return list(self._static_vars.keys()) + [arg[0] for arg in self.args]
+
+    def replace(self, str1, str2):
+        """Set verbatim code replacement
+
+        It is strongly recommended to use function['$foo'] = 'bar' where
+        possible because template variables are less likely to changed
+        than the code itself in future versions of vispy.
+
+        Parameters
+        ----------
+        str1 : str
+            String to replace
+        str2 : str
+            String to replace str1 with
+        """
+        if str2 != self._replacements.get(str1, None):
+            self._replacements[str1] = str2
+            self.changed(code_changed=True)
+            # self._last_changed = time.time()
+
+    # Private methods
+
+    def _parse_template_vars(self):
+        """Find all template variables in self._code, excluding the function name."""
+        template_vars = set()
+        for var in parsing.find_template_variables(self._code):
+            var = var.lstrip('$')
+            if var == self.name:
+                continue
+            if var in ('pre', 'post'):
+                raise ValueError('GLSL uses reserved template variable $%s' %
+                                 var)
+            template_vars.add(var)
+        return template_vars
+
+    def _get_replaced_code(self, names, version, shader):
+        """Return code, with new name, expressions, and replacements applied."""
+        code = self._code
+
+        # Modify name
+        fname = names[self]
+        code = code.replace(" " + self.name + "(", " " + fname + "(")
+
+        # Apply string replacements first -- these may contain $placeholders
+        for key, val in self._replacements.items():
+            code = code.replace(key, val)
+
+        # Apply assignments to the end of the function
+
+        # Collect post lines
+        post_lines = []
+        for key, val in self._assignments.items():
+            if isinstance(key, Variable):
+                key = names[key]
+            if isinstance(val, ShaderObject):
+                val = val.expression(names)
+            line = '    %s = %s;' % (key, val)
+            post_lines.append(line)
+
+        # Add a default $post placeholder if needed
+        if 'post' in self._expressions:
+            post_lines.append('    $post')
+
+        # Apply placeholders for hooks
+        post_text = '\n'.join(post_lines)
+        if post_text:
+            post_text = '\n' + post_text + '\n'
+        code = code.rpartition('}')
+        code = code[0] + post_text + code[1] + code[2]
+
+        # Add a default $pre placeholder if needed
+        if 'pre' in self._expressions:
+            m = re.search(fname + r'\s*\([^{]*\)\s*{', code)
+            if m is None:
+                raise RuntimeError("Cound not find beginning of function '%s'"
+                                   % fname)
+            ind = m.span()[1]
+            code = code[:ind] + "\n    $pre\n" + code[ind:]
+
+        # Apply template variables
+        for key, val in self._expressions.items():
+            val = val.expression(names)
+            search = r'\$' + key + r'($|[^a-zA-Z0-9_])'
+            code = re.sub(search, val+r'\1', code)
+
+        # Done
+        if '$' in code:
+            v = parsing.find_template_variables(code)
+            logger.warning('Unsubstituted placeholders in code: %s\n'
+                           '  replacements made: %s',
+                           v, list(self._expressions.keys()))
+
+        return code + '\n'
+
+    def definition(self, names, version, shader):
+        return self._get_replaced_code(names, version, shader)
+
+    def expression(self, names):
+        return names[self]
+
+    def _clean_code(self, code):
+        """Return *code* with indentation and leading/trailing blank lines removed."""
+        lines = code.split("\n")
+        min_indent = 100
+        for line in lines:
+            if line.strip() != "":
+                indent = len(line) - len(line.lstrip())
+                min_indent = min(indent, min_indent)
+        if min_indent > 0:
+            lines = [line[min_indent:] for line in lines]
+        code = "\n".join(lines)
+        return code
+
+    def __repr__(self):
+        try:
+            args = ', '.join([' '.join(arg) for arg in self.args])
+        except Exception:
+            return ('<%s (error parsing signature) at 0x%x>' %
+                    (self.__class__.__name__, id(self)))
+        return '<%s "%s %s(%s)" at 0x%x>' % (self.__class__.__name__,
+                                             self.rtype,
+                                             self.name,
+                                             args,
+                                             id(self))
+
+
+class MainFunction(Function):
+    """Subclass of Function that allows multiple functions and variables to
+    be defined in a single code string. The code must contain a main() function
+    definition.
+    """
+
+    def __init__(self, shader_type, *args, **kwargs):
+        self.shader_type = shader_type
+        self._chains = {}
+        Function.__init__(self, *args, **kwargs)
+
+    @property
+    def signature(self):
+        return ('main', [], 'void')
+
+    @property
+    def version_pragma(self):
+        """Return version number and extra qualifiers from pragma if present."""
+        m = re.search(parsing.re_version_pragma, self.code)
+        if m is None:
+            return None
+        return int(m.group(1)), m.group(2)
+
+    def definition(self, obj_names, version, shader):
+        code = Function.definition(self, obj_names, version, shader)
+        # strip out version pragma before returning code; this will be
+        # added to the final compiled code later.
+        code = re.sub(parsing.re_version_pragma, '', code)
+        return code
+
+    def static_names(self):
+        if self._static_vars is not None:
+            return self._static_vars
+
+        # parse static variables
+        names = Function.static_names(self)
+
+        # parse all function names + argument names
+        funcs = parsing.find_functions(self.code)
+        for f in funcs:
+            if f[0] == 'main':
+                continue
+            names.append(f[0])
+            for arg in f[1]:
+                names.append(arg[1])
+
+        self._static_vars = names
+        return names
+
+    def add_chain(self, var):
+        """Create a new ChainFunction and attach to $var."""
+        chain = FunctionChain(var, [])
+        self._chains[var] = chain
+        self[var] = chain
+
+    def add_callback(self, hook, func):
+        self._chains[hook].append(func)
+
+    def remove_callback(self, hook, func):
+        self._chains[hook].remove(func)
+
+
+class FunctionChain(Function):
+    """Subclass that generates GLSL code to call Function list in order
+
+    Functions may be called independently, or composed such that the
+    output of each function provides the input to the next.
+
+    Parameters
+    ----------
+    name : str
+        The name of the generated function
+    funcs : list of Functions
+        The list of Functions that will be called by the generated GLSL code.
+
+    Examples
+    --------
+    This creates a function chain::
+
+        >>> func1 = Function('void my_func_1() {}')
+        >>> func2 = Function('void my_func_2() {}')
+        >>> chain = FunctionChain('my_func_chain', [func1, func2])
+
+    If *chain* is included in a ModularProgram, it will generate the following
+    output::
+
+        void my_func_1() {}
+        void my_func_2() {}
+
+        void my_func_chain() {
+            my_func_1();
+            my_func_2();
+        }
+
+    The return type of the generated function is the same as the return type
+    of the last function in the chain. Likewise, the arguments for the
+    generated function are the same as the first function in the chain.
+
+    If the return type is not 'void', then the return value of each function
+    will be used to supply the first input argument of the next function in
+    the chain. For example::
+
+        vec3 my_func_1(vec3 input) {return input + vec3(1, 0, 0);}
+        void my_func_2(vec3 input) {return input + vec3(0, 1, 0);}
+
+        vec3 my_func_chain(vec3 input) {
+            return my_func_2(my_func_1(input));
+        }
+    """
+
+    def __init__(self, name=None, funcs=()):
+        # bypass Function.__init__ completely.
+        ShaderObject.__init__(self)
+        if not (name is None or isinstance(name, str)):
+            raise TypeError("Name argument must be string or None.")
+        self._funcs = []
+        self._code = None
+        self._name = name or "chain"
+        self._args = []
+        self._rtype = 'void'
+        self.functions = funcs
+
+    @property
+    def functions(self):
+        return self._funcs[:]
+
+    @functions.setter
+    def functions(self, funcs):
+        while self._funcs:
+            self.remove(self._funcs[0], update=False)
+        for f in funcs:
+            self.append(f, update=False)
+        self._update()
+
+    @property
+    def signature(self):
+        return self._name, self._args, self._rtype
+
+    def _update(self):
+        funcs = self._funcs
+        if len(funcs) > 0:
+            self._rtype = funcs[-1].rtype
+            self._args = funcs[0].args[:]
+        else:
+            self._rtype = 'void'
+            self._args = []
+
+        self.changed(code_changed=True)
+
+    @property
+    def code(self):
+        # Code is generated at compile time; hopefully it is not requested
+        # before then..
+        return None
+
+    @code.setter
+    def code(self, c):
+        raise TypeError("Cannot set code property on FunctionChain.")
+
+    @property
+    def template_vars(self):
+        return {}
+
+    def append(self, function, update=True):
+        """Append a new function to the end of this chain."""
+        self._funcs.append(function)
+        self._add_dep(function)
+        if update:
+            self._update()
+
+    def __setitem__(self, index, func):
+        self._remove_dep(self._funcs[index])
+        self._add_dep(func)
+        self._funcs[index] = func
+
+        self._update()
+
+    def __getitem__(self, k):
+        return self.functions[k]
+
+    def insert(self, index, function, update=True):
+        """Insert a new function into the chain at *index*."""
+        self._funcs.insert(index, function)
+        self._add_dep(function)
+        if update:
+            self._update()
+
+    def remove(self, function, update=True):
+        """Remove a function from the chain."""
+        self._funcs.remove(function)
+        self._remove_dep(function)
+        if update:
+            self._update()
+
+    def definition(self, obj_names, version, shader):
+        name = obj_names[self]
+
+        args = ", ".join(["%s %s" % arg for arg in self.args])
+        code = "%s %s(%s) {\n" % (self.rtype, name, args)
+
+        result_index = 0
+        if len(self.args) == 0:
+            last_rtype = 'void'
+            last_result = ''
+        else:
+            last_rtype, last_result = self.args[0][:2]
+
+        for fn in self._funcs:
+            # Use previous return value as an argument to the next function
+            if last_rtype == 'void':
+                args = ''
+            else:
+                args = last_result
+                if len(fn.args) != 1 or last_rtype != fn.args[0][0]:
+                    raise Exception("Cannot chain output '%s' of function to "
+                                    "input of '%s'" %
+                                    (last_rtype, fn.signature))
+            last_rtype = fn.rtype
+
+            # Store the return value of this function
+            if fn.rtype == 'void':
+                set_str = ''
+            else:
+                result_index += 1
+                result = 'result_%d' % result_index
+                set_str = '%s %s = ' % (fn.rtype, result)
+                last_result = result
+
+            code += "    %s%s(%s);\n" % (set_str, obj_names[fn], args)
+
+        # return the last function's output
+        if self.rtype != 'void':
+            code += "    return result_%d;\n" % result_index
+
+        code += "}\n"
+        return code
+
+    def static_names(self):
+        return []
+
+    def __repr__(self):
+        fn = ",\n                ".join(map(repr, self.functions))
+        return "<FunctionChain [%s] at 0x%x>" % (fn, id(self))
+
+
+class StatementList(ShaderObject):
+    """Represents a list of statements."""
+
+    def __init__(self):
+        self.items = {}
+        self.order = []
+        ShaderObject.__init__(self)
+
+    def add(self, item, position=5):
+        """Add an item to the list unless it is already present.
+
+        If the item is an expression, then a semicolon will be appended to it
+        in the final compiled code.
+        """
+        if item in self.items:
+            return
+        self.items[item] = position
+        self._add_dep(item)
+        self.order = None
+        self.changed(code_changed=True)
+
+    def remove(self, item):
+        """Remove an item from the list."""
+        self.items.pop(item)
+        self._remove_dep(item)
+        self.order = None
+        self.changed(code_changed=True)
+
+    def expression(self, obj_names):
+        if self.order is None:
+            self.order = list(self.items.items())
+            self.order.sort(key=lambda x: x[1])
+
+        code = ""
+        for item, pos in self.order:
+            code += item.expression(obj_names) + ';\n'
+        return code