def test_direct(self, s1, d1):

        def fn(a):
            a[:] = 1.0

        op2.par_loop(fn, s1, d1(op2.WRITE))
        assert np.allclose(d1.data, 1.0)

 def test_linalg_and_parloop(self, backend, x, y):
     """Linear algebra operators should force computation"""
     x._data = np.zeros(x.dataset.total_size, dtype=np.float64)
     k = op2.Kernel('void k(double *x) { *x = 1.0; }', 'k')
     op2.par_loop(k, x.dataset.set, x(op2.WRITE))
     z = x + y
     assert all(z.data == y.data + 1)

    def test_assemble_rhs(self, rhs, elements, b, coords, f, elem_node, expected_rhs):
        """Assemble a simple finite-element right-hand side and check result."""
        b.zero()
        op2.par_loop(rhs, elements, b(op2.INC, elem_node), coords(op2.READ, elem_node), f(op2.READ, elem_node))

        eps = 1.0e-12
        assert_allclose(b.data, expected_rhs, eps)

 def test_uninitialized_map(self, backend, iterset, indset, x):
     """Accessing a par_loop argument via an uninitialized Map should raise
     an exception."""
     kernel_wo = "void kernel_wo(unsigned int* x) { *x = 42; }\n"
     with pytest.raises(MapValueError):
         op2.par_loop(op2.Kernel(kernel_wo, "kernel_wo"), iterset,
                      x(op2.WRITE, op2.Map(iterset, indset, 1)))

    def test_onecolor_rw(self, backend, iterset, x, iterset2indset):
        """Increment each value of a Dat by one with op2.RW."""
        kernel_rw = "void kernel_rw(unsigned int* x) { (*x) = (*x) + 1; }\n"

        op2.par_loop(op2.Kernel(kernel_rw, "kernel_rw"),
                     iterset, x(op2.RW, iterset2indset[0]))
        assert sum(x.data) == nelems * (nelems + 1) / 2

    def test_invert_arg_similar_shape(self, iterset, iter2ind1, x, y):
        self.cache.clear()
        assert len(self.cache) == 0

        kernel_swap = """
static void swap(unsigned int* x, unsigned int* y)
{
  unsigned int t;
  t = *x;
  *x = *y;
  *y = t;
}
"""
        op2.par_loop(op2.Kernel(kernel_swap, "swap"),
                     iterset,
                     x(op2.RW, iter2ind1),
                     y(op2.RW, iter2ind1))

        base._trace.evaluate(set([x]), set())
        assert len(self.cache) == 1

        op2.par_loop(op2.Kernel(kernel_swap, "swap"),
                     iterset,
                     y(op2.RW, iter2ind1),
                     x(op2.RW, iter2ind1))

        base._trace.evaluate(set([y]), set())
        assert len(self.cache) == 1

    def test_vector_map(self, iterset, x2, iter2ind2):
        self.cache.clear()
        assert len(self.cache) == 0

        kernel_swap = """
static void swap(unsigned int* x)
{
  unsigned int t;
  t = x[0];
  x[0] = x[1];
  x[1] = t;
}
"""

        op2.par_loop(op2.Kernel(kernel_swap, "swap"),
                     iterset,
                     x2(op2.RW, iter2ind2))

        base._trace.evaluate(set([x2]), set())
        assert len(self.cache) == 1

        op2.par_loop(op2.Kernel(kernel_swap, "swap"),
                     iterset,
                     x2(op2.RW, iter2ind2))

        base._trace.evaluate(set([x2]), set())
        assert len(self.cache) == 1

    def test_sum_nodes_to_edges(self, backend):
        """Creates a 1D grid with edge values numbered consecutively.
        Iterates over edges, summing the node values."""

        nedges = nnodes - 1
        nodes = op2.Set(nnodes, "nodes")
        edges = op2.Set(nedges, "edges")

        node_vals = op2.Dat(nodes, numpy.arange(nnodes, dtype=numpy.uint32), numpy.uint32, "node_vals")
        edge_vals = op2.Dat(edges, numpy.zeros(nedges, dtype=numpy.uint32), numpy.uint32, "edge_vals")

        e_map = numpy.array([(i, i + 1) for i in range(nedges)], dtype=numpy.uint32)
        edge2node = op2.Map(edges, nodes, 2, e_map, "edge2node")

        kernel_sum = """
void kernel_sum(unsigned int* nodes, unsigned int *edge, int i)
{ *edge += nodes[0]; }
"""

        op2.par_loop(
            op2.Kernel(kernel_sum, "kernel_sum"), edges, node_vals(op2.READ, edge2node[op2.i[0]]), edge_vals(op2.INC)
        )

        expected = numpy.arange(1, nedges * 2 + 1, 2)
        assert all(expected == edge_vals.data)

 def test_assemble_mat(self, backend, mass, mat, coords, elements, elem_node,
                       elem_vnode, expected_matrix):
     op2.par_loop(mass, elements(3,3),
                  mat((elem_node[op2.i[0]], elem_node[op2.i[1]]), op2.INC),
                  coords(elem_vnode, op2.READ))
     eps=1.e-5
     assert_allclose(mat.values, expected_matrix, eps)

    def test_invert_arg_similar_shape(self, backend, iterset, iter2ind1, x, y):
        self.cache.clear()
        assert len(self.cache) == 0

        kernel_swap = """
void kernel_swap(unsigned int* x, unsigned int* y)
{
  unsigned int t;
  t = *x;
  *x = *y;
  *y = t;
}
"""
        op2.par_loop(op2.Kernel(kernel_swap, "kernel_swap"),
                     iterset,
                     x(iter2ind1[0], op2.RW),
                     y(iter2ind1[0], op2.RW))

        assert len(self.cache) == 1

        op2.par_loop(op2.Kernel(kernel_swap, "kernel_swap"),
                     iterset,
                     y(iter2ind1[0], op2.RW),
                     x(iter2ind1[0], op2.RW))

        assert len(self.cache) == 1

    def test_1d_inc_no_data(self, k1_inc_to_global, set, d1):
        g = op2.Global(1, dtype=numpy.uint32)
        op2.par_loop(k1_inc_to_global, set,
                     g(op2.INC),
                     d1(op2.READ))

        assert g.data == d1.data.sum()

    def test_idx_order(self, backend, iterset, iter2ind2, x):
        self.cache.clear()
        assert len(self.cache) == 0

        kernel_swap = """
void kernel_swap(unsigned int* x, unsigned int* y)
{
  unsigned int t;
  t = *x;
  *x = *y;
  *y = t;
}
"""
        op2.par_loop(op2.Kernel(kernel_swap, "kernel_swap"),
                     iterset,
                     x(iter2ind2[0], op2.RW),
                     x(iter2ind2[1], op2.RW))

        assert len(self.cache) == 1

        op2.par_loop(op2.Kernel(kernel_swap, "kernel_swap"),
                     iterset,
                     x(iter2ind2[1], op2.RW),
                     x(iter2ind2[0], op2.RW))

        assert len(self.cache) == 1

def get_restriction_weights(coarse, fine):
    mesh = coarse.mesh()
    assert hasattr(mesh, "_shared_data_cache")
    cache = mesh._shared_data_cache["hierarchy_restriction_weights"]
    key = entity_dofs_key(coarse.fiat_element.entity_dofs())
    try:
        return cache[key]
    except KeyError:
        # We hit each fine dof more than once since we loop
        # elementwise over the coarse cells.  So we need a count of
        # how many times we did this to weight the final contribution
        # appropriately.
        if not (coarse.ufl_element() == fine.ufl_element()):
            raise ValueError("Can't transfer between different spaces")
        if coarse.fiat_element.entity_dofs() == coarse.fiat_element.entity_closure_dofs():
            return cache.setdefault(key, None)
        ele = coarse.ufl_element()
        if isinstance(ele, ufl.VectorElement):
            ele = ele.sub_elements()[0]
            weights = firedrake.Function(firedrake.FunctionSpace(fine.mesh(), ele))
        else:
            weights = firedrake.Function(fine)
        c2f_map = coarse_to_fine_node_map(coarse, fine)
        kernel = get_count_kernel(c2f_map.arity)
        op2.par_loop(kernel, op2.LocalSet(mesh.cell_set),
                     weights.dat(op2.INC, c2f_map[op2.i[0]]))
        weights.assign(1/weights)
        return cache.setdefault(key, weights)

    def test_indirect(self, s1, d2, m12):

        def fn(a):
            a[0] = 1.0

        op2.par_loop(fn, s1, d2(op2.WRITE, m12))
        assert np.allclose(d2.data, 1.0)

    def test_2d_read(self, backend, k2_write_to_dat, set, d1):
        g = op2.Global(2, (1, 2), dtype=numpy.uint32)
        op2.par_loop(k2_write_to_dat, set,
                     d1(op2.IdentityMap, op2.WRITE),
                     g(op2.READ))

        assert all(d1.data == g.data.sum())

    def test_inc_2d_itspace_map(self, backend, d2, vd2, node2ele, node):
        vd2.data[:, 0] = 3
        vd2.data[:, 1] = 4
        d2.data[:] = numpy.arange(2 * nnodes).reshape(d2.data.shape)

        incs = Block([Incr(Symbol("vd", ("i",), ((1, 0),)), Symbol("d", (0,))),
                      Incr(
                          Symbol("vd", ("i",), ((1, 1),)), Symbol("d", (1,)))],
                     open_scope=True)
        k = FunDecl("void", "k",
                    [Decl("int*", c_sym("d")), Decl("int*", c_sym("vd"))],
                    c_for("i", 1, incs))

        op2.par_loop(op2.Kernel(k, 'k'), node,
                     d2(op2.READ),
                     vd2(op2.INC, node2ele[op2.i[0]]))

        expected = numpy.zeros_like(vd2.data)
        expected[:, 0] = 3
        expected[:, 1] = 4
        expected[:, 0] += numpy.arange(start=0, stop=2 * nnodes, step=4)
        expected[:, 0] += numpy.arange(start=2, stop=2 * nnodes, step=4)
        expected[:, 1] += numpy.arange(start=1, stop=2 * nnodes, step=4)
        expected[:, 1] += numpy.arange(start=3, stop=2 * nnodes, step=4)
        assert all(vd2.data[:, 0] == expected[:, 0])
        assert all(vd2.data[:, 1] == expected[:, 1])

    def test_1d_inc(self, backend, k1_inc_to_global, set, d1):
        g = op2.Global(1, 0, dtype=numpy.uint32)
        op2.par_loop(k1_inc_to_global, set,
                     d1(op2.IdentityMap, op2.READ),
                     g(op2.INC))

        assert g.data == d1.data.sum()

    def test_sum_nodes_to_edges(self, backend):
        """Creates a 1D grid with edge values numbered consecutively.
        Iterates over edges, summing the node values."""

        nedges = nnodes - 1
        nodes = op2.Set(nnodes, "nodes")
        edges = op2.Set(nedges, "edges")

        node_vals = op2.Dat(nodes, numpy.arange(
            nnodes, dtype=numpy.uint32), numpy.uint32, "node_vals")
        edge_vals = op2.Dat(
            edges, numpy.zeros(nedges, dtype=numpy.uint32), numpy.uint32, "edge_vals")

        e_map = numpy.array([(i, i + 1)
                            for i in range(nedges)], dtype=numpy.uint32)
        edge2node = op2.Map(edges, nodes, 2, e_map, "edge2node")

        kernel_sum = FunDecl("void", "kernel_sum",
                             [Decl(
                                 "int*", c_sym("nodes"), qualifiers=["unsigned"]),
                              Decl(
                                  "int*", c_sym("edge"), qualifiers=["unsigned"])],
                             c_for("i", 2, Incr(c_sym("*edge"), Symbol("nodes", ("i",)))))

        op2.par_loop(op2.Kernel(kernel_sum, "kernel_sum"), edges,
                     node_vals(op2.READ, edge2node[op2.i[0]]),
                     edge_vals(op2.INC))

        expected = numpy.arange(1, nedges * 2 + 1, 2)
        assert all(expected == edge_vals.data)

    def test_dloop_ignore_scalar(self, iterset, a, b):
        self.cache.clear()
        assert len(self.cache) == 0

        kernel_swap = """
static void swap(unsigned int* x, unsigned int* y)
{
  unsigned int t;
  t = *x;
  *x = *y;
  *y = t;
}
"""
        op2.par_loop(op2.Kernel(kernel_swap, "swap"),
                     iterset,
                     a(op2.RW),
                     b(op2.RW))

        base._trace.evaluate(set([a]), set())
        assert len(self.cache) == 1

        op2.par_loop(op2.Kernel(kernel_swap, "swap"),
                     iterset,
                     b(op2.RW),
                     a(op2.RW))

        base._trace.evaluate(set([b]), set())
        assert len(self.cache) == 1

    def test_1d_read(self, backend, k1_write_to_dat, set, d1):
        g = op2.Global(1, 1, dtype=numpy.uint32)
        op2.par_loop(k1_write_to_dat, set,
                     d1(op2.WRITE),
                     g(op2.READ))

        assert all(d1.data == g.data)