Source code for simudo.fem.function_space

from collections import OrderedDict

from cached_property import cached_property

import dolfin
import ufl

from ..util import SetattrInitMixin

__all__ = [
    'MixedFunctionSpace',
    'MixedFunctionHelper',
    'WithSubfunctionsMixin',
]

class MixedFunctionLikeMixin(object):
    def __init__(self, mixed_function_space, function=None):
        ''' pass `function` to use an existing dolfin function instead
        of allocating a new one '''
        self.mixed_function_space = mixed_function_space
        if function is not None:
            # TODO: check function space matches
            self.function = function

    @cached_property
    def function(self):
        return self._make_dolfin_function[0](
            self.mixed_function_space.function_space)

    def split(self):
        return self.mixed_function_space.split(self.function)

class MixedFunction(MixedFunctionLikeMixin):
    # the tuple is there to prevent methodification
    _make_dolfin_function = (dolfin.Function,)

class MixedTestFunction(MixedFunctionLikeMixin):
    _make_dolfin_function = (dolfin.TestFunction,)



class MixedFunctionSpace(object):
    '''Convenience class for dealing with mixed function spaces.
'''

    MixedFunction = MixedFunction
    MixedTestFunction = MixedTestFunction

    def __init__(self, mesh, subspace_descriptors,
                 function_space_cache):
        '''
        Parameters
        ----------
        mesh: dolfin.Mesh
        subspace_descriptors: iterable
            List of dict with keys:
            element, trial_key, trial_units, test_key, test_units
        function_space_cache: FunctionSpaceCache
            Used to recycle existing FunctionSpace.
        '''
        self.mesh = mesh
        self.subspace_descriptors = tuple(subspace_descriptors)
        self.function_space_cache = function_space_cache

    @cached_property
    def subspace_descriptors_dict(self):
        ''' Like `subspace_descriptors`, except it's a dictionary
        where the keys are "trial_key" as well as "test_key". '''
        r = {}
        for desc in self.subspace_descriptors:
            r[desc['trial_key']] = desc
            r[desc['test_key' ]] = desc
        return r



    def get_element(self):
        ''' use `element` property instead '''
        descs = self.subspace_descriptors
        elements = [desc['element'] for desc in descs]
        n = len(elements)
        if n == 0: raise AssertionError()
        elif n == 1:
            return elements[0]
        else:
            return dolfin.MixedElement(elements)




    def get_function_space(self):
        ''' use `function_space` property instead '''
        return self.function_space_cache.FunctionSpace(
            self.mesh, self.element)


    @cached_property
    def element(self):
        return self.get_element()

    @cached_property
    def function_space(self):
        return self.get_function_space()

    @cached_property
    def function_subspaces(self):
        descs = self.subspace_descriptors
        W = self.function_space
        keys = [desc['trial_key'] for desc in descs]
        n = len(keys)
        if n == 1:
            return {keys[0]: W}
        else:
            return {key: W.sub(i) for i, key in enumerate(keys)}



    def make_function(self, **kwargs):
        return self.MixedFunction(self, **kwargs)




    def make_test_function(self, **kwargs):
        return self.MixedTestFunction(self, **kwargs)




    def split(self, dolfin_function, kind=None):
        '''
        returns OrderedDict of (key: subfunction*unit)
        '''

        if kind is None:
            if isinstance(dolfin_function, dolfin.Function):
                kind = 'trial'
            elif isinstance(dolfin_function, ufl.Argument):
                kind = 'test'
            else:
                raise TypeError(
                    "failed to identify whether trial or test function")

        descs = self.subspace_descriptors

        n = len(descs)
        if n == 1:
            subfunctions = (dolfin_function,)
        else:
            subfunctions = dolfin.split(dolfin_function)

        return OrderedDict((inf[kind+'_key'], inf[kind+'_units']*sub)
                           for (inf, sub)
                           in zip(descs, subfunctions))





class MixedFunctionHelper(SetattrInitMixin):
    '''Utility mixin for objects that hold a mixed function made up of
many subfunctions that need to be gathered from different objects, for
example :py:class:`.poisson_drift_diffusion.PoissonDriftDiffusion`.

Parameters
----------
mesh_util: :py:class:`.mesh_util.MeshUtil`
    Instance of MeshUtil.
subspace_descriptors_for_solution_space: 
    Sequence of subspace descriptor dicts. This property must contain
    the subspace descriptors that will be used to build the solution
    space. Typically one would gather these by calling
    :py:attr:`WithSubfunctionsMixin.solution_subspace_descriptors`
    on all the objects that contain subspace descriptors.
'''

    @property
    def mesh(self):
        return self.mesh_util.mesh

    @property
    def function_space_cache(self):
        return self.mesh_util.function_space_cache

    @property
    def solution_function(self):
        '''Dolfin Function.'''
        return self.solution_mixed_function_data['trial'].function

    @property
    def solution_test_function(self):
        '''Dolfin TestFunction.'''
        return self.solution_mixed_function_data['test'].function

    @cached_property
    def solution_mixed_function_data(self):
        trial = self.make_solution_mixed_function()
        test = self.solution_mixed_space.make_test_function()
        split = OrderedDict(tuple(trial.split().items()) +
                            tuple(test.split().items()))
        return dict(trial=trial, test=test, split=split)



    def make_solution_mixed_function(self):
        ''' intercept this to use your own Function instead of
        allocating a new one '''
        return self.solution_mixed_space.make_function()


    @cached_property
    def solution_mixed_space(self):
        descs = self.subspace_descriptors_for_solution_space
        return MixedFunctionSpace(
            mesh=self.mesh,
            subspace_descriptors=descs,
            function_space_cache=self.function_space_cache)



    def make_tolerance_function(self):
        """
        Collect "trial_tolerance" from each of the subspaces and
        assign it as a constant in each subspace. This function is
        generally used in the Newton solver to determine if the
        iterative updates are less than an absolute tolerance.

        Returns
        -------
        function: MixedFunction
            The function. Use the ``.function`` attribute to get the
            dolfin Function.
        """
        space = self.solution_mixed_space
        mixedfunction = space.make_function()
        for k, subfunction in mixedfunction.split().items():
            tol = space.subspace_descriptors_dict[k]['trial_tolerance']
            # units = space.subspace_descriptors_dict[k]['trial_units']
            self.mesh_util.function_subspace_registry.assign_scalar(
                subfunction.magnitude, tol
            )
        return mixedfunction






class WithSubfunctionsMixin(SetattrInitMixin):
    '''Utility mixin for objects that need subfunctions
(:py:class:`.poisson_drift_diffusion.MixedPoisson`,
:py:class:`.poisson_drift_diffusion.MixedQflBand`, etc).

Parameters
----------
key: str
    Unique prefix to prevent clashes between trial/test function names
    across instances (for example, "CB" to distinguish between "CB/j"
    and "VB/j").
subfunctions_info:
    Sequence of dicts with keys "{trial,test}_{name,units}" and "element".
mixed_function_solution_object: :py:class:`MixedFunctionHelper`
    Where to register subfunctions.
'''

    @property
    def solution_subspace_descriptors(self):
        '''Creates properly namespaced subspace descriptors.

Returns
-------
:
    Returns properly namespaced subspace descriptors by prefixing the
    descriptors in :py:attr:`~subfunctions_info` with
    :code:`self.key + "/"`.
'''
        descs = []
        prefix = self.subfunction_prefix
        for inf in self.subfunctions_info:
            desc = inf.copy()
            desc['trial_key'] = prefix(desc['trial_key'])
            desc['test_key' ] = prefix(desc['test_key' ])
            descs.append(desc)
        return descs



    def get_subfunction(self, name):
        name = self.subfunction_prefix(name)
        return (self.mixed_function_solution_object
                .solution_mixed_function_data['split'][name])




    def subfunction_prefix(self, name):
        '''Apply prefix to this subfunction name.'''
        return '/'.join((self.key, name))