diff --git a/qmat/playgrounds/martin/diff_eqs/__init__.py b/qmat/playgrounds/martin/diff_eqs/__init__.py new file mode 100644 index 0000000..d33d182 --- /dev/null +++ b/qmat/playgrounds/martin/diff_eqs/__init__.py @@ -0,0 +1,10 @@ + +from qmat.playgrounds.martin.diff_eqs.de_solver import DESolver +from qmat.playgrounds.martin.diff_eqs.burgers import Burgers +from qmat.playgrounds.martin.diff_eqs.dahlquist2 import Dahlquist2 + +__all__ = [ + "DESolver", + "Burgers", + "Dahlquist2", +] diff --git a/qmat/playgrounds/martin/diff_eqs/burgers.py b/qmat/playgrounds/martin/diff_eqs/burgers.py new file mode 100644 index 0000000..f1771c4 --- /dev/null +++ b/qmat/playgrounds/martin/diff_eqs/burgers.py @@ -0,0 +1,159 @@ +import numpy as np +from qmat.playgrounds.martin.diff_eqs.de_solver import DESolver + + +class Burgers(DESolver): + """ + Class to handle the 1D viscous Burgers' equation. + """ + + def __init__(self, N: int, nu: float, domain_size: float = 2.0 * np.pi): + # Resolution + self._N: int = N + + # Viscosity + self._nu: float = nu + + # Domain size + self._domain_size: float = domain_size + + # Prepare spectral differentiation values + self._d_dx_ = 1j * np.fft.fftfreq(self._N, d=1.0 / self._N) * 2.0 * np.pi / self._domain_size + + def _d_dx(self, u: np.ndarray) -> np.ndarray: + """Compute the spatial derivative of `u` using spectral methods. + + Parameters + ---------- + u : np.ndarray + Array of shape (N,) representing the solution at the current time step. + + Returns + ------- + du_dx : np.ndarray + Array of shape (N,) representing the spatial derivative of `u`. + """ + u_hat = np.fft.fft(u) + du_dx_hat = u_hat * self._d_dx_ + du_dx = np.fft.ifft(du_dx_hat).real + return du_dx + + def initial_u0(self, mode: str) -> np.ndarray: + """Compute some initial conditions for the 1D viscous Burgers' equation.""" + + if mode == "sine": + x = np.linspace(0, self._domain_size, self._N, endpoint=False) + u0 = np.sin(x) + + elif mode == "hat": + x = np.linspace(0, self._domain_size, self._N, endpoint=False) + u0 = np.where((x >= np.pi / 2) & (x <= 3 * np.pi / 2), 1.0, 0.0) + + elif mode == "random": + np.random.seed(42) + u0 = np.random.rand(self._N) + + else: + raise ValueError(f"Unknown initial condition mode: {mode}") + + return u0 + + def evalF(self, u: np.ndarray, t: float) -> np.ndarray: + """ + Evaluate the right-hand side of the 1D viscous Burgers' equation. + + Parameters + ---------- + u : np.ndarray + Array of shape (N,) representing the solution at the current time step. + + Returns + ------- + f : np.ndarray + Array of shape (N,) representing the right-hand side evaluated at `u`. + """ + # Compute spatial derivatives using spectral methods + u_hat = np.fft.fft(u) + du_dx_hat = self._d_dx_ * u_hat + d2u_dx2_hat = (self._d_dx_**2) * u_hat + + du_dx = np.fft.ifft(du_dx_hat).real + d2u_dx2 = np.fft.ifft(d2u_dx2_hat).real + + f = -u * du_dx + self._nu * d2u_dx2 + return f + + def int_f(self, u0: np.ndarray, dt: float, t: float = 0.0) -> np.ndarray: + """ + Compute the analytical solution of the 1D viscous Burgers' equation at time `t`. + + See + https://gitlab.inria.fr/sweet/sweet/-/blob/6f20b19f246bf6fcc7ace1b69567326d1da78635/src/programs/_pde_burgers/time/Burgers_Cart2D_TS_ln_cole_hopf.cpp + + Parameters + ---------- + u0 : np.ndarray + Array of shape (N,) representing the initial condition. + t : float + Time at which to evaluate the analytical solution. + + Returns + ------- + u_analytical : np.ndarray + Array of shape (N,) representing the analytical solution at time `t`. + """ + + if self._nu < 0.05: + print("Viscosity is very small which can lead to errors in analytical solution!") + + u0_hat = np.fft.fft(u0) + + # Divide by d/dx operator in spectral space + tmp = np.zeros_like(u0_hat, dtype=complex) + tmp[1:] = u0_hat[1:] / self._d_dx_[1:] + + # Back to physical space + phi = np.fft.ifft(tmp).real + + # Apply exp(...) + phi = np.exp(-phi / (2 * self._nu)) + + phi_hat = np.fft.fft(phi) + + # Solve directly the heat equation in spectral space with exponential integration + phi_hat = phi_hat * np.exp(self._nu * self._d_dx_**2 * (t+dt)) + + phi = np.fft.ifft(phi_hat) + phi = np.log(phi) + + phi_hat = np.fft.fft(phi) + + u1_hat = -2.0 * self._nu * phi_hat * self._d_dx_ + return np.fft.ifft(u1_hat).real + + def test(self): + """ + Run test for currently set up Burgers instance. + """ + x = np.linspace(0, self._domain_size, self._N, endpoint=False) + w = 2.0 * np.pi / self._domain_size + u0 = np.sin(x * w) + + u1_analytical = np.cos(x * w) * w + u1_num = self._d_dx(u0) + + error: float = np.max(np.abs(u1_num - u1_analytical)) + + if error > 1e-10: + raise AssertionError(f"Test failed: error {error} too large for domain size {self._domain_size}.") + + def run_tests(self): + """ + Run basic tests to verify the correctness of the implementation. + + This doesn't change the current instance, but will create new instances. + """ + + for domain_size in [2.0 * np.pi, 1.0, 9.0]: + burgers = Burgers(N=128, nu=0.01, domain_size=domain_size) + burgers.test() diff --git a/qmat/playgrounds/martin/diff_eqs/dahlquist.py b/qmat/playgrounds/martin/diff_eqs/dahlquist.py new file mode 100644 index 0000000..8fc7447 --- /dev/null +++ b/qmat/playgrounds/martin/diff_eqs/dahlquist.py @@ -0,0 +1,121 @@ +import numpy as np +from qmat.playgrounds.martin.diff_eqs.de_solver import DESolver + + +class Dahlquist(DESolver): + """ + Standard Dahlquist test equation with two eigenvalues. + Optionally, some external frequency forcing can be added which is + configurable through `ext_scalar`. + + u(t) = exp(t*(lam1+lam2))*u(0) + s*sin(t) + + d/dt u(t) = (lam1 + lam2) * (u(t) - s*sin(t)) + s*cos(t) + """ + + def __init__(self, lam1: complex, lam2: complex, ext_scalar: float = 0.0): + # Lambda 1 + self.lam1: float = lam1 + + # Lambda 2 + self.lam2: float = lam2 + + # Scaling of external sin(t) frequency. Set to 0 to deactivate. + self.ext_scalar = ext_scalar + + def initial_u0(self, mode: str = None) -> np.ndarray: + return np.array([1.0 + 0.0j], dtype=np.complex128) + + def picardF(self, u: np.ndarray, dt: float, t: float) -> np.ndarray: + """ + Exactly integrate over one time step using the analytical solution + (=Picard). + """ + lam = self.lam1 + self.lam2 + s = self.ext_scalar + assert isinstance(t, float) + retval = np.exp(t * lam) * u + s * np.sin(t) + + assert retval.shape == u.shape + return retval + + def evalF(self, u: np.ndarray, t: float) -> np.ndarray: + lam = self.lam1 + self.lam2 + s = self.ext_scalar + retval = lam * (u - s * np.sin(t)) + s * np.cos(t) + + assert retval.shape == u.shape + return retval + + def evalF1(self, u: np.ndarray, t: float) -> np.ndarray: + lam = self.lam1 + retval = lam * u + + assert retval.shape == u.shape + return retval + + def evalF2(self, u: np.ndarray, t: float) -> np.ndarray: + lam = self.lam2 + s = self.ext_scalar + retval = lam * (u - s * np.sin(t)) + s * np.cos(t) + + assert retval.shape == u.shape + return retval + + def fSolve(self, u: np.ndarray, dt: float, t: float) -> np.ndarray: + t1 = t + dt + lam = self.lam1 + self.lam2 + s = self.ext_scalar + + rhs = u - s * dt * (lam * np.sin(t1) - np.cos(t1)) + retval = rhs / (1.0 - dt * lam) + + assert retval.shape == u.shape + return retval + + def fSolve1(self, u: np.ndarray, dt: float, t: float) -> np.ndarray: + retval = u / (1.0 - dt * self.lam1) + + assert retval.shape == u.shape + return retval + + def fSolve2(self, u: np.ndarray, dt: float, t: float) -> np.ndarray: + t1 = t + dt + lam = self.lam2 + s = self.ext_scalar + + rhs = u - s * dt * (lam * np.sin(t1) - np.cos(t1)) + retval = rhs / (1.0 - dt * lam) + + assert retval.shape == u.shape + return retval + + def int_f_t0(self, u0: np.ndarray, dt: float) -> np.ndarray: + lam = self.lam1 + self.lam2 + s = self.ext_scalar + assert isinstance(dt, float) + retval = np.exp(dt * lam) * u0 + s * np.sin(dt) + + assert retval.shape == u0.shape + return retval + + def int_f(self, u0: np.ndarray, dt: float, t: float = 0.0) -> np.ndarray: + """ + Integrate the solution from t0 to t1. + """ + + assert isinstance(t, (float, int)) + assert isinstance(dt, (float, int)) + + if t == 0: + return self.int_f_t0(u0, dt=dt) + + # Lambda + lam = self.lam1 + self.lam2 + + s = self.ext_scalar + + retval = np.exp(dt * lam) * (u0 - s*np.sin(t)) + s * np.sin(t+dt) + + assert retval.shape == u0.shape + return retval diff --git a/qmat/playgrounds/martin/diff_eqs/dahlquist2.py b/qmat/playgrounds/martin/diff_eqs/dahlquist2.py new file mode 100644 index 0000000..1fa42be --- /dev/null +++ b/qmat/playgrounds/martin/diff_eqs/dahlquist2.py @@ -0,0 +1,44 @@ +import numpy as np +from qmat.playgrounds.martin.diff_eqs.de_solver import DESolver + + +class Dahlquist2(DESolver): + """ + Modified Dahlquist test equation with superposition of two + frequencies in the solution u(t). + + u(t) = 0.5*(exp(lam1*t) + exp(lam2*t)) * u(0) + """ + + def __init__(self, lam1: complex, lam2: complex, s: float = 0.6): + """ + :param lam1: First eigenvalue (complex) + :param lam2: Second eigenvalue (complex) + :param s: Weighting between the two exponentials in the solution + """ + self.lam1: complex = lam1 + self.lam2: complex = lam2 + # Weighting between the two exponentials in the solution + # to avoid division by 0 + self.s: float = s + + assert 0 <= self.s <= 1, "s must be in [0,1]" + + def initial_u0(self, mode: str = None) -> np.ndarray: + return np.array([1.0 + 0.0j], dtype=np.complex128) + + def evalF(self, u: np.ndarray, t: float) -> np.ndarray: + retval = ( + (self.lam1 * self.s * np.exp(t * self.lam1) + self.lam2 * (1.0 - self.s) * np.exp(t * self.lam2)) + / (self.s * np.exp(t * self.lam1) + (1.0 - self.s) * np.exp(t * self.lam2)) + * u + ) + assert retval.shape == u.shape + return retval + + def int_f(self, u0: np.ndarray, dt: float, t: float = 0.0) -> np.ndarray: + assert isinstance(dt, float) + tfinal = t + dt + retval = (self.s * np.exp(tfinal * self.lam1) + (1.0 - self.s) * np.exp(tfinal * self.lam2)) * u0 + assert retval.shape == u0.shape + return retval diff --git a/qmat/playgrounds/martin/diff_eqs/de_solver.py b/qmat/playgrounds/martin/diff_eqs/de_solver.py new file mode 100644 index 0000000..e31b890 --- /dev/null +++ b/qmat/playgrounds/martin/diff_eqs/de_solver.py @@ -0,0 +1,133 @@ +from abc import ABC, abstractmethod +import numpy as np + + +class DESolver(ABC): + + @abstractmethod + def evalF(self, u: np.ndarray, t: float) -> np.ndarray: + """Evaluate the right-hand side of the equation. + + Parameters + ---------- + u : np.ndarray + Array of shape (N,) representing the solution at the current time step. + t : float + Current timestamp. + """ + pass + + def evalF1(self, u: np.ndarray, t: float) -> np.ndarray: + """Evaluate the 1st term of the equation + + Parameters + ---------- + u : np.ndarray + Array of shape (N,) representing the solution at the current time step. + t : float + Current timestamp. + """ + raise Exception("TODO: `evalF1` not implemented for this DE solver.") + + def evalF2(self, u: np.ndarray, t: float) -> np.ndarray: + """Evaluate the 2nd term of the equation + + Parameters + ---------- + u : np.ndarray + Array of shape (N,) representing the solution at the current time step. + t : float + Current timestamp. + """ + raise Exception("TODO: `evalF2` not implemented for this DE solver.") + + # This is optional since not every DE might have a solver for backward Euler + def fSolve(self, rhs: np.ndarray, dt: float, t: float) -> np.ndarray: + """Solve the right-hand side of an equation implicitly. + + # Solving this equation implicitly... + u_t = f(u, t) + + # ... means to u_new + u_new - dt * F(u_new, t + dt) = rhs + + Parameters + ---------- + rhs : np.ndarray + Right hand as given above. + t : float + Current timestamp. + Future one will be computed as `t + dt` + dt : float + Time step size. + """ + raise Exception("TODO: Implicit solver not implemented for this DE solver.") + + def fSolve1(self, rhs: np.ndarray, dt: float, t: float) -> np.ndarray: + """Solve the right-hand side of the 1st equation implicitly. + + # Solving this equation implicitly... + u_t = f(u, t) + + # ... means to u_new + u_new - dt * F(u_new, t + dt) = rhs + + Parameters + ---------- + rhs : np.ndarray + Right hand as given above. + t : float + Current timestamp. + Future one will be computed as `t + dt` + dt : float + Time step size. + """ + raise Exception("TODO: `fSolve1` not implemented for this DE solver.") + + def fSolve2(self, rhs: np.ndarray, dt: float, t: float) -> np.ndarray: + """Solve the right-hand side of the 1st equation implicitly. + + # Solving this equation implicitly... + u_t = f(u, t) + + # ... means to u_new + u_new - dt * F(u_new, t + dt) = rhs + + Parameters + ---------- + rhs : np.ndarray + Right hand as given above. + t : float + Current timestamp. + Future one will be computed as `t + dt` + dt : float + Time step size. + """ + raise Exception("TODO: `fSolve2` not implemented for this DE solver.") + + @abstractmethod + def initial_u0(self, mode: str) -> np.ndarray: + """Compute some initial conditions for the equation. + + Parameters + ---------- + mode : str + The type of initial condition to generate. + """ + pass + + @abstractmethod + def int_f(self, u0: np.ndarray, dt: float, t: float = 0.0) -> np.ndarray: + """ + Compute the (analytical) solution at time `dt` + t. + + Parameters + ---------- + u0 : np.ndarray + Array of shape (N,) representing the initial condition. + t : float + Time stamp of u0 + dt : float + Time step size + """ + pass diff --git a/qmat/playgrounds/martin/diff_eqs/two_freq.py b/qmat/playgrounds/martin/diff_eqs/two_freq.py new file mode 100644 index 0000000..0f045ee --- /dev/null +++ b/qmat/playgrounds/martin/diff_eqs/two_freq.py @@ -0,0 +1,99 @@ +import numpy as np +from qmat.playgrounds.martin.diff_eqs.de_solver import DESolver + + +class TwoFreq(DESolver): + """ + Test equation using a matrix to generate a + superposition of two frequencies in the solution. + """ + + def __init__(self, lam1: complex, lam2: complex, lam3: complex = None, lam4: complex = None): + """ + :param lam1: L[0,0] + :param lam2: L[0,1] + :param lam3: L[1,0] + :param lam4: L[1,1] + """ + + self.lam1: complex = lam1 + self.lam2: complex = lam2 + self.lam3: complex = lam3 if lam3 is not None else 0 + self.lam4: complex = lam4 if lam4 is not None else lam2 + + self.L = np.array([[self.lam1, self.lam2], [self.lam3, self.lam4]], dtype=np.complex128) + self.L1 = np.copy(self.L) + self.L1[:, 1] = 0 + self.L2 = np.copy(self.L) + self.L2[:, 0] = 0 + + # Compute eigenvalues and eigenvectors of L + self.eigvals, self.eigvecs = np.linalg.eig(self.L) + + if not np.all(np.isclose(np.real(self.eigvals), 0)): + raise Exception("Dahlquist3 matrix L must have purely imaginary eigenvalues") + + def initial_u0(self, mode: str = None) -> np.ndarray: + return np.array([1, 0.1], dtype=np.complex128) + + def evalF(self, u: np.ndarray, t: float) -> np.ndarray: + assert isinstance(t, float) + retval = self.L @ u + assert retval.shape == u.shape + return retval + + def evalF1(self, u: np.ndarray, t: float) -> np.ndarray: + assert isinstance(t, float) + retval = self.L1 @ u + assert retval.shape == u.shape + return retval + + def evalF2(self, u: np.ndarray, t: float) -> np.ndarray: + assert isinstance(t, float) + retval = self.L2 @ u + assert retval.shape == u.shape + return retval + + def fSolve(self, rhs: np.ndarray, dt: float, t: float) -> np.ndarray: + """ + Solve + `u_t = L u` + implicitly, i.e., return u_new such that + u_new = u + dt * L u_new + <=> u_new - dt * L u_new = u + <=> (I - dt * L) u_new = u + """ + assert isinstance(t, float) + retval = np.linalg.solve(np.eye(rhs.shape[0]) - dt*self.L, rhs) + assert retval.shape == rhs.shape + return retval + + def fSolve1(self, rhs: np.ndarray, dt: float, t: float) -> np.ndarray: + """ + Solve + `u_t = L1 u` + implicitly + """ + assert isinstance(t, float) + retval = np.linalg.solve(np.eye(rhs.shape[0]) - dt*self.L1, rhs) + assert retval.shape == rhs.shape + return retval + + def fSolve2(self, rhs: np.ndarray, dt: float, t: float) -> np.ndarray: + """ + Solve + `u_t = L2 u` + implicitly + """ + assert isinstance(t, float) + retval = np.linalg.solve(np.eye(rhs.shape[0]) - dt*self.L2, rhs) + assert retval.shape == rhs.shape + return retval + + def int_f(self, u0: np.ndarray, dt: float, t: float = 0.0) -> np.ndarray: + from scipy.linalg import expm + + assert isinstance(dt, float) + retval = expm(self.L * (t + dt)) @ u0 + assert retval.shape == u0.shape + return retval diff --git a/qmat/playgrounds/martin/ex_burgers_integrate.py b/qmat/playgrounds/martin/ex_burgers_integrate.py new file mode 100644 index 0000000..92d71ca --- /dev/null +++ b/qmat/playgrounds/martin/ex_burgers_integrate.py @@ -0,0 +1,77 @@ +# +# Example adopted from `test_4_sdc.py` +# + + +import numpy as np +from qmat.playgrounds.martin.diff_eqs.burgers import Burgers +from qmat.playgrounds.martin.time_integration.sdc_integration import SDCIntegration +from qmat.playgrounds.martin.time_integration.rk_integration import RKIntegration + + +from matplotlib import pyplot as plt + +N = 128 +nu = 0.2 +T: float = 4.0 +t: float = 0.0 # Starting point in time + +domain_size: float = 2.0 * np.pi +t_ = np.linspace(0, domain_size, N, endpoint=False) + +burgers: Burgers = Burgers(N=N, nu=nu, domain_size=domain_size) +u0 = burgers.initial_u0("sine") + +burgers.run_tests() + + +time_integration = "rk2" + +if 1: + results = [] + + u_analytical = burgers.int_f(u0, dt=T) + + for nt in range(4): + + num_timesteps = 2**nt * 1000 + print(f"Running simulation with num_timesteps={num_timesteps}") + + dt = T / num_timesteps + + burgers: Burgers = Burgers(N=N, nu=nu, domain_size=domain_size) + u0 = burgers.initial_u0("sine") + + u = u0.copy() + + if time_integration in RKIntegration.supported_methods: + rki = RKIntegration(method=time_integration) + u = rki.integrate_n(u, t, dt, num_timesteps, burgers) + + elif time_integration == "sdc": + sdci = SDCIntegration(num_nodes=3, node_type="LEGENDRE", quad_type="LOBATTO") + u = sdci.integrate_n(u, t, dt, num_timesteps, burgers) + + else: + raise Exception(f"Unsupported time integration method '{time_integration}'") + + plt.plot(t_, u, label=f"numerical {num_timesteps}", linestyle="dashed") + + error = np.max(np.abs(u - u_analytical)) + results.append({"N": num_timesteps, "dt": dt, "error": error}) + + prev_error = None + for r in results: + if prev_error is None: + conv = None + else: + conv = np.log2(prev_error / r["error"]) + + print(f"N={r["N"]}, dt={r["dt"]:.6e}, error={r["error"]:.6e}, conv={conv}") + prev_error = r["error"] + + if 1: + plt.plot(t_, u_analytical, label="analytical", linestyle="dotted") + + plt.legend() + plt.show() diff --git a/qmat/playgrounds/martin/ex_burgers_plot.py b/qmat/playgrounds/martin/ex_burgers_plot.py new file mode 100644 index 0000000..1f854b4 --- /dev/null +++ b/qmat/playgrounds/martin/ex_burgers_plot.py @@ -0,0 +1,40 @@ +# +# Example adopted from `test_4_sdc.py` +# + + +import numpy as np +from diff_eqs.burgers import Burgers + + +from matplotlib import pyplot as plt + +N = 128 +nu = 0.2 + +T: float = 4.0 +domain_size: float = 2.0 * np.pi +t = np.linspace(0, domain_size, N, endpoint=False) + +burgers: Burgers = Burgers(N=N, nu=nu, domain_size=domain_size) +u0 = burgers.initial_u0("sine") + +burgers.run_tests() + +if 1: + nt: int = 1000 + dt: float = T / nt + + if 0: + for t in [_ * dt for _ in range(nt)]: + ut = burgers.analytical_integration(u0, t=t) + print(f"t={t:.3f}, ut[0]={ut[0]:.6f}") + + plt.plot(t, ut) + else: + ut = burgers.int_f(u0, dt=T) + plt.plot(t, u0, label="u0") + plt.plot(t, ut, label="ut") + + plt.legend() + plt.show() diff --git a/qmat/playgrounds/martin/ex_dahlquist2_integrate.py b/qmat/playgrounds/martin/ex_dahlquist2_integrate.py new file mode 100644 index 0000000..b4291b3 --- /dev/null +++ b/qmat/playgrounds/martin/ex_dahlquist2_integrate.py @@ -0,0 +1,83 @@ +import numpy as np +from qmat.playgrounds.martin.diff_eqs.dahlquist2 import Dahlquist2 +#from qmat.playgrounds.martin.diff_eqs.dahlquist import Dahlquist +from time_integration.sdc_integration import SDCIntegration +from qmat.playgrounds.martin.time_integration.rk_integration import RKIntegration +from matplotlib import pyplot as plt + + +u0 = np.array([1.0]) # Initial condition +T: float = 2 * np.pi # Time interval +t: float = 0.0 # Starting time + +time_integration = "rk4" + +dahlquist2: Dahlquist2 = Dahlquist2(lam1=20.0j, lam2=1.0j, s=0.1) + +results = [] + +plt.close() + +for nt in range(1): + + num_timesteps = 2**nt * 300 + print(f"Running simulation with num_timesteps={num_timesteps}") + + dt = T / num_timesteps + + u_analytical_fin = dahlquist2.int_f(u0, dt=T) + u0 = dahlquist2.initial_u0() + + u = u0.copy() + + u_: np.ndarray = np.array([u]) + t_: np.ndarray = np.array([t]) + + if time_integration in RKIntegration.supported_methods: + rki = RKIntegration(method=time_integration) + + for n in range(num_timesteps): + u = rki.int_f(u, t + n * dt, dt, dahlquist2) + + u_ = np.concatenate((u_, np.expand_dims(u, axis=0))) + t_ = np.append(t_, t + (n + 1) * dt) + + elif time_integration == "sdc": + sdci = SDCIntegration(num_nodes=3, node_type="LEGENDRE", quad_type="LOBATTO") + + for n in range(num_timesteps): + u = sdci.int_f(u, t + n * dt, dt, dahlquist2) + + u_ = np.concatenate((u_, np.expand_dims(u, axis=0))) + t_ = np.append(t_, t + (n + 1) * dt) + + else: + raise Exception(f"Unsupported time integration method '{time_integration}'") + + plt.plot(t_, np.real(u_), label=f"ndt={num_timesteps}, real", linestyle="dashed") + plt.plot(t_, np.imag(u_), label=f"ndt={num_timesteps}, imag", linestyle="solid") + + error = np.max(np.abs(u - u_analytical_fin)) + results.append({"N": num_timesteps, "dt": dt, "error": error}) + +prev_error = None +for r in results: + if prev_error is None: + conv = None + else: + conv = np.log2(prev_error / r["error"]) + + print(f"N={r["N"]}, dt={r["dt"]:.6e}, error={r["error"]:.6e}, conv={conv}") + prev_error = r["error"] + +if 1: + t_ = np.linspace(0, T, 1000) + u_analytical_fin = np.array([dahlquist2.int_f(u0, t) for t in t_]) + plt.plot(t_, np.real(u_analytical_fin), linestyle="dotted", color="gray") + plt.plot(t_, np.imag(u_analytical_fin), linestyle="dotted", color="gray") + +if 1: + plt.ylim(-1.5, 1.5) + + plt.legend() + plt.show() diff --git a/qmat/playgrounds/martin/ex_dahlquist2_plot.py b/qmat/playgrounds/martin/ex_dahlquist2_plot.py new file mode 100644 index 0000000..6da0e78 --- /dev/null +++ b/qmat/playgrounds/martin/ex_dahlquist2_plot.py @@ -0,0 +1,18 @@ +import numpy as np +from matplotlib import pyplot as plt +from qmat.playgrounds.martin.diff_eqs.dahlquist2 import Dahlquist2 + + +N = 500 +u0 = np.array([1.0]) +t: np.array = np.linspace(0, 4*np.pi, N, endpoint=False) + +dahlquist2: Dahlquist2 = Dahlquist2(lam1=20.0j, lam2=1.0j, s=0.1) + +u_eval = np.array([dahlquist2.int_f(u0, _) for _ in t]) + + +plt.plot(t, np.real(u_eval), label="Re(u)") +plt.plot(t, np.imag(u_eval), label="Im(u)") +plt.legend() +plt.show() diff --git a/qmat/playgrounds/martin/ex_dahlquist_plot.py b/qmat/playgrounds/martin/ex_dahlquist_plot.py new file mode 100644 index 0000000..f9daa64 --- /dev/null +++ b/qmat/playgrounds/martin/ex_dahlquist_plot.py @@ -0,0 +1,18 @@ +import numpy as np +from matplotlib import pyplot as plt +from qmat.playgrounds.martin.diff_eqs.dahlquist import Dahlquist + + +N = 500 +u0 = np.array([1.0]) +t: np.array = np.linspace(0, 4*np.pi, N, endpoint=False) + +dahlquist: Dahlquist = Dahlquist(lam1=20.0j, lam2=1.0j) + +u_eval = np.array([dahlquist.int_f(u0, dt=_) for _ in t]) + + +plt.plot(t, np.real(u_eval), label="Re(u)") +plt.plot(t, np.imag(u_eval), label="Im(u)") +plt.legend() +plt.show() diff --git a/qmat/playgrounds/martin/ex_two_freq_integrate.py b/qmat/playgrounds/martin/ex_two_freq_integrate.py new file mode 100644 index 0000000..669dfa9 --- /dev/null +++ b/qmat/playgrounds/martin/ex_two_freq_integrate.py @@ -0,0 +1,96 @@ +import numpy as np +from qmat.playgrounds.martin.diff_eqs.two_freq import TwoFreq +from qmat.playgrounds.martin.time_integration.sdc_integration import SDCIntegration +from qmat.playgrounds.martin.time_integration.rk_integration import RKIntegration +from matplotlib import pyplot as plt + + +T: float = 2 * 2 * np.pi # Time interval +t: float = 0.0 # Starting time + +time_integration: str = "sdc" +sdc_micro_time_integration: str = "irk1" +sdc_micro_time_integration: str = "imex12" +sdc_num_sweeps: int = 4 + +two_freq: TwoFreq = TwoFreq(lam1=1.0j, lam2=20.0j, lam3=0.5j) +u0 = two_freq.initial_u0() + +results = [] + +plt.close() + +for nt in range(1): + + num_timesteps = 2**nt * 200 + print(f"Running simulation with num_timesteps={num_timesteps}") + + dt = T / num_timesteps + + u_analytical_fin = two_freq.int_f(u0, dt=T) + u0 = two_freq.initial_u0() + + u = u0.copy() + + u_: np.ndarray = np.array([u]) + t_: np.ndarray = np.array([t]) + + if time_integration in RKIntegration.supported_methods: + rki = RKIntegration(method=time_integration) + + for n in range(num_timesteps): + u = rki.int_f(u, t + n * dt, dt, two_freq) + + u_ = np.concatenate((u_, np.expand_dims(u, axis=0))) + t_ = np.append(t_, t + (n + 1) * dt) + + elif time_integration == "sdc": + sdci = SDCIntegration( + num_nodes=5, + node_type="LEGENDRE", + quad_type="LOBATTO", + num_sweeps=sdc_num_sweeps, + micro_time_integration=sdc_micro_time_integration, + ) + + for n in range(num_timesteps): + u = sdci.int_f(u, t + n * dt, dt, two_freq) + + u_ = np.concatenate((u_, np.expand_dims(u, axis=0))) + t_ = np.append(t_, t + (n + 1) * dt) + + else: + raise Exception(f"Unsupported time integration method '{time_integration}'") + + for i in range(u_.shape[1]): + plt.plot(t_, np.real(u_[:, i]), label=f"u[{i}].real", linestyle="dashed") + plt.plot(t_, np.imag(u_[:, i]), label=f"u[{i}].imag", linestyle="solid") + + error = np.max(np.abs(u - u_analytical_fin)) + results.append({"N": num_timesteps, "dt": dt, "error": error}) + + +if 1: + # Plot analytical solution + t_ = np.linspace(0, T, 1000) + u_analytical_fin = np.array([two_freq.int_f(u0, t) for t in t_]) + for i in range(u_analytical_fin.shape[1]): + plt.plot(t_, np.real(u_analytical_fin[:, i]), linestyle="dotted", color="black") + plt.plot(t_, np.imag(u_analytical_fin[:, i]), linestyle="dotted", color="black") + + +prev_error = None +for r in results: + if prev_error is None: + conv = None + else: + conv = np.log2(prev_error / r["error"]) + + print(f"N={r["N"]}, dt={r["dt"]:.6e}, error={r["error"]:.6e}, conv={conv}") + prev_error = r["error"] + +if 1: + # plt.ylim(-1.5, 1.5) + + plt.legend() + plt.show() diff --git a/qmat/playgrounds/martin/ex_two_freq_plot.py b/qmat/playgrounds/martin/ex_two_freq_plot.py new file mode 100644 index 0000000..c695e8c --- /dev/null +++ b/qmat/playgrounds/martin/ex_two_freq_plot.py @@ -0,0 +1,19 @@ +import numpy as np +from matplotlib import pyplot as plt +from qmat.playgrounds.martin.diff_eqs.two_freq import TwoFreq + + +N = 500 +t: np.array = np.linspace(0, 4*np.pi, N, endpoint=False) + +two_freq: TwoFreq = TwoFreq(lam1=1.0j, lam2=20.0j, lam3=0.5j) + +u0 = two_freq.initial_u0() +u_eval = np.array([two_freq.int_f(u0, dt=_) for _ in t]) + +for i in range(2): + plt.plot(t, np.real(u_eval[:, i]), label=f"Re(u[{i}])") + plt.plot(t, np.imag(u_eval[:, i]), label=f"Im(u[{i}])") + +plt.legend() +plt.show() diff --git a/qmat/playgrounds/martin/nyquist_test/ex_slow_fast.py b/qmat/playgrounds/martin/nyquist_test/ex_slow_fast.py new file mode 100644 index 0000000..d65aee1 --- /dev/null +++ b/qmat/playgrounds/martin/nyquist_test/ex_slow_fast.py @@ -0,0 +1,104 @@ +import numpy as np +from qmat.playgrounds.martin.diff_eqs.two_freq import TwoFreq +from qmat.playgrounds.martin.diff_eqs.dahlquist import Dahlquist +from qmat.playgrounds.martin.time_integration.sdc_integration import SDCIntegration +from qmat.playgrounds.martin.time_integration.rk_integration import RKIntegration +from matplotlib import pyplot as plt + + +s: int = 1 +T: float = 1.0 * np.pi * s # Time interval +t: float = 0.0 # Starting time + +time_integration: str = "sdc" +sdc_micro_time_integration: str = "irk1" +sdc_micro_time_integration: str = "erk1" +sdc_micro_time_integration: str = "imex21" +sdc_num_sweeps: int = 6 + +de_solver: Dahlquist = Dahlquist(lam1=1.0j, lam2=10.0j, ext_scalar=0.5) +u0 = de_solver.initial_u0() + +results = [] + +plt.close() + +if 1: + num_timesteps = 1 + print(f"Running simulation with num_timesteps={num_timesteps}") + + dt = T / num_timesteps + + u_analytical_fin = de_solver.int_f(u0, dt=T, t=0) + u0 = de_solver.initial_u0() + + u = u0.copy() + + u_: np.ndarray = np.array([u]) + t_: np.ndarray = np.array([t]) + + if time_integration in RKIntegration.supported_methods: + rki = RKIntegration(method=time_integration) + + for n in range(num_timesteps): + u = rki.int_f(u, t + n * dt, dt, de_solver) + + u_ = np.concatenate((u_, np.expand_dims(u, axis=0))) + t_ = np.append(t_, t + (n + 1) * dt) + + elif time_integration == "sdc": + sdci = SDCIntegration( + num_nodes=6, + # node_type="LEGENDRE", + node_type="EQUID", + quad_type="LOBATTO", + num_sweeps=sdc_num_sweeps, + micro_time_integration=sdc_micro_time_integration, + use_quadrature=False, + ) + + for n in range(num_timesteps): + u = sdci.int_f(u, t + n * dt, dt, de_solver) + + u_ = np.concatenate((u_, np.expand_dims(u, axis=0))) + t_ = np.append(t_, t + (n + 1) * dt) + + else: + raise Exception(f"Unsupported time integration method '{time_integration}'") + + # for i in range(u_.shape[1]): + for i in [0]: + # plt.plot(t_, np.real(u_[:, i]), label=f"u[{i}].real", linestyle="dashed") + plt.plot(t_, np.imag(u_[:, i]), label=f"u[{i}].imag", linestyle="solid") + + for t in t_: + plt.vlines(t, ymin=-1, ymax=1, colors="gray", linewidth=1) + + error = np.max(np.abs(u - u_analytical_fin)) + results.append({"N": num_timesteps, "dt": dt, "error": error}) + + +if 1: + # Plot analytical solution + t_ = np.linspace(0, T, 1000) + u_analytical_fin = np.array([de_solver.int_f(u0, dt=t, t=0) for t in t_]) + # for i in range(u_analytical_fin.shape[1]): + for i in [0]: + # plt.plot(t_, np.real(u_analytical_fin[:, i]), linestyle="dotted", color="black") + plt.plot(t_, np.imag(u_analytical_fin[:, i]), linestyle="dotted", color="black") + + +if 0: + prev_error = None + for r in results: + if prev_error is None: + conv = None + else: + conv = np.log2(prev_error / r["error"]) + + print(f"N={r["N"]}, dt={r["dt"]:.6e}, error={r["error"]:.6e}, conv={conv}") + prev_error = r["error"] + +if 1: + plt.legend() + plt.show() diff --git a/qmat/playgrounds/martin/tests/__init__.py b/qmat/playgrounds/martin/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/qmat/playgrounds/martin/tests/test_burgers.py b/qmat/playgrounds/martin/tests/test_burgers.py new file mode 100644 index 0000000..4517504 --- /dev/null +++ b/qmat/playgrounds/martin/tests/test_burgers.py @@ -0,0 +1,81 @@ + +import numpy as np +from qmat.playgrounds.martin.diff_eqs.burgers import Burgers +from qmat.playgrounds.martin.time_integration.sdc_integration import SDCIntegration +from qmat.playgrounds.martin.time_integration.rk_integration import RKIntegration + + +def test_burgers(): + N = 128 + nu = 0.2 + + T: float = 4.0 + t: float = 0.0 # Starting point in time + domain_size: float = 2.0 * np.pi + + burgers: Burgers = Burgers(N=N, nu=nu, domain_size=domain_size) + u0 = burgers.initial_u0("sine") + + burgers.run_tests() + + for time_integration_method in ["rk1", "rk2", "rk4", "sdc"]: + print("="*80) + print(f"Time integration method: {time_integration_method}") + print("="*80) + results = [] + + u_analytical = burgers.int_f(u0, dt=T) + + for nt in range(2, 4): + + num_timesteps = 2**nt * 500 + print(f"Running simulation with num_timesteps={num_timesteps}") + + dt = T / num_timesteps + + burgers: Burgers = Burgers(N=N, nu=nu, domain_size=domain_size) + u0 = burgers.initial_u0("sine") + + u = u0.copy() + + if time_integration_method in RKIntegration.supported_methods: + rki = RKIntegration(method=time_integration_method) + u = rki.integrate_n(u, t, dt, num_timesteps, burgers) + + elif time_integration_method == "sdc": + sdci = SDCIntegration(num_nodes=3, node_type="LEGENDRE", quad_type="LOBATTO") + u = sdci.integrate_n(u, t, dt, num_timesteps, burgers) + + else: + raise Exception("TODO") + + error = np.max(np.abs(u - u_analytical)) + results.append({"N": num_timesteps, "dt": dt, "error": error}) + + prev_error = None + for r in results: + if prev_error is None: + conv = None + else: + conv = np.log2(prev_error / r["error"]) + + print(f" - N={r["N"]}, dt={r["dt"]:.6e}, error={r["error"]:.6e}, conv={conv}") + prev_error = r["error"] + r["conv"] = conv + + if time_integration_method == "rk1": + assert results[-1]["error"] < 1e-4 + assert np.abs(results[-1]["conv"]-1.0) < 1e-3 + + elif time_integration_method == "rk2": + assert results[-1]["error"] < 1e-7 + assert np.abs(results[-1]["conv"]-2.0) < 1e-3 + + elif time_integration_method == "rk4": + assert results[-1]["error"] < 1e-14 + + elif time_integration_method == "sdc": + assert results[-1]["error"] < 1e-14 + + else: + raise Exception(f"TODO for {time_integration_method}") diff --git a/qmat/playgrounds/martin/tests/test_dahlquist.py b/qmat/playgrounds/martin/tests/test_dahlquist.py new file mode 100644 index 0000000..e8ec065 --- /dev/null +++ b/qmat/playgrounds/martin/tests/test_dahlquist.py @@ -0,0 +1,102 @@ +import numpy as np +from qmat.playgrounds.martin.diff_eqs.dahlquist import Dahlquist +from time_integration.sdc_integration import SDCIntegration +from qmat.playgrounds.martin.time_integration.rk_integration import RKIntegration + + +def test_dahlquist_integration(): + u0 = np.array([1.0]) # Initial condition + T: float = 4 * np.pi # Time interval + T: float = 0.5 # Time interval + t: float = 0.0 # Starting time + + dahlquist: Dahlquist = Dahlquist(lam1=1.0j, lam2=0.1j) + + for time_integration_method in ["rk1", "rk2", "rk4", "sdc"]: + if time_integration_method == "sdc": + micro_time_integration_ = ["erk1", "irk1", "imex12", "imex21"] + else: + micro_time_integration_ = ["-"] + + for micro_time_integration in micro_time_integration_: + print("=" * 80) + print(f"Time integration method: {time_integration_method} ({micro_time_integration})") + print("=" * 80) + results = [] + + u_analytical = dahlquist.int_f(u0, dt=T) + + for nt in range(4): + + num_timesteps = 2**nt * 10 + print(f"Running simulation with num_timesteps={num_timesteps}") + + dt = T / num_timesteps + u0 = dahlquist.initial_u0() + u = u0.copy() + + if time_integration_method in RKIntegration.supported_methods: + rki = RKIntegration(method=time_integration_method) + u = rki.integrate_n(u, t, dt, num_timesteps, dahlquist) + + elif time_integration_method == "sdc": + sdci = SDCIntegration( + num_nodes=3, + node_type="LEGENDRE", + quad_type="LOBATTO", + num_sweeps=3, + micro_time_integration=micro_time_integration, + ) + u = sdci.integrate_n(u, t, dt, num_timesteps, dahlquist) + + else: + raise Exception("TODO") + + error = np.max(np.abs(u - u_analytical)) + results.append({"N": num_timesteps, "dt": dt, "error": error, "mti": micro_time_integration}) + + prev_error = None + for r in results: + if prev_error is None: + conv = None + else: + conv = np.log2(prev_error / r["error"]) + + print(f" - N={r["N"]}, dt={r["dt"]:.6e}, error={r["error"]:.6e}, conv={conv}, mti={micro_time_integration}") + prev_error = r["error"] + r["conv"] = conv + + if time_integration_method == "rk1": + assert results[-1]["error"] < 1e-2 + assert np.abs(results[-1]["conv"] - 1.0) < 1e-2 + + elif time_integration_method == "rk2": + assert results[-1]["error"] < 1e-4 + assert np.abs(results[-1]["conv"] - 2.0) < 1e-3 + + elif time_integration_method == "rk4": + assert results[-1]["error"] < 1e-9 + assert np.abs(results[-1]["conv"] - 4.0) < 1e-4 + + elif time_integration_method == "sdc": + assert results[-1]["error"] < 1e-8 + assert np.abs(results[-1]["conv"] - 4.0) < 1e-3 + + else: + raise Exception(f"TODO for {time_integration_method}") + + +def test_dahlquist_integration_of_solution(): + u0 = np.array([1.0]) # Initial condition + T: float = 4 * np.pi # Time interval + + dahlquist: Dahlquist = Dahlquist(lam1=1.0j, lam2=0.1j) + + u0 = dahlquist.initial_u0() + + u_analytical = dahlquist.int_f_t0(u0, dt=T) + + u1 = dahlquist.int_f(u0, dt=T*0.5, t=0.0) + u2 = dahlquist.int_f(u1, dt=T*0.5, t=T*0.5) + + assert np.all(np.isclose(u_analytical, u2)) diff --git a/qmat/playgrounds/martin/tests/test_dahlquist2.py b/qmat/playgrounds/martin/tests/test_dahlquist2.py new file mode 100644 index 0000000..8936ee4 --- /dev/null +++ b/qmat/playgrounds/martin/tests/test_dahlquist2.py @@ -0,0 +1,77 @@ +import numpy as np +from qmat.playgrounds.martin.diff_eqs.dahlquist2 import Dahlquist2 +from time_integration.sdc_integration import SDCIntegration +from qmat.playgrounds.martin.time_integration.rk_integration import RKIntegration + + +def test_dahlquist2(): + u0 = np.array([1.0]) # Initial condition + T: float = 4 * np.pi # Time interval + T: float = 0.5 # Time interval + t: float = 0.0 # Starting time + + dahlquist2: Dahlquist2 = Dahlquist2(lam1=1.0j, lam2=0.1j) + + for time_integration_method in ["rk1", "rk2", "rk4", "sdc"]: + print("="*80) + print(f"Time integration method: {time_integration_method}") + print("="*80) + results = [] + + u_analytical = dahlquist2.int_f(u0, dt=T) + + for nt in range(4): + + num_timesteps = 2**nt * 10 + print(f"Running simulation with num_timesteps={num_timesteps}") + + dt = T / num_timesteps + + u0 = dahlquist2.initial_u0() + + u = u0.copy() + + if time_integration_method in RKIntegration.supported_methods: + rki = RKIntegration(method=time_integration_method) + + u = rki.integrate_n(u, t, dt, num_timesteps, dahlquist2) + + elif time_integration_method == "sdc": + sdci = SDCIntegration(num_nodes=3, node_type="LEGENDRE", quad_type="LOBATTO", num_sweeps=1) + u = sdci.integrate_n(u, t, dt, num_timesteps, dahlquist2) + + else: + raise Exception("TODO") + + error = np.max(np.abs(u - u_analytical)) + results.append({"N": num_timesteps, "dt": dt, "error": error}) + + prev_error = None + for r in results: + if prev_error is None: + conv = None + else: + conv = np.log2(prev_error / r["error"]) + + print(f" - N={r["N"]}, dt={r["dt"]:.6e}, error={r["error"]:.6e}, conv={conv}") + prev_error = r["error"] + r["conv"] = conv + + if time_integration_method == "rk1": + assert results[-1]["error"] < 1e-2 + assert np.abs(results[-1]["conv"]-1.0) < 1e-2 + + elif time_integration_method == "rk2": + assert results[-1]["error"] < 1e-5 + assert np.abs(results[-1]["conv"]-2.0) < 1e-3 + + elif time_integration_method == "rk4": + assert results[-1]["error"] < 1e-11 + assert np.abs(results[-1]["conv"]-4.0) < 1e-2 + + elif time_integration_method == "sdc": + assert results[-1]["error"] < 1e-5 + assert np.abs(results[-1]["conv"]-2.0) < 1e-3 + + else: + raise Exception(f"TODO for {time_integration_method}") diff --git a/qmat/playgrounds/martin/tests/test_dahlquist_ext_freq.py b/qmat/playgrounds/martin/tests/test_dahlquist_ext_freq.py new file mode 100644 index 0000000..403e19a --- /dev/null +++ b/qmat/playgrounds/martin/tests/test_dahlquist_ext_freq.py @@ -0,0 +1,97 @@ +import numpy as np +from qmat.playgrounds.martin.diff_eqs.dahlquist import Dahlquist +from time_integration.sdc_integration import SDCIntegration +from qmat.playgrounds.martin.time_integration.rk_integration import RKIntegration + + +def test_dahlquist(): + u0 = np.array([1.0]) # Initial condition + T: float = 4 * np.pi # Time interval + T: float = 0.5 # Time interval + t: float = 0.0 # Starting time + + dahlquist: Dahlquist = Dahlquist(lam1=1.0j, lam2=0.1j, ext_scalar=1.0) + + for time_integration_method in ["rk1", "rk2", "rk4", "irk1", "irk2", "sdc"]: + if time_integration_method == "sdc": + micro_time_integration_ = ["erk1", "irk1"] + else: + micro_time_integration_ = ["-"] + + for micro_time_integration in micro_time_integration_: + print("=" * 80) + print(f"Time integration method: {time_integration_method} ({micro_time_integration})") + print("=" * 80) + results = [] + + u_analytical = dahlquist.int_f(u0, dt=T) + + for nt in range(4): + + num_timesteps = 2**nt * 10 + print(f"Running simulation with num_timesteps={num_timesteps}") + + dt = T / num_timesteps + + u0 = dahlquist.initial_u0() + + u = u0.copy() + + if time_integration_method in RKIntegration.supported_methods: + rki = RKIntegration(method=time_integration_method) + + u = rki.integrate_n(u, t, dt, num_timesteps, dahlquist) + + elif time_integration_method == "sdc": + sdci = SDCIntegration( + num_nodes=3, + node_type="LEGENDRE", + quad_type="LOBATTO", + num_sweeps=3, + micro_time_integration=micro_time_integration, + ) + u = sdci.integrate_n(u, t, dt, num_timesteps, dahlquist) + + else: + raise Exception("TODO") + + error = np.max(np.abs(u - u_analytical)) + results.append({"N": num_timesteps, "dt": dt, "error": error, "mti": micro_time_integration}) + + prev_error = None + for r in results: + if prev_error is None: + conv = None + else: + conv = np.log2(prev_error / r["error"]) + + print(f" - N={r["N"]}, dt={r["dt"]:.6e}, error={r["error"]:.6e}, conv={conv}, mti={micro_time_integration}") + prev_error = r["error"] + r["conv"] = conv + + if time_integration_method == "rk1": + assert results[-1]["error"] < 1e-2 + assert np.abs(results[-1]["conv"] - 1.0) < 1e-2 + + elif time_integration_method == "rk2": + assert results[-1]["error"] < 1e-4 + assert np.abs(results[-1]["conv"] - 2.0) < 1e-2 + + elif time_integration_method == "rk4": + assert results[-1]["error"] < 1e-9 + assert np.abs(results[-1]["conv"] - 4.0) < 1e-2 + + elif time_integration_method == "irk1": + assert results[-1]["error"] < 1e-2 + assert np.abs(results[-1]["conv"] - 1.0) < 1e-2 + + elif time_integration_method == "irk2": + assert results[-1]["error"] < 1e-2 + assert np.abs(results[-1]["conv"] - 2.0) < 1e-2 + + elif time_integration_method == "sdc": + assert results[-1]["error"] < 1e-8 + assert np.abs(results[-1]["conv"] - 4.0) < 1e-1 + + else: + raise Exception(f"TODO for {time_integration_method}") diff --git a/qmat/playgrounds/martin/tests/test_two_freq.py b/qmat/playgrounds/martin/tests/test_two_freq.py new file mode 100644 index 0000000..58a056e --- /dev/null +++ b/qmat/playgrounds/martin/tests/test_two_freq.py @@ -0,0 +1,97 @@ +import numpy as np +from qmat.playgrounds.martin.diff_eqs.two_freq import TwoFreq +from time_integration.sdc_integration import SDCIntegration +from qmat.playgrounds.martin.time_integration.rk_integration import RKIntegration + + +def test_dahlquist2(): + T: float = 4 * np.pi # Time interval + T: float = 0.5 # Time interval + t: float = 0.0 # Starting time + + two_freq: TwoFreq = TwoFreq(lam1=1.0j, lam2=0.1j) + u0 = two_freq.initial_u0() + + for time_integration_method in ["rk1", "rk2", "rk4", "irk1", "irk2", "sdc"]: + if time_integration_method == "sdc": + micro_time_integration_ = ["erk1", "irk1", "imex12", "imex21"] + else: + micro_time_integration_ = ["-"] + + for micro_time_integration in micro_time_integration_: + print("=" * 80) + print(f"Time integration method: {time_integration_method} ({micro_time_integration})") + print("=" * 80) + results = [] + + u_analytical = two_freq.int_f(u0, dt=T) + + for nt in range(4): + + num_timesteps = 2**nt * 10 + print(f"Running simulation with num_timesteps={num_timesteps}") + + dt = T / num_timesteps + + u0 = two_freq.initial_u0() + + u = u0.copy() + + if time_integration_method in RKIntegration.supported_methods: + rki = RKIntegration(method=time_integration_method) + + u = rki.integrate_n(u, t, dt, num_timesteps, two_freq) + + elif time_integration_method == "sdc": + sdci = SDCIntegration( + num_nodes=3, + node_type="LEGENDRE", + quad_type="LOBATTO", + num_sweeps=3, + micro_time_integration=micro_time_integration, + ) + u = sdci.integrate_n(u, t, dt, num_timesteps, two_freq) + + else: + raise Exception("TODO") + + error = np.max(np.abs(u - u_analytical)) + results.append({"N": num_timesteps, "dt": dt, "error": error}) + + prev_error = None + for r in results: + if prev_error is None: + conv = None + else: + conv = np.log2(prev_error / r["error"]) + + print(f" - N={r["N"]}, dt={r["dt"]:.6e}, error={r["error"]:.6e}, conv={conv}, mti={micro_time_integration}") + prev_error = r["error"] + r["conv"] = conv + + if time_integration_method == "rk1": + assert results[-1]["error"] < 1e-2 + assert np.abs(results[-1]["conv"] - 1.0) < 1e-2 + + elif time_integration_method == "rk2": + assert results[-1]["error"] < 1e-5 + assert np.abs(results[-1]["conv"] - 2.0) < 1e-3 + + elif time_integration_method == "rk4": + assert results[-1]["error"] < 1e-11 + assert np.abs(results[-1]["conv"] - 4.0) < 1e-2 + + elif time_integration_method == "irk1": + assert results[-1]["error"] < 1e-2 + assert np.abs(results[-1]["conv"] - 1.0) < 1e-2 + + elif time_integration_method == "irk2": + assert results[-1]["error"] < 1e-2 + assert np.abs(results[-1]["conv"] - 2.0) < 1e-2 + + elif time_integration_method == "sdc": + assert results[-1]["error"] < 1e-8 + assert np.abs(results[-1]["conv"] - 4.0) < 1e-3 + + else: + raise Exception(f"TODO for {time_integration_method}") diff --git a/qmat/playgrounds/martin/time_integration/rk_integration.py b/qmat/playgrounds/martin/time_integration/rk_integration.py new file mode 100644 index 0000000..1cb2fd2 --- /dev/null +++ b/qmat/playgrounds/martin/time_integration/rk_integration.py @@ -0,0 +1,61 @@ +from typing import List +import numpy as np +from qmat.playgrounds.martin.diff_eqs.de_solver import DESolver + + +class RKIntegration: + + # List of supported time integration methods + supported_methods: List[str] = ["rk1", "rk2", "rk4", "irk1", "irk2"] + + def __init__(self, method: str): + assert method in self.supported_methods, "Unsupported RK method" + self.method = method + + def int_f(self, u0: np.array, t: float, dt: float, de_solver: DESolver) -> np.array: + u = u0 + + if self.method == "rk1": + # RK1 (Forward Euler) + return u + dt * de_solver.evalF(u, t) + + elif self.method == "rk2": + # RK2: Heun's method + k1 = de_solver.evalF(u, t) + k2 = de_solver.evalF(u + 0.5 * dt * k1, t + 0.5 * dt) + return u + dt * k2 + + elif self.method == "rk4": + # Classical RK4 method + k1 = de_solver.evalF(u, t) + k2 = de_solver.evalF(u + 0.5 * dt * k1, t + 0.5 * dt) + k3 = de_solver.evalF(u + 0.5 * dt * k2, t + 0.5 * dt) + k4 = de_solver.evalF(u + dt * k3, t + dt) + + # Update solution + return u + dt * (k1 + 2 * k2 + 2 * k3 + k4) / 6 + + elif self.method == "irk1": + # IRK1 (Implicit/backward Euler) + return de_solver.fSolve(u, dt, t) + + elif self.method == "irk2": + # IRK2 (Crank-Nicolson method) + # Implicit step + u = de_solver.fSolve(u, 0.5*dt, t) + # Forward step + u += 0.5 * dt * de_solver.evalF(u, t + 0.5*dt) + return u + + else: + raise Exception("TODO") + + return u + + def integrate_n(self, u0: np.array, t: float, dt: float, num_timesteps, de_solver: DESolver) -> np.array: + u_value = u0 + + for n in range(num_timesteps): + u_value = self.int_f(u_value, t + n * dt, dt, de_solver) + + return u_value diff --git a/qmat/playgrounds/martin/time_integration/sdc_integration.py b/qmat/playgrounds/martin/time_integration/sdc_integration.py new file mode 100644 index 0000000..a40f45d --- /dev/null +++ b/qmat/playgrounds/martin/time_integration/sdc_integration.py @@ -0,0 +1,398 @@ +import numpy as np +from qmat.playgrounds.martin.diff_eqs.de_solver import DESolver + + +class SDCIntegration: + def __init__( + self, + num_nodes: int = 3, + node_type: str = "LOBATTO", # Basic node types to generate SDC nodes + quad_type: str = "LOBATTO", # 'LOBATTO': Always include {0, 1} in quadrature points. Add them if they don't exist. + num_sweeps: int = None, + micro_time_integration: str = None, # 'erk1' = explicit Euler, 'irk1' = implicit Euler, 'imex' = implicit-explicit + use_quadrature: bool = True, # Use final quadrature to compute final solution + ): + from qmat.qcoeff.collocation import Collocation + import qmat.qdelta.timestepping as module + + coll = Collocation(nNodes=num_nodes, nodeType=node_type, quadType=quad_type) + + self.gen: module.FE = module.FE(coll.nodes) + + self.nodes, self.weights, self.q = coll.genCoeffs(form="N2N") + + self.q_delta: np.ndarray = self.gen.getQDelta() + # Deltas are the \tau + self.deltas: np.ndarray = self.gen.deltas + + # Number of nodes + self.N = len(self.nodes) + + print(f"self.nodes: {self.nodes}") + print(f"self.deltas: {self.deltas}") + + if num_sweeps is None: + self.num_sweeps = len(self.nodes) + else: + self.num_sweeps = num_sweeps + + # Time integration to be used within SDC sweeps + # 'erk1' = explicit Euler + # 'irk1' = implicit Euler + # 'imex' = implicit-explicit => 'imex12' + # 'imex12' = implicit-explicit: 1st term treated implicitly, 2nd term explicitly + # 'imex21' = implicit-explicit: 2nd term treated implicitly, 1st term explicitly + self.time_integration_method = micro_time_integration if micro_time_integration is not None else "erk1" + + if self.time_integration_method == "imex": + self.time_integration_method = "imex12" + + self.use_quadrature: bool = use_quadrature + + assert self.num_sweeps >= 1 + + def integrate_erk1(self, u0: np.array, t: float, dt: float, de_solver: DESolver) -> np.array: + if not np.isclose(self.nodes[0], 0.0): + raise Exception("SDC nodes must include the left interval boundary.") + + if not np.isclose(self.nodes[-1], 1.0): + raise Exception("SDC nodes must include the right interval boundary.") + + assert self.N == self.deltas.shape[0] + + shape = (self.N,) + u0.shape + u = np.zeros_like(u0, shape=shape) + + u[0, :] = u0 + evalF_k0 = np.empty_like(u) + evalF_k1 = np.empty_like(u) + + # + # Propagate initial condition to all nodes + # + for m in range(0, self.N): + if m > 0: + u[m] = u[m - 1] + dt * self.deltas[m] * evalF_k0[m - 1] + evalF_k0[m] = de_solver.evalF(u[m], t + dt * self.nodes[m]) + + # + # Iteratively sweep over SDC nodes + # + for _ in range(1, self.num_sweeps): + for m in range(0, self.N): + + if m > 0: + qeval = self.q[m] @ evalF_k0 + u[m] = u[m - 1] + dt * (self.deltas[m] * (evalF_k1[m - 1] - evalF_k0[m - 1]) + qeval) + + evalF_k1[m] = de_solver.evalF(u[m], t + dt * self.nodes[m]) + + # Copy tendency arrays + evalF_k0[:] = evalF_k1[:] + + if self.use_quadrature: + # Compute new starting value with quadrature on tendencies + u[0] = u[0] + dt * self.weights @ evalF_k0 + else: + # If we're using Radau-right, we can just use the last value + u[0] = u[-1] + + assert u0.shape == u[0].shape + return u[0] + + def integrate_irk1(self, u0: np.array, t: float, dt: float, de_solver: DESolver) -> np.array: + if not np.isclose(self.nodes[0], 0.0): + raise Exception("SDC nodes must include the left interval boundary.") + + if not np.isclose(self.nodes[-1], 1.0): + raise Exception("SDC nodes must include the right interval boundary.") + + assert self.N == self.deltas.shape[0] + + shape = (self.N,) + u0.shape + u = np.zeros_like(u0, shape=shape) + + u[0, :] = u0 + evalF_k0 = np.empty_like(u) + evalF_k1 = np.empty_like(u) + + # Backup integrator contribution I[...] of previous iteration + ISolves = np.empty_like(u) + + # + # Propagate initial condition to all nodes + # + for m in range(0, self.N): + if m > 0: + # + # Solve implicit step: + # + # u^n+1 = u^n + dt * delta * F(u^n+1) + # <=> u^n+1 - dt * delta * F(u^n+1) = u^n + # <=> (I - dt * delta * F) * u^n+1 = u^n + # + rhs = u[m - 1] + u[m] = de_solver.fSolve(rhs, dt * self.deltas[m], t + dt * self.nodes[m]) + # Compute I[...] term + # u^n+1 = u^n + dt * delta * F(u^n+1) + # dt * delta * F(u^n+1) = u^n+1 - u^n + ISolves[m] = u[m] - u[m - 1] + + evalF_k0[m] = de_solver.evalF(u[m], t + dt * self.nodes[m]) + + # + # Iteratively sweep over SDC nodes + # + for _ in range(1, self.num_sweeps): + for m in range(0, self.N): + if m > 0: + # + # Solve implicit step: + # + # u^n+1 = u^n + dt * delta * (F(u^n+1)) - I(u^n) + dt * Q * F(u^n) + # <=> u^n+1 - dt * delta * F(u^n+1) = u^n - I(u^n) + dt * Q * F(u^n) + # <=> (I - dt * delta * F) * u^n+1 = u^n - I(u^n) + dt * Q * F(u^n) + # + # rhs = u^n - I(u^n) + dt * Q * F(u^n) + # + qeval = self.q[m] @ evalF_k0 + rhs = u[m - 1] - ISolves[m] + dt * qeval + + u[m] = de_solver.fSolve(rhs, dt * self.deltas[m], t + dt * self.nodes[m]) + + # Update I[...] term for next correction + # <=> - dt * delta * F(u^n+1) = u^n - I(u^n) + dt * Q * F(u^n) - u^n+1 + # <=> dt * delta * F(u^n+1) = u^n+1 - rhs + ISolves[m] = u[m] - rhs + + evalF_k1[m] = de_solver.evalF(u[m], t + dt * self.nodes[m]) + + # Copy tendency arrays + evalF_k0[:] = evalF_k1[:] + + if self.use_quadrature: + # Compute new starting value with quadrature on tendencies + u[0] = u[0] + dt * self.weights @ evalF_k0 + + else: + # If we're using Radau-right, we can just use the last value + u[0] = u[-1] + + assert u0.shape == u[0].shape + return u[0] + + def integrate_imex21(self, u0: np.array, t: float, dt: float, de_solver: DESolver) -> np.array: + """ + IMEX SDC where the first term is treated implicitly and the second term explicitly. + """ + if not np.isclose(self.nodes[0], 0.0): + raise Exception("SDC nodes must include the left interval boundary.") + + if not np.isclose(self.nodes[-1], 1.0): + raise Exception("SDC nodes must include the right interval boundary.") + + assert self.N == self.deltas.shape[0] + + shape = (self.N,) + u0.shape + u = np.zeros_like(u0, shape=shape) + + u[0, :] = u0 + evalF_k0 = np.empty_like(u) + evalF_k1 = np.empty_like(u) + + # Backup integrator contribution I[...] of previous iteration + ISolves = np.empty_like(u) + + print("="*80) + print("INITIAL") + # + # Propagate initial condition to all nodes + # + for m in range(0, self.N): + if m > 0: + # + # Solve explicit step first + # u* = u^n + dt * delta * F1(u^n) + # + rhs = u[m - 1] + dt * self.deltas[m] * de_solver.evalF1(u[m - 1], t + dt * self.nodes[m]) + + # + # Solve implicit step next (see integrate_irk1) + # u^{n+1} = u* + dt * delta * F2(u^{n+1}) + # + u[m] = de_solver.fSolve2(rhs, dt * self.deltas[m], t + dt * self.nodes[m]) + + # + # Compute I[...] term for implicit and explicit parts + # + # u^n+1 = u^n + dt * delta * F1(u^n) + dt * delta * F2(u^n+1) + # dt * delta * F1(u^n) + dt * delta * F2(u^n+1) = u^n+1 - u^n + # + ISolves[m] = u[m] - u[m - 1] + + evalF_k0[m] = de_solver.evalF(u[m], t + dt * self.nodes[m]) + print(m, "eval_f", evalF_k0[m]) + print(" ", "u", u[m]) + + # + # Iteratively sweep over SDC nodes + # + for _ in range(1, self.num_sweeps): + print("="*80) + print(f"SWEET {_}") + for m in range(0, self.N): + if m > 0: + # + # Solve IMEX step: + # + # u^n+1 = u^n + dt * delta * (F1(u^n) + F2(u^n+1)) - I(u^n) + dt * Q * F(u^n) + # <=> u^n+1 - dt * delta * F2(u^n+1) = u^n + dt * delta * F1(u^n) - I(u^n) + dt * Q * F(u^n) + # <=> (I - dt * delta * F) * u^n+1 = u^n + dt * delta * F1(u^n) - I(u^n) + dt * Q * F(u^n) + # + # rhs = u^n + dt * delta * F1(u^n) - I(u^n) + dt * Q * F(u^n) + # + qeval = self.q[m] @ evalF_k0 + euler = dt * self.deltas[m] * de_solver.evalF1(u[m - 1], t + dt * self.nodes[m]) + rhs = u[m - 1] + euler - ISolves[m] + dt * qeval + + u[m] = de_solver.fSolve2(rhs, dt * self.deltas[m], t + dt * self.nodes[m]) + + # + # Update I[...] term for next correction + # <=> dt * delta * (F1(u^n) + F2(u^n+1)) = u^n+1 - u^n + I(u^n) - dt * Q * F(u^n) + # + ISolves[m] = u[m] - u[m - 1] + ISolves[m] - dt * qeval + + evalF_k1[m] = de_solver.evalF(u[m], t + dt * self.nodes[m]) + print(m, "eval_f", evalF_k1[m]) + print(" ", "u", u[m]) + + # Copy tendency arrays + evalF_k0[:] = evalF_k1[:] + + if self.use_quadrature: + # Compute new starting value with quadrature on tendencies + u[0] = u[0] + dt * self.weights @ evalF_k0 + else: + # If we're using Radau-right, we can just use the last value + u[0] = u[-1] + + assert u0.shape == u[0].shape + return u[0] + + def integrate_imex12(self, u0: np.array, t: float, dt: float, de_solver: DESolver) -> np.array: + """ + IMEX SDC where the first term is treated implicitly and the second term explicitly. + """ + if not np.isclose(self.nodes[0], 0.0): + raise Exception("SDC nodes must include the left interval boundary.") + + if not np.isclose(self.nodes[-1], 1.0): + raise Exception("SDC nodes must include the right interval boundary.") + + assert self.N == self.deltas.shape[0] + + shape = (self.N,) + u0.shape + u = np.zeros_like(u0, shape=shape) + + u[0, :] = u0 + evalF_k0 = np.empty_like(u) + evalF_k1 = np.empty_like(u) + + # Backup integrator contribution I[...] of previous iteration + ISolves = np.empty_like(u) + + # + # Propagate initial condition to all nodes + # + for m in range(0, self.N): + if m > 0: + # + # Solve explicit step first + # u* = u^n + dt * delta * F1(u^n) + # + rhs = u[m - 1] + dt * self.deltas[m] * de_solver.evalF2(u[m - 1], t + dt * self.nodes[m]) + + # + # Solve implicit step next (see integrate_irk1) + # u^{n+1} = u* + dt * delta * F2(u^{n+1}) + # + u[m] = de_solver.fSolve1(rhs, dt * self.deltas[m], t + dt * self.nodes[m]) + + # + # Compute I[...] term for implicit and explicit parts + # + # u^n+1 = u^n + dt * delta * F1(u^n) + dt * delta * F2(u^n+1) + # dt * delta * F1(u^n) + dt * delta * F2(u^n+1) = u^n+1 - u^n + # + ISolves[m] = u[m] - u[m - 1] + + evalF_k0[m] = de_solver.evalF(u[m], t + dt * self.nodes[m]) + + # + # Iteratively sweep over SDC nodes + # + for _ in range(1, self.num_sweeps): + for m in range(0, self.N): + if m > 0: + # + # Solve IMEX step: + # + # u^n+1 = u^n + dt * delta * (F1(u^n) + F2(u^n+1)) - I(u^n) + dt * Q * F(u^n) + # <=> u^n+1 - dt * delta * F2(u^n+1) = u^n + dt * delta * F1(u^n) - I(u^n) + dt * Q * F(u^n) + # <=> (I - dt * delta * F) * u^n+1 = u^n + dt * delta * F1(u^n) - I(u^n) + dt * Q * F(u^n) + # + # rhs = u^n + dt * delta * F1(u^n) - I(u^n) + dt * Q * F(u^n) + # + qeval = self.q[m] @ evalF_k0 + euler = dt * self.deltas[m] * de_solver.evalF2(u[m - 1], t + dt * self.nodes[m]) + rhs = u[m - 1] + euler - ISolves[m] + dt * qeval + + u[m] = de_solver.fSolve1(rhs, dt * self.deltas[m], t + dt * self.nodes[m]) + + # + # Update I[...] term for next correction + # <=> dt * delta * (F1(u^n) + F2(u^n+1)) = u^n+1 - u^n + I(u^n) - dt * Q * F(u^n) + # + ISolves[m] = u[m] - u[m - 1] + ISolves[m] - dt * qeval + + evalF_k1[m] = de_solver.evalF(u[m], t + dt * self.nodes[m]) + + # Copy tendency arrays + evalF_k0[:] = evalF_k1[:] + + if self.use_quadrature: + # Compute new starting value with quadrature on tendencies + u[0] = u[0] + dt * self.weights @ evalF_k0 + else: + # If we're using Radau-right, we can just use the last value + u[0] = u[-1] + + assert u0.shape == u[0].shape + return u[0] + + def int_f(self, u0: np.array, t: float, dt: float, de_solver: DESolver) -> np.array: + if self.time_integration_method == "erk1": + return self.integrate_erk1(u0, t, dt, de_solver) + elif self.time_integration_method == "irk1": + return self.integrate_irk1(u0, t, dt, de_solver) + elif self.time_integration_method == "imex12": + return self.integrate_imex12(u0, t, dt, de_solver) + elif self.time_integration_method == "imex21": + return self.integrate_imex21(u0, t, dt, de_solver) + else: + raise Exception(f"Unsupported time integration within SDC: '{self.time_integration_method}'") + + def integrate_n(self, u0: np.array, t: float, dt: float, num_timesteps, de_solver: DESolver) -> np.array: + if not np.isclose(self.nodes[0], 0.0): + raise Exception("SDC nodes must include the left interval boundary.") + + if not np.isclose(self.nodes[-1], 1.0): + raise Exception("SDC nodes must include the right interval boundary.") + + u_value = u0 + + for n in range(num_timesteps): + u_value = self.int_f(u_value, t + n * dt, dt, de_solver) + + return u_value diff --git a/tests/test_qdelta/test_timestepping.py b/tests/test_qdelta/test_timestepping.py index d568eba..5b056f8 100644 --- a/tests/test_qdelta/test_timestepping.py +++ b/tests/test_qdelta/test_timestepping.py @@ -18,7 +18,7 @@ def testBE(nNodes, nodeType, quadType): coll = Collocation(nNodes, nodeType, quadType) nodes = coll.nodes - gen = module.BE(nodes) + gen = module.BE(nodes) QDelta = gen.getQDelta() assert np.allclose(np.tril(QDelta), QDelta), \