Coverage for pySDC/helpers/problem_helper.py: 100%
103 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
1import numpy as np
2import scipy.sparse
3from scipy.special import factorial
6def get_steps(derivative, order, stencil_type):
7 """
8 Get the offsets for the FD stencil.
10 Args:
11 derivative (int): Order of the derivative
12 order (int): Order of accuracy
13 stencil_type (str): Type of the stencil
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
109 bc_params (dict or list of dicts): Parameters `val`, `neumann_bc_order` and `reduce` of Dirichlet and Neumann
110 BCs, for both sides or one dict per side
112 Returns:
113 Sparse matrix: Finite difference matrix
114 numpy.ndarray: Vector containing information about the boundary conditions
115 """
116 if cupy:
117 import cupyx.scipy.sparse as sp
118 else:
119 import scipy.sparse as sp
121 # get stencil
122 coeff, steps = get_finite_difference_stencil(
123 derivative=derivative, order=order, stencil_type=stencil_type, steps=steps
124 )
126 if type(bc) is not tuple:
127 assert type(bc) == str, 'Please pass BCs as string or tuple of strings'
128 bc = (bc, bc)
129 bc_params = bc_params if bc_params is not None else {}
130 if type(bc_params) is not list:
131 bc_params = [bc_params, bc_params]
133 b = np.zeros(size**dim)
135 if bc[0] == 'periodic':
136 assert bc[1] == 'periodic'
137 A_1d = 0 * sp.eye(size, format='csc')
138 for i in steps:
139 A_1d += coeff[i] * sp.eye(size, k=steps[i])
140 if steps[i] > 0:
141 A_1d += coeff[i] * sp.eye(size, k=-size + steps[i])
142 if steps[i] < 0:
143 A_1d += coeff[i] * sp.eye(size, k=size + steps[i])
144 else:
145 # Built on the host whatever the target: the boundary rows are edited in place below, which
146 # needs `lil`, and CuPy has no `lil`. The matrix moves to the device once it is finished.
147 A_1d = scipy.sparse.diags(coeff, steps, shape=(size, size), format='lil')
149 # Default parameters for Dirichlet and Neumann BCs
150 bc_params_defaults = {
151 'val': 0.0,
152 'neumann_bc_order': order,
153 'reduce': False,
154 }
156 # Loop over each side (0 for left, 1 for right)
157 for iS in [0, 1]:
158 # -- check Boundary condition types
159 assert "neumann" in bc[iS] or "dirichlet" in bc[iS], f"unknown BC type : {bc[iS]}"
161 # -- boundary condition parameters
162 bc_params[iS] = {**bc_params_defaults, **bc_params[iS]}
163 par = bc_params[iS].copy()
165 # -- extract parameters and raise an error if additionals
166 val = par.pop('val')
167 reduce = par.pop('reduce')
168 neumann_bc_order = par.pop('neumann_bc_order')
169 assert len(par) == 0, f"unused BCs parameters : {par}"
171 # -- half stencil width
172 sWidth = -min(steps) if iS == 0 else max(steps)
174 # -- loop over lines of A that have to be modified
175 for i in range(sWidth):
176 # -- index of the line
177 iLine = i if iS == 0 else -i - 1
178 # -- slice for coefficients used in the A matrix
179 sCoeff = slice(1, None) if iS == 0 else slice(None, -1)
180 # -- index of coefficient used in the b vector
181 iCoeff = 0 if iS == 0 else -1
183 if reduce:
184 # -- reduce order close to boundary
185 b_coeff, b_steps = get_finite_difference_stencil(
186 derivative=derivative,
187 order=2 * (i + 1),
188 stencil_type='center',
189 )
190 else:
191 # -- shift stencil close to boundary
192 b_steps = (
193 np.arange(-(i + 1), order + derivative - (i + 1))
194 if iS == 0
195 else np.arange(-(order + derivative) + (i + 2), (i + 2))
196 )
198 b_coeff, b_steps = get_finite_difference_stencil(derivative=derivative, steps=b_steps)
200 # -- column slice where to put coefficients in the A matrix
201 colSlice = slice(None, len(b_coeff) - 1) if iS == 0 else slice(-len(b_coeff) + 1, None)
203 # -- modify A
204 A_1d[iLine, :] = 0
205 A_1d[iLine, colSlice] = b_coeff[sCoeff]
207 if "dirichlet" in bc[iS]:
208 # -- modify b
209 b[iLine] = val * b_coeff[iCoeff]
211 elif "neumann" in bc[iS]:
212 nOrder = neumann_bc_order
214 # -- generate the first derivative stencil
215 n_coeff, n_steps = get_finite_difference_stencil(
216 derivative=1, order=nOrder, stencil_type="forward" if iS == 0 else "backward"
217 )
219 # -- column slice where to put coefficients in the A matrix
220 colSlice = slice(None, len(n_coeff) - 1) if iS == 0 else slice(-len(n_coeff) + 1, None)
222 # -- additional modification to A
223 A_1d[iLine, colSlice] -= b_coeff[iCoeff] / n_coeff[iCoeff] * n_coeff[sCoeff]
225 # -- modify B
226 b[iLine] = val * b_coeff[iCoeff] / n_coeff[iCoeff] * dx
228 # TODO: extend the BCs to higher dimensions
229 A_1d = sp.csc_matrix(A_1d)
230 if dim == 1:
231 A = A_1d
232 elif dim == 2:
233 A = sp.kron(A_1d, sp.eye(size)) + sp.kron(sp.eye(size), A_1d)
234 elif dim == 3:
235 A = (
236 sp.kron(A_1d, sp.eye(size**2))
237 + sp.kron(sp.eye(size**2), A_1d)
238 + sp.kron(sp.kron(sp.eye(size), A_1d), sp.eye(size))
239 )
240 else:
241 raise NotImplementedError(f'Dimension {dim} not implemented.')
243 A /= dx**derivative
244 b /= dx**derivative
246 return A, b
249def get_1d_grid(size, bc, left_boundary=0.0, right_boundary=1.0):
250 """
251 Generate a grid in one dimension and obtain mesh spacing for finite difference discretization.
253 Args:
254 size (int): Number of degrees of freedom per dimension
255 bc (str): Boundary conditions for both sides
256 left_boundary (float): x value at the left boundary
257 right_boundary (float): x value at the right boundary
259 Returns:
260 float: mesh spacing
261 numpy.ndarray: 1d mesh
262 """
263 L = right_boundary - left_boundary
264 if bc == 'periodic':
265 dx = L / size
266 xvalues = np.array([left_boundary + dx * i for i in range(size)])
267 elif "dirichlet" in bc or "neumann" in bc:
268 dx = L / (size + 1)
269 xvalues = np.array([left_boundary + dx * (i + 1) for i in range(size)])
270 else:
271 raise NotImplementedError(f'Boundary conditions \"{bc}\" not implemented.')
273 return dx, xvalues