update

Source/integrator.py

@@ -34,54 +34,65 @@ along with this program. If not, see https://www.gnu.org/licenses/.
 """
 import numpy as np
 
-def euler(x0, y0, h, n, func): # FIXME cannot be used with vectors
-    """DEPRECIATED - DO NOT USE
-    Euler method adapted for state vector [[x, y], [u, v]]
-    :param x0: initial time value
-    :param y0: initial state vector [[x, y], [u, v]]
-    :param h: step size (time step)
-    :param n: number of steps
-    :param func: RHS of differential equation
-    :returns: x array, solution array
+def euler(t0: float, 
+          W0: np.ndarray, 
+          h: float, 
+          n: int, 
+          func):
+    """Euler method adapted for state vector [[x, y], [u, v]]
+    @ params
+        - t0: initial time value
+        - W0: initial state vector [[x, y], [u, v]]
+        - h: step size (time step)
+        - n: number of steps
+        - func: RHS of differential equation
+    @returns: 
+        - t, W: time and state (solution) arrays
     """
-    x_values = np.zeros(n)
-    y_values = np.zeros((n, 2, 2))  #  to accommodate the state vector
+    time = np.zeros(n)
+    W = np.zeros((n,) + np.shape(W0))
 
+    t = t0
+    w = W0
     for i in range(n):
-        dydt = func(x0, y0)
-        y0 = y0 + h * dydt
-        x0 = x0 + h
+        k1 = func(t, w)
+        w = w + h*k1
+        t = t + h
 
-        x_values[i] = x0
-        y_values[i, :, :] = y0
+        time[i] = t
+        W[i] = w
+    return time, W
 
-    return x_values, y_values
-
-# Updated RK2 integrator
-def rk2(x0, y0, h, n, func): # FIXME cannot be used with vectors
-    """ DEPRECIATED - DO NOT USE
-    RK2 method adapted for state vector [[x, y], [u, v]]
-    :param x0: initial time value
-    :param y0: initial state vector [[x, y], [u, v]]
-    :param h: step size (time step)
-    :param n: number of steps
-    :param func: RHS of differential equation
-    :returns: x array, solution array
+def rk2(t0: float, 
+        W0: np.ndarray, 
+        h: float, 
+        n: int, 
+        func):
+    """RK2 method adapted for state vector [[x, y], [u, v]]
+    @ params
+        - t0: initial time value
+        - W0: initial state vector [[x, y], [u, v]]
+        - h: step size (time step)
+        - n: number of steps
+        - func: RHS of differential equation
+    @returns: 
+        - t, W: time and state (solution) arrays
     """
-    x_values = np.zeros(n)
-    y_values = np.zeros((n, 2, 2))  #  to accommodate the state vector
+    time = np.zeros(n)
+    W = np.zeros((n,) + np.shape(W0))
 
+    t = t0
+    w = W0
     for i in range(n):
-        k1 = func(x0, y0)
-        k2 = func(x0 + h / 2., y0 + h / 2. * k1)
+        k1 = func(t, w)
+        k2 = func(t + h/2, w + h/2*k1)
 
-        y0 = y0 + h * (k1 / 2. + k2 / 2.)
-        x0 = x0 + h
+        w = w + h*k2
+        t = t + h
 
-        x_values[i] = x0
-        y_values[i, :, :] = y0
-
-    return x_values, y_values
+        time[i] = t
+        W[i] = w
+    return time, W
 
 def rk4(t0: float, 
         W0: np.ndarray, 
@@ -90,13 +101,13 @@ def rk4(t0: float,
         func):
     """RK4 method adapted for state vector [[x, y], [u, v]]
     @ params
-        - x0: initial time value
-        - y0: initial state vector [[x, y], [u, v]]
+        - t0: initial time
+        - W0: initial state vector [[x, y], [u, v]]
         - h: step size (time step)
         - n: number of steps
         - func: RHS of differential equation
     @returns: 
-        - t, W: time and state (solution) arrays, 
+        - t, W: time and state (solution) arrays
     """
     time = np.zeros(n)
     W = np.zeros((n,) + np.shape(W0))
@@ -105,19 +116,50 @@ def rk4(t0: float,
     w = W0
     for i in range(n):
         k1 = func(t, w)
-        k2 = func(t + h / 2., w + h / 2. * k1)
-        k3 = func(t + h / 2., w + h / 2. * k2)
-        k4 = func(t + h, w + h * k3)
+        k2 = func(t + h/2, w + h/2*k1)
+        k3 = func(t + h/2, w + h/2*k2)
+        k4 = func(t + h, w + h*k3)
 
-        w = w + h * (k1 / 6. + k2 / 3. + k3 / 3. + k4 / 6.)
+        w = w + h*(k1/6 + k2/3 + k3/3 + k4/6)
         t = t + h
 
         time[i] = t
         W[i] = w
+    return time, W
 
-    return t, W
+def integrator_type(t0, W0, h, n, func, integrator):
+    return integrator(t0, W0, h, n, func)
 
+def kepler_analytical(t0: float, 
+                      W0: np.ndarray, 
+                      h: float, 
+                      n: int):
+    """Computes the evolution from the Kepler potential derivative
+    @ params
+        - t0: initial time value
+        - W0: initial state vector [[x, y], [u, v]]
+        - h: step size (time step)
+        - n: number of steps
+    @returns: 
+        - t, W: time and state (solution) arrays  
+    """
+    X0 = W0[0 ,0]
+    Y0 = W0[0, 1]
+    U0 = W0[1, 0]
+    V0 = W0[1, 1]
 
-def integrator_type(x0, y0, h, n, func, int_type):
-    """DEPRECIATED - DO NOT USE"""
-    return int_type(x0, y0, h, n, func)
+    time = np.arange(t0, t0 + n*h, h)
+    W = np.zeros((n,) + np.shape(W0))
+
+    R0 = np.sqrt(X0**2 + Y0**2)
+    Omega0 = np.sqrt(U0**2 + V0**2)/R0
+
+    X = R0 * np.cos(Omega0 * time)
+    Y = R0 * np.sin(Omega0 * time)
+    U = -R0 * Omega0 * np.sin(Omega0 * time)
+    V = R0 * Omega0 * np.cos(Omega0 * time)
+
+    W = np.array([[X, Y], [U, V]])
+    W = np.swapaxes(W, 0, 2)
+    W = np.swapaxes(W, 1, 2)
+    return time, W