Coverage for pySDC/helpers/problem_helper.py: 97%
102 statements
« prev ^ index » next coverage.py v7.5.0, created at 2024-04-29 09:02 +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