From 9d99538a60980557032c9c05e2a06cf397d506ff Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Sat, 19 Oct 2024 14:01:07 +0200 Subject: [PATCH] Simplify time stepping schemes and add type hints and documentation (#572) * Improve naming in time stepping schemes. * Add documentation and type hints. Remove support for monolithic problem to simplify code. * Tidy up and improve inheritance and documentation. * Do not instantiate objects from problemDefinition --- oscillator/solver-python/mypy.ini | 7 + oscillator/solver-python/oscillator.py | 64 +++-- oscillator/solver-python/problemDefinition.py | 27 ++- oscillator/solver-python/timeSteppers.py | 226 ++++++++++-------- 4 files changed, 182 insertions(+), 142 deletions(-) create mode 100644 oscillator/solver-python/mypy.ini diff --git a/oscillator/solver-python/mypy.ini b/oscillator/solver-python/mypy.ini new file mode 100644 index 000000000..ff162b0b1 --- /dev/null +++ b/oscillator/solver-python/mypy.ini @@ -0,0 +1,7 @@ +[mypy] + +[mypy-scipy.*] +ignore_missing_imports = True + +[mypy-precice.*] +ignore_missing_imports = True \ No newline at end of file diff --git a/oscillator/solver-python/oscillator.py b/oscillator/solver-python/oscillator.py index 8ea30042f..0720cb6bc 100644 --- a/oscillator/solver-python/oscillator.py +++ b/oscillator/solver-python/oscillator.py @@ -7,9 +7,10 @@ from enum import Enum import csv import os +from typing import Type import problemDefinition -import timeSteppers +from timeSteppers import TimeStepper, TimeSteppingSchemes, GeneralizedAlpha, RungeKutta4, RadauIIA class Participant(Enum): @@ -25,12 +26,17 @@ class Participant(Enum): help="Time stepping scheme being used.", type=str, choices=[ - s.value for s in timeSteppers.TimeSteppingSchemes], - default=timeSteppers.TimeSteppingSchemes.NEWMARK_BETA.value) + s.value for s in TimeSteppingSchemes], + default=TimeSteppingSchemes.NEWMARK_BETA.value) args = parser.parse_args() participant_name = args.participantName +this_mass: Type[problemDefinition.Mass] +other_mass: Type[problemDefinition.Mass] +this_spring: Type[problemDefinition.Spring] +connecting_spring = problemDefinition.SpringMiddle + if participant_name == Participant.MASS_LEFT.value: write_data_name = 'Force-Left' read_data_name = 'Force-Right' @@ -38,7 +44,6 @@ class Participant(Enum): this_mass = problemDefinition.MassLeft this_spring = problemDefinition.SpringLeft - connecting_spring = problemDefinition.SpringMiddle other_mass = problemDefinition.MassRight elif participant_name == Participant.MASS_RIGHT.value: @@ -48,7 +53,6 @@ class Participant(Enum): this_mass = problemDefinition.MassRight this_spring = problemDefinition.SpringRight - connecting_spring = problemDefinition.SpringMiddle other_mass = problemDefinition.MassLeft else: @@ -89,25 +93,19 @@ class Participant(Enum): a = a0 t = 0 -if args.time_stepping == timeSteppers.TimeSteppingSchemes.GENERALIZED_ALPHA.value: - time_stepper = timeSteppers.GeneralizedAlpha(stiffness=stiffness, mass=mass, alpha_f=0.4, alpha_m=0.2) -elif args.time_stepping == timeSteppers.TimeSteppingSchemes.NEWMARK_BETA.value: - time_stepper = timeSteppers.GeneralizedAlpha(stiffness=stiffness, mass=mass, alpha_f=0.0, alpha_m=0.0) -elif args.time_stepping == timeSteppers.TimeSteppingSchemes.RUNGE_KUTTA_4.value: - ode_system = np.array([ - [0, 1], # du - [-stiffness / mass, 0], # dv - ]) - time_stepper = timeSteppers.RungeKutta4(ode_system=ode_system) -elif args.time_stepping == timeSteppers.TimeSteppingSchemes.Radau_IIA.value: - ode_system = np.array([ - [0, 1], # du - [-stiffness / mass, 0], # dv - ]) - time_stepper = timeSteppers.RadauIIA(ode_system=ode_system) +time_stepper: TimeStepper + +if args.time_stepping == TimeSteppingSchemes.GENERALIZED_ALPHA.value: + time_stepper = GeneralizedAlpha(stiffness=stiffness, mass=mass, alpha_f=0.4, alpha_m=0.2) +elif args.time_stepping == TimeSteppingSchemes.NEWMARK_BETA.value: + time_stepper = GeneralizedAlpha(stiffness=stiffness, mass=mass, alpha_f=0.0, alpha_m=0.0) +elif args.time_stepping == TimeSteppingSchemes.RUNGE_KUTTA_4.value: + time_stepper = RungeKutta4(stiffness=stiffness, mass=mass) +elif args.time_stepping == TimeSteppingSchemes.Radau_IIA.value: + time_stepper = RadauIIA(stiffness=stiffness, mass=mass) else: raise Exception( - f"Invalid time stepping scheme {args.time_stepping}. Please use one of {[ts.value for ts in timeSteppers.TimeSteppingSchemes]}") + f"Invalid time stepping scheme {args.time_stepping}. Please use one of {[ts.value for ts in TimeSteppingSchemes]}") positions = [] @@ -133,34 +131,32 @@ class Participant(Enum): precice_dt = participant.get_max_time_step_size() dt = np.min([precice_dt, my_dt]) - def f(t): return participant.read_data(mesh_name, read_data_name, vertex_ids, t)[0] + def f(t: float) -> float: return participant.read_data(mesh_name, read_data_name, vertex_ids, t)[0] # do time step, write data, and advance - # performs adaptive time stepping with dense output; multiple write calls per time step - if args.time_stepping == timeSteppers.TimeSteppingSchemes.Radau_IIA.value: - u_new, v_new, a_new, sol = time_stepper.do_step(u, v, a, f, dt) - t_new = t + dt + u_new, v_new, a_new = time_stepper.do_step(u, v, a, f, dt) + + t_new = t + dt + + # RadauIIA time stepper provides dense output. Do multiple write calls per time step. + if isinstance(time_stepper, RadauIIA): # create n samples_per_step of time stepping scheme. Degree of dense # interpolating function is usually larger 1 and, therefore, we need # multiple samples per step. samples_per_step = 5 - n_time_steps = len(sol.ts) # number of time steps performed by adaptive time stepper + n_time_steps = len(time_stepper.dense_output.ts) # number of time steps performed by adaptive time stepper n_pseudo = samples_per_step * n_time_steps # samples_per_step times no. time steps per window. - t_pseudo = 0 for i in range(n_pseudo): # perform n_pseudo pseudosteps dt_pseudo = dt / n_pseudo t_pseudo += dt_pseudo - write_data = [connecting_spring.k * sol(t_pseudo)[0]] + write_data = np.array([connecting_spring.k * time_stepper.dense_output(t_pseudo)[0]]) participant.write_data(mesh_name, write_data_name, vertex_ids, write_data) participant.advance(dt_pseudo) else: # simple time stepping without dense output; only a single write call per time step - u_new, v_new, a_new = time_stepper.do_step(u, v, a, f, dt) - t_new = t + dt - - write_data = [connecting_spring.k * u_new] + write_data = np.array([connecting_spring.k * u_new]) participant.write_data(mesh_name, write_data_name, vertex_ids, write_data) participant.advance(dt) diff --git a/oscillator/solver-python/problemDefinition.py b/oscillator/solver-python/problemDefinition.py index b0ac2aea3..65ca35d63 100644 --- a/oscillator/solver-python/problemDefinition.py +++ b/oscillator/solver-python/problemDefinition.py @@ -1,20 +1,33 @@ import numpy as np from numpy.linalg import eig +from typing import Callable -class SpringLeft: +class Spring: + k: float + + +class SpringLeft(Spring): k = 4 * np.pi**2 -class SpringMiddle: +class SpringMiddle(Spring): k = 16 * (np.pi**2) -class SpringRight: +class SpringRight(Spring): k = 4 * np.pi**2 -class MassLeft: +class Mass: + m: float + u0: float + v0: float + u_analytical: Callable[[float | np.ndarray], float | np.ndarray] + v_analytical: Callable[[float | np.ndarray], float | np.ndarray] + + +class MassLeft(Mass): # mass m = 1 @@ -23,10 +36,8 @@ class MassLeft: u0 = 1.0 v0 = 0.0 - u_analytical, v_analytical = None, None # will be defined below - -class MassRight: +class MassRight(Mass): # mass m = 1 @@ -35,8 +46,6 @@ class MassRight: u0 = 0.0 v0 = 0.0 - u_analytical, v_analytical = None, None # will be defined below - # Mass matrix M = np.array([ diff --git a/oscillator/solver-python/timeSteppers.py b/oscillator/solver-python/timeSteppers.py index da737a192..4935f20df 100644 --- a/oscillator/solver-python/timeSteppers.py +++ b/oscillator/solver-python/timeSteppers.py @@ -1,7 +1,8 @@ -from typing import Tuple, List +from typing import Tuple, Callable import numpy as np import numbers import scipy as sp +from scipy.integrate import OdeSolution from enum import Enum @@ -12,15 +13,50 @@ class TimeSteppingSchemes(Enum): Radau_IIA = "radauIIA" # 5th order implicit -class GeneralizedAlpha(): - alpha_f = None - alpha_m = None - gamma = None - beta = None - mass = None - stiffness = None - - def __init__(self, stiffness, mass, alpha_f=0.4, alpha_m=0.2) -> None: +class TimeStepper: + """Time stepper for a mass spring system with given mass and spring stiffness + """ + + def __init__(self, + stiffness: float, + mass: float) -> None: + """Initializes time stepper with parameters describing mass spring system + + Args: + stiffness (float): the stiffness of the spring connected to the mass + mass (float): the mass + """ + raise NotImplementedError() + + def do_step(self, + u0: float, + v0: float, + a0: float, + rhs: Callable[[float], float], + dt: float + ) -> Tuple[float, float, float]: + """Perform a time step of size dt with given time stepper + + Args: + u0 (float): displacement at time t0 + v0 (float): velocity at time t0 + a0 (float): acceleration at time t0 + rhs (Callable[[float],float]): time dependent right-hand side + dt (float): time step size + + Returns: + Tuple[float, float, float]: returns computed displacement, velocity, and acceleration at time t1 = t0 + dt. + """ + raise NotImplementedError() + + +class GeneralizedAlpha(TimeStepper): + """TimeStepper implementing generalized Alpha or Newmark Beta scheme (depends on parameters alpha_f and alpha_m set in constructor) + """ + alpha_f: float + alpha_m: float + + def __init__(self, stiffness: float, mass: float, alpha_f: float = 0.4, alpha_m: float = 0.2) -> None: self.alpha_f = alpha_f self.alpha_m = alpha_m @@ -30,32 +66,33 @@ def __init__(self, stiffness, mass, alpha_f=0.4, alpha_m=0.2) -> None: self.stiffness = stiffness self.mass = mass - def do_step(self, u, v, a, rhs, dt) -> Tuple[float, float, float]: - f = rhs((1 - self.alpha_f) * dt) + def do_step(self, + u0: float, + v0: float, + a0: float, + rhs: Callable[[float], float], + dt: float + ) -> Tuple[float, float, float]: + f = rhs((1.0 - self.alpha_f) * dt) - m = 3 * [None] - m[0] = (1 - self.alpha_m) / (self.beta * dt**2) - m[1] = (1 - self.alpha_m) / (self.beta * dt) - m[2] = (1 - self.alpha_m - 2 * self.beta) / (2 * self.beta) + m = [(1 - self.alpha_m) / (self.beta * dt**2), + (1 - self.alpha_m) / (self.beta * dt), + (1 - self.alpha_m - 2 * self.beta) / (2 * self.beta)] k_bar = self.stiffness * (1 - self.alpha_f) + m[0] * self.mass # do generalized alpha step - if (type(self.stiffness)) is np.ndarray: - u_new = np.linalg.solve( - k_bar, - (f - self.alpha_f * self.stiffness.dot(u) + self.mass.dot((m[0] * u + m[1] * v + m[2] * a))) - ) - else: - u_new = (f - self.alpha_f * self.stiffness * u + self.mass * (m[0] * u + m[1] * v + m[2] * a)) / k_bar - - a_new = 1.0 / (self.beta * dt**2) * (u_new - u - dt * v) - (1 - 2 * self.beta) / (2 * self.beta) * a - v_new = v + dt * ((1 - self.gamma) * a + self.gamma * a_new) + u1 = (f - self.alpha_f * self.stiffness * u0 + self.mass * (m[0] * u0 + m[1] * v0 + m[2] * a0)) / k_bar + a1 = 1.0 / (self.beta * dt**2) * (u1 - u0 - dt * v0) - (1 - 2 * self.beta) / (2 * self.beta) * a0 + v1 = v0 + dt * ((1 - self.gamma) * a0 + self.gamma * a1) - return u_new, v_new, a_new + return u1, v1, a1 -class RungeKutta4(): +class RungeKutta4(TimeStepper): + """TimeStepper implementing classic Runge Kutta scheme (4 stages) + """ + # parameters from Butcher tableau of classic Runge Kutta scheme (RK4) a = np.array([[0, 0, 0, 0], [0.5, 0, 0, 0], [0, 0.5, 0, 0], @@ -63,79 +100,70 @@ class RungeKutta4(): b = np.array([1 / 6, 1 / 3, 1 / 3, 1 / 6]) c = np.array([0, 0.5, 0.5, 1]) - def __init__(self, ode_system) -> None: - self.ode_system = ode_system - pass - - def do_step(self, u, v, a, rhs, dt) -> Tuple[float, float, float]: - assert (isinstance(u, type(v))) - - n_stages = 4 - - if isinstance(u, np.ndarray): - x = np.concatenate([u, v]) - def f(t): return np.concatenate([np.array([0, 0]), rhs(t)]) - elif isinstance(u, numbers.Number): - x = np.array([u, v]) - def f(t): return np.array([0, rhs(t)]) - else: - raise Exception(f"Cannot handle input type {type(u)} of u and v") - - s = n_stages * [None] - s[0] = self.ode_system.dot(x) + f(self.c[0] * dt) - s[1] = self.ode_system.dot(x + self.a[1, 0] * s[0] * dt) + f(self.c[1] * dt) - s[2] = self.ode_system.dot(x + self.a[2, 1] * s[1] * dt) + f(self.c[2] * dt) - s[3] = self.ode_system.dot(x + self.a[3, 2] * s[2] * dt) + f(self.c[3] * dt) - - x_new = x - - for i in range(n_stages): - x_new += dt * self.b[i] * s[i] - - if isinstance(u, np.ndarray): - u_new = x_new[0:2] - v_new = x_new[2:4] - elif isinstance(u, numbers.Number): - u_new = x_new[0] - v_new = x_new[1] - - a_new = None - - return u_new, v_new, a_new - - -class RadauIIA(): - def __init__(self, ode_system) -> None: - self.ode_system = ode_system - pass - - def do_step(self, u, v, a, rhs, dt) -> Tuple[float, float, float]: - t0 = 0 - - assert (isinstance(u, type(v))) - - if isinstance(u, np.ndarray): - x0 = np.concatenate([u, v]) - f = np.array(f) - assert (u.shape[0] == f.shape[1]) - def rhs_fun(t): return np.concatenate([np.array([np.zeros_like(t), np.zeros_like(t)]), rhs(t)]) - elif isinstance(u, numbers.Number): - x0 = np.array([u, v]) - def rhs_fun(t): return np.array([np.zeros_like(t), rhs(t)]) - else: - raise Exception(f"Cannot handle input type {type(u)} of u and v") - - def fun(t, x): - return self.ode_system.dot(x) + rhs_fun(t) + def __init__(self, stiffness: float, mass: float) -> None: + self.ode_system = np.array([ + [0, 1], # du + [-stiffness / mass, 0], # dv + ]) + + def do_step(self, + u0: float, + v0: float, + _: float, + rhs: Callable[[float], float], + dt: float + ) -> Tuple[float, float, float]: + k = 4 * [None] # store stages in list + + x0 = np.array([u0, v0]) + def f(t, x): return self.ode_system.dot(x) + np.array([0, rhs(t)]) + + k[0] = f(self.c[0] * dt, x0) + k[1] = f(self.c[1] * dt, x0 + self.a[1, 0] * k[0] * dt) + k[2] = f(self.c[2] * dt, x0 + self.a[2, 1] * k[1] * dt) + k[3] = f(self.c[3] * dt, x0 + self.a[3, 2] * k[2] * dt) + + x1 = x0 + dt * sum(b_i * k_i for k_i, b_i in zip(k, self.b)) + + return x1[0], x1[1], f(dt, x1)[1] + + +class RadauIIA(TimeStepper): + """Perform a step with the adaptive RadauIIA time stepper of scipy.integrate + """ + _dense_output: OdeSolution + + def __init__(self, stiffness: float, mass: float) -> None: + self.ode_system = np.array([ + [0, 1], # du + [-stiffness / mass, 0], # dv + ]) + + def do_step(self, + u0: float, + v0: float, + _: float, + rhs: Callable[[float], float], + dt: float + ) -> Tuple[float, float, float]: + t0, t1 = 0, dt + + x0 = np.array([u0, v0]) + def f(t, x): return self.ode_system.dot(x) + np.array([np.zeros_like(t), rhs(t)]) # use adaptive time stepping; dense_output=True allows us to sample from continuous function later - ret = sp.integrate.solve_ivp(fun, [t0, t0 + dt], x0, method="Radau", + ret = sp.integrate.solve_ivp(f, [t0, t1], x0, method="Radau", dense_output=True, rtol=10e-5, atol=10e-9) - a_new = None - if isinstance(u, np.ndarray): - u_new, v_new = ret.y[0:2, -1], ret.y[2:4, -1] - elif isinstance(u, numbers.Number): - u_new, v_new = ret.y[:, -1] + self._dense_output = ret.sol # store dense output in class + + return ret.y[0, -1], ret.y[1, -1], f(dt, ret.y[:, -1])[1] + + @property + def dense_output(self) -> OdeSolution: + """Returns dense output created for the last call of do_step - return u_new, v_new, a_new, ret.sol + Returns: + OdeSolution: dense output + """ + return self._dense_output