Coverage for pySDC/implementations/transfer_classes/TransferFiredrakeMesh.py: 100%

1from firedrake import assemble, prolong, inject 

2from firedrake.__future__ import interpolate 

3 

4from pySDC.core.errors import TransferError 

5from pySDC.core.space_transfer import SpaceTransfer 

6from pySDC.implementations.datatype_classes.firedrake_mesh import firedrake_mesh, IMEX_firedrake_mesh 

7 

8 

9class MeshToMeshFiredrake(SpaceTransfer): 

10 """ 

11 This implementation can restrict and prolong between Firedrake meshes 

12 """ 

13 

14 def restrict(self, F): 

15 """ 

16 Restrict from fine to coarse grid 

17 

18 Args: 

19 F: the fine level data 

20 

21 Returns: 

22 Coarse level data 

23 """ 

24 if isinstance(F, firedrake_mesh): 

25 u_coarse = self.coarse_prob.dtype_u(assemble(interpolate(F.functionspace, self.coarse_prob.init))) 

26 elif isinstance(F, IMEX_firedrake_mesh): 

27 u_coarse = IMEX_firedrake_mesh(self.coarse_prob.init) 

28 u_coarse.impl.functionspace.assign(assemble(interpolate(F.impl.functionspace, self.coarse_prob.init))) 

29 u_coarse.expl.functionspace.assign(assemble(interpolate(F.expl.functionspace, self.coarse_prob.init))) 

30 else: 

31 raise TransferError('Unknown type of fine data, got %s' % type(F)) 

32 

33 return u_coarse 

34 

35 def prolong(self, G): 

36 """ 

37 Prolongate from coarse to fine grid 

38 

39 Args: 

40 G: the coarse level data 

41 

42 Returns: 

43 fine level data 

44 """ 

45 if isinstance(G, firedrake_mesh): 

46 u_fine = self.fine_prob.dtype_u(assemble(interpolate(G.functionspace, self.fine_prob.init))) 

47 elif isinstance(G, IMEX_firedrake_mesh): 

48 u_fine = IMEX_firedrake_mesh(self.fine_prob.init) 

49 u_fine.impl.functionspace.assign(assemble(interpolate(G.impl.functionspace, self.fine_prob.init))) 

50 u_fine.expl.functionspace.assign(assemble(interpolate(G.expl.functionspace, self.fine_prob.init))) 

51 else: 

52 raise TransferError('Unknown type of coarse data, got %s' % type(G)) 

53 

54 return u_fine 

55 

56 

57class MeshToMeshFiredrakeHierarchy(SpaceTransfer): 

58 """ 

59 This implementation can restrict and prolong between Firedrake meshes that are generated from a hierarchy. 

60 Example: 

61 

62 .. code-block:: python 

63 

64 from firedrake import * 

65 

66 mesh = UnitSquareMesh(8, 8) 

67 hierarchy = MeshHierarchy(mesh, 4) 

68 

69 mesh = hierarchy[-1] 

70 

71 """ 

72 

73 @staticmethod 

74 def _restrict(u_fine, u_coarse): 

75 """Perform restriction in Firedrake""" 

76 inject(u_fine.functionspace, u_coarse.functionspace) 

77 

78 @staticmethod 

79 def _prolong(u_coarse, u_fine): 

80 """Perform prolongation in Firedrake""" 

81 prolong(u_coarse.functionspace, u_fine.functionspace) 

82 

83 def restrict(self, F): 

84 """ 

85 Restrict from fine to coarse grid 

86 

87 Args: 

88 F: the fine level data 

89 

90 Returns: 

91 Coarse level data 

92 """ 

93 if isinstance(F, firedrake_mesh): 

94 G = self.coarse_prob.u_init 

95 self._restrict(u_fine=F, u_coarse=G) 

96 elif isinstance(F, IMEX_firedrake_mesh): 

97 G = IMEX_firedrake_mesh(self.coarse_prob.init) 

98 self._restrict(u_fine=F.impl, u_coarse=G.impl) 

99 self._restrict(u_fine=F.expl, u_coarse=G.expl) 

100 else: 

101 raise TransferError('Unknown type of fine data, got %s' % type(F)) 

102 

103 return G 

104 

105 def prolong(self, G): 

106 """ 

107 Prolongate from coarse to fine grid 

108 

109 Args: 

110 G: the coarse level data 

111 

112 Returns: 

113 fine level data 

114 """ 

115 if isinstance(G, firedrake_mesh): 

116 F = self.fine_prob.u_init 

117 self._prolong(u_coarse=G, u_fine=F) 

118 elif isinstance(G, IMEX_firedrake_mesh): 

119 F = self.fine_prob.f_init 

120 self._prolong(u_coarse=G.impl, u_fine=F.impl) 

121 self._prolong(u_coarse=G.expl, u_fine=F.expl) 

122 else: 

123 raise TransferError('Unknown type of coarse data, got %s' % type(G)) 

124 

125 return F