import numpy as np
import matplotlib.pyplot as plt

import scipy.integrate as spi

L = 1e-3       # inductance: 1mH
C = 10e-9      # capacitance: 10nF
R = 250        # resistance: 250ohm

# function that returns dz/dt=[Vcdot, Idot]
def model(z,t):
  Vc, I = z
  Vcdot = I/C
  if (t>25e-6 and t<=100e-6):
    Vin = 5
  else:
    Vin = 0
  Idot = (Vin-Vc-I*R)/L
  return np.array([Vcdot, Idot])

# initial condition
ic = [0, 0]
# time points
t = np.linspace(0,200e-6,1000)
u = np.zeros(1000)
u[125:500] = 5;
# solve ODE
states = spi.odeint(model,ic,t)

# plot results
plt.figure(figsize=(8,6), dpi=100)
plt.subplot(211)
plt.plot(t*10**6, u, '--', label=r'$V_{in}$')
plt.plot(t*10**6, states[:,0], label=r'$V_C$')
plt.legend()
plt.grid()
plt.ylabel(r'$V_c (V)$')
plt.subplot(212)
plt.plot(t*10**6, states[:,1])
plt.xlabel(r'time ($\mu$s)')
plt.ylabel(r'$I (A)$')
plt.grid()
plt.show()

Cd = 1.0  # drag coefficient

# function that returns dy/dt=[hdot, vdot]
def model(z,t):
  h, v = z                # h (altitude) and v (altitude rate)
  hdot = v                # dh/dt
  vdot = -10              # dv/dt
  return np.array([hdot, vdot])

# initial condition
ic = [2000, 0]
# time points
t = np.linspace(0,20,100)

# solve ODE
states = spi.odeint(model,ic,t)

# plot results
plt.figure(figsize=(8,6), dpi=100)
plt.subplot(211)
plt.plot(t, states[:,0], label='Numerical solution')
plt.plot(t, 2000-0.5*10*t*t, '--', label='Analytic solution')
plt.ylabel('h (m)')
plt.legend()
plt.grid()
plt.subplot(212)
plt.plot(t, states[:,1], label='Numerical solution')
plt.plot(t, -10*t, '--', label='Analytic solution')
plt.xlabel('time (s)')
plt.ylabel('v (m/s)')
plt.legend()
plt.grid()
plt.show()

Cd = 1.0  # drag coefficient

# function that returns dy/dt=[hdot, vdot]
def model(z,t,d):
  S = np.pi*d**2/4                  # cross-section area of the droplet
  m = np.pi*d**3/6*1000             # mass of the droplet
  h, v = z                          # h (altitude) and v (altitude rate)
  hdot = v                          # dh/dt
  rho = 1.225*(1-2.256e-5*h)**5.256 # air density
  vdot = -10 + 0.5*rho*v*v*S*Cd/m   # dv/dt
  return np.array([hdot, vdot])

# initial condition
ic = [2000, 0]
# time points
t = np.linspace(0,2,100)

# solve ODEs
d = 0.01e-3
states1 = spi.odeint(model,ic,t,args=(d,))
d = 0.1e-3
states2 = spi.odeint(model,ic,t,args=(d,))
d = 1e-3
states3 = spi.odeint(model,ic,t,args=(d,))

# plot results
plt.figure(figsize=(8,6), dpi=100)
plt.subplot(211)
plt.plot(t, states1[:,0], label='d=0.01mm')
plt.plot(t, states2[:,0], label='d=0.1mm')
plt.plot(t, states3[:,0], label='d=1mm')
plt.ylabel('h (m)')
plt.legend()
plt.grid()
plt.subplot(212)
plt.plot(t, states1[:,1], label='d=0.01mm')
plt.plot(t, states2[:,1], label='d=0.1mm')
plt.plot(t, states3[:,1], label='d=1mm')
plt.xlabel('time (s)')
plt.ylabel('v (m/s)')
plt.legend()
plt.grid()
plt.show()


T = len(t)
dt = t[-1]/T

# solve ODEs using Euler integration
states3_e = np.zeros((T,2))
states3_e[0,:] = ic

for k in range(T-1):
  states3_e[k+1,:] = states3_e[k,:] + dt*model(states3_e[k,:],t[k],1.00e-3)

# plot results
plt.figure(figsize=(8,6), dpi=100)
plt.plot(t, states3[:,1], label='scipy.integrate.odeint')
plt.plot(t, states3_e[:,1], '-.', label='Euler')
plt.xlabel('time (s)')
plt.ylabel('v (m/s)')
plt.legend()
plt.grid()
plt.show()

# solve ODEs using Adams-Bashforth integration
states3_ab = np.zeros((T,2))
states3_ab[0,:] = ic
deriv_p = model(states3_ab[0,:],t[0],1.00e-3)
for k in range(T-1):
  deriv = model(states3_ab[k,:],t[k],1.00e-3)
  states3_ab[k+1,:] = states3_ab[k,:] + 0.5*dt*(3*deriv - deriv_p)
  deriv_p = deriv

# plot results
plt.figure(figsize=(8,6), dpi=100)
plt.plot(t, states3[:,1], label='scipy.integrate.odeint')
plt.plot(t, states3_e[:,1], '-.', label='Euler')
plt.plot(t, states3_ab[:,1], ':', label='Adams-Bashforth')
plt.xlabel('time (s)')
plt.ylabel('v (m/s)')
plt.legend()
plt.grid()
plt.show()