# Coverage for pySDC/helpers/problem_helper.py: 97%

## 102 statements

, created at 2024-09-20 17:10 +0000

1import numpy as np

2from scipy.special import factorial

5def get_steps(derivative, order, stencil_type):

6 """

7 Get the offsets for the FD stencil.

9 Args:

10 derivative (int): Order of the derivative

11 order (int): Order of accuracy

12 stencil_type (str): Type of the stencil

13 steps (list): Provide specific steps, overrides `stencil_type`

15 Returns:

16 int: The number of elements in the stencil

17 numpy.ndarray: The offsets for the stencil

18 """

19 if stencil_type == 'center':

20 n = order + derivative - (derivative + 1) % 2 // 1

21 steps = np.arange(n) - n // 2

22 elif stencil_type == 'forward':

23 n = order + derivative

24 steps = np.arange(n)

25 elif stencil_type == 'backward':

26 n = order + derivative

27 steps = -np.arange(n)

28 elif stencil_type == 'upwind':

29 n = order + derivative

31 if n <= 3:

32 n, steps = get_steps(derivative, order, 'backward')

33 else:

34 steps = np.append(-np.arange(n - 1)[::-1], [1])

35 else:

36 raise ValueError(

37 f'Stencil must be of type "center", "forward", "backward" or "upwind", not {stencil_type}. If you want something else you can also give specific steps.'

38 )

39 return n, steps

42def get_finite_difference_stencil(derivative, order=None, stencil_type=None, steps=None):

43 """

44 Derive general finite difference stencils from Taylor expansions

46 Args:

47 derivative (int): Order of the derivative

48 order (int): Order of accuracy

49 stencil_type (str): Type of the stencil

50 steps (list): Provide specific steps, overrides `stencil_type`

52 Returns:

53 numpy.ndarray: The weights of the stencil

54 numpy.ndarray: The offsets for the stencil

55 """

57 if steps is not None:

58 n = len(steps)

59 else:

60 n, steps = get_steps(derivative, order, stencil_type)

62 # make a matrix that contains the Taylor coefficients

63 A = np.zeros((n, n))

64 idx = np.arange(n)

65 inv_facs = 1.0 / factorial(idx)

66 for i in range(0, n):

67 A[i, :] = steps ** idx[i] * inv_facs[i]

69 # make a right hand side vector that is zero everywhere except at the position of the desired derivative

70 sol = np.zeros(n)

71 sol[derivative] = 1.0

73 # solve the linear system for the finite difference coefficients

74 coeff = np.linalg.solve(A, sol)

76 # sort coefficients and steps

77 coeff = coeff[np.argsort(steps)]

78 steps = np.sort(steps)

80 return coeff, steps

83def get_finite_difference_matrix(

84 derivative,

85 order,

86 stencil_type=None,

87 steps=None,

88 dx=None,

89 size=None,

90 dim=None,

91 bc=None,

92 cupy=False,

93 bc_params=None,

94):

95 """

96 Build FD matrix from stencils, with boundary conditions.

97 Keep in mind that the boundary conditions may require further modification of the right hand side.

99 Args:

100 derivative (int): Order of the spatial derivative

101 order (int): Order of accuracy

102 stencil_type (str): Type of stencil

103 steps (list): Provide specific steps, overrides `stencil_type`

104 dx (float): Mesh width

105 size (int): Number of degrees of freedom per dimension

106 dim (int): Number of dimensions

107 bc (str): Boundary conditions for both sides

108 cupy (bool): Construct a GPU ready matrix if yes

110 Returns:

111 Sparse matrix: Finite difference matrix

112 numpy.ndarray: Vector containing information about the boundary conditions

113 """

114 if cupy:

115 import cupyx.scipy.sparse as sp

116 else:

117 import scipy.sparse as sp

119 # get stencil

120 coeff, steps = get_finite_difference_stencil(

121 derivative=derivative, order=order, stencil_type=stencil_type, steps=steps

122 )

124 if type(bc) is not tuple:

125 assert type(bc) == str, 'Please pass BCs as string or tuple of strings'

126 bc = (bc, bc)

127 bc_params = bc_params if bc_params is not None else {}

128 if type(bc_params) is not list:

129 bc_params = [bc_params, bc_params]

131 b = np.zeros(size**dim)

133 if bc[0] == 'periodic':

134 assert bc[1] == 'periodic'

135 A_1d = 0 * sp.eye(size, format='csc')

136 for i in steps:

137 A_1d += coeff[i] * sp.eye(size, k=steps[i])

138 if steps[i] > 0:

139 A_1d += coeff[i] * sp.eye(size, k=-size + steps[i])

140 if steps[i] < 0:

141 A_1d += coeff[i] * sp.eye(size, k=size + steps[i])

142 else:

143 A_1d = sp.diags(coeff, steps, shape=(size, size), format='lil')

145 # Default parameters for Dirichlet and Neumann BCs

146 bc_params_defaults = {

147 'val': 0.0,

148 'neumann_bc_order': order,

149 'reduce': False,

150 }

152 # Loop over each side (0 for left, 1 for right)

153 for iS in [0, 1]:

154 # -- check Boundary condition types

155 assert "neumann" in bc[iS] or "dirichlet" in bc[iS], f"unknown BC type : {bc[iS]}"

157 # -- boundary condition parameters

158 bc_params[iS] = {**bc_params_defaults, **bc_params[iS]}

159 par = bc_params[iS].copy()

161 # -- extract parameters and raise an error if additionals

162 val = par.pop('val')

163 reduce = par.pop('reduce')

164 neumann_bc_order = par.pop('neumann_bc_order')

165 assert len(par) == 0, f"unused BCs parameters : {par}"

167 # -- half stencil width

168 sWidth = -min(steps) if iS == 0 else max(steps)

170 # -- loop over lines of A that have to be modified

171 for i in range(sWidth):

172 # -- index of the line

173 iLine = i if iS == 0 else -i - 1

174 # -- slice for coefficients used in the A matrix

175 sCoeff = slice(1, None) if iS == 0 else slice(None, -1)

176 # -- index of coefficient used in the b vector

177 iCoeff = 0 if iS == 0 else -1

179 if reduce:

180 # -- reduce order close to boundary

181 b_coeff, b_steps = get_finite_difference_stencil(

182 derivative=derivative,

183 order=2 * (i + 1),

184 stencil_type='center',

185 )

186 else:

187 # -- shift stencil close to boundary

188 b_steps = (

189 np.arange(-(i + 1), order + derivative - (i + 1))

190 if iS == 0

191 else np.arange(-(order + derivative) + (i + 2), (i + 2))

192 )

194 b_coeff, b_steps = get_finite_difference_stencil(derivative=derivative, steps=b_steps)

196 # -- column slice where to put coefficients in the A matrix

197 colSlice = slice(None, len(b_coeff) - 1) if iS == 0 else slice(-len(b_coeff) + 1, None)

199 # -- modify A

200 A_1d[iLine, :] = 0

201 A_1d[iLine, colSlice] = b_coeff[sCoeff]

203 if "dirichlet" in bc[iS]:

204 # -- modify b

205 b[iLine] = val * b_coeff[iCoeff]

207 elif "neumann" in bc[iS]:

208 nOrder = neumann_bc_order

210 # -- generate the first derivative stencil

211 n_coeff, n_steps = get_finite_difference_stencil(

212 derivative=1, order=nOrder, stencil_type="forward" if iS == 0 else "backward"

213 )

215 # -- column slice where to put coefficients in the A matrix

216 colSlice = slice(None, len(n_coeff) - 1) if iS == 0 else slice(-len(n_coeff) + 1, None)

218 # -- additional modification to A

219 A_1d[iLine, colSlice] -= b_coeff[iCoeff] / n_coeff[iCoeff] * n_coeff[sCoeff]

221 # -- modify B

222 b[iLine] = val * b_coeff[iCoeff] / n_coeff[iCoeff] * dx

224 # TODO: extend the BCs to higher dimensions

225 A_1d = A_1d.tocsc()

226 if dim == 1:

227 A = A_1d

228 elif dim == 2:

229 A = sp.kron(A_1d, sp.eye(size)) + sp.kron(sp.eye(size), A_1d)

230 elif dim == 3:

231 A = (

232 sp.kron(A_1d, sp.eye(size**2))

233 + sp.kron(sp.eye(size**2), A_1d)

234 + sp.kron(sp.kron(sp.eye(size), A_1d), sp.eye(size))

235 )

236 else:

237 raise NotImplementedError(f'Dimension {dim} not implemented.')

239 A /= dx**derivative

240 b /= dx**derivative

242 return A, b

245def get_1d_grid(size, bc, left_boundary=0.0, right_boundary=1.0):

246 """

247 Generate a grid in one dimension and obtain mesh spacing for finite difference discretization.

249 Args:

250 size (int): Number of degrees of freedom per dimension

251 bc (str): Boundary conditions for both sides

252 left_boundary (float): x value at the left boundary

253 right_boundary (float): x value at the right boundary

255 Returns:

256 float: mesh spacing

257 numpy.ndarray: 1d mesh

258 """

259 L = right_boundary - left_boundary

260 if bc == 'periodic':

261 dx = L / size

262 xvalues = np.array([left_boundary + dx * i for i in range(size)])

263 elif "dirichlet" in bc or "neumann" in bc:

264 dx = L / (size + 1)

265 xvalues = np.array([left_boundary + dx * (i + 1) for i in range(size)])

266 else:

267 raise NotImplementedError(f'Boundary conditions \"{bc}\" not implemented.')

269 return dx, xvalues