Skip to content

Commit 0ea97b4

Browse files
committed
update
1 parent 45ac740 commit 0ea97b4

File tree

2 files changed

+91
-0
lines changed

2 files changed

+91
-0
lines changed
25.1 KB
Loading

doc/src/week46/programs/rk4.py

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
import numpy as np
2+
import pandas as pd
3+
from math import *
4+
import matplotlib.pyplot as plt
5+
import os
6+
7+
# Where to save the figures and data files
8+
PROJECT_ROOT_DIR = "Results"
9+
FIGURE_ID = "Results/FigureFiles"
10+
DATA_ID = "DataFiles/"
11+
12+
if not os.path.exists(PROJECT_ROOT_DIR):
13+
os.mkdir(PROJECT_ROOT_DIR)
14+
15+
if not os.path.exists(FIGURE_ID):
16+
os.makedirs(FIGURE_ID)
17+
18+
if not os.path.exists(DATA_ID):
19+
os.makedirs(DATA_ID)
20+
21+
def image_path(fig_id):
22+
return os.path.join(FIGURE_ID, fig_id)
23+
24+
def data_path(dat_id):
25+
return os.path.join(DATA_ID, dat_id)
26+
27+
def save_fig(fig_id):
28+
plt.savefig(image_path(fig_id) + ".png", format='png')
29+
30+
31+
def SpringForce(v,x,t):
32+
# note here that we have divided by mass and we return the acceleration
33+
return -2*gamma*v-x+Ftilde*cos(t*Omegatilde)
34+
35+
36+
def RK4(v,x,t,n,Force):
37+
for i in range(n-1):
38+
# Setting up k1
39+
k1x = DeltaT*v[i]
40+
k1v = DeltaT*Force(v[i],x[i],t[i])
41+
# Setting up k2
42+
vv = v[i]+k1v*0.5
43+
xx = x[i]+k1x*0.5
44+
k2x = DeltaT*vv
45+
k2v = DeltaT*Force(vv,xx,t[i]+DeltaT*0.5)
46+
# Setting up k3
47+
vv = v[i]+k2v*0.5
48+
xx = x[i]+k2x*0.5
49+
k3x = DeltaT*vv
50+
k3v = DeltaT*Force(vv,xx,t[i]+DeltaT*0.5)
51+
# Setting up k4
52+
vv = v[i]+k3v
53+
xx = x[i]+k3x
54+
k4x = DeltaT*vv
55+
k4v = DeltaT*Force(vv,xx,t[i]+DeltaT)
56+
# Final result
57+
x[i+1] = x[i]+(k1x+2*k2x+2*k3x+k4x)/6.
58+
v[i+1] = v[i]+(k1v+2*k2v+2*k3v+k4v)/6.
59+
t[i+1] = t[i] + DeltaT
60+
61+
62+
# Main part begins here
63+
64+
DeltaT = 0.001
65+
#set up arrays
66+
tfinal = 20 # in dimensionless time
67+
n = ceil(tfinal/DeltaT)
68+
# set up arrays for t, v, and x
69+
t = np.zeros(n)
70+
v = np.zeros(n)
71+
x = np.zeros(n)
72+
# Initial conditions (can change to more than one dim)
73+
x0 = 1.0
74+
v0 = 0.0
75+
x[0] = x0
76+
v[0] = v0
77+
gamma = 0.2
78+
Omegatilde = 0.5
79+
Ftilde = 1.0
80+
# Start integrating using Euler's method
81+
# Note that we define the force function as a SpringForce
82+
RK4(v,x,t,n,SpringForce)
83+
84+
# Plot position as function of time
85+
fig, ax = plt.subplots()
86+
ax.set_ylabel('x[m]')
87+
ax.set_xlabel('t[s]')
88+
ax.plot(t, x)
89+
fig.tight_layout()
90+
save_fig("ForcedBlockRK4")
91+
plt.show()

0 commit comments

Comments
 (0)