#  Python 3
#  File: midpoint_rule.py
#  Demo Python code for Differential Equations class, Spring 2026
#  Authors: Jonathan Goodman, goodman@cims.nyu.edu
#  class web site: https://math.nyu.edu/~goodman/teaching/DifferentialEquations2026/DifferentialEquations.html

"""  
    Demonstrate the midpoint rule time stepping method for solving the
    initial value problem for an ODE.  Use a 1D ODE.
    This code is identical to the one using the forward Euler method
    except that the this uses the second order accurate midpoint rule.
    1. Compute the and plot the (approximate) solution for various dt.
    2. Plot the error as a function of dt on a log-log plot
"""

import numpy             as np    # load the nympy library, call it np
import matplotlib.pyplot as plt   # the plotting package

print("Demo of the midpoint rule.")

#  Set parameters

x_0 = 1.05    # initial condition
T   = 1.5    # integrate to this time
n_steps   = [5, 10, 20, 50, 100, 200, 300, 10000]  # different computational experiments
n_runs    = len(n_steps)             # the number of curves to plot
last_n    = n_steps[n_runs-1]        # the largest n gets a fatter line
x_T_vals  = np.zeros(n_runs)         # will be the x values at time T
                                     # for each value of n (dt = T/n)
dt_vals   = np.zeros(n_runs)         # for plotting error vs. dt


def f(x,t):
   """
      The ODE is (dx/dt) = f(x,t)
      This one is f = x^2-10*t^3 , made up to illustrate the process and accuracy
      Inputs:
         x (float): the state
         t (float): the time
      Return value
         (float), computed value of f
   """
   return x**2 - 10*t**3

fig,ax = plt.subplots()            # create the figure, empty at first, where plots will go
ax.set_title(r"midpoint ODE approximate solution, $\dot{x}=x^2-10t^3$")

#    Compute the function curves one by one

line_width      = 1                # make thin lines for all but the last plot
last_line_width = 2                # a fat line for the last n value
i               = 0                # index that determines which n is being used
for n in n_steps:                  # take each n value in turn

   dt = T/n                        # n time steps of size dt go from 0 to T
   
   t_vals = np.zeros(n+1)          # hold t_0=0 up to T = t_n
   x_vals = np.zeros(n+1)          # hold x_n from x_0 (initial val) to x_n
   
   t = 0.                          # starting time
   x = x_0                         # initial condition
   
   t_vals[0] = t                   # not necessary because it's already zero
   x_vals[0] = x_0
   
   for k in range(n):              # the time step loop
      xt = x + .5*dt*f(x,t)        # predict the midpoint value
      x = x + dt*f(xt,t+.5*dt)     # update using the midpoint value
      t = t + dt                   # replace t_k with t_{k+1}
      
      t_vals[k+1] = t              # these are the values at the next time
      x_vals[k+1] = x
   
   x_T_vals[i] = x                 # record for error analysis
   dt_vals[i]  = dt
   i           = i+1
      
   curve_label = "time step = {dt:11.2e}".format(dt=dt)
   if (n == last_n):          # the last plot has i = n_plots-1
      line_width = last_line_width
   ax.plot( t_vals, x_vals, label = curve_label, linewidth = line_width)

ax.legend()
ax.grid()
ax.set_xlabel(r"$t$")
ax.set_ylabel(r"$x$")
plt.savefig("midpoint_rule_curves.pdf")

#       The log log plot of the error as a function of delta t

fig,ax = plt.subplots()
ax.set_title(r"midpoint rule error at time $T$, as a function of $\Delta t$")

x_ground_truth = x_T_vals[n_runs-1]      # use a teeny tiny dt for "truth"
errors         = np.zeros(n_runs-1)      # omit the smallest dt in the plot
for i in range(n_runs-1):
   error     = np.abs( x_T_vals[i] - x_ground_truth)
   errors[i] = error
ax.loglog(dt_vals[0:n_runs-1], errors, "o-")

ax.set_xlabel(r"$\Delta t$")
ax.set_ylabel(r"error")
ax.grid()
plt.savefig("midpoint_rule_error.pdf")