# -*- coding: utf-8 -*-
"""
Created on Thu Jan 25 17:22:01 2021
@author: neil
"""

# a 'short' program to demonstrate using a ODE solver
# in this first example we assume we know the 'span' of the problem
# calls the ODE function np.odeint, which is a wrapper around a basic FORTRAN
# ODE solver 

from scipy.integrate import odeint  # odeint is the basic solver, handles arrays of equations or one
import numpy as np
import matplotlib.pyplot as plt

# **********************header*************************************
# we know that the settling velocity of a rock in a low viscousity fluid
# is related to many variables, but that the settling velocity is proportional to the 
# square root of gravity.  But you probably don't know how the velocity changes from
# when it starts falling, to when it reaches terminal velocity(?)
# We ignore the details, and just use the fact that the drag force on a falling object
# is proportional to the square of the velocity. Using F = ma, a = dv/dt, and Sum F = Fgrav - Fdrag,
# we want to solve the eqn  dv/dt = g - c v**2, where v is the velocity downwards.
# 'c' would be all the other variables in the problem (density, size, shape etc.)

# So solve the velocity trajectory of a human body or a rock falling straightdown
# in different gravity fields and varying atmospheres or liquids
# NOTE the entire program is actually only 1 line, the rest is header or plotting

# **************************header*********************************************************
grav  = 9.8            # gravity
rhoAir= 1              # density of air
rhoRoc= 2700           # density of rock
DRoc  = .01            # diameter of rock
Cd    = .75            # drag coeff for a rough round rock (approx)
cons  = 3*Cd*rhoAir/(4*rhoRoc*DRoc)   # all the 'constants' piled together

vz0   = 0              # vertical velocity at t = 0, positive downwards(!)
z0    = 0              # starting point 
t0    = 0              # starting time

tmax  = 7.             # this is a kludge, we should really solve for when v reaches terminal velocity 
tarray= np.linspace(t0,tmax,100)  # make an array of times to send to the ODE solver

# ********************************************************************************************
#  minimal ODE usage x = odeint(func, y0, t), func is a function that we have to write that uses
#  the eqn of the derivative to return the value of the derivative at a specific point.
#  As well as the derivative function, odeint also requires arguements: y0 is the initial conditions 
#  (can be an array), t an array of time points to solve for. Returned is an array(len(t), len(y0)). 

# *******************************************************************************************
# the 'myderivative' function file looks like: dv/dt = g - c * v*v
def myderivative(vel,time):  # odeint requires the input to be [initial conditions], soln t points.
    return grav-cons*vel*vel # note that you can use variables from the main program ('namespace') 
                             # in the function.  It is usually considered poor programming practise.
# *******************************************************************************************

# ****************** Make our own ODE method, note it doesn't have any checks, nor is it very accurate
def myodeint(f,vstart,t): # make our own ode method, this is Euler stepping in time (discuss in class)
    """ primative ODE (forward difference), takes: a derivative function, start point, and
    evenly spaced vector of solution points"""
    v = np.empty_like(t)  # make an array to output our results of velocities
    delt = t[1] - t[0]    # assume evenly spaced solution points, so delt is the spacing
    v[0] = vstart         # first output is the starting point
    for n in np.arange(1,len(t)):       #fill in all the other points after the first
        v[n] = v[n-1]+f(v[n-1],1)*delt  #this ia a standard forward difference representation of the next value
    return v              # return our calculated solution vector at the 't' points

Vz  = myodeint(myderivative,vz0,tarray)     # myderivative is our function, 
                                          # Vz is an array of velocities vs time 

# ************************above is the program, below is plotting*******************************************
# for plotting, convert t to x using dz/dt = v, or dz = v*dt
dt = tarray[1]-tarray[0]
z = np.empty_like(tarray)  # create an array the same length as tarray
z[0] = 0                   # set the first value to 0
for n in np.arange(1,len(tarray)): # step thru Vz, summing the distance traveled at each time step
    z[n] = z[n-1] + dt*Vz[n]
print('position at tmax ',z[-1])
# some illustrations of plotting, using dx/dt = v to convert v or t to x
# this shows making a multi-panel plot to show various aspects of the solution
# make our basic plotting canvas, with a title in the header bar
fig = plt.figure()
fig.canvas.set_window_title ("Plots for falling bodies, using a coefficent of {:6.3f}".format(cons))

# Make a 2x2 array of plots, and start filling the first one
ax1 = fig.add_subplot(2,2,1)
ax1.grid(True)
ax1.plot(tarray,Vz, 'r-+')  
ax1.set_title("Velocity vs Time for a falling body")
ax1.set_xlabel("time")  
ax1.set_ylabel("velocity")    

# now we have x as well as t, we can plot the other 2 views
ax2 = fig.add_subplot(2,2,2)
ax2.grid(True)
ax2.plot(-z,Vz, 'b-+')  
ax2.set_title("Velocity vs Elevation for a falling body")
ax2.set_xlabel("Elevation (m)")  
ax2.set_ylabel("velocity (m/s)") 
ax2.invert_xaxis()

#  making 2 curves on one plot (reusing one of the axes [x])
ax2t = ax2.twinx()    
ax2t.plot(-z,Vz[-1]-Vz, 'g-')  
ax2t.set_ylabel("velocity below terminal velocity") 

# illustrating adding labels (including a dummy label)
ax3 = fig.add_subplot(2,2,3)
ax3.grid(True)
ax3.plot(tarray,-z, 'k-+', label = 'Elevation (m)')  
ax3.plot(np.nan, 'k+', label = 'soln points')
ax3.set_title("Distance vs Time for a falling body")
ax3.set_xlabel("time")  
ax3.set_ylabel("distance") 
ax3.legend(loc = 0, shadow=True, title = "Plot Labels" )             

plt.show()