Coverage for pySDC / projects / parallelSDC_reloaded / utils.py: 100%
168 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-23 17:16 +0000
1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3"""
4Created on Sun Nov 12 18:50:39 2023
6Utility functions to investigate parallel SDC on non-linear problems
7"""
9import os
10import json
11import numpy as np
12from time import time
13import warnings
15from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
17from pySDC.implementations.problem_classes.Van_der_Pol_implicit import vanderpol
18from pySDC.implementations.problem_classes.Lorenz import LorenzAttractor
19from pySDC.implementations.problem_classes.odeScalar import ProtheroRobinson
20from pySDC.implementations.problem_classes.odeSystem import (
21 ProtheroRobinsonAutonomous,
22 Kaps,
23 ChemicalReaction3Var,
24 JacobiElliptic,
25)
26from pySDC.implementations.problem_classes.AllenCahn_1D_FD import (
27 allencahn_front_fullyimplicit,
28 allencahn_periodic_fullyimplicit,
29)
30from pySDC.implementations.problem_classes.TestEquation_0D import testequation0d
32from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit
33import pySDC.implementations.sweeper_classes.Runge_Kutta as rk
35import matplotlib.pyplot as plt
37PATH = '/' + os.path.join(*__file__.split('/')[:-1])
39# General matplotlib settings
40plt.rc('font', size=12)
41plt.rcParams['lines.linewidth'] = 2
42plt.rcParams['axes.titlesize'] = 16
43plt.rcParams['axes.labelsize'] = 16
44plt.rcParams['xtick.labelsize'] = 15
45plt.rcParams['ytick.labelsize'] = 15
46plt.rcParams['xtick.major.pad'] = 3
47plt.rcParams['ytick.major.pad'] = 2
48plt.rcParams['axes.labelpad'] = 6
49plt.rcParams['markers.fillstyle'] = 'none'
50plt.rcParams['lines.markersize'] = 7.0
51plt.rcParams['lines.markeredgewidth'] = 1.5
52plt.rcParams['mathtext.fontset'] = 'cm'
53plt.rcParams['mathtext.rm'] = 'serif'
54plt.rcParams['figure.max_open_warning'] = 100
57def getParamsSDC(
58 quadType="RADAU-RIGHT", numNodes=4, qDeltaI="IE", nSweeps=3, nodeType="LEGENDRE", collUpdate=False, initType="copy"
59):
60 description = {
61 # Sweeper and its parameters
62 "sweeper_class": generic_implicit,
63 "sweeper_params": {
64 "quad_type": quadType,
65 "num_nodes": numNodes,
66 "node_type": nodeType,
67 "initial_guess": initType,
68 "do_coll_update": collUpdate,
69 "QI": qDeltaI,
70 'skip_residual_computation': ('IT_CHECK', 'IT_DOWN', 'IT_UP', 'IT_FINE', 'IT_COARSE'),
71 },
72 # Step parameters
73 "step_params": {
74 "maxiter": 1,
75 },
76 # Level parameters
77 "level_params": {
78 "restol": -1,
79 "nsweeps": nSweeps,
80 },
81 }
83 return description
86RK_SWEEPERS = {
87 "BE": rk.BackwardEuler,
88 "FE": rk.ForwardEuler,
89 "RK4": rk.RK4,
90 "DIRK43": rk.DIRK43,
91 "ESDIRK53": rk.ESDIRK53,
92 "ESDIRK43": rk.ESDIRK43,
93}
96def getParamsRK(method="RK4"):
97 description = {
98 # Sweeper and its parameters
99 "sweeper_class": RK_SWEEPERS[method],
100 "sweeper_params": {'skip_residual_computation': ('IT_CHECK', 'IT_DOWN', 'IT_UP', 'IT_FINE', 'IT_COARSE')},
101 # Step parameters
102 "step_params": {
103 "maxiter": 1,
104 },
105 # Level parameters
106 "level_params": {
107 "restol": -1,
108 "nsweeps": 1,
109 },
110 }
112 return description
115def setupProblem(name, description, dt, **kwargs):
116 """Add problem settings to pySDC description parameters"""
118 # Common newton tolerance and max number of iterations
119 description["problem_params"] = {
120 'newton_tol': 1e-8,
121 'newton_maxiter': 300,
122 }
123 # Level parameters
124 description["level_params"]["dt"] = dt
126 if name == "VANDERPOL":
127 description["problem_class"] = vanderpol
128 description["problem_params"].update(
129 {
130 'newton_tol': 1e-12,
131 'mu': kwargs.get("mu", 10), # vanderpol parameter
132 'u0': np.array([2.0, 0]),
133 }
134 )
135 elif name == "LORENZ":
136 description["problem_class"] = LorenzAttractor
137 description["problem_params"].update(
138 {
139 'newton_tol': 1e-12,
140 'u0': kwargs.get("u0", (1, 1, 1)),
141 }
142 )
143 elif name == "PROTHERO-ROBINSON":
144 description["problem_class"] = ProtheroRobinson
145 description["problem_params"].update(
146 {
147 'newton_tol': 1e-12,
148 'epsilon': kwargs.get("epsilon", 1e-3),
149 }
150 )
151 elif name == "PROTHERO-ROBINSON-NL":
152 description["problem_class"] = ProtheroRobinson
153 description["problem_params"].update(
154 {
155 'newton_tol': 1e-12,
156 'epsilon': kwargs.get("epsilon", 1e-3),
157 'nonLinear': True,
158 }
159 )
160 elif name == "PROTHERO-ROBINSON-A":
161 description["problem_class"] = ProtheroRobinsonAutonomous
162 description["problem_params"].update(
163 {
164 'newton_tol': 1e-12,
165 'epsilon': kwargs.get("epsilon", 1e-3),
166 }
167 )
168 elif name == "PROTHERO-ROBINSON-A-NL":
169 description["problem_class"] = ProtheroRobinsonAutonomous
170 description["problem_params"].update(
171 {
172 'newton_tol': 1e-12,
173 'epsilon': kwargs.get("epsilon", 1e-3),
174 'nonLinear': True,
175 }
176 )
177 elif name == "KAPS":
178 description["problem_class"] = Kaps
179 description["problem_params"].update(
180 {
181 'epsilon': kwargs.get("epsilon", 1e-3),
182 }
183 )
184 elif name == "CHEMREC":
185 description["problem_class"] = ChemicalReaction3Var
186 elif name == "ALLEN-CAHN":
187 periodic = kwargs.get("periodic", False)
188 description["problem_class"] = allencahn_periodic_fullyimplicit if periodic else allencahn_front_fullyimplicit
189 description["problem_params"].update(
190 {
191 'newton_tol': 1e-8,
192 'nvars': kwargs.get("nvars", 128 if periodic else 127),
193 'eps': kwargs.get("epsilon", 0.04),
194 'stop_at_maxiter': True,
195 }
196 )
197 elif name == "JACELL":
198 description["problem_class"] = JacobiElliptic
199 elif name == "DAHLQUIST":
200 lambdas = kwargs.get("lambdas", None)
201 description["problem_class"] = testequation0d
202 description["problem_params"].update(
203 {
204 "lambdas": lambdas,
205 "u0": 1.0,
206 }
207 )
208 description["problem_params"].pop("newton_tol")
209 description["problem_params"].pop("newton_maxiter")
210 else:
211 raise NotImplementedError(f"problem {name} not implemented")
214def solutionSDC(tEnd, nSteps, params, probName, verbose=True, noExcept=False, **kwargs):
215 dt = tEnd / nSteps
216 setupProblem(probName, params, dt, **kwargs)
217 if noExcept:
218 params["problem_params"]["stop_at_nan"] = False
219 params["problem_params"]["stop_at_maxiter"] = False
221 controller = controller_nonMPI(num_procs=1, controller_params={'logger_level': 30}, description=params)
223 prob = controller.MS[0].levels[0].prob
225 uInit = prob.u_exact(0)
226 uTmp = uInit.copy()
228 uSDC = np.zeros((nSteps + 1, uInit.size), dtype=uInit.dtype)
229 uSDC[0] = uInit
230 tVals = np.linspace(0, tEnd, nSteps + 1)
231 tBeg = time()
232 if verbose:
233 print(" -- computing numerical solution with pySDC ...")
234 warnings.filterwarnings("ignore")
235 for i in range(nSteps):
236 uTmp[:] = uSDC[i]
237 try:
238 uSDC[i + 1], _ = controller.run(u0=uTmp, t0=tVals[i], Tend=tVals[i + 1])
239 except Exception as e:
240 print(f" -- exception when running controller : {e}")
241 warnings.resetwarnings()
242 return None, (0, 0, 0), False
243 warnings.resetwarnings()
244 tComp = time() - tBeg
246 try:
247 nNewton = prob.work_counters["newton"].niter
248 except KeyError:
249 nNewton = 0
250 nRHS = prob.work_counters["rhs"].niter
251 if verbose:
252 print(f" done, newton:{nNewton}, rhs:{nRHS}, tComp:{tComp}")
253 try:
254 parallel = controller.MS[0].levels[0].sweep.parallelizable
255 except AttributeError: # pragma: no cover
256 parallel = False
258 return uSDC, (nNewton, nRHS, tComp), parallel
261def solutionExact(tEnd, nSteps, probName, **kwargs):
262 """Return the exact solution of the Van-der-Pol problem at tEnd"""
264 if probName == "VANDERPOL":
265 mu = kwargs.get('mu', 10)
266 key = f"{tEnd}_{nSteps}_{mu}"
267 cacheFile = '_solVanderpolExact.json'
268 elif probName == "LORENZ":
269 u0 = kwargs.get('u0', (1, 1, 1))
270 key = f"{tEnd}_{nSteps}_{u0}"
271 cacheFile = '_solLorenzExact.json'
272 elif probName == "CHEMREC":
273 key = f"{tEnd}_{nSteps}"
274 cacheFile = '_solChemicalReactionExact.json'
275 elif probName == "JACELL":
276 key = f"{tEnd}_{nSteps}"
277 cacheFile = '_solJacobiEllipticExact.json'
279 # Potentially load already computed solution from local cache
280 try:
281 with open(f"{PATH}/{cacheFile}", "r") as f:
282 cache = json.load(f)
283 if key in cache:
284 return np.array(cache[key])
285 except (FileNotFoundError, json.JSONDecodeError, UnboundLocalError):
286 cache = {}
288 # Compute solution
289 params = getParamsSDC()
290 dt = tEnd / nSteps
291 setupProblem(probName, params, dt, **kwargs)
293 controller = controller_nonMPI(num_procs=1, controller_params={'logger_level': 30}, description=params)
294 solver = controller.MS[0].levels[0].prob.u_exact
296 print(" -- computing analytical solution with P.u_exact ...")
297 tBeg = time()
298 tVals = np.linspace(0, tEnd, nSteps + 1)
299 uExact = [solver(0)]
300 for i in range(nSteps):
301 try:
302 uExact.append(solver(tVals[i + 1], uExact[-1], tVals[i]))
303 except TypeError:
304 uExact.append(solver(tVals[i + 1]))
305 uExact = np.array(uExact)
306 print(f" done in {time()-tBeg:1.2f}s")
308 try:
309 # Save solution in local cache
310 cache[key] = uExact.tolist()
311 with open(f"{PATH}/{cacheFile}", "w") as f:
312 json.dump(cache, f)
313 except UnboundLocalError:
314 pass
316 return uExact
319# Plotting functions
320def plotStabContour(reVals, imVals, stab, ax=None):
321 if ax is None:
322 ax = plt.gca()
323 ax.contour(reVals, imVals, stab, levels=[1.0], colors='black', linewidths=1)
324 ax.contourf(reVals, imVals, stab, levels=[1.0, np.inf], colors='gainsboro')
325 ax.hlines(0, min(reVals), max(reVals), linestyles='--', colors='black', linewidth=0.5)
326 ax.vlines(0, min(imVals), max(imVals), linestyles='--', colors='black', linewidth=0.5)
327 ax.set_aspect('equal', 'box')
328 ax.set_xlabel(r"$Re(z)$", labelpad=0, fontsize=10)
329 ax.set_ylabel(r"$Im(z)$", labelpad=0, fontsize=10)
330 return ax