Improve naming in time stepping schemes.

oscillator/solver-python/oscillator.py

@@ -136,17 +136,18 @@ class Participant(Enum):
     def f(t): 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, *other = time_stepper.do_step(u, v, a, f, dt)
+    t_new = t + dt
+
+    if other:
+        # if dense output is available, performed adaptive time stepping. Do multiple write calls per time step
+        sol, *_ = other
         # 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_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
@@ -157,9 +158,6 @@ def f(t): return participant.read_data(mesh_name, read_data_name, vertex_ids, t)
             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]
         participant.write_data(mesh_name, write_data_name, vertex_ids, write_data)
         participant.advance(dt)

oscillator/solver-python/timeSteppers.py

@@ -30,7 +30,7 @@ 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]:
+    def do_step(self, u0, v0, a0, rhs, dt) -> Tuple[float, float, float]:
         f = rhs((1 - self.alpha_f) * dt)
 
         m = 3 * [None]
@@ -42,17 +42,17 @@ def do_step(self, u, v, a, rhs, dt) -> Tuple[float, float, float]:
 
         # do generalized alpha step
         if (type(self.stiffness)) is np.ndarray:
-            u_new = np.linalg.solve(
+            u1 = 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)))
+                (f - self.alpha_f * self.stiffness.dot(u0) + self.mass.dot((m[0] * u0 + m[1] * v0 + m[2] * a0)))
             )
         else:
-            u_new = (f - self.alpha_f * self.stiffness * u + self.mass * (m[0] * u + m[1] * v + m[2] * a)) / k_bar
+            u1 = (f - self.alpha_f * self.stiffness * u0 + self.mass * (m[0] * u0 + m[1] * v0 + m[2] * a0)) / 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)
+        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():
@@ -67,75 +67,63 @@ 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)))
+    def do_step(self, u0, v0, _, rhs, dt) -> Tuple[float, float, float]:
+        assert (isinstance(u0, type(v0)))
 
-        n_stages = 4
+        k = 4 * [None]  # store stages in list
 
-        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)])
+        if isinstance(u0, np.ndarray):
+            x0 = np.concatenate([u0, v0])
+            def f(t,x): return self.ode_system.dot(x) + np.concatenate([np.array([0, 0]), rhs(t)])
+        elif isinstance(u0, numbers.Number):
+            x0 = np.array([u0, v0])
+            def f(t,x): return self.ode_system.dot(x) + np.array([0, rhs(t)])
         else:
-            raise Exception(f"Cannot handle input type {type(u)} of u and v")
+            raise Exception(f"Cannot handle input type {type(u0)} 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)
+        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)
 
-        x_new = x
+        x1 = x0 + dt * sum(b_i * k_i for k_i, b_i in zip(k, self.b))
 
-        for i in range(n_stages):
-            x_new += dt * self.b[i] * s[i]
+        if isinstance(u0, np.ndarray):
+            u1 = x1[0:2]
+            v1 = x1[2:4]
+        elif isinstance(u0, numbers.Number):
+            u1 = x1[0]
+            v1 = x1[1]
 
-        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
+        return u1, v1, None
 
 
 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
+    def do_step(self, u0, v0, _, rhs, dt) -> Tuple[float, float, float]:
+        assert (isinstance(u0, type(v0)))
 
-        assert (isinstance(u, type(v)))
+        t0 = 0
 
-        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)])
+        if isinstance(u0, np.ndarray):
+            x0 = np.concatenate([u0, v0])
+            def f(t,x): return self.ode_system.dot(x) + np.concatenate([np.array([np.zeros_like(t), np.zeros_like(t)]), rhs(t)])
+        elif isinstance(u0, numbers.Number):
+            x0 = np.array([u0, v0])
+            def f(t,x): return self.ode_system.dot(x) + 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)
+            raise Exception(f"Cannot handle input type {type(u0)} of u and v")
 
         # 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, t0 + dt], 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]
+        if isinstance(u0, np.ndarray):
+            u1, v1 = ret.y[0:2, -1], ret.y[2:4, -1]
+        elif isinstance(u0, numbers.Number):
+            u1, v1 = ret.y[:, -1]
 
-        return u_new, v_new, a_new, ret.sol
+        return u1, v1, None, ret.sol