def evaluate_sparse_function_at_dofs(input_function, dofs_list, output_V=None, reduced_dofs_list=None):
assert (output_V is None) == (reduced_dofs_list is None)
if output_V is None:
return evaluate_sparse_vector_at_dofs(input_function.vector(), dofs_list)
else:
vec = to_petsc4py(input_function.vector())
output_function = Function(output_V)
out = to_petsc4py(output_function.vector())
_evaluate_sparse_function_at_dofs(vec, dofs_list, out, reduced_dofs_list)
return output_function
def evaluate_sparse_vector_at_dofs(sparse_vector, dofs_list):
vec = to_petsc4py(sparse_vector)
row_start, row_end = vec.getOwnershipRange()
out_size = len(dofs_list)
out = OnlineVector(out_size)
mpi_comm = vec.comm.tompi4py()
for (index, dofs) in enumerate(dofs_list):
assert len(dofs) == 1
i = dofs[0]
out_index = None
vec_i_processor = -1
if i >= row_start and i < row_end:
out_index = vec.getValue(i)
vec_i_processor = mpi_comm.rank
vec_i_processor = mpi_comm.allreduce(vec_i_processor, op=MAX)
assert vec_i_processor >= 0
out[index] = mpi_comm.bcast(out_index, root=vec_i_processor)
return out
def evaluate_and_vectorize_sparse_matrix_at_dofs(sparse_matrix, dofs_list):
mat = to_petsc4py(sparse_matrix)
row_start, row_end = mat.getOwnershipRange()
out_size = len(dofs_list)
out = OnlineVector(out_size)
mpi_comm = mat.comm.tompi4py()
for (index, dofs) in enumerate(dofs_list):
assert len(dofs) == 2
i = dofs[0]
out_index = None
mat_ij_processor = -1
if i >= row_start and i < row_end:
j = dofs[1]
out_index = mat.getValue(i, j)
mat_ij_processor = mpi_comm.rank
mat_ij_processor = mpi_comm.allreduce(mat_ij_processor, op=MAX)
assert mat_ij_processor >= 0
out[index] = mpi_comm.bcast(out_index, root=mat_ij_processor)
return out
def vectorized_matrix_inner_vectorized_matrix(matrix, other_matrix):
matrix = to_petsc4py(matrix)
other_matrix = to_petsc4py(other_matrix)
# petsc4py does not expose MatGetTrace, we do this by hand
return matrix.transposeMatMult(other_matrix).getDiagonal().sum()
