Skip to content
Snippets Groups Projects

Optimize the method of leapfrog

Closed Optimize the method of leapfrog
integrators into main
6 unresolved threads
Closed jim92 requested to merge integrators into main
6 unresolved threads
Compare
and
Show latest version
4 files
+ 15
− 178
Compare changes
  • Side-by-side
  • Inline
Files
4
tasks/src/integrators.py
+ 11
− 115
@@ -2,137 +2,33 @@ import numpy as np
def leapfrog(initial_states, num_steps, dt, gravitational_constant):
position, velocity, masses = initial_states
trajectory, _ = leapfrog_core(initial_states, num_steps, dt, gravitational_constant, masses)
return trajectory
positions, velocities = initial_states
traj = [positions]
def leapfrog_core(initial_states, num_steps, dt, gravitational_constant, masses):
# calculate the initial accelerations for all particles
position, velocity, _ = initial_states
trajectory = np.zeros((len(position), 3, num_steps + 1))
trajectory[:, :, 0] = position
velocities = np.zeros((len(position), 3, num_steps + 1))
velocities[:, :, 0] = velocity
accelerations = getAcceleration(position, gravitational_constant, masses)
for index in range(num_steps):
for _ in range(num_steps):
# Leapfrog integration
'''
# using leapfrog integration. https://en.wikipedia.org/wiki/Leapfrog_integration
The basic idea of the Leapfrog method is to update the position
and momentum of a particle or system at half-integer time steps.
The algorithm proceeds in discrete steps, and the update scheme looks like this:
1. Update Velocity at Half Step:
2. Update Position:
3. Update Velocity at Full Step:
'''
# 1. Update Momentum at Half Step:
velocity += dt / 2 * accelerations
# 2. Update Position:
position += velocity * dt
# 3. Update Velocity at Full Step:
# step 1: update accelerations
accelerations = getAcceleration(position, gravitational_constant, masses)
# step 2: update velocity
velocity += accelerations * dt / 2
# save the velocity and trajectory for plotting
velocities[:, :, index + 1] = velocity
trajectory[:, :, index + 1] = position
return trajectory, velocities
accelerations = -gravitational_constant * positions / np.linalg.norm(positions, axis=1)[:, None]**3
velocities += accelerations * dt
positions += velocities * dt
traj.append(positions.copy())
# TODO: maybe this method need to be optimized to improve efficiency, i am not sure. 19.01.2024
def getAcceleration(positions, gravitational_constant, masses):
accelerations = np.zeros_like(positions)
for i in range(len(positions)):
for j in range(len(positions)):
if i != j:
r_ij = positions[i] - positions[j]
distance = np.linalg.norm(r_ij)
accelerations[i] += gravitational_constant * masses[j] * r_ij / (distance ** 3)
# force = gravitational_constant * masses[i] * masses[j] * r_ij / (distance ** 3)
# accelerations[i] += force / masses[i]
return accelerations
return np.array(traj)
def stormer_verlet(initial_states, num_steps, dt, gravitational_constant):
positions, velocities, masses = initial_states
positions, velocities = initial_states
traj = [positions]
for _ in range(num_steps):
# Stormer-Verlet integration
half_velocities = velocities + 0.5 * dt * (-gravitational_constant * positions /
np.linalg.norm(positions, axis=1)[:, None] ** 3)
np.linalg.norm(positions, axis=1)[:, None]**3)
positions += dt * half_velocities
velocities = half_velocities + 0.5 * dt * (-gravitational_constant * positions /
np.linalg.norm(positions, axis=1)[:, None] ** 3)
np.linalg.norm(positions, axis=1)[:, None]**3)
traj.append(positions.copy())
return np.array(traj)
# NOTE: this method maybe is not applicable to the problem in the project.
# The Tutor will inform me of the result tomorrow(17.01.2024)
def rkn_matrix(f, y0, t_span, dt):
"""
- t_span: Time span [t0, tn].
- dt: Step size.
- y0: Initial conditions [Y0, Y'_0].
"""
# Initial conditions
t0, tn = t_span
y = np.array(y0)
t = np.arange(t0, tn + dt, dt)
# Arrays to store solutions
y_solution = np.zeros((len(t), *y0.shape), dtype=y0.dtype)
y_solution[0] = y0
# RKN integration loop
for i in range(1, len(t)):
j = 0
for each_y in y:
y_solution_each_y_i = rkn_matrix_core(f, each_y, dt, (i - 1) * dt, y_solution[i - 1])
y_solution[i][j] = y_solution_each_y_i
j += 1
return t, y_solution
def rkn_matrix_core(f, y, dt, base_t, y_all):
"""
Solve a system of second-order matrix ODEs using the Runge-Kutta-Nystrom method.
Parameters:
- f: A function representing the matrix second-order ODE (Y'' = F(t, Y, Y')).
- y: Initial conditions [Y_[i-1], Y'_[i-1]] for every loop.
- dt: Step size.
- y_all: the latest solution, namely the last updated solution before call this function
Returns:
- y: Array of solutions [Y, Y'].
NOTE: the N-Bodies problems are not second-order equations. Maybe this method can not be employed for it.
How to use this method:
1. convert all second-order equations into a set of one-second equations
2. for every group one-second equations generated from #1, translate them into a matrix format
3. traverse the results generated from #2
above steps are only for one step, if you want to integrate the method for many times,
please explicitly call for it outside the method
"""
# RKN integration
k1 = dt * np.array([y[1], f(base_t, y[0], y[1], y_all)])
k2 = dt * np.array([y[1] + 0.5 * k1[1], f(base_t + 0.5 * dt, y[0] + 0.5 * k1[0], y[1] + 0.5 * k1[1], y_all)])
k3 = dt * np.array([y[1] + 0.5 * k2[1], f(base_t + 0.5 * dt, y[0] + 0.5 * k2[0], y[1] + 0.5 * k2[1], y_all)])
k4 = dt * np.array([y[1] + k3[1], f(base_t + dt, y[0] + k3[0], y[1] + k3[1], y_all)])
return (k1 + 2 * k2 + 2 * k3 + k4) / 6.0