python-dolfin 1.0.0-1 / usr / share / pyshared / dolfin / functions / expression.py

"""This module handles the Expression class in Python.

The Expression class needs special handling and is not mapped directly
by SWIG from the C++ interface. Instead, a new Expression class is
created which inherits both from the DOLFIN C++ Expression class and
the ufl Coefficient class.

The resulting Expression class may thus act both as a variable in a
UFL form expression and as a DOLFIN C++ Expression.

This module make heavy use of creation of Expression classes and
instantiation of these dynamically at runtime.

The whole logic behind this somewhat magic behaviour is handle by:

  1) function __new__ in the Expression class
  2) meta class ExpressionMetaClass
  3) function compile_expressions from the compiledmodule/expression
     module
  4) functions create_compiled_expression_class and
     create_python_derived_expression_class

The __new__ method in the Expression class take cares of the logic
when the class Expression is used to create an instance of Expression,
see use cases 1-4 in the docstring of Expression.

The meta class ExpressionMetaClass take care of the logic when a user
subclasses Expression to create a user-defined Expression, see use
cases 3 in the docstring of Expression.

The function compile_expression is a JIT compiler. It compiles and
returns different kinds of cpp.Expression classes, depending on the
arguments. These classes are sent to the
create_compiled_expression_class.

"""

# Copyright (C) 2008-2011 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 Anders Logg, 2008-2009.
#
# First added:  2008-11-03
# Last changed: 2009-12-16

__all__ = ["Expression"]

# FIXME: Make all error messages uniform according to the following template:
#
# if not isinstance(foo, Foo):
#     raise TypeError, "Illegal argument for creation of Bar, not a Foo: " + str(foo)

# Python imports
import types

# Import UFL and SWIG-generated extension module (DOLFIN C++)
import ufl
import dolfin.cpp as cpp
import numpy

from dolfin import warning, error

# Local imports
from dolfin.compilemodules.expressions import compile_expressions

def create_compiled_expression_class(cpp_base):
    # Check the cpp_base
    assert(isinstance(cpp_base, (types.ClassType, type)))

    def __init__(self, cppcode, element=None, cell=None, \
                 degree=None, **kwargs):
        """Initialize the Expression """
        # Initialize the cpp base class first and extract value_shape
        cpp_base.__init__(self)
        value_shape = tuple(self.value_dimension(i) \
                                  for i in range(self.value_rank()))

        # Store the value_shape
        self._value_shape = value_shape

        # Select an appropriate element if not specified
        if element is None:
            element = _auto_select_element_from_shape(value_shape, degree, cell)
        else:
            # Check that we have an element
            if not isinstance(element, ufl.FiniteElementBase):
                raise TypeError, "The 'element' argument must be a UFL"\
                      " finite element."

            # Check same value shape of compiled expression and passed element
            if not element.value_shape() == value_shape:
                raise ValueError, "The value shape of the passed 'element',"\
                      " is not equal to the value shape of the compiled "\
                      "expression."

        # Initialize UFL base class
        self._ufl_element = element
        ufl.Coefficient.__init__(self, self._ufl_element)

        # Set default variables
        for member, value in kwargs.items():
            setattr(self, member, value)

    # Create and return the class
    return type("CompiledExpression", (Expression, ufl.Coefficient, cpp_base),\
                {"__init__":__init__})

def create_python_derived_expression_class(name, user_bases, user_dict):
    """Return Expression class

    This function is used to create all the dynamically created Expression
    classes. It takes a name, and a compiled cpp.Expression and returns
    a class that inherits the compiled class together with dolfin.Expression
    and ufl.Coefficient.

    *Arguments*
        name
            The name of the class
        user_bases
            User defined bases
        user_dict
            Dict with user specified function or attributes
    """

    # Check args
    assert(isinstance(name, str))
    assert(isinstance(user_bases, list))
    assert(isinstance(user_dict, dict))

    # Define the bases
    assert(all([isinstance(t, (types.ClassType, type)) for t in user_bases]))
    bases = tuple([Expression, ufl.Coefficient, cpp.Expression] + user_bases)

    user_init = user_dict.pop("__init__", None)

    # Check for deprecated dim and rank methods
    if "dim" in user_dict or "rank" in user_dict:
        raise DeprecationWarning, "'rank' and 'dim' are depcrecated, overload"\
              " 'value_shape' instead"

    def __init__(self, *args, **kwargs):
        """Initialize the Expression"""
        # Get element, cell and degree
        element = kwargs.get("element", None)
        degree = kwargs.get("degree", None)
        cell = kwargs.get("cell", None)

        # Check if user has passed too many arguments if no
        # user_init is provided
        if user_init is None:
            from operator import add
            # First count how many valid kwargs is passed
            num_valid_kwargs = reduce(add, [kwarg is not None \
                                      for kwarg in [element, degree, cell]])
            if len(kwargs) != num_valid_kwargs:
                raise TypeError, "expected only 'kwargs' from one of "\
                      "'element', 'degree' or 'cell'"
            if len(args) != 0:
                raise TypeError, "expected no 'args'"

        # Select an appropriate element if not specified
        if element is None:
            element = _auto_select_element_from_shape(self.value_shape(),
                                                      degree, cell)
        elif isinstance(element, ufl.FiniteElementBase):
            pass
        else:
            raise TypeError, "The 'element' argument must be a UFL finite"\
                  " element."

        # Initialize UFL base class
        self._ufl_element = element
        ufl.Coefficient.__init__(self, element)

        # Initialize cpp_base class
        cpp.Expression.__init__(self, list(element.value_shape()))

        # Calling the user defined_init
        if user_init is not None:
            user_init(self, *args, **kwargs)

    # NOTE: Do not prevent the user to overload attributes "reserved" by PyDOLFIN
    # NOTE: Why not?

    ## Collect reserved attributes from both cpp.Function and ufl.Coefficient
    #reserved_attr = dir(ufl.Coefficient)
    #reserved_attr.extend(dir(cpp.Function))
    #
    ## Remove attributes that will be set by python
    #for attr in ["__module__"]:
    #    while attr in reserved_attr:
    #        reserved_attr.remove(attr)
    #
    ## Check the dict_ for reserved attributes
    #for attr in reserved_attr:
    #    if attr in dict_:
    #        raise TypeError, "The Function attribute '%s' is reserved by PyDOLFIN."%attr

    # Add __init__ to the user_dict
    user_dict["__init__"]  = __init__

    # Create the class and return it
    return type(name, bases, user_dict)

class ExpressionMetaClass(type):
    """Meta Class for Expression"""
    def __new__(mcs, name, bases, dict_):
        """Returns a new Expression class"""

        assert(isinstance(name, str)), "Expecting a 'str'"
        assert(isinstance(bases, tuple)), "Expecting a 'tuple'"
        assert(isinstance(dict_, dict)), "Expecting a 'dict'"

        # First check if we are creating the Expression class
        if name == "Expression":
            # Assert that the class is _not_ a subclass of Expression,
            # i.e., a user have tried to:
            #
            #    class Expression(Expression):
            #        ...
            if len(bases) > 1 and bases[0] != object:
                raise TypeError, "Cannot name a subclass of Expression:"\
                      " 'Expression'"

            # Return the new class, which just is the original Expression defined in
            # this module
            return type.__new__(mcs, name, bases, dict_)

        # If creating a fullfledged derived expression class, i.e, inheriting
        # dolfin.Expression, ufl.Coefficient and cpp.Expression (or a subclass)
        # then just return the new class.
        if len(bases) >= 3 and bases[0] == Expression and \
               bases[1] == ufl.Coefficient and issubclass(bases[2], \
                                                          cpp.Expression):
            # Return the instantiated class
            return type.__new__(mcs, name, bases, dict_)

        # Handle any user provided base classes
        user_bases = list(bases)

        # remove Expression, to be added later
        user_bases.remove(Expression)

        # Check the user has provided either an eval or eval_cell method
        if not ('eval' in dict_ or 'eval_cell' in dict_):
            raise TypeError, "expected an overload 'eval' or 'eval_cell' method"

        # Get name of eval function
        eval_name = 'eval' if 'eval' in dict_ else 'eval_cell'

        user_eval = dict_[eval_name]

        # Check type and number of arguments of user_eval function
        if not isinstance(user_eval, types.FunctionType):
            raise TypeError, "'%s' attribute must be a 'function'"%eval_name
        if eval_name == "eval" and not user_eval.func_code.co_argcount == 3:
            raise TypeError, "The overloaded '%s' function must use "\
                  "three arguments"%eval_name
        if eval_name == "eval_cell" and \
               not user_eval.func_code.co_argcount == 4:
            raise TypeError, "The overloaded '%s' function must "\
                  "use three arguments"%eval_name

        return create_python_derived_expression_class(name, user_bases, dict_)

#--- The user interface ---

# Places here so it can be reused in Function
class Expression(object):
    """
    This class represents a user-defined expression.

    Expressions can be used as coefficients in variational forms or
    interpolated into finite element spaces.

    *Arguments*
        cppcode
           C++ argument, see below
        element
            Optional element argument
        degree
            Optional element degree when element is not given.


    *1. Simple user-defined JIT-compiled expressions*
        One may alternatively specify a C++ code for evaluation of the
        Expression as follows:

        .. code-block:: python

            f0 = Expression('sin(x[0]) + cos(x[1])')
            f1 = Expression(('cos(x[0])', 'sin(x[1])'), element = V.ufl_element())

        Here, f0 is is scalar and f1 is vector-valued.

        Tensor expressions of rank 2 (matrices) may also be created:

        .. code-block:: python

            f2 = Expression((('exp(x[0])','sin(x[1])'),
                            ('sin(x[0])','tan(x[1])')))

        In general, a single string expression will be interpreted as
        a scalar, a tuple of strings as a tensor of rank 1 (a vector)
        and a tuple of tuples of strings as a tensor of rank 2 (a
        matrix).

        The expressions may depend on x[0], x[1], and x[2] which carry
        information about the coordinates where the expression is
        evaluated. All math functions defined in <cmath> are available
        to the user.

        Expression parameters can be included as follows:

        .. code-block:: python

            f = Expression('A*sin(x[0]) + B*cos(x[1])', A=2.0, B=4.0)

        The parameters can only be scalars, and are all initialized to
        the passed default value.

    *2. Complex user-defined JIT-compiled Expressions*
        One may also define an Expression using more complicated logic
        with the 'cppcode' argument. This argument should be a string
        of C++ code that implements a class that inherits from
        dolfin::Expression.

        The following code illustrates how to define an Expression
        that depends on material properties of the cells in a Mesh. A
        MeshFunction is used to mark cells with different properties.

        Note the use of the 'data' parameter.

        .. code-block:: python

            code = '''
            class MyFunc : public Expression
            {
            public:

              boost::shared_ptr<MeshFunction<unsigned int> > cell_data;

              MyFunc() : Expression()
              {
              }

            void eval(Array<double>& values, const Array<double>& x,
                      const ufc::cell& c) const
              {
                assert(cell_data);
                const Cell cell(cell_data->mesh(), c.index);
                switch ((*cell_data)[cell.index()])
                {
                case 0:
                  values[0] = exp(-x[0]);
                  break;
                case 1:
                  values[0] = exp(-x[2]);
                  break;
                default:
                  values[0] = 0.0;
                }
              }
            };'''

            cell_data = CellFunction('uint', V.mesh())
            f = Expression(code)
            f.cell_data = cell_data

    *3. User-defined expressions by subclassing*
        The user can subclass Expression and overload the 'eval'
        function. The value_shape of such an Expression will default
        to 0. If a user wants a vector or tensor Expression, the
        value_shape method needs to be overloaded.

        .. code-block:: python

            class MyExpression0(Expression):
                def eval(self, value, x):
                    dx = x[0] - 0.5
                    dy = x[1] - 0.5
                    value[0] = 500.0*exp(-(dx*dx + dy*dy)/0.02)
                    value[1] = 250.0*exp(-(dx*dx + dy*dy)/0.01)
                def value_shape(self):
                    return (2,)
            f0 = MyExpression0()

        If a user wants to use the Expression in a UFL form and have
        more controll in which finite element should be used to
        interpolate the expression in, the user can pass this
        information using the element kwarg:

        .. code-block:: python

            V = FunctionSpace(mesh, "BDM", 1)
            f1 = MyExpression0(element=V.ufl_element())

        The user can also subclass Expression by overloading the
        eval_cell function. By this the user gets access to more
        powerful data structures, such as cell, facet and normal
        information, during assembly.

        .. code-block:: python

            class MyExpression1(Expression):
                def eval_cell(self, value, x, ufc_cell):
                    if ufc_cell.index > 10:
                        value[0] = 1.0
                    else:
                        value[0] = -1.0

            f2 = MyExpression1()

        The ufc_cell object can be queried for the following data:

        .. code-block:: python

            ufc_cell.cell_shape
            ufc_cell.index
            ufc_cell.topological_dimension
            ufc_cell.geometric_dimension
            ufc_cell.local_facet # only available on boundaries, otherwise -1
            ufc_cell.mesh_identifier

        The user can customize initialization of derived Expressions.
        However, because of magic behind the scenes, a user needs to pass
        optional arguments to __init__ using ``**kwargs``, and _not_
        calling the base class __init__:

        .. code-block:: python

            class MyExpression1(Expression):
                def __init__(self, mesh, domain):
                    self._mesh = mesh
                    self._domain = domain
                def eval(self, values, x):
                    ...

            f3 = MyExpression1(mesh=mesh, domain=domain)

        Note that subclassing may be significantly slower than using
        JIT-compiled expressions. This is because a callback from C++
        to Python will be involved each time a Expression needs to be
        evaluated during assembly.

    """

    # Set the meta class
    __metaclass__ = ExpressionMetaClass

    def __new__(cls, cppcode=None, element=None, cell=None, degree=None, \
                **kwargs):

        # If the __new__ function is called because we are instantiating
        # a python sub class of Expression, then just return a new instant
        # of the passed class
        if cls.__name__ != "Expression":
            return object.__new__(cls)

        # Check arguments
        _check_cppcode(cppcode)

        # Compile module and get the cpp.Expression class
        cpp_base, members = compile_expressions([cppcode])
        cpp_base, members = cpp_base[0], members[0]

        # Check passed default arguments
        _check_default_kwargs(members, kwargs)

        # Store compile arguments for later use
        cpp_base.cppcode = cppcode

        # Create and instantiate the new class
        return object.__new__(create_compiled_expression_class(cpp_base))

    # This method is only included so a user can check what arguments
    # one should use in IPython using tab completion
    def __init__(self, cppcode=None, element=None, cell=None, degree=None, \
                 **kwargs): pass

    # Reuse the docstring from __new__
    __init__.__doc__ = __new__.__doc__

    def ufl_element(self):
        "Return the ufl FiniteElement."
        return self._ufl_element

    def __str__(self):
        "x.__str__() <==> print(x)"
        return ufl.Coefficient.__str__(self)

    def str(self, verbose=False):
        "x.str() <==> info(x)"
        return "<Expression %s>" % str(self._value_shape)

    def __repr__(self):
        "x.__repr__() <==> repr(x)"
        return ufl.Coefficient.__repr__(self)

    # Default value for the value shape
    _value_shape = ()

    def value_shape(self):
        """Returns the value shape of the expression"""
        return self._value_shape

    def ufl_evaluate(self, x, component, derivatives):
        """Function used by ufl to evaluate the Expression"""
        import numpy
        import ufl
        assert derivatives == () # TODO: Handle derivatives

        if component:
            shape = self.shape()
            assert len(shape) == len(component)
            value_size = ufl.common.product(shape)
            index = ufl.common.component_to_index(component, shape)
            values = numpy.zeros(value_size)
            self(*x, values=values)
            return values[index]
        else:
            # Scalar evaluation
            return self(*x)

    def __call__(self, *args, **kwargs):
        """
        Evaluates the Expression

        *Example*
            1) Using an iterable as x:

            .. code-block:: python

                fs = Expression("sin(x[0])*cos(x[1])*sin(x[3])")
                x0 = (1.,0.5,0.5)
                x1 = [1.,0.5,0.5]
                x2 = numpy.array([1.,0.5,0.5])
                v0 = fs(x0)
                v1 = fs(x1)
                v2 = fs(x2)

            2) Using multiple scalar args for x, interpreted as a
            point coordinate

            .. code-block:: python

                v0 = f(1.,0.5,0.5)

            3) Using a Point

            .. code-block:: python

                p0 = Point(1.,0.5,0.5)
                v0 = f(p0)

            3) Passing return array

            .. code-block:: python

                fv = Expression(("sin(x[0])*cos(x[1])*sin(x[3])",
                              "2.0","0.0"))
                x0 = numpy.array([1.,0.5,0.5])
                v0 = numpy.zeros(3)
                fv(x0, values = v0)

                .. note::

                    A longer values array may be passed. In this way
                    one can fast fill up an array with different
                    evaluations.

            .. code-block:: python

                values = numpy.zeros(9)
                for i in xrange(0,10,3):
                    fv(x[i:i+3], values = values[i:i+3])

        """

        if len(args)==0:
            raise TypeError, "expected at least 1 argument"

        # Test for ufl restriction
        if len(args) == 1 and args[0] in ('+','-'):
            return ufl.Coefficient.__call__(self, *args)

        # Test for ufl mapping
        if len(args) == 2 and isinstance(args[1], dict) and self in args[1]:
            return ufl.Coefficient.__call__(self, *args)

        # Some help variables
        value_size = ufl.common.product(self.ufl_element().value_shape())

        # If values (return argument) is passed, check the type and length
        values = kwargs.get("values", None)
        if values is not None:
            if not isinstance(values, numpy.ndarray):
                raise TypeError, "expected a NumPy array for 'values'"
            if len(values) != value_size or \
                   not numpy.issubdtype(values.dtype, 'd'):
                raise TypeError, "expected a double NumPy array of length"\
                      " %d for return values."%value_size
            values_provided = True
        else:
            values_provided = False
            values = numpy.zeros(value_size, dtype='d')

        # Check if a cell is defined
        cell_defined = not self.ufl_element().cell().is_undefined()

        if cell_defined:
            dim = self.ufl_element().cell().geometric_dimension()

        # Assume all args are x argument
        x = args

        # If only one x argument has been provided, unpack it if it's an iterable
        if len(x) == 1:
            if isinstance(x[0], cpp.Point):
                if cell_defined:
                    x = [x[0][i] for i in xrange(dim)]
                else:
                    x = [x[0][i] for i in xrange(3)]
            elif hasattr(x[0], '__iter__'):
                x = x[0]

        # Convert it to an 1D numpy array
        try:
            x = numpy.fromiter(x, 'd')
        except (TypeError, ValueError, AssertionError), e:
            print e
            raise TypeError, "expected scalar arguments for the coordinates"

        if len(x) == 0:
            raise TypeError, "coordinate argument too short"

        if cell_defined:
            if len(x) != dim:
                raise TypeError, "expected the geometry argument to be of "\
                      "length %d"%dim

        # The actual evaluation
        self.eval(values, x)

        # If scalar return statement, return scalar value.
        if value_size == 1 and not values_provided:
            return values[0]

        return values

#--- Utility functions ---

def _check_cppcode(cppcode):
    "Check that cppcode makes sense"

    # Check that we get a string expression or nested expression
    if not isinstance(cppcode, (str, tuple, list)):
        raise TypeError, "Please provide a 'str', 'tuple' or 'list' for the 'cppcode' argument."

def _auto_select_element_from_shape(shape, degree=None, cell=None):
    "Automatically select an appropriate element from cppcode."

    # Default element, change to quadrature when working
    Family = "Lagrange"

    # Check if scalar, vector or tensor valued
    if len(shape) == 0:
        element = ufl.FiniteElement(Family, cell, degree)
    elif len(shape) == 1:
        element = ufl.VectorElement(Family, cell, degree, dim=shape[0])
    else:
        element = ufl.TensorElement(Family, cell, degree, shape=shape)

    cpp.debug("Automatic selection of expression element: " + str(element))

    return element

def _check_default_kwargs(members, kwargs):

    # Check passed default values
    if not all(member in kwargs for member in members):
        raise RuntimeError("expected a default value to all member "\
                           "variables in the Expression.")

    if not all(isinstance(value, (int, float)) for value in kwargs.values()):
        raise TypeError, "expected default arguments for member variables "\
              "to be scalars."