Coverage for python/src/dolfinx_mpc/rbe.py: 89%

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

2# 

3# This file is part of DOLFINX_MPC 

4# 

5# SPDX-License-Identifier: MIT 

6"""Constraints tying dofs to spiders (RBE2, RBE3): points of a spider mesh, see :mod:`dolfinx_mpc.spider`.""" 

7 

8from __future__ import annotations 

9 

10import dolfinx.fem as _fem 

11import numpy as np 

12import numpy.typing as npt 

13 

14from .container import _cpp_function, _mpc_data_classes, _scalar_type 

15 

16__all__ = ["create_rbe2", "create_rbe3"] 

17 

18 

19def create_rbe2( 

20 V: _fem.FunctionSpace, 

21 dofs: npt.NDArray[np.int32], 

22 spiders: npt.NDArray[np.int64], 

23 W: _fem.FunctionSpace, 

24 dtype: npt.DTypeLike | None = None, 

25 x: npt.NDArray[np.floating] | None = None, 

26) -> _mpc_data_classes: 

27 r"""The constraint tying every component of the blocked `dofs` to the rigid-body motion of a spider. 

28 

29 Foot `i`, at `x`, follows spider `spiders[i]` at :math:`x_c`: 

30 :math:`u_j = t_j + (\theta \times (x - x_c))_j`. The translation :math:`t` is the first 

31 `gdim` components of the spider's block of dofs in `W`, the rotation :math:`\theta` the rest 

32 (components 3, 4, 5 in 3D, component 2 in 2D), and :math:`x_c` the coordinate of that block. 

33 Every rotation term is kept, also where its coefficient is zero, so that the masters do not 

34 depend on the configuration. Wraps `dolfinx_mpc::create_rbe2`. 

35 

36 Args: 

37 V: The space of the feet, with one component per dimension 

38 dofs: The feet, blocked dofs of `V` local to the process, ghosts included 

39 spiders: The spider of each foot, its input index (see :func:`create_spider_mesh`) 

40 W: The space on the spider mesh 

41 dtype: The scalar type of the coefficients. Defaults to the default scalar type of 

42 DOLFINx, real or complex, at the precision of the meshes. 

43 x: The coordinates of all dofs of `V` local to the process, from 

44 `V.tabulate_dof_coordinates()`, if already at hand 

45 

46 Returns: 

47 The slaves, masters (global dofs of `W`), coefficients, owners and offsets. 

48 

49 Note: 

50 Collective. 

51 """ 

52 real_type = V.mesh.geometry.x.dtype 

53 if W.mesh.geometry.x.dtype != real_type: 

54 raise ValueError("The mesh of the feet and the spider mesh must have the same coordinate type") 

55 create = _cpp_function("create_rbe2", _scalar_type(real_type, dtype)) 

56 return create( 

57 V._cpp_object, 

58 np.ascontiguousarray(dofs, dtype=np.int32), 

59 np.ascontiguousarray(spiders, dtype=np.int64), 

60 W._cpp_object, 

61 None if x is None else np.ascontiguousarray(x, dtype=real_type), 

62 ) 

63 

64 

65def create_rbe3( 

66 W: _fem.FunctionSpace, 

67 V: list[_fem.FunctionSpace], 

68 dofs: list[npt.NDArray[np.int32]], 

69 spiders: list[npt.NDArray[np.int64]], 

70 weights: list[npt.NDArray[np.floating]], 

71 dtype: npt.DTypeLike | None = None, 

72) -> tuple[_mpc_data_classes, npt.NDArray[np.int32]]: 

73 r"""The constraint tying the dofs of spiders to the motion of their feet (RBE3). 

74 

75 Each spider moves with the rigid motion that best fits its feet, 

76 :math:`\min_{t, \theta} \sum_i w_i |u_i - t - \theta \times (x_i - x_c)|^2`. Without rotations, 

77 :math:`t` is the weighted mean of the feet. Wraps `dolfinx_mpc::create_rbe3`. 

78 

79 Args: 

80 W: The space on the spider mesh, holding the slaves 

81 V: The spaces of the feet, with one component per dimension 

82 dofs: The feet in each space, blocked dofs local to the process 

83 spiders: The spider of each foot, its input index (see :func:`create_spider_mesh`) 

84 weights: The weight of each foot, non-negative 

85 dtype: The scalar type of the coefficients. Defaults to the default scalar type of 

86 DOLFINx, real or complex, at the precision of the meshes. 

87 

88 Returns: 

89 The slaves, masters, coefficients, owners and offsets, and the position in `V` of the 

90 space of each master. 

91 

92 Note: 

93 Collective. 

94 """ 

95 real_type = W.mesh.geometry.x.dtype 

96 if any(V_s.mesh.geometry.x.dtype != real_type for V_s in V): 

97 raise ValueError("The meshes of the feet and the spider mesh must have the same coordinate type") 

98 create = _cpp_function("create_rbe3", _scalar_type(real_type, dtype)) 

99 return create( 

100 W._cpp_object, 

101 [V_s._cpp_object for V_s in V], 

102 [np.ascontiguousarray(d, dtype=np.int32) for d in dofs], 

103 [np.ascontiguousarray(k, dtype=np.int64) for k in spiders], 

104 [np.ascontiguousarray(w, dtype=real_type) for w in weights], 

105 )