compile_functions.py

利用C

"This module provides functionality for compilation of strings as dolfin Functions."

__author__ = "Martin Sandve Alnes ([e-mail])"
__date__ = "2008-06-04 -- 2008-06-04"
__copyright__ = "Copyright (C) 2008-2008 Martin Sandve Alnes"
__license__  = "GNU LGPL Version 2.1"

import re
import numpy
import instant

__all__ = ["compile_functions",]

# FIXME: Extend this list, needed to autodetect variable names that are not builtins
_builtins = [
             # local symbols:
             "v", "x", "_dim", "pi",
             # cmath funcions:
             "cos", "sin", "tan", "acos", "asin", "atan", "atan2",
             "cosh", "sinh", "tanh",
             "exp", "frexp", "ldexp", "log", "log10", "modf",
             "pow", "sqrt", "ceil", "fabs", "floor", "fmod",
            ]

# Add utility code here
_header = """
// cmath functions
using std::cos;
using std::sin;
using std::tan;
using std::acos;
using std::asin;
using std::atan;
using std::atan2;
using std::cosh;
using std::sinh;
using std::tanh;
using std::exp;
using std::frexp;
using std::ldexp;
using std::log;
using std::log10;
using std::modf;
using std::pow;
using std::sqrt;
using std::ceil;
using std::fabs;
using std::floor;
using std::fmod;

const double pi = acos(-1.0);
"""


_function_template = """
class %(classname)s: public dolfin::Function
{
public:
  unsigned _dim;
%(members)s

  %(classname)s(dolfin::Mesh & mesh):
    dolfin::Function(mesh)
  {
    _dim = mesh.geometry().dim();
%(constructor)s
  }

  void eval(real* v, const real* x) const
  {
%(eval)s
  }
%(rank)s%(dim)s
};
"""

_ranktemplate = """\
  dolfin::uint rank() const
  {
    return %d;
  }
"""

_dimtemplate = """\
  dolfin::uint dim(dolfin::uint i) const
  {
    switch(i) {
%s
    }
    throw std::runtime_error("Invalid dimension i in dim(i).");
  }
"""

def expression_to_function(e, classname):
    "Generates code for a dolfin::Function subclass for a single expression."
    # Get shape from e and make e a flat tuple of strings
    if isinstance(e, str):
        e = (e,)
        shape = ()
    elif isinstance(e, tuple):
        if isinstance(e[0], str):
            shape = (len(e),)
        elif isinstance(e[0], tuple):
            shape = (len(e),len(e[0]))
            assert isinstance(e[0][0], str)
            e = sum(e, ())
    else:
        raise RuntimeError("Invalid expression %s" % e)

    # Autodetect variables from function strings
    variables = set()
    for c in e:
        # Find groups of connected alphanumeric letters
        variables.update(re.findall(r"([\w]+)", c))
    variables.difference_update(_builtins)
    numerals = [v for v in variables if v[0] in "0123456789"]
    variables.difference_update(numerals)
    
    # Generate code for member variables
    memberscode = "\n".join("  double %s;" % name for name in variables)
    
    # Generate constructor code for initialization of member variables
    constructorcode = "\n".join("    %s = 0.0;" % name for name in variables)
    
    # Generate code for rank and dim
    if len(shape) > 0:
        rankcode = _ranktemplate % len(shape)
        cases = "\n".join("      case %d: return %d;" % (d,n) for (d,n) in enumerate(shape))
        dimcode = _dimtemplate % cases
    else:
        dimcode  = ""
        rankcode = ""
    
    # Generate code for the actual function evaluation
    evalcode = "\n".join("    v[%d] = %s;" % (i, c) for (i,c) in enumerate(e))
    
    # Connect the code fragments using the function template code
    fragments = {}
    fragments["classname"] = classname
    fragments["members"]   = memberscode
    fragments["rank"]      = rankcode
    fragments["dim"]       = dimcode
    fragments["eval"]      = evalcode
    fragments["constructor"] = constructorcode
    code = _function_template % fragments
    return code


_function_count = 0
def generate_functions(expressions):
    "Generates code for dolfin::Function subclasses for a list of expressions."
    global _function_count
    assert isinstance(expressions, list)
    code = ""
    classnames = []
    for e in expressions:
        classname = "function_%d" % _function_count
        code += expression_to_function(e, classname)
        classnames.append(classname)
        _function_count += 1
    return code, classnames

# NB! This code is highly dependent on the dolfin swig setup!
_additional_declarations = r"""
%rename(cpp_Function) Function;
%include exception.i
%include cpointer.i
%pointer_class(int, intp);
%pointer_class(double, doublep);
%include std_vector.i
%template(STLVectorFunctionPtr) std::vector<dolfin::Function *>;
%import dolfin.i
"""
    
_additional_definitions  = """
#include <dolfin.h>
using namespace dolfin;
#include <numpy/arrayobject.h>
"""

_module_count = 0
def compile_function_code(code, mesh, classnames=None, module_name=None):
    # Create unique module name for this application run
    global _module_count, _header, _additional_definitions, _additional_declarations
    if module_name is None:
        module_name = "functions_%d" % _module_count
        _module_count += 1

    # Autodetect classnames:
    _classnames = re.findall(r"class[ ]+([\w]+).*", code)
    # Just a little assertion for safety:
    if classnames is None:
        classnames = _classnames
    else:
        assert all(a == b for (a,b) in zip(classnames, _classnames))
    
    # Get system configuration   
    (includes, flags, libraries, libdirs) = instant.header_and_libs_from_pkgconfig("dolfin")
    dolfin_include_dir = includes[0] # FIXME: is this safe?
    numpy_dir = numpy.get_include()
    includes.append(numpy_dir)
    
    sysheaders = ["cmath", "iostream", "stdexcept",
                  "dolfin.h", "dolfin/mesh/Mesh.h", "dolfin/function/Function.h"]
    
    # FIXME: use dolfin flags?
    cppargs = flags
    cppargs = " ".join(flags)
    cppargs = ""
    
    # Let swig see the installed dolfin swig files
    swigopts = "-c++ -I%s -I%s/dolfin/swig" % (dolfin_include_dir, dolfin_include_dir)
    
    # Compile extension module with instant
    compiled_module = instant.create_extension(\
             code           = _header + code,
             module         = module_name,
             swigopts       = swigopts,
             system_headers = sysheaders,
             include_dirs   = includes,
             cppargs        = cppargs,
             libraries      = libraries,
             library_dirs   = libdirs,
             additional_definitions  = _additional_definitions,
             additional_declarations = _additional_declarations
             )
    compiled_module = __import__(module_name)
    
    # Construct instances of the compiled functor classes
    functions = [eval("compiled_module.%s(mesh)" % name) for name in classnames]
    return functions


def compile_functions(expressions, mesh):
    """Compiles C++ string expressions into dolfin::Function instances.
    
    The variable 'expressions' can either be a str or a list.
    
    If 'expressions' is a str, it is interpreted as a C++ string
    with complete implementations of subclasses of dolfin::Function.
    The compiled functions returned will be in the same order 
    as they are defined in this code.

    If it is a list, each item of the list is interpreted as
    a function expression, and the compiled functions returned
    will be in the same order as they occur in this list.
    
    Each expression item in the list can either be a str,
    in which case it is interpreted as a scalar expression
    and a scalar Function is generated, or it can be a tuple.
    
    A tuple of str objects is interpreted as a vector expression,
    and a rank 1 Function is generated.
    
    A tuple of tuples of str objects is interpreted as a matrix
    expression, and a rank 2 Function is generated.

    If an expression string contains a name, it is assumed to
    be a scalar variable name, and is added as a public member
    of the generated function. The exceptions are set in the
    variable dolfin.compile_functions._builtins."""
    #, which contains:
    #    %s
    #""" % "\n".join("        " + b for b in _builtins)
    if isinstance(expressions, list):
        code, classnames = generate_functions(expressions)
        functions = compile_function_code(code, mesh, classnames)
    else:
        functions = compile_function_code(expressions, mesh)
    return functions


if __name__ == "__main__":
    code, cn = generate_functions([
                              "exp(alpha)",
                              ("sin(x[0])", "cos(x[1])", "0.0"),
                              (("sin(x[0])", "cos(x[1])"), ("0.0", "1.0"))
                             ])
    print code
    print cn