Exemplo n.º 1
0
def get_aspect_ratios2d(mesh, python=False):
    """
    Computes the aspect ratio of each cell in a 2D triangular mesh

    :arg mesh: the input mesh to do computations on
    :kwarg python: compute the measure using Python?

    :rtype: firedrake.function.Function aspect_ratios with
        aspect ratio data
    """
    P0 = firedrake.FunctionSpace(mesh, "DG", 0)
    if python:
        P0_ten = firedrake.TensorFunctionSpace(mesh, "DG", 0)
        J = firedrake.interpolate(ufl.Jacobian(mesh), P0_ten)
        edge1 = ufl.as_vector([J[0, 0], J[1, 0]])
        edge2 = ufl.as_vector([J[0, 1], J[1, 1]])
        edge3 = edge1 - edge2
        a = ufl.sqrt(ufl.dot(edge1, edge1))
        b = ufl.sqrt(ufl.dot(edge2, edge2))
        c = ufl.sqrt(ufl.dot(edge3, edge3))
        aspect_ratios = firedrake.interpolate(
            a * b * c / ((a + b - c) * (b + c - a) * (c + a - b)), P0
        )
    else:
        coords = mesh.coordinates
        aspect_ratios = firedrake.Function(P0)
        op2.par_loop(
            get_pyop2_kernel("get_aspect_ratio", 2),
            mesh.cell_set,
            aspect_ratios.dat(op2.WRITE, aspect_ratios.cell_node_map()),
            coords.dat(op2.READ, coords.cell_node_map()),
        )
    return aspect_ratios
Exemplo n.º 2
0
def test_de_rahm_2D(order):
    mesh = create_unit_square(MPI.COMM_WORLD, 3, 4)
    W = FunctionSpace(mesh, ("Lagrange", order))
    w = Function(W)
    w.interpolate(lambda x: x[0] + x[0] * x[1] + 2 * x[1]**2)

    g = ufl.grad(w)
    Q = FunctionSpace(mesh, ("N2curl", order - 1))
    q = Function(Q)
    q.interpolate(Expression(g, Q.element.interpolation_points))

    x = ufl.SpatialCoordinate(mesh)
    g_ex = ufl.as_vector((1 + x[1], 4 * x[1] + x[0]))
    assert np.isclose(
        np.abs(assemble_scalar(form(ufl.inner(q - g_ex, q - g_ex) * ufl.dx))),
        0)

    V = FunctionSpace(mesh, ("BDM", order - 1))
    v = Function(V)

    def curl2D(u):
        return ufl.as_vector((ufl.Dx(u[1], 0), -ufl.Dx(u[0], 1)))

    v.interpolate(
        Expression(curl2D(ufl.grad(w)), V.element.interpolation_points))
    h_ex = ufl.as_vector((1, -1))
    assert np.isclose(
        np.abs(assemble_scalar(form(ufl.inner(v - h_ex, v - h_ex) * ufl.dx))),
        0)
Exemplo n.º 3
0
 def grad(self, o, expr):
     shape = o.ufl_shape
     if len(shape) == 1:
         return as_vector([expr] * shape[0])
     elif len(shape) == 2:
         return as_vector([[expr[i]] * shape[1] for i in range(shape[0])])
     else:
         raise NotImplementedError()
Exemplo n.º 4
0
def UFLFunction(grid, name, order, expr, renumbering=None, virtualize=True, tempVars=True,
                predefined=None, **kwargs):
    scalar = False
    if type(expr) == list or type(expr) == tuple:
        expr = ufl.as_vector(expr)
    elif type(expr) == int or type(expr) == float:
        expr = ufl.as_vector( [expr] )
        scalar = True
    try:
        if expr.ufl_shape == ():
            expr = ufl.as_vector([expr])
            scalar = True
    except:
        return None
    _, coeff_ = ufl.algorithms.analysis.extract_arguments_and_coefficients(expr)
    coeff   = {c : c.toVectorCoefficient()[0] for c in coeff_ if len(c.ufl_shape) == 0 and not c.is_cellwise_constant()}
    expr = replace(expr,coeff)

    if len(expr.ufl_shape) > 1:
        raise AttributeError("can only generate grid functions from vector values UFL expressions not from expressions with shape=",expr.ufl_shape)

    # set up the source class
    source = UFLFunctionSource(grid, expr,
            name,order,
            tempVars=tempVars,virtualize=virtualize,
            predefined=predefined)

    coefficients = source.coefficientList
    numCoefficients = len(coefficients)
    if renumbering is None:
        renumbering = dict()
        renumbering.update((c, i) for i, c in enumerate(sorted((c for c in coefficients if not c.is_cellwise_constant()), key=lambda c: c.count())))
        renumbering.update((c, i) for i, c in enumerate(c for c in coefficients if c.is_cellwise_constant()))
    coefficientNames = ['coefficient' + str(i) if n is None else n for i, n in enumerate(getattr(c, 'name', None) for c in coefficients if not c.is_cellwise_constant())]

    # call code generator
    from dune.generator import builder
    module = builder.load(source.name(), source, "UFLLocalFunction")

    assert hasattr(module,"UFLLocalFunction"),\
          "GridViews of coefficients need to be compatible with the grid view of the ufl local functions"

    class LocalFunction(module.UFLLocalFunction):
        def __init__(self, gridView, name, order, *args, **kwargs):
            self.base = module.UFLLocalFunction
            self._coefficientNames = {n: i for i, n in enumerate(source.coefficientNames)}
            if renumbering is not None:
                self._renumbering = renumbering
                self._setConstant = self.setConstant # module.UFLLocalFunction.__dict__['setConstant']
                self.setConstant = lambda *args: setConstant(self,*args)
            self.constantShape = source._constantShapes
            self._constants = [c for c in source.constantList if isinstance(c,Constant)]
            self.scalar = scalar
            init(self, gridView, name, order, *args, **kwargs)

    return LocalFunction
Exemplo n.º 5
0
 def argument(self, o):
     shape = o.ufl_shape
     if len(shape) == 0:
         return 1
     elif len(shape) == 1:
         return as_vector([1] * shape[0])
     elif len(shape) == 2:
         return as_vector([as_vector([1] * shape[1]) for _ in range(shape[0])])
     else:
         raise NotImplementedError()
Exemplo n.º 6
0
 def list_tensor(self, o):
     shape = o.ufl_shape
     if len(shape) == 1:
         return as_vector([self.visit(op) for op in o.ufl_operands])
     elif len(shape) == 2:
         return as_vector(
             [as_vector([self.visit(op) for op in row.ufl_operands]) for row in o.ufl_operands]
         )
     else:
         raise NotImplementedError()
Exemplo n.º 7
0
def codeDG(self):
    code = self._code()

    u = self.trialFunction
    ubar = Coefficient(u.ufl_function_space())
    penalty = self.penalty
    if penalty is None:
        penalty = 1
    if isinstance(penalty, Expr):
        if penalty.ufl_shape == ():
            penalty = as_vector([penalty])
        try:
            penalty = expand_indices(
                expand_derivatives(expand_compounds(penalty)))
        except:
            pass
        assert penalty.ufl_shape == u.ufl_shape
        dmPenalty = as_vector([
            replace(
                expand_derivatives(diff(replace(penalty, {u: ubar}),
                                        ubar))[i, i], {ubar: u})
            for i in range(u.ufl_shape[0])
        ])
    else:
        dmPenalty = None

    code.append(AccessModifier("public"))
    x = SpatialCoordinate(self.space.cell())
    predefined = {}
    self.predefineCoefficients(predefined, x)
    spatial = Variable('const auto', 'y')
    predefined.update({
        x:
        UnformattedExpression(
            'auto', 'entity().geometry().global( Dune::Fem::coordinate( x ) )')
    })
    generateMethod(code,
                   penalty,
                   'RRangeType',
                   'penalty',
                   args=['const Point &x', 'const DRangeType &u'],
                   targs=['class Point', 'class DRangeType'],
                   static=False,
                   const=True,
                   predefined=predefined)
    generateMethod(code,
                   dmPenalty,
                   'RRangeType',
                   'linPenalty',
                   args=['const Point &x', 'const DRangeType &u'],
                   targs=['class Point', 'class DRangeType'],
                   static=False,
                   const=True,
                   predefined=predefined)
    return code
Exemplo n.º 8
0
 def __init__(self, functionSpace, value, subDomain=None):
     if functionSpace.scalar:
         self.functionSpace = functionSpace.toVectorSpace()
         value = ufl.as_vector([value])
     else:
         self.functionSpace = functionSpace
     self.value = value
     self.subDomain = subDomain
     if type(value) is list:
         self.ufl_value = value  # DirichletBC.flatten(value)
         self.ufl_value = [0 if v is None else v for v in self.ufl_value]
         self.ufl_value = ufl.as_vector(self.ufl_value)
     else:
         self.ufl_value = value
     assert self.ufl_value.ufl_shape[0] == functionSpace.dimRange
Exemplo n.º 9
0
def mark(indicator,
         refineTolerance,
         coarsenTolerance=0,
         minLevel=0,
         maxLevel=None,
         gridView=None):
    if ufl and (isinstance(indicator, list) or isinstance(indicator, tuple)):
        indicator = ufl.as_vector(indicator)
    if ufl and isinstance(indicator, ufl.core.expr.Expr):  #
        if gridView is None:
            _, coeff_ = extract_arguments_and_coefficients(indicator)
            gridView = [c.grid for c in coeff_ if hasattr(c, "grid")]
            gridView += [c.gridView for c in coeff_ if hasattr(c, "gridView")]
            if len(gridView) == 0:
                raise ValueError(
                    "if a ufl expression is passed as indicator then the 'gridView' must also be provided"
                )
            gridView = gridView[0]
        indicator = expression2GF(gridView, indicator, 0)
    if gridView is None:
        gridView = indicator.gridView
    try:
        if not gridView.canAdapt:
            raise AttributeError(
                "indicator function must be over grid view that supports adaptation"
            )
    except AttributeError:
        raise AttributeError(
            "indicator function must be over grid view that supports adaptation"
        )
    if maxLevel == None:
        maxLevel = -1
    return gridView.mark(indicator, refineTolerance, coarsenTolerance,
                         minLevel, maxLevel)
Exemplo n.º 10
0
 def argument(self, o):
     from ufl import as_vector
     from ufl.constantvalue import Zero
     from dolfin import Argument
     from numpy import ndindex
     V = o.function_space()
     if V.num_sub_spaces() == 0:
         # Not on a mixed space, just return ourselves.
         return o
     # Already seen this argument, return the cached version.
     if o in self._args:
         return self._args[o]
     args = []
     for i, V_i_sub in enumerate(V.split()):
         # Walk over the subspaces and build a vector that is zero
         # where the argument does not match the one we're looking
         # for and is just the non-mixed argument when we do want
         # it.
         V_i = V_i_sub.collapse()
         a = Argument(V_i, o.number(), part=o.part())
         indices = ndindex(a.ufl_shape)
         if self.idx[o.number()] == i:
             args += [a[j] for j in indices]
         else:
             args += [Zero() for j in indices]
     self._args[o] = as_vector(args)
     return self._args[o]
Exemplo n.º 11
0
    def argument(self, arg):
        """Split an argument into its constituent spaces."""
        from itertools import product
        if isinstance(arg.function_space(), functionspace.MixedFunctionSpace):
            # Look up the split argument in cache since we want it unique
            if arg in self._args:
                return self._args[arg]
            args = []
            for i, fs in enumerate(arg.function_space().split()):
                # Build the sub-space Argument (not part of the mixed
                # space).
                a = Argument(fs, arg.number(), part=arg.part())

                # Produce indexing iterator.
                # For scalar-valued spaces this results in the empty
                # tuple (), which returns the Argument when indexing.
                # For vector-valued spaces, this is just
                # range(a.ufl_shape[0])
                # For tensor-valued spaces, it's the nested loop of
                # range(x for x in a.ufl_shape).
                #
                # Each of the indexed things is then scalar-valued
                # which we turn into a ufl Vector.
                indices = product(*map(range, a.ufl_shape))
                if self._idx[arg.number()] == i:
                    args += [a[idx] for idx in indices]
                else:
                    args += [Zero() for _ in indices]
            self._args[arg] = as_vector(args)
            return self._args[arg]
        return arg
Exemplo n.º 12
0
    def __init__(self, grid, expr,
            name,order,
            tempVars=True, virtualize=True,
            predefined=None):
        gridType = grid.cppTypeName
        gridIncludes = grid.cppIncludes
        if len(expr.ufl_shape) == 0:
            expr = as_vector( [ expr ] )
        dimRange = expr.ufl_shape[0]
        predefined = {} if predefined is None else predefined
        codegen.ModelClass.__init__(self,"UFLLocalFunction",[expr],
          virtualize,dimRange=dimRange, predefined=predefined)

        self.evalCode = []
        self.jacCode = []
        self.hessCode = []
        self.gridType = gridType
        self.gridIncludes = gridIncludes

        self.functionName = name
        self.functionOrder = order
        self.expr = expr
        self.tempVars = tempVars
        self.virtualize = virtualize

        self.codeString = self.toString()
Exemplo n.º 13
0
    def argument(self, obj):
        Q = obj.ufl_function_space()
        dom = Q.ufl_domain()
        sub_elements = obj.ufl_element().sub_elements()

        # If not a mixed element, do nothing
        if not isinstance(obj.ufl_element(), MixedElement):
            return obj

        # Split into sub-elements, creating appropriate space for each
        args = []
        for i, sub_elem in enumerate(sub_elements):
            Q_i = FunctionSpace(dom, sub_elem)
            a = Argument(Q_i, obj.number(), part=obj.part())

            indices = [()]
            for m in a.ufl_shape:
                indices = [(k + (j, )) for k in indices for j in range(m)]

            if (i == self.idx[obj.number()]):
                args += [a[j] for j in indices]
            else:
                args += [Zero() for j in indices]

        return as_vector(args)
Exemplo n.º 14
0
    def amg_solve(N, method):
        # Elasticity parameters
        E = 1.0e9
        nu = 0.3
        mu = E / (2.0 * (1.0 + nu))
        lmbda = E * nu / ((1.0 + nu) * (1.0 - 2.0 * nu))

        # Stress computation
        def sigma(v):
            return 2.0 * mu * sym(grad(v)) + lmbda * tr(sym(
                grad(v))) * Identity(2)

        # Define problem
        mesh = UnitSquareMesh(MPI.comm_world, N, N)
        V = VectorFunctionSpace(mesh, 'Lagrange', 1)
        bc0 = Function(V)
        with bc0.vector().localForm() as bc_local:
            bc_local.set(0.0)

        def boundary(x, only_boundary):
            return [only_boundary] * x.shape(0)

        bc = DirichletBC(V.sub(0), bc0, boundary)
        u = TrialFunction(V)
        v = TestFunction(V)

        # Forms
        a, L = inner(sigma(u), grad(v)) * dx, dot(ufl.as_vector(
            (1.0, 1.0)), v) * dx

        # Assemble linear algebra objects
        A = assemble_matrix(a, [bc])
        A.assemble()
        b = assemble_vector(L)
        apply_lifting(b, [a], [[bc]])
        b.ghostUpdate(addv=PETSc.InsertMode.ADD,
                      mode=PETSc.ScatterMode.REVERSE)
        set_bc(b, [bc])

        # Create solution function
        u = Function(V)

        # Create near null space basis and orthonormalize
        null_space = build_nullspace(V, u.vector())

        # Attached near-null space to matrix
        A.set_near_nullspace(null_space)

        # Test that basis is orthonormal
        assert null_space.is_orthonormal()

        # Create PETSC smoothed aggregation AMG preconditioner, and
        # create CG solver
        solver = PETScKrylovSolver("cg", method)

        # Set matrix operator
        solver.set_operator(A)

        # Compute solution and return number of iterations
        return solver.solve(u.vector(), b)
Exemplo n.º 15
0
    def argument(self, arg):
        """Split an argument into its constituent spaces."""
        from itertools import product
        if isinstance(arg.function_space(), functionspace.MixedFunctionSpace):
            # Look up the split argument in cache since we want it unique
            if arg in self._args:
                return self._args[arg]
            args = []
            for i, fs in enumerate(arg.function_space().split()):
                # Build the sub-space Argument (not part of the mixed
                # space).
                a = Argument(fs, arg.number(), part=arg.part())

                # Produce indexing iterator.
                # For scalar-valued spaces this results in the empty
                # tuple (), which returns the Argument when indexing.
                # For vector-valued spaces, this is just
                # range(a.ufl_shape[0])
                # For tensor-valued spaces, it's the nested loop of
                # range(x for x in a.ufl_shape).
                #
                # Each of the indexed things is then scalar-valued
                # which we turn into a ufl Vector.
                indices = product(*map(range, a.ufl_shape))
                if self._idx[arg.number()] == i:
                    args += [a[idx] for idx in indices]
                else:
                    args += [Zero() for _ in indices]
            self._args[arg] = as_vector(args)
            return self._args[arg]
        return arg
Exemplo n.º 16
0
    def argument(self, arg):
        """
        Argument is UFL lingo for test (num=0) and trial (num=1) functions
        """
        # Do not make several new args for the same original arg
        if arg in self._cache:
            return self._cache[arg]

        # Get some data about the argument and get our target index
        V = arg.function_space()
        N = V.num_sub_spaces()
        num = arg.number()
        idx_wanted = [self._index_v, self._index_u][num]

        new_arg = []
        for idx in range(N):
            # Construct non-coupled argument
            Vi = V.sub(idx).collapse()
            a = dolfin.function.argument.Argument(Vi, num, part=arg.part())
            indices = numpy.ndindex(a.ufl_shape)

            if idx == idx_wanted:
                # This is a wanted index
                new_arg.extend(a[I] for I in indices)
            else:
                # This index should be pruned
                new_arg.extend(Zero() for _ in indices)

        new_arg = as_vector(new_arg)
        self._cache[arg] = new_arg
        return new_arg
Exemplo n.º 17
0
    def amg_solve(N, method):
        # Elasticity parameters
        E = 1.0e9
        nu = 0.3
        mu = E / (2.0 * (1.0 + nu))
        lmbda = E * nu / ((1.0 + nu) * (1.0 - 2.0 * nu))

        # Stress computation
        def sigma(v):
            return 2.0 * mu * sym(grad(v)) + lmbda * tr(sym(
                grad(v))) * Identity(2)

        # Define problem
        mesh = create_unit_square(MPI.COMM_WORLD, N, N)
        V = VectorFunctionSpace(mesh, 'Lagrange', 1)
        u = TrialFunction(V)
        v = TestFunction(V)

        facetdim = mesh.topology.dim - 1
        bndry_facets = locate_entities_boundary(
            mesh, facetdim, lambda x: np.full(x.shape[1], True))
        bdofs = locate_dofs_topological(V.sub(0), V, facetdim, bndry_facets)
        bc = dirichletbc(PETSc.ScalarType(0), bdofs, V.sub(0))

        # Forms
        a, L = inner(sigma(u), grad(v)) * dx, dot(ufl.as_vector(
            (1.0, 1.0)), v) * dx

        # Assemble linear algebra objects
        A = assemble_matrix(a, [bc])
        A.assemble()
        b = assemble_vector(L)
        apply_lifting(b, [a], [[bc]])
        b.ghostUpdate(addv=PETSc.InsertMode.ADD,
                      mode=PETSc.ScatterMode.REVERSE)
        set_bc(b, [bc])

        # Create solution function
        u = Function(V)

        # Create near null space basis and orthonormalize
        null_space = build_nullspace(V, u.vector)

        # Attached near-null space to matrix
        A.set_near_nullspace(null_space)

        # Test that basis is orthonormal
        assert null_space.is_orthonormal()

        # Create PETSC smoothed aggregation AMG preconditioner, and
        # create CG solver
        solver = PETSc.KSP().create(mesh.comm)
        solver.setType("cg")

        # Set matrix operator
        solver.setOperators(A)

        # Compute solution and return number of iterations
        return solver.solve(b, u.vector)
Exemplo n.º 18
0
def rebuild_expression_from_graph(G):
    "This is currently only used by tests."
    w = rebuild_with_scalar_subexpressions(G)

    # Find expressions of final v
    if len(w) == 1:
        return w[0]
    else:
        return as_vector(w)  # TODO: Consider shape of initial v
Exemplo n.º 19
0
    def split(self, fields):
        from ufl import as_vector, replace
        from firedrake import NonlinearVariationalProblem as NLVP, FunctionSpace
        splits = self._splits.get(tuple(fields))
        if splits is not None:
            return splits

        splits = []
        problem = self._problem
        splitter = ExtractSubBlock()
        for field in fields:
            try:
                if len(field) > 1:
                    raise NotImplementedError("Can't split into subblock")
            except TypeError:
                # Just a single field, we can handle that
                pass
            F = splitter.split(problem.F, argument_indices=(field, ))
            J = splitter.split(problem.J, argument_indices=(field, field))
            us = problem.u.split()
            subu = us[field]
            vec = []
            for i, u in enumerate(us):
                for idx in numpy.ndindex(u.ufl_shape):
                    vec.append(u[idx])
            u = as_vector(vec)
            F = replace(F, {problem.u: u})
            J = replace(J, {problem.u: u})
            if problem.Jp is not None:
                Jp = splitter.split(problem.Jp,
                                    argument_indices=(field, field))
                Jp = replace(Jp, {problem.u: u})
            else:
                Jp = None
            bcs = []
            for bc in problem.bcs:
                if bc.function_space().index == field:
                    V = FunctionSpace(subu.ufl_domain(), subu.ufl_element())
                    bcs.append(
                        type(bc)(V,
                                 bc.function_arg,
                                 bc.sub_domain,
                                 method=bc.method))
            new_problem = NLVP(
                F,
                subu,
                bcs=bcs,
                J=J,
                Jp=None,
                form_compiler_parameters=problem.form_compiler_parameters)
            new_problem._constant_jacobian = problem._constant_jacobian
            splits.append(
                type(self)(new_problem,
                           mat_type=self.mat_type,
                           pmat_type=self.pmat_type,
                           appctx=self.appctx))
        return self._splits.setdefault(tuple(fields), splits)
Exemplo n.º 20
0
def test_2D_lagrange_to_curl(order):
    mesh = create_unit_square(MPI.COMM_WORLD, 3, 4)
    V = FunctionSpace(mesh, ("N1curl", order))
    u = Function(V)

    W = FunctionSpace(mesh, ("Lagrange", order))
    u0 = Function(W)
    u0.interpolate(lambda x: -x[1])
    u1 = Function(W)
    u1.interpolate(lambda x: x[0])

    f = ufl.as_vector((u0, u1))
    f_expr = Expression(f, V.element.interpolation_points)
    u.interpolate(f_expr)
    x = ufl.SpatialCoordinate(mesh)
    f_ex = ufl.as_vector((-x[1], x[0]))
    assert np.isclose(
        np.abs(assemble_scalar(form(ufl.inner(u - f_ex, u - f_ex) * ufl.dx))),
        0)
Exemplo n.º 21
0
def test_ufl_only_to_contravariant_piola():
    mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1))
    V = ufl.FunctionSpace(mesh, ufl.FiniteElement("P", ufl.triangle, 2))
    v = ufl.Coefficient(V)
    expr = ufl.as_vector([v, v])
    W = ufl.FunctionSpace(mesh, ufl.FiniteElement("RT", ufl.triangle, 1))
    to_element = create_element(W.ufl_element())
    ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, *_ = compile_expression_dual_evaluation(
        expr, to_element, coffee=False)
    assert first_coeff_fake_coords is True
Exemplo n.º 22
0
def test_rank1(compile_args):
    """Tests evaluation of rank-1 form.

    Builds a linear operator which takes vector-valued functions in P1 space
    and evaluates expression [u_y, u_x] + grad(u_x) at specified points.

    """
    e = ufl.VectorElement("P", "triangle", 1)
    mesh = ufl.Mesh(e)

    V = ufl.FunctionSpace(mesh, e)
    u = ufl.TrialFunction(V)

    expr = ufl.as_vector([u[1], u[0]]) + ufl.grad(u[0])

    points = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]])
    obj, module = ffcx.codegeneration.jit.compile_expressions(
        [(expr, points)], cffi_extra_compile_args=compile_args)

    ffi = cffi.FFI()
    kernel = obj[0][0]

    c_type, np_type = float_to_type("double")

    # 2 components for vector components of expression
    # 3 points of evaluation
    # 6 degrees of freedom for rank1 form
    A = np.zeros((3, 2, 6), dtype=np_type)

    # Coefficient storage XYXYXY
    w = np.array([0.0], dtype=np_type)
    c = np.array([0.0], dtype=np_type)

    # Coords storage XYXYXY
    coords = np.array(points.flatten(), dtype=np.float64)
    kernel.tabulate_expression(
        ffi.cast('{type} *'.format(type=c_type), A.ctypes.data),
        ffi.cast('{type} *'.format(type=c_type), w.ctypes.data),
        ffi.cast('{type} *'.format(type=c_type), c.ctypes.data),
        ffi.cast('double *', coords.ctypes.data))

    f = np.array([[1.0, 2.0, 3.0], [-4.0, -5.0, 6.0]])

    # Apply the operator on some test input data
    u_ffcx = np.einsum("ijk,k", A, f.T.flatten())

    # Compute the correct values using NumPy
    # Gradf0 is gradient of f[0], each component of the gradient is constant
    gradf0 = np.array(
        [[f[0, 1] - f[0, 0], f[0, 1] - f[0, 0], f[0, 1] - f[0, 0]],
         [f[0, 2] - f[0, 0], f[0, 2] - f[0, 0], f[0, 2] - f[0, 0]]])

    u_correct = np.array([f[1], f[0]]) + gradf0

    assert np.allclose(u_ffcx, u_correct.T)
Exemplo n.º 23
0
def _uflToExpr(grid, order, f):
    if not ufl: return f
    if isinstance(f, list) or isinstance(f, tuple):
        if isinstance(f[0], ufl.core.expr.Expr):
            f = ufl.as_vector(f)
    if isinstance(f, GridFunction):
        return f
    if isinstance(f, ufl.core.expr.Expr):
        return expression2GF(grid, f, order).as_ufl()
    else:
        return f
Exemplo n.º 24
0
def a(phi):
    # phi = 1e5*phi
    n = ufl.grad(phi) / ufl.sqrt(ufl.dot(ufl.grad(phi), ufl.grad(phi)) + 1e-5)

    theta = ufl.atan_2(n[1], n[0])

    n = 1e5 * n + ufl.as_vector([1e-5, 1e-5])
    xx = n[1] / n[0]
    # theta = ufl.asin(xx/ ufl.sqrt(1 + xx**2))
    theta = ufl.atan(xx)

    return 1.0 + epsilon_4 * ufl.cos(m * (theta - theta_0))
Exemplo n.º 25
0
def test_curl(space_type, order):
    """Test that curl is consistent for different cell permutations of a tetrahedron."""

    tdim = cpp.mesh.cell_dim(CellType.tetrahedron)
    points = unit_cell_points(CellType.tetrahedron)

    spaces = []
    results = []
    cell = list(range(len(points)))
    # Assemble vector on 5 randomly numbered cells
    for i in range(5):
        shuffle(cell)

        domain = ufl.Mesh(
            ufl.VectorElement("Lagrange",
                              cpp.mesh.to_string(CellType.tetrahedron), 1))
        mesh = create_mesh(MPI.COMM_WORLD, [cell], points, domain)
        mesh.topology.create_connectivity_all()

        V = FunctionSpace(mesh, (space_type, order))
        v = ufl.TestFunction(V)

        f = ufl.as_vector(tuple(1 if i == 0 else 0 for i in range(tdim)))
        form = ufl.inner(f, ufl.curl(v)) * ufl.dx
        result = fem.assemble_vector(form)
        spaces.append(V)
        results.append(result.array)

    # Check that all DOFs on edges agree
    V = spaces[0]
    result = results[0]
    connectivity = V.mesh.topology.connectivity(1, 0)
    for i, edge in enumerate(V.mesh.topology.connectivity(tdim, 1).links(0)):
        vertices = connectivity.links(edge)
        values = sorted([
            result[V.dofmap.cell_dofs(0)[a]]
            for a in V.dofmap.dof_layout.entity_dofs(1, i)
        ])

        for s, r in zip(spaces[1:], results[1:]):
            c = s.mesh.topology.connectivity(1, 0)
            for j, e in enumerate(
                    s.mesh.topology.connectivity(tdim, 1).links(0)):
                if sorted(c.links(e)) == sorted(vertices):
                    v = sorted([
                        r[s.dofmap.cell_dofs(0)[a]]
                        for a in s.dofmap.dof_layout.entity_dofs(1, j)
                    ])
                    assert np.allclose(values, v)
                    break
            else:
                continue
            break
Exemplo n.º 26
0
def generateMethodBody(cppType, expr, returnResult, default, predefined):
    if expr is not None and not expr == [None]:
        try:
            dimR = expr.ufl_shape[0]
        except:
            if isinstance(expr,list) or isinstance(expr,tuple):
                expr = as_vector(expr)
            else:
                expr = as_vector([expr])
            dimR = expr.ufl_shape[0]

        _, coeff = extract_arguments_and_coefficients(expr)
        coeff = {c : c.toVectorCoefficient()[0] for c in coeff if len(c.ufl_shape) == 0 and not c.is_cellwise_constant()}
        expr = replace(expr, coeff)

        t = ExprTensor(expr.ufl_shape) # , exprTensorApply(lambda u: u, expr.ufl_shape, expr))
        expression = [expr[i] for i in t.keys()]
        u = extract_arguments_and_coefficients(expr)[0]
        if u != []:
            u = u[0]
            du = Grad(u)
            d2u = Grad(du)
            arg_u = Variable("const RangeType &", "u")
            arg_du = Variable("const JacobianRangeType &", "du")
            arg_d2u = Variable("const HessianRangeType &", "d2u")
            predefined.update( {u: arg_u, du: arg_du, d2u: arg_d2u} )
        code, results = generateCode(predefined, expression, tempVars=False)
        result = Variable(cppType, 'result')
        if cppType == 'double':
            code = code + [assign(result, results[0])]
        else:
            code = code + [assign(result[i], r) for i, r in zip(t.keys(), results)]
        if returnResult:
            code = [Declaration(result)] + code + [return_(result)]
    else:
        result = Variable(cppType, 'result')
        code = [assign(result, construct(cppType,default) )]
        if returnResult:
            code = [Declaration(result)] + code + [return_(result)]
    return code
Exemplo n.º 27
0
def dfInterpolate(self, f):
    if isinstance(f, list) or isinstance(f, tuple):
        if isinstance(f[0], ufl.core.expr.Expr):
            f = ufl.as_vector(f)
    dimExpr = 0
    if isinstance(f, GridFunction):
        func = f.gf
        dimExpr = func.dimRange
    elif isinstance(f, ufl.core.expr.Expr):
        func = expression2GF(self.space.gridView, f, self.space.order).as_ufl()
        if func.ufl_shape == ():
            dimExpr = 1
        else:
            dimExpr = func.ufl_shape[0]
    else:
        try:
            gl = len(inspect.getargspec(f)[0])
            func = None
        except TypeError:
            func = f
            if isinstance(func, int) or isinstance(func, float):
                dimExpr = 1
            else:
                dimExpr = len(func)
        if func is None:
            if gl == 1:  # global function
                func = function.globalFunction(self.space.gridView, "tmp",
                                               self.space.order, f).gf
            elif gl == 2:  # local function
                func = function.localFunction(self.space.gridView, "tmp",
                                              self.space.order, f).gf
            elif gl == 3:  # local function with self argument (i.e. from @gridFunction)
                func = function.localFunction(self.space.gridView, "tmp",
                                              self.space.order,
                                              lambda en, x: f(en, x)).gf
            dimExpr = func.dimRange

    if dimExpr == 0:
        raise AttributeError("can not determine if expression shape"\
                " fits the space's range dimension")
    elif dimExpr != self.space.dimRange:
        raise AttributeError("trying to interpolate an expression"\
            " of size "+str(dimExpr)+" into a space with range dimension = "\
            + str(self.space.dimRange))

    try:
        # API GridView: assert func.gridView == self.gridView, "can only interpolate with same grid views"
        assert func.gridView == self.gridView, "can only interpolate with same grid views"
        assert func.dimRange == self.dimRange, "range dimension mismatch"
    except AttributeError:
        pass
    return self._interpolate(func)
Exemplo n.º 28
0
Arquivo: fem.py Projeto: peide/tsfc
    def physical_edge_lengths(self):
        expr = ufl.classes.CellEdgeVectors(self.mt.terminal.ufl_domain())
        if self.mt.restriction == '+':
            expr = PositiveRestricted(expr)
        elif self.mt.restriction == '-':
            expr = NegativeRestricted(expr)

        expr = ufl.as_vector([ufl.sqrt(ufl.dot(expr[i, :], expr[i, :])) for i in range(3)])
        expr = preprocess_expression(expr)
        config = {"point_set": PointSingleton([1/3, 1/3])}
        config.update(self.config)
        context = PointSetContext(**config)
        return map_expr_dag(context.translator, expr)
Exemplo n.º 29
0
def rotate(v, angle, origin=None):
    """
    Rotate a UFL :class:`as_vector` ``v``
    by ``angle`` about an ``origin``.
    """
    dim = len(v)
    origin = origin or ufl.as_vector(np.zeros(dim))
    assert len(origin) == dim, "Origin does not match dimension"
    if dim == 2:
        R = rotation_matrix_2d(angle)
    else:
        raise NotImplementedError
    return ufl.dot(R, v - origin) + origin
Exemplo n.º 30
0
 def argument(self, arg):
     """Split an argument into its constituent spaces."""
     if isinstance(arg.function_space(), functionspace.MixedFunctionSpace):
         if arg in self._args:
             return self._args[arg]
         args = []
         for i, fs in enumerate(arg.function_space().split()):
             # Look up the split argument in cache since we want it unique
             a = Argument(fs, arg.number(), part=arg.part())
             if a.shape():
                 if self._idx[arg.number()] == i:
                     args += [a[j] for j in range(a.shape()[0])]
                 else:
                     args += [Zero() for j in range(a.shape()[0])]
             else:
                 if self._idx[arg.number()] == i:
                     args.append(a)
                 else:
                     args.append(Zero())
         self._args[arg] = as_vector(args)
         return as_vector(args)
     return arg
Exemplo n.º 31
0
 def argument(self, arg):
     """Split an argument into its constituent spaces."""
     if isinstance(arg.function_space(), functionspace.MixedFunctionSpace):
         if arg in self._args:
             return self._args[arg]
         args = []
         for i, fs in enumerate(arg.function_space().split()):
             # Look up the split argument in cache since we want it unique
             a = Argument(fs.ufl_element(), fs, arg.number())
             if a.shape():
                 if self._idx[arg.number()] == i:
                     args += [a[j] for j in range(a.shape()[0])]
                 else:
                     args += [Zero() for j in range(a.shape()[0])]
             else:
                 if self._idx[arg.number()] == i:
                     args.append(a)
                 else:
                     args.append(Zero())
         self._args[arg] = as_vector(args)
         return as_vector(args)
     return arg
Exemplo n.º 32
0
def get_scaled_jacobians2d(mesh, python=False):
    """
    Computes the scaled Jacobian of each cell in a 2D triangular mesh

    :arg mesh: the input mesh to do computations on
    :kwarg python: compute the measure using Python?

    :rtype: firedrake.function.Function scaled_jacobians with scaled
        jacobian data.
    """
    P0 = firedrake.FunctionSpace(mesh, "DG", 0)
    if python:
        P0_ten = firedrake.TensorFunctionSpace(mesh, "DG", 0)
        J = firedrake.interpolate(ufl.Jacobian(mesh), P0_ten)
        edge1 = ufl.as_vector([J[0, 0], J[1, 0]])
        edge2 = ufl.as_vector([J[0, 1], J[1, 1]])
        edge3 = edge1 - edge2
        a = ufl.sqrt(ufl.dot(edge1, edge1))
        b = ufl.sqrt(ufl.dot(edge2, edge2))
        c = ufl.sqrt(ufl.dot(edge3, edge3))
        detJ = ufl.JacobianDeterminant(mesh)
        jacobian_sign = ufl.sign(detJ)
        max_product = ufl.Max(
            ufl.Max(ufl.Max(a * b, a * c), ufl.Max(b * c, b * a)), ufl.Max(c * a, c * b)
        )
        scaled_jacobians = firedrake.interpolate(detJ / max_product * jacobian_sign, P0)
    else:
        coords = mesh.coordinates
        scaled_jacobians = firedrake.Function(P0)
        op2.par_loop(
            get_pyop2_kernel("get_scaled_jacobian", 2),
            mesh.cell_set,
            scaled_jacobians.dat(op2.WRITE, scaled_jacobians.cell_node_map()),
            coords.dat(op2.READ, coords.cell_node_map()),
        )
    return scaled_jacobians
Exemplo n.º 33
0
    def __init__(self, angular_quad, L):
      from transport_data import angular_tensors_ext_module

      i,j,k1,k2,p,q = ufl.indices(6)

      tensors = angular_tensors_ext_module.AngularTensors(angular_quad, L)
      self.Y = ufl.as_tensor( numpy.reshape(tensors.Y(), tensors.shape_Y()) )
      self.Q = ufl.as_tensor( numpy.reshape(tensors.Q(), tensors.shape_Q()) )
      self.QT = ufl.transpose(self.Q)
      self.Qt = ufl.as_tensor( numpy.reshape(tensors.Qt(), tensors.shape_Qt()) )
      self.QtT = ufl.as_tensor( self.Qt[k1,p,i], (p,i,k1) )
      self.G = ufl.as_tensor( numpy.reshape(tensors.G(), tensors.shape_G()) )
      self.T = ufl.as_tensor( numpy.reshape(tensors.T(), tensors.shape_T()) )
      self.Wp = ufl.as_vector( angular_quad.get_pw() )
      self.W = ufl.diag(self.Wp)
Exemplo n.º 34
0
    def split(self, fields):
        from ufl import as_vector, replace
        from firedrake import NonlinearVariationalProblem as NLVP, FunctionSpace
        splits = self._splits.get(tuple(fields))
        if splits is not None:
            return splits

        splits = []
        problem = self._problem
        splitter = ExtractSubBlock()
        for field in fields:
            try:
                if len(field) > 1:
                    raise NotImplementedError("Can't split into subblock")
            except TypeError:
                # Just a single field, we can handle that
                pass
            F = splitter.split(problem.F, argument_indices=(field, ))
            J = splitter.split(problem.J, argument_indices=(field, field))
            us = problem.u.split()
            subu = us[field]
            vec = []
            for i, u in enumerate(us):
                for idx in numpy.ndindex(u.ufl_shape):
                    vec.append(u[idx])
            u = as_vector(vec)
            F = replace(F, {problem.u: u})
            J = replace(J, {problem.u: u})
            if problem.Jp is not None:
                Jp = splitter.split(problem.Jp, argument_indices=(field, field))
                Jp = replace(Jp, {problem.u: u})
            else:
                Jp = None
            bcs = []
            for bc in problem.bcs:
                if bc.function_space().index == field:
                    V = FunctionSpace(subu.ufl_domain(), subu.ufl_element())
                    bcs.append(type(bc)(V,
                                        bc.function_arg,
                                        bc.sub_domain,
                                        method=bc.method))
            new_problem = NLVP(F, subu, bcs=bcs, J=J, Jp=None,
                               form_compiler_parameters=problem.form_compiler_parameters)
            new_problem._constant_jacobian = problem._constant_jacobian
            splits.append(type(self)(new_problem, mat_type=self.mat_type, pmat_type=self.pmat_type,
                                     appctx=self.appctx))
        return self._splits.setdefault(tuple(fields), splits)
Exemplo n.º 35
0
  def solve(self, it=0):
    if self.group_GS:
      self.solve_group_GS(it)
    else:
      super(EllipticSNFluxModule, self).solve(it)
      self.slns_mg = split(self.sln)

    i,p,q,k1,k2 = ufl.indices(5)

    sol_timer = Timer("-- Complete solution")
    aux_timer = Timer("---- SOL: Computing angular flux + adjoint")

    # TODO: Move to Discretization
    V11 = FunctionSpace(self.DD.mesh, "CG", self.DD.parameters["p"])

    for gto in range(self.DD.G):
      self.PD.get_xs('D', self.D, gto)

      form = self.D * ufl.diag_vector(ufl.as_matrix(self.DD.ordinates_matrix[i,p]*self.slns_mg[gto][q].dx(i), (p,q)))

      for gfrom in range(self.DD.G):
        pres_Ss = False

        # TODO: Enlarge self.S and self.C to (L+1)^2 (or 1./2.*(L+1)*(L+2) in 2D) to accomodate for anisotropic
        # scattering (lines below using CC, SS are correct only for L = 0, when the following inner loop runs only
        # once.
        for l in range(self.L+1):
          for m in range(-l, l+1):
            if self.DD.angular_quad.get_D() == 2 and divmod(l+m, 2)[1] == 0:
              continue

            pres_Ss |= self.PD.get_xs('Ss', self.S[l], gto, gfrom, l)
            self.PD.get_xs('C', self.C[l], gto, gfrom, l)

        if pres_Ss:
          Cd = ufl.diag(self.C)
          CC = self.tensors.Y[p,k1] * Cd[k1,k2] * self.tensors.Qt[k2,q,i]

          form += ufl.as_vector(CC[p,q,i] * self.slns_mg[gfrom][q].dx(i), p)

      # project(form, self.DD.Vpsi1, function=self.aux_slng, preconditioner_type="petsc_amg")
      # FASTER, but requires form compilation for each dir.:
      for pp in range(self.DD.M):
        assign(self.aux_slng.sub(pp), project(form[pp], V11, preconditioner_type="petsc_amg"))

      self.psi_mg[gto].assign(self.slns_mg[gto] + self.aux_slng)
      self.adj_psi_mg[gto].assign(self.slns_mg[gto] - self.aux_slng)
def test_dolfin_expression_compilation_of_vector_literal():
    # Define some literal ufl expression
    uexpr = ufl.as_vector((1.23, 7.89))

    # Define expected output from compilation
    expected_lines = ['double s[2];',
                      's[0] = 1.23;',
                      's[1] = 7.89;',
                      'values[0] = s[0];',
                      'values[1] = s[1];']

    # Define expected evaluation values: [(x,value), (x,value), ...]
    expected_values = [((0.0, 0.0), (1.23, 7.89)),
                       ((0.6, 0.7), (1.23, 7.89)),
                       ]

    # Execute all tests
    check_dolfin_expression_compilation(uexpr, expected_lines, expected_values)
Exemplo n.º 37
0
    def argument(self, o):
        from ufl import split
        from firedrake import MixedFunctionSpace, FunctionSpace
        V = o.function_space()
        if len(V) == 1:
            # Not on a mixed space, just return ourselves.
            return o

        if o in self._arg_cache:
            return self._arg_cache[o]

        V_is = V.split()
        indices = self.blocks[o.number()]

        try:
            indices = tuple(indices)
            nidx = len(indices)
        except TypeError:
            # Only one index provided.
            indices = (indices, )
            nidx = 1

        if nidx == 1:
            W = V_is[indices[0]]
            W = FunctionSpace(W.mesh(), W.ufl_element())
            a = (Argument(W, o.number(), part=o.part()), )
        else:
            W = MixedFunctionSpace([V_is[i] for i in indices])
            a = split(Argument(W, o.number(), part=o.part()))
        args = []
        for i in range(len(V_is)):
            if i in indices:
                c = indices.index(i)
                a_ = a[c]
                if len(a_.ufl_shape) == 0:
                    args += [a_]
                else:
                    args += [a_[j] for j in numpy.ndindex(a_.ufl_shape)]
            else:
                args += [Zero()
                         for j in numpy.ndindex(
                         V_is[i].ufl_element().value_shape())]
        return self._arg_cache.setdefault(o, as_vector(args))
Exemplo n.º 38
0
def _find_linear_terms(rhs_exprs, u):
    """Help function that takes a list of rhs expressions and return a
    list of bools determining what component, rhs_exprs[i], is linear
    wrt u[i].

    """

    uu = [Constant(1.0) for _ in rhs_exprs]
    if len(rhs_exprs) > 1:
        repl = {u: ufl.as_vector(uu)}
    else:
        repl = {u: uu[0]}

    linear_terms = []
    for i, ui in enumerate(uu):
        comp_i_s = expand_indices(ufl.replace(rhs_exprs[i], repl))
        linear_terms.append(ui in extract_coefficients(comp_i_s) and
                            ui not in extract_coefficients(
                                expand_derivatives(ufl.diff(comp_i_s, ui))))
    return linear_terms
Exemplo n.º 39
0
 def argument(self, o):
     V = o.function_space()
     if len(V) == 1:
         # Not on a mixed space, just return ourselves.
         return o
     # Already seen this argument, return the cached version.
     if o in self._args:
         return self._args[o]
     args = []
     for i, V_i in enumerate(V.split()):
         # Walk over the subspaces and build a vector that is zero
         # where the argument does not match the one we're looking
         # for and is just the non-mixed argument when we do want
         # it.
         a = Argument(V_i, o.number(), part=o.part())
         indices = numpy.ndindex(a.ufl_shape)
         if self.idx[o.number()] == i:
             args += [a[j] for j in indices]
         else:
             args += [Zero() for j in indices]
     self._args[o] = as_vector(args)
     return self._args[o]
Exemplo n.º 40
0
def _rush_larsen_scheme_generator(rhs_form, solution, time, order, generalized):
    """Generates a list of forms and solutions for a given Butcher
    tableau

    *Arguments*
        rhs_form (ufl.Form)
            A UFL form representing the rhs for a time differentiated equation
        solution (_Function_)
            The prognostic variable
        time (_Constant_)
            A Constant holding the time at the start of the time step
        order (int)
            The order of the scheme
        generalized (bool)
            If True generate a generalized Rush Larsen scheme, linearizing all
            components.

    """

    DX = _check_form(rhs_form)

    if DX != ufl.dP:
        raise TypeError("Expected a form with a Pointintegral.")

    # Create time step
    dt = Constant(0.1)

    # Get test function
    #    arguments = rhs_form.arguments()
    #    coefficients = rhs_form.coefficients()

    # Get time dependent expressions
    time_dep_expressions = _time_dependent_expressions(rhs_form, time)

    # Extract rhs expressions from form
    rhs_integrand = rhs_form.integrals()[0].integrand()
    rhs_exprs, v = extract_tested_expressions(rhs_integrand)
    vector_rhs = len(v.ufl_shape) > 0 and v.ufl_shape[0] > 1

    system_size = v.ufl_shape[0] if vector_rhs else 1

    # Fix for indexing of v for scalar expressions
    v = v if vector_rhs else [v]

    # Extract linear terms if not using generalized Rush Larsen
    if not generalized:
        linear_terms = _find_linear_terms(rhs_exprs, solution)
    else:
        linear_terms = [True for _ in range(system_size)]

    # Wrap the rhs expressions into a ufl vector type
    rhs_exprs = ufl.as_vector([rhs_exprs[i] for i in range(system_size)])
    rhs_jac = ufl.diff(rhs_exprs, solution)

    # Takes time!
    if vector_rhs:
        diff_rhs_exprs = [expand_indices(expand_derivatives(rhs_jac[ind, ind]))
                          for ind in range(system_size)]
    else:
        diff_rhs_exprs = [expand_indices(expand_derivatives(rhs_jac[0]))]
        solution = [solution]

    ufl_stage_forms = []
    dolfin_stage_forms = []
    dt_stage_offsets = []

    # Stage solutions (3 per order rhs, linearized, and final step)
    # If 2nd order the final step for 1 step is a stage
    if order == 1:
        stage_solutions = []
        rl_ufl_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs, linear_terms,
                                        system_size, solution, None, dt, time, 1.0,
                                        0.0, v, DX, time_dep_expressions)
    elif order == 2:

        # Stage solution for order 2
        if vector_rhs:
            stage_solutions = [Function(solution.function_space(), name="y_1/2")]
        else:
            stage_solutions = [Function(solution[0].function_space(), name="y_1/2")]

        stage_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
                                       linear_terms, system_size,
                                       solution, None, dt, time, 0.5,
                                       0.0, v, DX,
                                       time_dep_expressions)

        rl_ufl_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
                                        linear_terms, system_size,
                                        solution, stage_solutions[0],
                                        dt, time, 1.0, 0.5, v, DX,
                                        time_dep_expressions)

        ufl_stage_forms.append([stage_form])
        dolfin_stage_forms.append([Form(stage_form)])

    # Get last stage form
    last_stage = Form(rl_ufl_form)

    human_form = "%srush larsen %s" % ("generalized " if generalized else "",
                                       str(order))

    return rhs_form, linear_terms, ufl_stage_forms, dolfin_stage_forms, last_stage, \
        stage_solutions, dt, dt_stage_offsets, human_form, None
Exemplo n.º 41
0
def _rush_larsen_scheme_generator_adm(rhs_form, solution, time, order,
                                      generalized, perturbation):
    """Generates a list of forms and solutions for the adjoint
    linearisation of the Rush-Larsen scheme

    *Arguments*
        rhs_form (ufl.Form)
            A UFL form representing the rhs for a time differentiated equation
        solution (_Function_)
            The prognostic variable
        time (_Constant_)
            A Constant holding the time at the start of the time step
        order (int)
            The order of the scheme
        generalized (bool)
            If True generate a generalized Rush Larsen scheme, linearizing all
            components.
        perturbation (Function)
            The vector on which we compute the adjoint action.

    """

    DX = _check_form(rhs_form)

    if DX != ufl.dP:
        raise TypeError("Expected a form with a Pointintegral.")

    # Create time step
    dt = Constant(0.1)

    # Get test function
    #    arguments = rhs_form.arguments()
    #    coefficients = rhs_form.coefficients()

    # Get time dependent expressions
    time_dep_expressions = _time_dependent_expressions(rhs_form, time)

    # Extract rhs expressions from form
    rhs_integrand = rhs_form.integrals()[0].integrand()
    rhs_exprs, v = extract_tested_expressions(rhs_integrand)
    vector_rhs = len(v.ufl_shape) > 0 and v.ufl_shape[0] > 1

    system_size = v.ufl_shape[0] if vector_rhs else 1

    # Fix for indexing of v for scalar expressions
    v = v if vector_rhs else [v]

    # Extract linear terms if not using generalized Rush Larsen
    if not generalized:
        linear_terms = _find_linear_terms(rhs_exprs, solution)
    else:
        linear_terms = [True for _ in range(system_size)]

    # Wrap the rhs expressions into a ufl vector type
    rhs_exprs = ufl.as_vector([rhs_exprs[i] for i in range(system_size)])
    rhs_jac = ufl.diff(rhs_exprs, solution)

    # Takes time!
    if vector_rhs:
        diff_rhs_exprs = [expand_indices(expand_derivatives(rhs_jac[ind, ind]))
                          for ind in range(system_size)]
        soln = solution
    else:
        diff_rhs_exprs = [expand_indices(expand_derivatives(rhs_jac[0]))]
        solution = [solution]
        soln = solution[0]

    ufl_stage_forms = []
    dolfin_stage_forms = []
    dt_stage_offsets = []
    trial = TrialFunction(soln.function_space())

    # Stage solutions (3 per order rhs, linearized, and final step)
    # If 2nd order the final step for 1 step is a stage
    if order == 1:
        stage_solutions = []

        # Fetch the original step
        rl_ufl_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
                                        linear_terms, system_size,
                                        solution, None, dt, time, 1.0,
                                        0.0, v, DX,
                                        time_dep_expressions)

        # If this is commented out, we don't get NaNs.  Yhy is
        # solution a list of length zero anyway?
        rl_ufl_form = safe_action(safe_adjoint(derivative(rl_ufl_form, soln, trial)),
                                  perturbation)

    elif order == 2:
        # Stage solution for order 2
        fn_space = soln.function_space()

        stage_solutions = [Function(fn_space, name="y_1/2"),
                           Function(fn_space, name="y_1"),
                           Function(fn_space, name="y_bar_1/2")]

        y_half_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
                                        linear_terms, system_size,
                                        solution, None, dt, time, 0.5,
                                        0.0, v, DX,
                                        time_dep_expressions)

        y_one_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
                                       linear_terms, system_size,
                                       solution, stage_solutions[0],
                                       dt, time, 1.0, 0.5, v, DX,
                                       time_dep_expressions)

        y_bar_half_form = safe_action(safe_adjoint(derivative(y_one_form,
                                                              stage_solutions[0], trial)), perturbation)

        rl_ufl_form = safe_action(safe_adjoint(derivative(y_one_form, soln, trial)), perturbation) + \
            safe_action(safe_adjoint(derivative(y_half_form, soln, trial)), stage_solutions[2])

        ufl_stage_forms.append([y_half_form])
        ufl_stage_forms.append([y_one_form])
        ufl_stage_forms.append([y_bar_half_form])
        dolfin_stage_forms.append([Form(y_half_form)])
        dolfin_stage_forms.append([Form(y_one_form)])
        dolfin_stage_forms.append([Form(y_bar_half_form)])

    # Get last stage form
    last_stage = Form(rl_ufl_form)

    human_form = "%srush larsen %s" % ("generalized " if generalized else "",
                                       str(order))

    return rhs_form, linear_terms, ufl_stage_forms, dolfin_stage_forms, last_stage, \
        stage_solutions, dt, dt_stage_offsets, human_form, perturbation