1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
|
# Copyright (C) 2021 Chris Richardson and Igor Baratta
#
# This file is part of DOLFINx (https://www.fenicsproject.org)
#
# SPDX-License-Identifier: LGPL-3.0-or-later
"""Unit tests for the KrylovSolver interface"""
from mpi4py import MPI
import numpy as np
import pytest
from basix.ufl import element
from dolfinx import default_real_type, la
from dolfinx.fem import Function, functionspace
from dolfinx.mesh import create_unit_square
@pytest.mark.parametrize(
"e",
[
element("Lagrange", "triangle", 1, dtype=default_real_type),
element("Lagrange", "triangle", 1, shape=(2,), dtype=default_real_type),
],
)
def test_scatter_forward(e):
mesh = create_unit_square(MPI.COMM_WORLD, 5, 5)
V = functionspace(mesh, e)
u = Function(V)
bs = V.dofmap.bs
u.interpolate(lambda x: [x[i] for i in range(bs)])
# Forward scatter should have no effect
w0 = u.x.array.copy()
u.x.scatter_forward()
assert np.allclose(w0, u.x.array)
# Fill local array with the mpi rank
u.x.array.fill(MPI.COMM_WORLD.rank)
w0 = u.x.array.copy()
u.x.scatter_forward()
# Now the ghosts should have the value of the rank of the owning
# process
ghost_owners = u.function_space.dofmap.index_map.owners
ghost_owners = np.repeat(ghost_owners, bs)
local_size = u.function_space.dofmap.index_map.size_local * bs
assert np.allclose(u.x.array[local_size:], ghost_owners)
@pytest.mark.parametrize(
"e",
[
element("Lagrange", "triangle", 1, dtype=default_real_type),
element("Lagrange", "triangle", 1, shape=(2,), dtype=default_real_type),
],
)
def test_scatter_reverse(e):
comm = MPI.COMM_WORLD
mesh = create_unit_square(MPI.COMM_WORLD, 5, 5)
V = functionspace(mesh, e)
u = Function(V)
bs = V.dofmap.bs
u.interpolate(lambda x: [x[i] for i in range(bs)])
# Reverse scatter (insert) should have no effect
w0 = u.x.array.copy()
u.x.scatter_reverse(la.InsertMode.insert)
assert np.allclose(w0, u.x.array)
# Fill with MPI rank, and sum all entries in the vector (including
# ghosts)
u.x.array.fill(comm.rank)
all_count0 = MPI.COMM_WORLD.allreduce(u.x.array.sum(), op=MPI.SUM)
# Reverse scatter (add)
u.x.scatter_reverse(la.InsertMode.add)
num_ghosts = V.dofmap.index_map.num_ghosts
ghost_count = MPI.COMM_WORLD.allreduce(num_ghosts * comm.rank, op=MPI.SUM)
# New count should have gone up by the number of ghosts times their
# rank on all processes
all_count1 = MPI.COMM_WORLD.allreduce(u.x.array.sum(), op=MPI.SUM)
assert all_count1 == (all_count0 + bs * ghost_count)
@pytest.mark.parametrize(
"dtype",
[
np.float32,
np.float64,
np.complex64,
np.complex128,
np.int8,
np.int32,
np.int64,
],
)
def test_vector_from_index_map_scatter_forward(dtype):
comm = MPI.COMM_WORLD
mesh = create_unit_square(comm, 5, 5)
for d in range(mesh.topology.dim):
mesh.topology.create_entities(d)
im = mesh.topology.index_map(d)
vector = la.vector(im, dtype=dtype)
vector.array[: im.size_local] = np.arange(*im.local_range, dtype=dtype)
vector.scatter_forward()
global_idxs = im.local_to_global(np.arange(im.size_local + im.num_ghosts, dtype=np.int32))
global_idxs = np.asarray(global_idxs, dtype)
assert np.all(vector.array == global_idxs)
|
