assemble.py

利用C

"""This module provides functionality for form assembly in Python,
corresponding to the C++ assembly and PDE classes.

The C++ assemble function (renamed to cpp_assemble) is wrapped with
an additional preprocessing step where code is generated using the
FFC JIT compiler.

The C++ PDE classes are reimplemented in Python since the C++ classes
rely on the dolfin::Form class which is not used on the Python side."""

__author__ = "Anders Logg ([e-mail])"
__date__ = "2007-08-15 -- 2008-06-21"
__copyright__ = "Copyright (C) 2007-2008 Anders Logg"
__license__  = "GNU LGPL Version 2.1"

__all__ = ["assemble",
           "Function",
           "FacetNormal",
           "MeshSize",
           "AvgMeshSize",
           "FacetArea",
           "InvFacetArea",
           "LinearPDE",
           "DirichletBC",
           "PeriodicBC",
           "compile_functions"]

from ffc import *
from dolfin import *
from compile_functions import compile_functions

# Cache for dof maps
_dof_map_cache = {}

# Cache for tensors
_tensor_cache = {}

# JIT assembler
def assemble(form, mesh, coefficients=None, dof_maps=None,
    cell_domains=None, exterior_facet_domains=None, interior_facet_domains=None, reset_tensor=None,
    tensor=None, backend=None, return_dofmaps=False):
    "Assemble form over mesh and return tensor"

    # Create empty list of coefficients, filled below
    _coefficients = ArrayFunctionPtr()

    # Extract coefficients
    if not coefficients is None: 
        # Compile all strings as dolfin::Function
        string_expressions = []
        for c in coefficients:
            # Note: To allow tuples of floats or ints below, this logic becomes more involved...
            if isinstance(c, (tuple, str)):
                string_expressions.append(c)
        if string_expressions:
            compiled_functions = compile_functions(string_expressions, mesh)
            compiled_functions.reverse()
        
        # Build list of coefficients
        for c in coefficients:
            # Note: We could generalize this to support more objects 
            # like sympy expressions, tuples for constant vectors, etc...
            if isinstance(c, (float, int)):
                c = cpp_Function(mesh, float(c))
            elif isinstance(c, (tuple, str)):
                c = compiled_functions.pop()
            _coefficients.push_back(c)

    # Check if we need to compile the form (JIT)
    if not hasattr(form, "create_cell_integral"):
        # FFC form, call JIT compile
        optimize = dolfin_get("optimize form") or dolfin_get("optimize")
        (compiled_form, module, form_data) = jit(form, optimize=optimize)
        
        # Extract coefficients from form if no coefficients are provided
        if coefficients is None:
            for c in form_data.coefficients:
                _coefficients.push_back(c.f)
        
    else:    
        # UFC form, no need to compile
        compiled_form = form
        

    # FIXME: do we need these lines now? None works fine with Assembler
    # Create dummy arguments for domains (not yet supported in Python)
    if cell_domains is None:
        cell_domains = MeshFunction("uint")
    if exterior_facet_domains is None:
        exterior_facet_domains = MeshFunction("uint")
    if interior_facet_domains is None:
        interior_facet_domains = MeshFunction("uint")

    # Create dof map set
    dof_maps = _create_dof_map_set(form, compiled_form, mesh, dof_maps)

    # Create tensor
    rank = compiled_form.rank()
    (tensor, reset_tensor) = _create_tensor(form, rank, backend, tensor, reset_tensor)

    # Set default value for reset_tensor if not specified
    if reset_tensor is None:
        reset_tensor = True

    # Assemble tensor from compiled form
    cpp_assemble(tensor, compiled_form, mesh, _coefficients, dof_maps,
                 cell_domains, exterior_facet_domains, interior_facet_domains, reset_tensor)
    
    # Convert to float for scalars
    if rank == 0:
        tensor = tensor.getval()
    
    # Return value
    if return_dofmaps:
        return tensor, dof_maps
    else:
        return tensor

def _create_dof_map_set(form, compiled_form, mesh, dof_maps):
    "Create dof map set for form"

    # Check if dof map set is supplied by user
    if dof_maps:
        return dof_maps

    # Check if we should use the cache
    use_cache = dolfin_get("optimize use dof map cache") or dolfin_get("optimize")
    if use_cache and form in _dof_map_cache:
        return _dof_map_cache[form]

    # Create dof map set
    dof_maps = DofMapSet(compiled_form, mesh)

    # Store in cache
    if use_cache:
        _dof_map_cache[form] = dof_maps

    return dof_maps

def _create_tensor(form, rank, backend, tensor, reset_tensor):
    "Create tensor for form"

    # Check if tensor is supplied by user
    if tensor:
        return (tensor, reset_tensor)

    # Decide if we should reset the tensor
    use_cache = dolfin_get("optimize use tensor cache") or dolfin_get("optimize")
    if use_cache and reset_tensor is None:
        reset_tensor = not form in _tensor_cache

    # Check if we should use the cache
    if use_cache and form in _tensor_cache:
        return (_tensor_cache[form], reset_tensor)

    # Create tensor
    if rank == 0:
        tensor = Scalar()
    elif rank == 1:
        if backend: tensor = backend.createVector()
        else:       tensor = Vector()
    elif rank == 2:
        if backend: tensor = backend.createMatrix()
        else:       tensor = Matrix()
    else:
        raise RuntimeError, "Unable to create tensors of rank %d." % rank

    # Store in cache
    if use_cache:
        _tensor_cache[form] = tensor

    return (tensor, reset_tensor)

# Rename FFC Function
ffc_Function = Function

# Create new class inheriting from both FFC and DOLFIN Function
class Function(ffc_Function, cpp_Function):

    def __init__(self, element, *others):
        "Create Function"
        # Special case, Function(element, mesh, x), need to create simple form to get arguments
        if (isinstance(element, FiniteElement) or isinstance(element, MixedElement)) and \
               len(others) == 2 and isinstance(others[0], Mesh) and \
               ( isinstance(others[1],Vector) or  isinstance(others[1], GenericVector) ):
            mesh = others[0]
            self.x = others[1]
            # Create simplest possible form
            if element.value_dimension(0) > 1:
                form = TestFunction(element)[0]*dx
            else:
                form = TestFunction(element)*dx
            # Compile form and create dof map
            (compiled_form, module, form_data) = jit(form)
            self.dof_maps = DofMapSet(compiled_form, mesh)
            # Initialize FFC and DOLFIN Function
            ffc_Function.__init__(self, element)
            cpp_Function.__init__(self, mesh, self.x, self.dof_maps.sub(0), compiled_form, 0)
        # If we have an element, then give element to FFC and the rest to DOLFIN
        elif isinstance(element, FiniteElement) or isinstance(element, MixedElement):
            ffc_Function.__init__(self, element)
            cpp_Function.__init__(self, *others)
        # Otherwise give all to DOLFIN
        else:
            cpp_Function.__init__(self, *((element,) + others))

    def split(self):
        "Extract subfunctions"
        return tuple([Function(self.e0.sub_element(i), self.sub(i)) for i in range(self.numSubFunctions())])

# Create new class inheriting from both FFC and DOLFIN FacetNormal
# (FFC FacetNormal is a function that returns a FFC Function object)
class FacetNormal(ffc_Function, cpp_FacetNormal):

    def __init__(self, shape, mesh):
        "Create FacetNormal"

        element = VectorElement("Discontinuous Lagrange", shape, 0)
        ffc_Function.__init__(self, element)
        cpp_FacetNormal.__init__(self, mesh)

# Create new class inheriting from FFC MeshSize and DOLFIN MeshSize
# (FFC MeshSize is a function that returns a FFC Function object)
class MeshSize(ffc_Function, cpp_MeshSize):

    def __init__(self, shape, mesh):
        "Create MeshSize"

        element = FiniteElement("Discontinuous Lagrange", shape, 0)
        ffc_Function.__init__(self, element)
        cpp_MeshSize.__init__(self, mesh)

# Create new class inheriting from FFC MeshSize and DOLFIN AvgMeshSize
# (FFC MeshSize is a function that returns a FFC Function object)
class AvgMeshSize(ffc_Function, cpp_AvgMeshSize):

    def __init__(self, shape, mesh):
        "Create AvgMeshSize"

        element = FiniteElement("Discontinuous Lagrange", shape, 0)
        ffc_Function.__init__(self, element)
        cpp_AvgMeshSize.__init__(self, mesh)

# Create new class inheriting from FFC FacetArea and DOLFIN FacetArea
# (FFC FacetArea is a function that returns a FFC Function object)
class FacetArea(ffc_Function, cpp_FacetArea):

    def __init__(self, shape, mesh):
        "Create FacetArea"

        element = FiniteElement("Discontinuous Lagrange", shape, 0)
        ffc_Function.__init__(self, element)
        cpp_FacetArea.__init__(self, mesh)

# Create new class inheriting from FFC InvFacetArea and DOLFIN InvFacetArea
# (FFC InvFacetArea is a function that returns a FFC Function object)
class InvFacetArea(ffc_Function, cpp_InvFacetArea):

    def __init__(self, shape, mesh):
        "Create InvFacetArea"

        element = FiniteElement("Discontinuous Lagrange", shape, 0)
        ffc_Function.__init__(self, element)
        cpp_InvFacetArea.__init__(self, mesh)

# LinearPDE class
class LinearPDE:
    """A LinearPDE represents a (system of) linear partial differential
    equation(s) in variational form: Find u in V such that
    
        a(v, u) = L(v) for all v in V',

    where a is a bilinear form and L is a linear form."""

    def __init__(self, a, L, mesh, bcs=[]):
        "Create LinearPDE"

        self.a = a
        self.L = L
        self.mesh = mesh
        self.bcs = bcs
        self.x = Vector()
        self.dof_maps = DofMapSet()

        # Make sure we have a list
        if not isinstance(self.bcs, list):
            self.bcs = [self.bcs]

    def solve(self):
        "Solve PDE and return solution"

        begin("Solving linear PDE.");
        # Assemble linear system
        (A, self.dof_maps) = assemble(self.a, self.mesh, return_dofmaps=True)
        (b, dof_maps_L)    = assemble(self.L, self.mesh, return_dofmaps=True)

        # FIXME: Maybe there is a better solution?
        # Compile form, needed to create discrete function
        (compiled_form, module, form_data) = jit(self.a)

        # Apply boundary conditions
        for bc in self.bcs:
            if isinstance(bc, DirichletBC):
                cpp_DirichletBC.apply(bc, A, b, self.dof_maps.sub(1), compiled_form)
            elif isinstance(bc, PeriodicBC):
                cpp_PeriodicBC.apply(bc, A, b, self.dof_maps.sub(1), compiled_form)
            else:
                raise RuntimeError("Unable to apply boundary conditions, unknown type: " + str(bc))

        #message("Matrix:")
        #A.disp()

        #message("Vector:")
        #b.disp()

        # Choose linear solver
        solver_type = dolfin_get("PDE linear solver")
        if solver_type == "direct":
            solver = LUSolver()
            #solver.set("parent", self)
        elif solver_type == "iterative":
            solver = KrylovSolver(gmres)
            #solver.set("parent", self)
        else:
            error("Unknown solver type \"%s\"." % solver_type)

        # Solver linear system
        solver.solve(A, self.x, b)
        
        #message("Solution vector:")
        #self.x.disp()

        # Get trial element
        element = form_data.elements[1]
  
        # Create Function
        u = Function(element, self.mesh, self.x, self.dof_maps.sub(1), compiled_form)

        end()

        return u

# DirichletBC class (need to compile form before calling constructor)
class DirichletBC(cpp_DirichletBC):

    def __init__(self, *args):
        "Create Dirichlet boundary condition"
        cpp_DirichletBC.__init__(self, *args)

    def apply(self, A, b, form):
        "Apply boundary condition to linear system"
        
        # Compile form
        (compiled_form, module, form_data) = jit(form)        

        # Create dof maps
        dof_maps = DofMapSet(compiled_form, self.mesh())
        
        # Apply boundary condition
        cpp_DirichletBC.apply(self, A, b, dof_maps.sub(1), compiled_form)

    def zero(self, A, form):
        "Apply boundary condition to linear system"
        
        # Compile form
        (compiled_form, module, form_data) = jit(form)        

        # Create dof maps
        dof_maps = DofMapSet(compiled_form, self.mesh())
        
        # Apply boundary condition
        cpp_DirichletBC.zero(self, A, dof_maps.sub(1), compiled_form)

# PeriodicBC class (need to compile form before calling constructor)
class PeriodicBC(cpp_PeriodicBC):

    def __init__(self, *args):
        "Create periodic boundary condition"
        cpp_PeriodicBC.__init__(self, *args)

    def apply(self, A, b, form):
        "Apply boundary condition to linear system"
        
        # Compile form
        (compiled_form, module, form_data) = jit(form)        

        # Create dof maps
        dof_maps = DofMapSet(compiled_form, self.mesh())
        
        # Apply boundary condition
        cpp_PeriodicBC.apply(self, A, b, dof_maps.sub(1), compiled_form)