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
|
# Copyright (C) 2020-2021 Jørgen S. Dokken
#
# This file is part of DOLFINX_MPC
#
# SPDX-License-Identifier: MIT
import numba
import numpy
import numpy.typing as npt
from typing import Union
import dolfinx.cpp as _cpp
_forms = Union[_cpp.fem.Form_float32, _cpp.fem.Form_float64, _cpp.fem.Form_complex128]
_bcs = Union[_cpp.fem.DirichletBC_float32, _cpp.fem.DirichletBC_float64,
_cpp.fem.DirichletBC_complex64, _cpp.fem.DirichletBC_complex128]
@numba.njit(fastmath=True, cache=True)
def extract_slave_cells(cell_offset: npt.NDArray[numpy.int32]) -> npt.NDArray[numpy.int32]:
""" From an offset determine which entries are nonzero"""
slave_cells = numpy.zeros(len(cell_offset) - 1, dtype=numpy.int32)
c = 0
for cell in range(len(cell_offset) - 1):
num_cells = cell_offset[cell + 1] - cell_offset[cell]
if num_cells > 0:
slave_cells[c] = cell
c += 1
return slave_cells[:c]
@numba.njit(fastmath=True, cache=True)
def pack_slave_facet_info(facets: npt.NDArray[numpy.int32],
slave_cells: npt.NDArray[numpy.int32]) -> npt.NDArray[numpy.int32]:
"""
Given an MPC and a set of facets (cell index, local_facet_index),
compress the set to those that only contain slave cells
"""
facet_info = numpy.zeros((len(facets), 2), dtype=numpy.int32)
i = 0
for facet in facets:
if sum(slave_cells == facet[0]) > 0:
facet_info[i, :] = [facet[0], facet[1]]
i += 1
return facet_info[:i, :]
|
