simple_edo_solver_and_plotter.py

import matplotlib.pyplot as plt
from matplotlib.ticker import (MultipleLocator, AutoMinorLocator, FormatStrFormatter)
import math
import numpy as np

# Funciton that returns the next differential.
# In the case of the example:
# dy/dt = cos(t)
def f(t, y):
    return math.cos(t)

# Using Runge-Kutta 4 method
def rk4(f, t, h, s):
    k1 = h * f(t, s)
    k2 = h * f(t + h / 2, s + k1 / 2)
    k3 = h * f(t + h / 2, s + k2 / 2)
    k4 = h * f(t + h, s + k3)
    return s + (k1 + 2 * k2 + 2 * k3 + k4) / 6
    
 # Initial condition:
s = 0

# Up to value of t:
t_max = 30

t = np.linspace(0, t_max, t_max)

## If discretization is 2, it goes like: 0, 0.5, 1, 1,5, 2...
discretization = 4
dt = 1/discretization
ts = [ t / discretization  for t in range(0, t_max * discretization) ]
S = []

for t in ts:
    S.append(s)
    # Optional print
    # print(s)
    s = rk4(f, t, dt, s)
    
def myPlot(x, S, x0, x1, y0, y1, dis_x, dis_y, mytitle):
    # Size of the plot (should be proportional axis scale)
    plot, ax = plt.subplots(figsize = (12, 12));

    # Axis
    plt.xlabel("t", fontsize = 16);
    plt.ylabel("y", fontsize = 16);

    # Title
    plt.title(mytitle, fontsize = 25)

    # Axis names
    plt.tick_params(axis='both', which='major', labelsize = 10)
    plt.tick_params(axis='both', which='minor', labelsize = 10)

    # Axis discretization
    ax.xaxis.set_major_locator( MultipleLocator(dis_x) )
    ax.yaxis.set_major_locator( MultipleLocator(dis_y) )

    # Axis bounds
    plt.xlim([x0, x1]);
    plt.ylim([y0, y1]);

    plt.plot(x, S, color = "blue", label='solution')
    plt.legend()
    ax.legend(fontsize="x-large")
    
# The function is similar to sin(t) (using the 'f' example)
myPlot(ts, S, 0, t_max, -2, 2, 1, 1, "title")