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

"""
This module contains functionality for function spaces in particular
discrete function spaces defined over meshes in terms of finite
elements.
"""

# Copyright (C) 2008 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
# Modified by Martin Alnes 2008
# Modified by Kent-Andre Mardal 2009
#
# First added:  2008-11-03
# Last changed: 2011-11-15

__all__ = ["FunctionSpace", "MixedFunctionSpace",
           "VectorFunctionSpace", "TensorFunctionSpace",
           "FunctionSpaceBase","FunctionSpaceFromCpp",
           "create_ufc_function_spaces"]

import copy

# Import UFL and SWIG-generated extension module (DOLFIN C++)
import ufl

import dolfin.cpp as cpp
from dolfin.compilemodules.jit import jit

# Mapping from dimension to domain
dim2domain = {1: "interval", 2: "triangle", 3: "tetrahedron"}

class FunctionSpaceBase(cpp.FunctionSpace):
    "Base class for all function spaces."

    def __init__(self, mesh, element):
        """Create function space on given mesh for given finite element.

        *Arguments*
            mesh
                A :py:class:`Mesh <dolfin.cpp.Mesh>`
            element
                A :py:class:`(UFL) FiniteElement <ufl.FiniteElementBase>`

        """

        # Check arguments
        if not isinstance(mesh, cpp.Mesh):
            cpp.dolfin_error("functionspace.py",
                             "create function space",
                             "Illegal argument, not a mesh: " + str(mesh))
        if not isinstance(element, (ufl.FiniteElementBase)):
            cpp.dolfin_error("functionspace.py",
                             "create function space",
                             "Illegal argument, not a finite element: " + str(element))

        # Store element
        # Note: self._element cannot be a private attribute as we want to be able to
        #       set the element from a derived class.
        self._ufl_element = element

        # JIT-compile element to get ufc_element and ufc_dofmap
        ufc_element, ufc_dofmap = jit(self._ufl_element)

        # Instantiate DOLFIN FiniteElement and DofMap
        self._dolfin_element = cpp.FiniteElement(ufc_element)
        dolfin_dofmap  = cpp.DofMap(ufc_dofmap, mesh)

        # Initialize the cpp_FunctionSpace
        cpp.FunctionSpace.__init__(self, mesh,
                                   self._dolfin_element,
                                   dolfin_dofmap)

    # FIXME: Sort out consistent interface for access to DOLFIN, UFL and UFC objects...

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

    def ufl_element(self):
        "Return the UFL element."
        return self._ufl_element

    def dolfin_element(self):
        "Return the DOLFIN element."
        return self._dolfin_element

    def __add__(self, other):
        "Create enriched function space."
        return EnrichedFunctionSpace((self, other))

    def __mul__(self, other):
        "Create mixed function space."
        return MixedFunctionSpace((self, other))

    def __str__(self):
        "Pretty-print."
        return "<Function space of dimension %d (%s)>" % \
               (self.dofmap().global_dimension(), str(self._ufl_element))

    def num_sub_spaces(self):
        "Return the number of sub spaces"
        return self._dolfin_element.num_sub_elements()

    def sub(self, i):
        "Return the i:th cpp.SubSpace"
        # Fixme: Should we have a more extensive check other than whats included in
        # the cpp code?
        #if i not in self._cpp_sub_spaces.keys():
            # Store the created subspace to prevent swig garbage collection
            # Should not be needed as SubSpace is shared_ptr stored
        #    self._cpp_sub_spaces[i] = FunctionSpaceFromCpp(cpp.SubSpace(self,i))
        if not isinstance(i, int):
            raise TypeError, "expected an int for 'i'"
        if self.num_sub_spaces() == 1:
            raise ValueError, "no SubSpaces to extract"
        if i >= self.num_sub_spaces():
            raise ValueError, "Can only extract SubSpaces with i = 0 ... %d" % \
                  (self.num_sub_spaces() - 1)
        assert(hasattr(self._ufl_element,"sub_elements"))
        element = self._ufl_element.sub_elements()[i]
        return FunctionSpaceFromCpp(cpp.SubSpace(self,i), element)

    def split(self):
        """
        Split a mixed functionspace into its sub spaces
        """
        S = []
        for i in range(self.num_sub_spaces()):
            S.append(self.sub(i))
        return S

    def collapse(self, collapsed_dofs=False):
        """
        Collapse a subspace and return a new function space and a map
        from new to old dofs

        *Arguments*
            collapsed_dofs (bool)
                Return the map from new to old dofs

       *Returns*
           _FunctionSpace_
                The new function space.
           dict
                The map from new to old dofs (optional)
        """
        cpp_space, dofs = cpp.FunctionSpace.collapse(self)
        if collapsed_dofs:
            return FunctionSpaceFromCpp(cpp_space), dofs
        return FunctionSpaceFromCpp(cpp_space)

class FunctionSpaceFromCpp(FunctionSpaceBase):
    " FunctionSpace represents a finite element function space."
    def __init__(self, cppV, element=None):
        """ Initialize a FunctionSpace from an already excisting
            cpp.FunctionSpace.
        """
        if not isinstance(cppV, (cpp.FunctionSpace)):
            cpp.dolfin_error("functionspace.py",
                             "create function space",
                             "Illegal argument, not a cpp.FunctionSpace: " + str(cppV))

        if element is None:
            # Get the ufl.FiniteElement from the compiled element sigature
            self._ufl_element = eval(cppV.element().signature(), ufl.__dict__)
        else:
            if not isinstance(element, ufl.FiniteElementBase):
                raise TypeError, "expected a ufl.FiniteElementBase"
            self._ufl_element = element

        # Init the cpp.FunctionSpace with the provided cppV
        cpp.FunctionSpace.__init__(self, cppV)

        # Assign the dolfin element
        self._dolfin_element = cpp.FunctionSpace.element(self)

def create_ufc_function_spaces(mesh, ufc_form, cache=None):
    """
    Instantiate cpp.FunctionSpaces from compiled ufc form.

    *Arguments*
        mesh
            a :py:class:`Mesh <dolfin.cpp.Mesh>`.
        ufc_form
            compiled ufc form
        cache
            a 'dict' with already instantiated
            :py:class:`cpp.FunctionSpaces <dolfin.cpp.FunctionSpace>`.

    *Examples of usage*

        .. code-block:: python

            fs0, c = create_ufc_function_spaces(mesh, ufc_form0)
            fs1, c = create_ufc_function_spaces(mesh, ufc_form1, c)

    """

    # Initialize return arguments
    functionspaces = []
    if cache is None:
        cache = {}

    # Iterate over all known ufc_finite_elements
    for i in range(ufc_form.rank() + ufc_form.num_coefficients()):
        # Create a ufc_finite_elements
        fe = ufc_form.create_finite_element(i)

        # Use the signature of the element as key in the cache dict
        fesig = fe.signature()

        # Try to access the cache
        V = cache.get(fesig)

        # If the cpp.FunctionSpace does not excists in the cache
        if V is None:
            # Instantiate a dofmap
            dm = ufc_form.create_dofmap(i)

            # Instantiate the UFCFunctionSpace
            V = UFCFunctionSpace(mesh, fe, dm)
            cache[fesig] = V

        functionspaces.append(V)

    return functionspaces, cache

class UFCFunctionSpace(cpp.FunctionSpace):
    "FunctionSpace represents a finite element function space."
    def __init__(self, mesh, ufc_finite_element, ufc_dofmap):
        " Initialize a FunctionSpace from ufc data "
        self._mesh = mesh
        self._finite_element     = cpp.FiniteElement(ufc_finite_element)
        self._dofmap            = cpp.DofMap(ufc_dofmap, mesh)
        self._ufc_finite_element = ufc_finite_element
        self._ufc_dofmap        = ufc_dofmap
        cpp.FunctionSpace.__init__(self, self._mesh, self._finite_element, \
                                   self._dofmap)

class FunctionSpace(FunctionSpaceBase):
    "FunctionSpace represents a finite element function space."

    def __init__(self, mesh, family, degree, form_degree=None,
                 restriction=None):
        """
        Create finite element function space.

        *Arguments*
            mesh (:py:class:`Mesh <dolfin.cpp.Mesh>`)
                the mesh
            family (string)
                specification of the element family, see below for
                alternatives.
            degree (int)
                the degree of the element.
            form_degree (int)
                the form degree (FEEC notation, used when field is
                viewed as k-form)
            restriction
                restriction of the element (e.g. to cell facets).

        Which families and degrees that are supported is determined by the
        form compiler used to generate the element, but typical families
        include

        =================================  =========
        Name                               Usage
        =================================  =========
        Argyris*                           "ARG"
        Arnold-Winther*                    "AW"
        Brezzi-Douglas-Fortin-Marini*      "BDFM"
        Brezzi-Douglas-Marini              "BDM"
        Bubble                             "B"
        Crouzeix-Raviart                   "CR"
        Discontinuous Lagrange             "DG"
        Hermite*                           "HER"
        Lagrange                           "CG"
        Mardal-Tai-Winther*                "MTW"
        Morley*                            "MOR"
        Nedelec 1st kind H(curl)           "N1curl"
        Nedelec 2nd kind H(curl)           "N2curl"
        Quadrature                         "Q"
        Raviart-Thomas                     "RT"
        =================================  =========

        *only partly supported.

        *Examples of usage*
            To define a discrete function space over e.g. the unit square:

            .. code-block:: python

                mesh = UnitSquare(32,32)
                V = FunctionSpace(mesh, "CG", 1)

            Here, ``"CG"`` stands for Continuous Galerkin, implying the
            standard Lagrange family of elements. Instead of ``"CG"``, we
            could have written ``"Lagrange"``. With degree 1, we get the
            linear Lagrange element. Other examples include:

            .. code-block:: python

                # Discontinuous element of degree 0
                V = FunctionSpace(mesh, "DG", 0)

                # Brezzi-Douglas-Marini element of degree 2
                W = FunctionSpace(mesh, "BDM", 2)

        """

        # Check arguments
        if not isinstance(mesh, cpp.Mesh):
            cpp.dolfin_error("functionspace.py",
                             "create function space",
                             "Illegal argument, not a mesh: " + str(mesh))
        if not isinstance(family, str):
            cpp.dolfin_error("functionspace.py",
                             "create function space",
                             "Illegal argument for finite element family, not a string: " + str(family))
        if not isinstance(degree, int):
            cpp.dolfin_error("functionspace.py",
                             "create function space",
                             "Illegal argument for degree, not an integer: " + str(degree))
        dim = mesh.topology().dim()
        if not dim in dim2domain:
            cpp.dolfin_error("functionspace.py",
                             "create function space",
                             "Invalid mesh dimension (%s)" % str(dim))

        # Create element
        domain = dim2domain[mesh.topology().dim()]
        element = ufl.FiniteElement(family, domain, degree,
                                    form_degree=form_degree)
        if restriction is not None:
            element = element[restriction]

        # Initialize base class
        FunctionSpaceBase.__init__(self, mesh, element)

        self.___degree = degree
        self.___family = family
        self.___mesh = mesh
        self.___form_degree = form_degree

    def restriction(self, meshfunction):
        space = FunctionSpace(self.___mesh, self.___family, self.___degree,
                              form_degree=self.___form_degree)
        space.attach(meshfunction)
        return space

class EnrichedFunctionSpace(FunctionSpaceBase):
    "EnrichedFunctionSpace represents an enriched finite element function space."

    def __init__(self, spaces):
        """
        Create enriched finite element function space.

        *Arguments*
            spaces
                a list (or tuple) of :py:class:`FunctionSpaces
                <dolfin.functions.functionspace.FunctionSpace>`.

        *Usage*
            The function space may be created by

            .. code-block:: python

                V = EnrichedFunctionSpace(spaces)

        """

        # Check arguments
        if not len(spaces) > 0:
            cpp.dolfin_error("functionspace.py",
                             "create enriched function space",
                             "Need at least one subspace")
        if not all(isinstance(V, FunctionSpaceBase) for V in spaces):
            cpp.dolfin_error("functionspace.py",
                             "create enriched function space",
                             "Invalid subspaces: " + str(spaces))
        #if not all(V.mesh() == spaces[0].mesh() for V in spaces):
        #    cpp.dolfin_error("functionspace.py",
        # "Nonmatching meshes for mixed function space: " + str([V.mesh() for V in spaces]))

        # Create element
        element = ufl.EnrichedElement(*[V.ufl_element() for V in spaces])

        # Initialize base class
        FunctionSpaceBase.__init__(self, spaces[0].mesh(), element)

class MixedFunctionSpace(FunctionSpaceBase):
    "MixedFunctionSpace represents a mixed finite element function space."

    def __init__(self, spaces):
        """
        Create mixed finite element function space.

        *Arguments*
            spaces
                a list (or tuple) of :py:class:`FunctionSpaces
                <dolfin.functions.functionspace.FunctionSpace>`.

        *Examples of usage*
            The function space may be created by

            .. code-block:: python

                V = MixedFunctionSpace(spaces)

            ``spaces`` may consist of multiple occurances of the same space:

            .. code-block:: python

                P1  = FunctionSpace(mesh, "CG", 1)
                P2v = VectorFunctionSpace(mesh, "Lagrange", 2)

                ME  = MixedFunctionSpace([P2v, P1, P1, P1])

        """

        # Check arguments
        if not len(spaces) > 0:
            cpp.dolfin_error("functionspace.py",
                             "create mixed function space",
                             "Need at least one subspace")
        if not all(isinstance(V, FunctionSpaceBase) for V in spaces):
            cpp.dolfin_error("functionspace.py",
                             "create mixed function space",
                             "Invalid subspaces: " + str(spaces))
        #if not all(V.mesh() == spaces[0].mesh() for V in spaces):
        #    cpp.dolfin_error("functionspace.py", "Nonmatching meshes for mixed function space: " \
        #              + str([V.mesh() for V in spaces]))

        # Create element
        element = ufl.MixedElement(*[V.ufl_element() for V in spaces])

        # Initialize base class
        FunctionSpaceBase.__init__(self, spaces[0].mesh(), element)

class VectorFunctionSpace(MixedFunctionSpace):
    "VectorFunctionSpace represents a vector-valued finite element function space."

    def __init__(self, mesh, family, degree, dim=None, form_degree=None,
                 restriction=None):
        """
        Create vector-valued finite element function space.

        Use VectorFunctionSpace if the unknown is a vector field,
        instead of a :py:class:`FunctionSpace
        <dolfin.functions.functionspace.FunctionSpace>` object for
        scalar fields.

        *Arguments*
            mesh (:py:class:`Mesh <dolfin.cpp.Mesh>`)
                the mesh
            family (string)
                a string specifying the element family, see
                :py:class:`FunctionSpace
                <dolfin.functions.functionspace.FunctionSpace>` for
                alternatives.
            degree (int)
                the (polynomial) degree of the element.
            dim (int)
                an optional argument specifying the number of components.
            form_degree (int)
                an optional argument specifying the degree of the
                k-form (used for FEEC notation)

        If the dim argument is not provided, the dimension will be deduced from
        the dimension of the mesh.

        *Example of usage*

            .. code-block:: python

                V = VectorFunctionSpace(mesh, "CG", 1)

        """

        # Create element
        dim = dim or mesh.geometry().dim()
        cell = dim2domain[mesh.topology().dim()]
        element = ufl.VectorElement(family, cell, degree, dim=dim,
                                    form_degree=form_degree)

        if restriction is not None:
            element = element[restriction]

        # Initialize base class
        FunctionSpaceBase.__init__(self, mesh, element)

class TensorFunctionSpace(FunctionSpaceBase):
    "TensorFunctionSpace represents a tensor-valued finite element function space."

    def __init__(self, mesh, family, degree, shape=None, symmetry=None, \
                 restriction=None):
        """
        Create tensor-valued finite element function space.

        *Arguments*
            mesh
                a :py:class:`Mesh <dolfin.cpp.Mesh>`.
            family
                a string specifying the element family,
                see :py:class:`FunctionSpace
                <dolfin.functions.functionspace.FunctionSpace>`
                for alternatives.
            degree
                the degree of the element.
            shape
                an optional argument specifying the shape of the tensor.
            symmetry
                optional argument specifying whether the tensor is symmetric.

        If the shape argument is not provided, the dimension will be deduced from
        the dimension of the mesh.

        *Example of usage*

            .. code-block:: python

                V = TensorFunctionSpace(mesh, "CG", 1)

        """

        # Create subspaces
        if shape == None:
            dim = mesh.topology().dim()
            shape = (dim,dim)
        cell = dim2domain[mesh.topology().dim()]
        element = ufl.TensorElement(family, cell, degree, shape, symmetry)
        if restriction is not None:
            element = element[restriction]

        # Initialize base class
        FunctionSpaceBase.__init__(self, mesh, element)