Skip to content
Snippets Groups Projects
Select Git revision
  • 08edcc179af0edc604e32e4a32d4681f4f237fc1
  • main default protected
2 results

week4_3.py

Blame
    • week4_3.py 2.74 KiB
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74 
    

    

    

    import numpy as np
import matplotlib.pyplot as plt

'''
6.3)
Write a fourth-order Runge Kutta adaptive stepper for the preceding problem,
and check how the average step size that it finds depends on the desired local error.
'''

# Adaptive Runge-Kutta method (rewritten to be a step call)
def rk4_step(step, y, dy):
    k1, k2, k3, k4 = [0, 0], [0, 0], [0, 0], [0, 0] # make position and velocity vector
    k1[0] = step * dy # normal euler step for velocity
    k1[1] = - step * y  # normal euler step for position
    k2[0] = step * (dy + k1[1]*0.5) # half euler step in the direction of k1
    k2[1] = -step * (y + k1[0]*0.5)
    k3[0] = step * (dy + k2[1]*0.5) # half euler step in the direction of k2
    k3[1] = -step * (y + k2[0]*0.5)
    k4[0] = step * (dy + k3[1]) # full euler step in the direction of k3
    k4[1] = -step * (y + k3[0])
    y_new = y + (k1[0] + 2*k2[0] + 2*k3[0] + k4[0])/6 # divide by 6 (not 4) because we assigned 2x value to k2 and k3
    dy_new = dy + (k1[1] + 2*k2[1] + 2*k3[1] + k4[1])/6
    return y_new, dy_new

def rk4_adaptive(step_0, step_nudge, error_min, error_max):
    ''' In this case we don't know how many timesteps we will take, so it's best to
    start with lists and convert to numpy arrays at the end...
    '''
    t_final = 100*np.pi
    t = [0]
    y = [1]
    dy = [0]
    steps = [0.1]
    
    step = steps[0]
    time = step
    count = 1

    while time < t_final:
        y_full, dy_full = rk4_step(step, y[count-1], dy[count-1])
        y_half1, dy_half1 = rk4_step(step*0.5, y[count-1], dy[count-1])
        y_half2, dy_half2 = rk4_step(step*0.5, y_half1, dy_half1)

        error = np.abs(y_full - y_half2) # error is how closely the full and half steps agree
        # if we are too high or too low, nudge the step size and try again
        if error > error_max:
            step *= step_nudge
        elif error < error_min:
            step *= 1 - (step_nudge-1)
        # otherwise go ahead and append and increment time and count
        else:
            y.append(y_full)
            dy.append(dy_full)
            steps.append(step)
            time += step
            t.append(time)
            count += 1
    return t, y, dy, steps

step_0 = 0.1
nudge = 0.99 # increase or decrease the step by this percentage to nudge the step size
error_min =  1e-6 # easier to choose if theses are informed by something physical
error_max = 1e-4

t, y, dy, steps = rk4_adaptive(step_0, nudge, error_min, error_max)

fig, ax = plt.subplots(2)
ax[0].plot(t, y)
ax[0].title.set_text('\n Max Error: ' + str(round(error_max, 10)) + '\n Min Error: ' + str(round(error_min, 10)))
ax[1].plot(t, steps)
ax[1].title.set_text('Total number of steps: ' + str(len(t)) + '\n Average step size: ' + str(round(np.mean(steps), 6)))
plt.tight_layout()
plt.show()