python-dolfin 1.0.0-7 / usr / share / pyshared / 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

# Import PyDOLFIN
import dolfin

from dolfin.compilemodules.jit import mpi_jit_decorator
from dolfin.compilemodules.sharedptrclasses import shared_ptr_classes


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

# A list of supported math builtins
_math_builtins = [
    # cmath functions:
    "cos", "sin", "tan", "acos", "asin", "atan", "atan2",
    "cosh", "sinh", "tanh",
    "exp", "frexp", "ldexp", "log", "log10", "modf",
    "pow", "sqrt", "ceil", "fabs", "floor", "fmod"]

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

math_header = """
// cmath functions

%s

const double pi = DOLFIN_PI;
""" % "\n".join("using std::%s;"%mf for mf in _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"]

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

    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_][\w]*)", c)
        assert isinstance(symbols, list)
        variables.update(symbols)
        for sym in symbols:
            if sym in _cpp_keywords:
                raise TypeError, "The C++ keyword '%s' was detected in C++ expression." % sym
        # 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(_math_builtins+_math_dolfin)

    # 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"])

    # Generate code for member variables
    members_code = "\n".join("  double %s;" % name for name in variables)

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

    # 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 get_dolfin_include_dir():
    # Get system configuration
    (dolfin_include_dirs,
     dummy,
     dummy,
     dummy) = instant.header_and_libs_from_pkgconfig("dolfin")

    # Let swig see the installed dolfin swig files
    dolfin_include_found = False
    for inc_dir in dolfin_include_dirs:
        # Check if dolfin.h is installed in dolfin_include_dir
        if os.path.exists(os.path.join(inc_dir, "dolfin.h")):
            return os.path.abspath(inc_dir)
    raise OSError, "Could not find 'dolfin.h', make sure it is installed in a"\
          " directory that 'pkg-config --cflags dolfin' returns."

def configure_instant():
    """Check system requirements

    Returns a dict with kwargs that can be passed to instant.build_module.
    """

    instant_kwargs = {}
    # Get system configuration
    (instant_kwargs['include_dirs'],
     instant_kwargs['cppargs'],
     instant_kwargs['libraries'],
     instant_kwargs['library_dirs'],
     instant_kwargs['lddargs']) = instant.header_and_libs_from_pkgconfig("dolfin", returnLinkFlags=True)

    # Let swig see the installed dolfin swig files
    swig_include_dirs = []
    dolfin_include_found = False
    ufc_include_found = False
    for inc_dir in instant_kwargs['include_dirs']:
        # Check if dolfin is installed in inc_dir
        if os.path.exists(os.path.join(inc_dir, "dolfin", "swig", "dolfin.i")):
            dolfin_include_found = True
            if inc_dir not in swig_include_dirs:
                swig_include_dirs.append(inc_dir)

        # Check if ufc is installed in inc_dir
        if os.path.exists(os.path.join(inc_dir, "swig", "ufc.i")):
            ufc_include_found = True
            if inc_dir not in swig_include_dirs:
                swig_include_dirs.append(inc_dir)

    if not dolfin_include_found:
        raise OSError, """
Didn't find dolfin.i in include paths returned from pkg-config.
Please make sure that your DOLFIN installation corresponds with the
one returned by pkg-config"""

    if not ufc_include_found:
        raise OSError, """Didn't find ufc.i in include paths returned from pkg-config.
Please make sure that your UFC installation corresponds with the one
returned by pkg-config"""

    # Check if UFC is importable and what version of swig was used to
    # create the UFC extension module
    try: import ufc
    except: raise OSError, "Please install the python extenstion module of UFC "\
            "on your system.\n"

    # Check that the form compiler will use the same swig version
    # that UFC was compiled with
    if not instant.check_swig_version(ufc.__swigversion__, same=True):
        raise OSError, """The python extension module of UFC was not compiled with the present version of swig.
Install swig version %s or recompile UFC with present swig
"""%ufc.__swigversion__

    # Check that the form compiler will use the same swig version
    # that PyDOLFIN was compiled with
    if not instant.check_swig_version(dolfin.__swigversion__,same=True):
        raise OSError, """PyDOLFIN was not compiled with the present version of swig.
Install swig version %s or recompile PyDOLFIN with present swig
"""%dolfin.__swigversion__

    # Check for boost installation
    # Set a default directory for the boost installation
    if sys.platform == "darwin":
        # use MacPorts as default
        default = os.path.join(os.path.sep,"opt","local")
    else:
        default = os.path.join(os.path.sep,"usr")

    # If BOOST_DIR is not set use default directory
    boost_dir = os.getenv("BOOST_DIR", default)
    boost_is_found = False
    for inc_dir in ["", "include"]:
        if os.path.isfile(os.path.join(boost_dir, inc_dir, "boost", "version.hpp")):
            boost_include_dir = os.path.join(boost_dir, inc_dir)
            boost_is_found = True
            break

    if not boost_is_found:
        raise OSError, """The Boost library was not found.
If Boost is installed in a nonstandard location,
set the environment variable BOOST_DIR.
"""

    instant_kwargs['swig_include_dirs'] = swig_include_dirs
    instant_kwargs['include_dirs'].append(numpy.get_include())
    instant_kwargs['include_dirs'].append(boost_include_dir)
    instant_kwargs['system_headers'] = ["cmath", "iostream","complex",
                                        "stdexcept","numpy/arrayobject.h",
                                        "dolfin.h","boost/shared_ptr.hpp",
                                        "dolfin/math/basic.h"]
    instant_kwargs['swigargs'] =['-c++','-I.']

    return instant_kwargs


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

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

// DOLFIN typemaps
%%include "dolfin/swig/typemaps.i"
%%include "dolfin/swig/std_pair_typemaps.i"
%%include "dolfin/swig/numpy_typemaps.i"
%%include "dolfin/swig/array_typemaps.i"
%%include "dolfin/swig/std_vector_typemaps.i"
%%include "dolfin/swig/std_set_typemaps.i"

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

%(additional_declarations)s

// Shared_ptr declarations
%%include "dolfin/swig/shared_ptr_classes.i"

%(shared_ptr_declarations)s

// Import statements
%(dolfin_import_statement)s

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

def expect_list_of(argtype,arg,argname):
    assert(isinstance(argtype,type))
    error_str = "Provide a 'tuple' or 'list' with '%s', for the '%s' argument"%(argtype.__name__,argname)
    if arg is None:
        return []
    if not isinstance(arg,(types.NoneType,list,tuple)):
        raise TypeError, error_str
    if not all(isinstance(s, str) for s in arg):
        raise TypeError, error_str
    return arg

def expect_arg(argtype,arg,argname):
    if isinstance(argtype,(tuple,list)):
        all(isinstance(s, type) for s in argtype)
    else:
        assert(isinstance(argtype,type))
        type_str = "'%s'"%argtype.__name__

    # Check the type of the argument
    if not isinstance(arg,argtype):
        if isinstance(argtype,(tuple,list)):
            if len(argtype) > 2:
                argtype_str = ", ".join("'%s'"%s.__name__ for s in argtype[:-1]) + " or " + "'%s'"%argtype[-1]
            else:
                argtype_str = " or ".join("'%s'"%s.__name__ for s in argtype)
        else:
            argtype_str = "'%s'"%s.__name__
        raise TypeError, "Expected a '%s' for the '%s' argument"%argname


@mpi_jit_decorator
def compile_extension_module(code, module_name="",
                             recompile=False,
                             dolfin_module_import=None,
                             additional_declarations="",
                             additional_system_headers=None):
    """
    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 code need to be wrapped in
            a dolfin namespace.
  
        module_name
            Force a name of the module. If not set a name based on the hex
            representation of the code will be used.

        recompile
            Force recompilation of a module, also if it is found in cache.

        dolfin_module_import
            If only a few types from the original DOLFIN is used in the code
            the amount of code generated can be limited by only importing
            certain DOLFIN kernel modules. If the only DOLFIN type used in the
            code is from the Mesh class, this argument can be set to:
    
            .. code-block:: python

                dolfin_module_import = ["mesh"],
  
            and if in addition to the Mesh class a GenericVector is used the
            "la" module can also be included:
            
            .. code-block:: python

                dolfin_module_import = ["mesh", "la"],

        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, 
                         dolfin_module_import=["la"], 
                         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")
    expect_arg(bool,recompile,"recompile")
    additional_system_headers = \
                expect_list_of(str, additional_system_headers,"additional_system_headers")
    dolfin_module_import = \
                expect_list_of(str, dolfin_module_import, "dolfin_module_import")

    # Check that the code does not use 'using namespace dolfin'
    if re.search("using[ \n]+namespace[ \n]+dolfin",code):
        raise ValueError, "Do not use 'using namespace dolfin'. Include the code in namespace dolfin {...} instead."

    # Check that the code is included in dolfin namespace
    if not re.search("namespace[ \n]+dolfin[ \n]{",code):
        raise ValueError, "Include the C++ code in the 'dolfin namspace', otherwise SWIG cannot figure out the right PyDOLFIN types."


    # Check and set swig binary
    if not instant.check_and_set_swig_binary(dolfin.parameters["swig_binary"], \
                                             dolfin.parameters["swig_path"]):
        raise OSError, "Could not find swig installation. Pass an existing "\
              "swig binary or install SWIG version 2.0 or higher.\n"

    # Create unique module name for this application run
    if module_name is "":
        module_name = "dolfin_compile_code_%s" % hashlib.md5(repr(code) +\
                                            instant.get_swig_version() +\
                                            dolfin.__version__).hexdigest()
    # Check the handed import files
    if dolfin_module_import:
        interface_files = []
        dolfin_include_dir = get_dolfin_include_dir()
        for m in dolfin_module_import:
            interface_file = os.path.join("dolfin", "swig", "import", m + ".i")
            if not os.path.isfile(os.path.join(dolfin_include_dir, interface_file)):
                raise OSError, "'%s' is not a module of DOLFIN"%m
            interface_files.append(interface_file)

        dolfin_module_import = interface_files

    # Check cache
    if not recompile:
        compiled_module = instant.import_module(module_name)

        if 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
    instant_kwargs = configure_instant()
    instant_kwargs["system_headers"] += additional_system_headers

    declaration_strs = {"additional_declarations":""}
    if dolfin_module_import:
        using_reduced_import = True
        declaration_strs["dolfin_import_statement"] = \
                        "\n".join(['%%include %s'%f
                                   for f in dolfin_module_import])
    else:
        using_reduced_import = False
        declaration_strs["dolfin_import_statement"] = r"%import dolfin/swig/dolfin.i"

    # If using numpy typemaps add  %include "numpy.i"
    # FIXME: Need to fix some standard way to apply the types maps...
    #if use_numpy_typemaps:
    #    raise NotImplementedError
    #    declaration_strs["additional_declarations"] += '\n%include "numpy.i"\n'

    # 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(code, shared_ptr_classes)

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

    sys.stdout.flush()
    return compiled_module

def extract_shared_ptr_declaration(code, 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)"
    derived_and_base = [(derived, base) for derived, base in \
                        re.findall(r"class[ ]+([\w]+)[ ]*: public[ ]+([\w:]+)", code)\
                        if base.split("::")[-1] in shared_ptr_classes]

    derived_shared_ptr_classes = [derived for derived, base in \
                                  derived_and_base if base in shared_ptr_classes]

    shared_ptr_declarations = "\n".join([shared_ptr_format%derived \
                                         for derived in derived_shared_ptr_classes])

    return shared_ptr_declarations