update

Source/potentials.py

@@ -0,0 +1,65 @@
+import numpy as np
+
+MAX_VAL = 1e3
+
+def kepler_potential(W_grid: np.ndarray, 
+                     position_only: bool = False) -> np.ndarray:
+    """Computes the Kepler potential: V(R) = -G*m1*m2/R 
+    (assuming G = 1, m1 = 1, m2 = 1)
+    assuming the point mass at (x = 0, y = 0).
+    @params:
+        - W: Phase-space vector
+        - position_only: True if W is np.array([X, Y])
+    @returns:
+        - V: computed potential
+    """
+    if position_only: 
+        X = W_grid[0]
+        Y = W_grid[1]
+    else:
+        X = W_grid[0,0]
+        Y = W_grid[0,1]
+    # If X or Y is not an array (or a list), but rather a scalar, then we
+    # create a list of one element so that it can work either way
+    if np.ndim(X) == 0: X = np.array([X])
+    if np.ndim(Y) == 0: Y = np.array([Y])
+    R = np.sqrt(X**2 + Y**2)
+    return -1/R
+
+def hh_potential(W_grid: np.ndarray, 
+                 position_only=False) -> np.ndarray:
+    """Computes the Hénon-Heiles potential.
+    :param W: Phase-space vector
+    :output V: Potential
+    """
+    if position_only:
+        X = W_grid[0]
+        Y = W_grid[1]
+    else:
+        X = W_grid[0, 0]
+        Y = W_grid[0, 1]
+        
+    # If X or Y is not an array (or a list), but rather a scalar, then we
+    # create a list of one element so that it can work either way
+    if np.ndim(X) == 0: X = np.array([X])
+    if np.ndim(Y) == 0: Y = np.array([Y])
+
+    POT = (X**2 + Y**2 + 2*X**2*Y - 2*Y**3/3)/2
+    return POT        
+
+def hh_evolution(t: np.ndarray, W: np.ndarray):
+    """Computes the evolution from the HH potential
+    :param t: Time (not used)
+    :param W: Phase space vector
+    :returns dot W: Time derivative of the phase space vector
+    """
+    X = W[0 ,0]
+    Y = W[0, 1]
+    U = W[1, 0]
+    V = W[1, 1]
+    DX = U 
+    DY = V
+    DU = -(2*X*Y + X)
+    DV = -(X**2 - Y**2 + Y)
+
+    return np.array([[DX, DY], [DU, DV]])