python-dolfin 1.3.0+dfsg-2 / usr / lib / python2.7 / dist-packages / dolfin / compilemodules / compilemodule.py

"""This module provides functionality to compile PyDOLFIN compatible
extension modules."""

# Copyright (C) 2009 Johan Hake
#
# This file is part of DOLFIN.
#
# DOLFIN is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# DOLFIN is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with DOLFIN. If not, see <http://www.gnu.org/licenses/>.
#
# Modified by Johannes Ring, 2011
#
# First added:  2009-08-15
# Last changed: 2011-08-31

import sys
import os
import re
import numpy
import instant
import hashlib
import types
from dolfin_utils.cppparser import *

# Import PyDOLFIN
import dolfin
import dolfin.cpp as cpp

from dolfin.compilemodules.jit import mpi_jit_decorator
from dolfin.compilemodules.swigimportinfo import *

__all__ = ["compile_extension_module",
           "expression_to_code_fragments",
           "math_header"]

# Bump the interface version if anything changes that invalidates cached
# modules (not required for change in generated code, swig version or dolfin
# version)
_interface_version = 1

# A list of supported math builtins
_cpp_math_builtins = [
    # <cmath> functions: from http://www.cplusplus.com/reference/cmath/
    "cos", "sin", "tan", "acos", "asin", "atan", "atan2",
    "cosh", "sinh", "tanh", "exp", "frexp", "ldexp", "log", "log10", "modf",
    "pow", "sqrt", "ceil", "fabs", "floor", "fmod",
    "max", "min"]

_math_builtins = [
    # math.h functions: http://en.wikibooks.org/wiki/C_Programming/C_Reference/math.h
    "acos", "asin", "atan", "atan2", "ceil", "cos", "cosh", "exp",
    "fabs", "floor", "fmod", "frexp", "ldexp", "log", "log10", "modf",
    "pow", "sin", "sinh", "sqrt", "tan", "tanh", "acosh", "asinh", "atanh",
    "cbrt", "copysign", "erf", "erfc", "exp2", "expm1", "fdim", "fma", "fmax",
    "fmin", "hypot", "ilogb", "lgamma", "llrint", "lrint", "llround", "lround",
    "log1p", "log2", "logb", "nan", "nearbyint", "nextafter", "nexttoward",
    "remainder", "remquo", "rint", "round", "scalbln", "scalbn", "tgamma", "trunc"]

_math_dolfin = [
    # functions from dolfin::math:
    "sqr", "ipow", "rand", "near", "DOLFIN_EPS", "DOLFIN_PI", "pi"]

_all_math = list(set(_math_builtins).difference(_cpp_math_builtins)) + \
            _math_dolfin + _cpp_math_builtins

math_header = """
// cmath functions
%s

const double pi = DOLFIN_PI;
""" % "\n".join("using std::%s;" % mf for mf in _cpp_math_builtins)


_cpp_keywords = ["auto","const","double","float","int","short","struct","unsigned",
                 "break","continue","else","for","long","signed","switch","void",
                 "case","default","enum","goto","register","sizeof","typedef",
                 "char","do","extern","if","return","static","union","while",
                 "asm","dynamic_cast","namespace","reinterpret_cast","try",
                 "bool","explicit","new","static_cast","typeid","volatile",
                 "catch","operator","template","typename",
                 "class","friend","private","this","using",
                 "const_cast","inline","public","throw","virtual",
                 "delete","mutable","protected","wchar_t",
                 "or","and","xor","not"]

_additional_declarations = r"""
%%init%%{
import_array();
%%}

// Include global SWIG interface files:
// Typemaps, shared_ptr declarations, exceptions, version
%%include <boost_shared_ptr.i>

%%{
#define SWIG_SHARED_PTR_QNAMESPACE boost
%%}

// Global typemaps and forward declarations
%%include "dolfin/swig/typemaps/includes.i"
%%include "dolfin/swig/forwarddeclarations.i"

// Global exceptions
%%include <exception.i>
%%include "dolfin/swig/exceptions.i"

// Do not expand default arguments in C++ by generating two an extra 
// function in the SWIG layer. This reduces code bloat.
%%feature("compactdefaultargs");

// STL SWIG string class
%%include <std_string.i>

// Manually import ufc:
%%shared_ptr(ufc::cell_integral)
%%shared_ptr(ufc::dofmap)
%%shared_ptr(ufc::finite_element)
%%shared_ptr(ufc::function)
%%shared_ptr(ufc::form)
%%shared_ptr(ufc::exterior_facet_integral)
%%shared_ptr(ufc::interior_facet_integral)
%%import(module="ufc") "ufc.h"

// Local shared_ptr declarations
%(shared_ptr_declarations)s

%(additional_declarations)s

// Import statements
%(dolfin_import_statement)s

%%feature("autodoc", "1");

%%inline %%{
int get_swigversion() { return  SWIGVERSION; }
%%}

%%pythoncode %%{
tmp = hex(get_swigversion())
swigversion = "%%d.%%d.%%d"%%(tuple(map(int, [tmp[-5], tmp[-3], tmp[-2:]])))
del tmp, get_swigversion
%%}

"""

def expression_to_code_fragments(expr, arguments, generic_function_members=None):
    "A help function which extract a dict with code snippets from an expression"

    generic_function_members = generic_function_members or []
    expr = list(expr)

    # Autodetect variables from function strings
    variables = set()

    for i, c in enumerate(expr):
        # Find groups of connected alphanumeric letters
        symbols = re.findall(r"([a-zA-Z_]*:{0,2}[\w]+)", c)
        assert isinstance(symbols, list)
        variables.update(symbols)
        for sym in symbols:
            if sym in _cpp_keywords:
                cpp.dolfin_error("compilemodule.py",
                                 "parse expression string",
                                 "Detected C++ keyword (%s) in C++ expression" % sym)
        # FIXME: Remove outcommented code?
        # NOTE: A hack to get around an ambiguous overloading of
        #       dolfin::pow(double,int)

        #if "pow" in symbols:
        #    c = c.replace("pow","std::pow")
        expr[i] = c

    # Remove any variables defined in the arguments list
    variables.difference_update(arguments)

    # Remove the builtin math functions from the variables
    variables.difference_update(_all_math)

    # Remove the numerals from the variables
    numerals = [v for v in variables if v[0] in "0123456789"]
    variables.difference_update(numerals)

    # Remove any exponential representation
    exponentials = [v for v in variables if bool(re.search("e[0-9,\-,\.,\+]+", v))]
    variables.difference_update(exponentials+["e"])

    # Remove namespace-specified variables
    namespacevars = [v for v in variables if "::" in v]
    variables.difference_update(namespacevars)

    # Generate code for member variables
    members_code = ["  double %s;" % name for name in variables \
                    if name not in generic_function_members]
    members_code.extend("  boost::shared_ptr<dolfin::GenericFunction> shared_%s;" % name \
                        for name in generic_function_members)
    members_code = "\n".join(members_code)

    # Generate constructor code for initialization of member variables
    constructor_code = "\n".join("    %s = 0;" % name for name in variables\
                                 if name not in generic_function_members)

    # Connect the code fragments using the function template code
    fragments = {}
    fragments["members"]        = members_code
    fragments["constructor"]    = constructor_code

    # Return the code fragments
    return fragments, variables

def expect_list_of(argtype, arg, argname):
    if arg is None:
        return []
    
    if isinstance(arg, (types.NoneType, list, tuple)):
        if all(isinstance(s, argtype) for s in arg):
            return arg
    
    cpp.dolfin_error("compilemodule.py",
                     "ensure correct argument for compile_extension_module",
                     "Provide a 'tuple' or 'list' with '%s', for the "\
                     "'%s' argument" % (argtype.__name__, argname))

def expect_arg(argtype, arg, argname):
    # Check the type of the argument
    if isinstance(arg, argtype):
        return 
    
    cpp.dolfin_error("compilemodule.py",
                     "ensure correct argument for compile_extension_module",
                     "Provide a '%s', for the '%s' argument" % \
                     (argtype.__name__, argname))
    
def check_swig_version(compiled_module):
    
    # Check swig version of compiled module
    if compiled_module and compiled_module.swigversion != cpp.__swigversion__:
        cpp.dolfin_error("compilemodule.py",
                         "compiling extension module",
                         "Incompatible swig versions detected. DOLFIN swig "\
                         "version is not the same as extension module swig "\
                         "version: '%s' != '%s' " % \
                         (cpp.__swigversion__, compiled_module.swigversion))

@mpi_jit_decorator
def compile_extension_module(code, module_name="",
                             additional_declarations="",
                             additional_system_headers=None,
                             **instant_kwargs):
    """
    Just In Time compile DOLFIN C++ code into a Python module.

    *Arguments*
        code
            C++ code which implements any function or C++ class. Any function
            or class available in the C++ DOLFIN namespace can be used and/or
            subclassed. All typemaps from the original Python interface are
            available, making it possible to interface with for example NumPy
            for Array<double/int> arguments. Source which is not wrapped in
            a dolfin namespace will be automatically wrapped.
  
        module_name
            Force a name of the module. If not set a name based on the hex
            representation of the code will be used.

        additional_declarations
            Additional SWIG declarations can be passed using this argument.

        additional_system_headers :
            System headers needed to compile the generated can be included
            using this argument. The headers are passed using a list of 'str'
        
    *Returns*

        The JIT compiled extension module

    *Examples of usage*

        The following toy example shows how one can use compiled extension
        modules to access low level PETSc routines:
        
        .. code-block:: python

            from numpy import arange
            code = '''
            namespace dolfin {
            
              void PETSc_exp(boost::shared_ptr<dolfin::PETScVector> vec) 
              {
                boost::shared_ptr<Vec> x = vec->vec();
                assert(x);
                VecExp(*x);
              }
            }
            '''
            ext_module = compile_extension_module(code, 
                         additional_system_headers=["petscvec.h"])
            vec = PETScVector(10)
            vec[:] = arange(10)
            print vec[-1]
            ext_module.PETSc_exp(vec)
            print vec[-1]

    """
    # Check the provided arguments
    expect_arg(str, code, "first")
    expect_arg(str, module_name, "module_name")
    expect_arg(str, additional_declarations, "additional_declarations")
    additional_system_headers = \
                expect_list_of(str, additional_system_headers, "additional_system_headers")

    # Check that the code does not use 'using namespace dolfin'
    if re.search("using\s+namespace\s+dolfin",code):
        cpp.dolfin_error("compilemodule.py",
                         "ensure correct argument to compile_extension_module",
                         "Do not use 'using namespace dolfin'. "\
                         "Include the code in namespace dolfin {...} instead")

    # Check if the code does not use namespace dolfin {...}
    if not re.search("namespace\s+dolfin\s*\{[\s\S]+\}", code):

        # Wrap and indet code in namespace dolfin
        codelines = ["namespace dolfin","{"]
        codelines += ["  " + line for line in code.split("\n")]
        codelines += ["}"]
        code = "\n".join(codelines)

    # Create unique module name for this application run
    if module_name is "":
        module_name = "dolfin_compile_code_%s" % \
                      hashlib.md5(repr(code) + dolfin.__version__ + \
                                  str(_interface_version)+\
                                  additional_declarations +\
                                  str(additional_system_headers)).hexdigest()

    # Extract dolfin dependencies and class names
    used_types, declared_types = parse_and_extract_type_info(code)

    # Add any bases of the declared types to used_types
    for declared_type, bases in declared_types.items():
        used_types.update(bases)

    # Filter out dolfin types and add derived and bases for each type
    used_dolfin_types = []
    for dolfin_type in dolfin_type_def:
        for used_type in used_types:
            if dolfin_type in used_type:
                
                # Add bases and derived types
                used_dolfin_types.extend(\
                    dolfin_type_def[dolfin_type]["bases"])
                
                # Add dolfin type
                used_dolfin_types.append(dolfin_type)

                break
    
    # Generate dependency info
    dependencies = {}
    for dolfin_type in used_dolfin_types:
        if dolfin_type_def[dolfin_type]["submodule"] not in dependencies:
            dependencies[dolfin_type_def[dolfin_type]["submodule"]] = []
        
        dependencies[dolfin_type_def[dolfin_type]["submodule"]].append(\
            dolfin_type_def[dolfin_type]["header"])

    # Need special treatment for template definitions in function/pre.i
    if "function" in dependencies:
        for dolfin_type in ["FunctionSpace", "Function"]:
            dependencies["function"].append(dolfin_type_def[dolfin_type]["header"])

    # Add uint type
    if "common" in dependencies:
        dependencies["common"].append("dolfin/common/types.h")
    else:
        dependencies["common"] = ["dolfin/common/types.h"]
    
    # Sort the dependencies
    dependencies = sort_submodule_dependencies(dependencies, submodule_info)

    import_lines, headers_includes, file_dependencies = \
                  build_swig_import_info(dependencies, submodule_info, "dolfin.cpp.")
    
    # Extract header info
    dolfin_system_headers = [header for header in file_dependencies \
                             if not "pre.i" in header]

    # Check the handed import files
    interface_import_files = []

    # Check cache
    compiled_module = instant.import_module(module_name)

    if compiled_module:
        # Check that the swig version of the compiled module is the same as
        # dolfin was compiled with
        check_swig_version(compiled_module)
        return compiled_module

    sys.stdout.flush()
    dolfin.info("Calling DOLFIN just-in-time (JIT) compiler, this may take some time.")
    
    # Configure instant and add additional system headers
    
    # Add dolfin system headers
    instant_kwargs["system_headers"] = ["cmath", "iostream","complex",
                                        "stdexcept","numpy/arrayobject.h",
                                        "boost/shared_ptr.hpp",
                                        "dolfin/common/types.h",
                                        "dolfin/math/basic.h"] + \
                                        instant_kwargs.get("system_headers", [])
    instant_kwargs["system_headers"] += dolfin_system_headers

    # Add user specified system headers
    instant_kwargs["system_headers"] += additional_system_headers

    # Add cmake packages
    instant_kwargs["cmake_packages"]  = ["DOLFIN"] + \
                                        instant_kwargs.get("cmake_packages", [])
    instant_kwargs["signature"] = module_name

    declaration_strs = {"additional_declarations":""}
    declaration_strs["dolfin_import_statement"] = \
                    "\n".join(import_lines)
    
    # Add any provided additional declarations
    if additional_declarations is not None:
        declaration_strs["additional_declarations"] += additional_declarations

    # Add any shared_ptr declarations
    declaration_strs["shared_ptr_declarations"] = \
        extract_shared_ptr_declaration(declared_types, used_dolfin_types, \
                                       shared_ptr_classes)

    # Compile extension module with instant
    compiled_module = instant.build_module(\
        code              = code,
        additional_declarations = _additional_declarations % declaration_strs,
        **instant_kwargs)

    sys.stdout.flush()

    # Check that the swig version of the compiled module is the same as
    # dolfin was compiled with
    check_swig_version(compiled_module)
    
    return compiled_module

def extract_shared_ptr_declaration(declared_types, used_dolfin_types, shared_ptr_classes):
    " Extract any declaration for shared_ptr"
    # Check if there are any classes that is derived from any of the
    # shared_ptr classes in PyDOLFIN and declare if any
    shared_ptr_format = "%%shared_ptr(dolfin::%s)"

    used_shared_ptr_types = []

    # Collect used types which should be shared_ptr declared
    for dolfin_type in used_dolfin_types:
        if dolfin_type in shared_ptr_classes:
            used_shared_ptr_types.append(dolfin_type)

    used_shared_ptr_types.extend(derived for derived, bases in declared_types.items() 
                                 if any(base in shared_ptr_classes for base in bases))
            
    shared_ptr_declarations = "\n".join([shared_ptr_format % derived \
                                         for derived in used_shared_ptr_types])

    return shared_ptr_declarations