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

1from scipy.signal import resample 

2 

3from pySDC.core.errors import TransferError 

4from pySDC.core.space_transfer import SpaceTransfer 

5 

6 

7class mesh_to_mesh_fft2d(SpaceTransfer): 

8 """ 

9 Custon base_transfer class, implements Transfer.py 

10 

11 This implementation can restrict and prolong between 2d meshes with FFT for periodic boundaries 

12 

13 Attributes: 

14 ratio: refinement factor between the two meshes 

15 """ 

16 

17 def __init__(self, fine_prob, coarse_prob, params): 

18 """ 

19 Initialization routine 

20 

21 Args: 

22 fine_prob: fine problem 

23 coarse_prob: coarse problem 

24 params: parameters for the transfer operators 

25 """ 

26 # invoke super initialization 

27 super(mesh_to_mesh_fft2d, self).__init__(fine_prob, coarse_prob, params) 

28 

29 assert len(self.fine_prob.nvars) == 2 

30 assert len(self.coarse_prob.nvars) == 2 

31 assert self.fine_prob.nvars[0] == self.fine_prob.nvars[1] 

32 assert self.coarse_prob.nvars[0] == self.coarse_prob.nvars[1] 

33 

34 self.ratio = int(self.fine_prob.nvars[0] / self.coarse_prob.nvars[0]) 

35 

36 def restrict(self, F): 

37 """ 

38 Restriction implementation 

39 

40 Args: 

41 F: the fine level data (easier to access than via the fine attribute) 

42 """ 

43 G = type(F)(self.coarse_prob.init, val=0.0) 

44 

45 def _restrict(fine, coarse): 

46 coarse[:] = fine[:: self.ratio, :: self.ratio] 

47 

48 # note that a `MultiComponentMesh` is also an instance of `mesh`, so ask for the components 

49 # rather than for the type 

50 if hasattr(type(F), 'components'): 

51 for comp in F.components: 

52 _restrict(getattr(F, comp), getattr(G, comp)) 

53 elif type(F).__name__ == 'mesh': 

54 _restrict(F, G) 

55 else: 

56 raise TransferError('Unknown data type, got %s' % type(F)) 

57 return G 

58 

59 def prolong(self, G): 

60 """ 

61 Prolongation implementation 

62 

63 Args: 

64 G: the coarse level data (easier to access than via the coarse attribute) 

65 """ 

66 F = type(G)(self.fine_prob.init, val=0.0) 

67 

68 def _prolong(coarse, fine): 

69 # Fourier interpolation along both axes. `resample` also gets the normalisation and the 

70 # splitting of the Nyquist mode right, which hand-rolled zero padding of the spectrum 

71 # only did for a refinement factor of two. 

72 fine[:] = resample(resample(coarse, fine.shape[0], axis=0), fine.shape[1], axis=1) 

73 

74 if hasattr(type(G), 'components'): 

75 for comp in G.components: 

76 _prolong(getattr(G, comp), getattr(F, comp)) 

77 elif type(G).__name__ == 'mesh': 

78 _prolong(G, F) 

79 else: 

80 raise TransferError('Unknown data type, got %s' % type(G)) 

81 return F