Coverage for python/src/dolfinx_mpc/dirichletbc.py: 100%

1# Copyright (C) 2026 Jørgen S. Dokken 

2# 

3# This file is part of DOLFINX_MPC 

4# 

5# SPDX-License-Identifier: MIT 

6"""Dirichlet condition data shared by the assemblers. 

7 

8To avoid re-computing the marked degrees of freedom that 

9gets its rows and columns zeroed in assembly. 

10""" 

11 

12from __future__ import annotations 

13 

14from collections.abc import Sequence 

15from typing import Optional 

16 

17import dolfinx.fem as _fem 

18import numpy as np 

19import numpy.typing as npt 

20 

21__all__ = ["BCData"] 

22 

23 

24class BCData: 

25 """Dof markers, diagonal rows and lifting values for a set of Dirichlet conditions. 

26 

27 Markers and rows depend only on a function space and the conditions, not on 

28 the condition values. We cache these computations so that multiple 

29 assembly calls inside {py:class}`dolfinx_mpc.LinearProblem` or 

30 {py:class}`dolfinx_mpc.NonlinearProblem` do not repeat them. 

31 """ 

32 

33 def __init__(self, bcs: Optional[Sequence[_fem.DirichletBC]] = None): 

34 self._bcs = list(bcs) if bcs else [] 

35 self._markers: dict[int, npt.NDArray[np.int8]] = {} 

36 self._rows: dict[int, npt.NDArray[np.int32]] = {} 

37 # The constraints, by the ids of their C++ objects, checked against these conditions 

38 self._checked: set[tuple[int, ...]] = set() 

39 

40 def markers( 

41 self, V0: _fem.FunctionSpace, V1: _fem.FunctionSpace 

42 ) -> tuple[npt.NDArray[np.int8], npt.NDArray[np.int8]]: 

43 """Constrained dof markers for the test and trial spaces of a form.""" 

44 return self._marker(V0), self._marker(V1) 

45 

46 def rows(self, V: _fem.FunctionSpace) -> npt.NDArray[np.int32]: 

47 """Locally owned rows of `V` carrying a Dirichlet condition.""" 

48 key = id(V._cpp_object) 

49 if key not in self._rows: 

50 owned = [bc.dof_indices()[0][: bc.dof_indices()[1]] for bc in self._bcs if V.contains(bc.function_space)] 

51 self._rows[key] = np.concatenate(owned) if owned else np.empty(0, dtype=np.int32) 

52 return self._rows[key] 

53 

54 def lifting( 

55 self, V: _fem.FunctionSpace, dtype: npt.DTypeLike 

56 ) -> tuple[npt.NDArray[np.int8], npt.NDArray[np.generic]]: 

57 """Constrained dof markers and Dirichlet values on the trial space `V`. 

58 

59 The markers are cached; the values are re-read from the conditions on 

60 every call, so a time-dependent condition is picked up. Both are empty 

61 when no condition applies to `V`. 

62 """ 

63 markers = self._marker(V) 

64 values = np.zeros(markers.size, dtype=dtype) 

65 if markers.size > 0: 

66 for bc in self._bcs: 

67 if V.contains(bc.function_space): 

68 bc.set(values, None, 1) 

69 return markers, values 

70 

71 def _marker(self, V: _fem.FunctionSpace) -> npt.NDArray[np.int8]: 

72 key = id(V._cpp_object) 

73 if key not in self._markers: 

74 self._markers[key] = _fem.assemble._bc_dof_markers(V, self._bcs) 

75 return self._markers[key]