python-dolfin 1.3.0+dfsg-2 / usr / lib / python2.7 / dist-packages / 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

# Local help class
class UserDefinedParameters(dict):
    """
    Help class to set user defined parameters in Expressions
    """
    def __init__(self, expr, **kwargs):
        assert(isinstance(expr, cpp.Expression))
        self._expr = expr
        dict.__init__(self, **kwargs)
        
    def __setitem__(self, name, value):
        __doc__ = dict.__setitem__.__doc__
        if name not in self:
            raise KeyError("'%s' is not a parameter name" % name)
        setattr(self._expr, name, value)
        dict.__setattr__(self, name, value)

    def update(self, other):
        __doc__ = dict.update.__doc__
        if hasattr(other, "keys"):
            for name in other.keys():
                self[name] = other[name]
        else:
            for name, value in other:
                self[name] = value
            
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, name=None, label=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)

        name = name or "f_" + str(ufl.Coefficient.count(self))
        label = label or "User defined expression"

        self.rename(name, label)

        # Store the rest of the kwargs (parameters)
        self.user_parameters = UserDefinedParameters(self, **kwargs)

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

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

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

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

    # Check args
    assert(isinstance(class_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)
        name = kwargs.get("name", None)
        label = kwargs.get("label", 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, name, label]])
            if len(kwargs) != num_valid_kwargs:
                raise TypeError, "expected only 'kwargs' from one of "\
                      "'element', 'degree', 'cell', 'name' or 'label'"
            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)

        name = name or "f_" + str(ufl.Coefficient.count(self))
        label = label or "User defined expression"

        self.rename(name, label)

    # 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(class_name, bases, user_dict)

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

        assert(isinstance(class_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 class_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, class_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, class_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(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.

        User defined parameters can be included as follows:

        .. code-block:: python

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

        The parameters can be scalars and any scalar valued
        GenericFunction and are all initialized to the passed default
        value. The user defined parameters can be accessed and set
        as attributes or via the dict-like user_parameters attribute:

        .. code-block:: python

            f.A = 5.0
            f.B = Expression("value", value=6.0)
            f.user_parameters["A"] = 1.0
            f.user_parameters["B"] = Constant(5.0)

        A parameter can only be updated with its original
        value-type. So if a parameter is a float, it can only be
        updated with float.

    *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 'cell' 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

        While JIT compiling an Expression the public interface is
        scanned for dependencies. A user therefore need to include
        header files (also dolfin header files) declaring types that
        is used inside a method. The following example illustrates
        this. Note the inclusion of the dolfin namespace.

        ... code-block:: python

            code = '''
            #include "dolfin/fem/GenericDofMap.h"
            namespace dolfin
            {
              class Delta : public Expression
              {
              public:
              
                Delta() : Expression() {}
              
                void eval(Array<double>& values, const Array<double>& data,
                          const ufc::cell& cell) const
                { }
              
                void update(const boost::shared_ptr<const Function> u, 
                            double nu, double dt, double C1,
                            double U_infty, double chord)
                {
                  const boost::shared_ptr<const Mesh> mesh = u->function_space()->mesh();
                  const boost::shared_ptr<const GenericDofMap> dofmap = u->function_space()->dofmap();
                  const uint ncells = mesh->num_cells();
                  uint ndofs_per_cell;
                  if (ncells > 0) 
                  {
                    CellIterator cell(*mesh);
                    ndofs_per_cell = dofmap->cell_dimension(cell->index());
                  } 
                  else 
                  {
                     return;
                  }
                }
              };
            }'''
            
            e = Expression(code)

    *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, \
                name=None, label=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)

        # Extract any Constants from the kwargs
        generic_function_members = [arg_name for arg_name, value in kwargs.items() \
                                    if isinstance(value, cpp.GenericFunction)]

        # Compile module and get the cpp.Expression class
        cpp_base, members = compile_expressions([cppcode],
                                                [generic_function_members])
        cpp_base, members = cpp_base[0], members[0]
        
        # Check passed default arguments 
        not_allowed = [n for n in dir(cls) if n[0] !="_"]
        not_allowed += ["cppcode", "user_parameters"]
        _check_kwargs(members, kwargs, not_allowed)

        # 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, \
                 name=None, label=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 self.name()

    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 = self.ufl_element().cell()
        cell_defined = cell is not None

        if cell_defined:
            dim = 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_kwargs(members, kwargs, not_allowed):
    """
    Check that all kwargs passed is either scalars or scalar Constants,
    and that the name is allowed
    """
    
    # Check passed default values
    if not all(member in kwargs for member in members):
        missing = []
        for member in members:
            if member not in kwargs:
                missing.append(member)
        missing = ", ".join("'%s'" % miss for miss in missing)
        raise RuntimeError("expected a default value to all member "\
                           "variables in the Expression. Missing: %s." % missing)

    for name in kwargs.keys():
        if name in not_allowed:
            raise RuntimeError("Parameter name: '%s' is not allowed. It is "\
                               "part of the interface of Expression" % name)
            
    if not (all(isinstance(value, (int, float)) or \
                (isinstance(value, cpp.GenericFunction) and value.value_size()==1) \
                for value in kwargs.values())):
        raise TypeError, "expected default arguments for member variables "\
              "to be scalars or a scalar GenericFunctions."