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

# a 'short' program to demonstrate using a ODE solver
# in this example we assume we know the 'span' of the problem
# BUT we call our own ODE solver

# this version deals with the non-linearity of the underlying abstraction of the problem

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

# **************************header*********************************************************
g     = 9.8            # gravity
rho   = 1              # density of air
rhor  = 2700           # density of rock
D     = .01            # diameter of rock
Cd    = .75            # drag coeff for a rough round rock (approx)
c     = 3*Cd*rho/(4*rhor*D)   # all the 'constants' piled together
#c = 0

vz0   = -200              # vertical velocity at t = 0
z0    = 0              # starting point 
t0    = 0              # starting time

#tmax  = 1/(np.sqrt(c)) # this is a kludge, we should really solve for when v reaches terminal velocity
tmax = 10
tarray= np.linspace(t0,tmax,400)  # make an array of times to send to the ODE solver
delt = tarray[1] - tarray[0]

# ********************************************************************************************
#  minimal ODE usage x = odeint(func, y0, t), func is the eqn of the derivative as a function,
#  y0 is the initial conditions (can be an array), t an array of time points to solve for,
#  The solution is an array with shape (len(t), len(y0)). 

# *******************************************************************************************
# the 'derivative' function file looks like: dv/dt = g - c * v*v
# (the name can be anything)
def derivative(v,t):      # odeint requires the input order to the func be [initial conditions], soln t points.
    if v < 0 :
        return g + c*v*v  # this is subtle, since the underlying eqn requires the air resistance to slow
                          # the particule, we need to change the direction (sign) of the drag, depending on +/-V
    else :
        return g - c*v*v 
        # note that you can use variables in the main program ('namespace'), except the 
                          # names of the input variables are overwritten (here I have kept them the same
                          # but you could call them anything
# *******************************************************************************************

# ****************** 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
    """ 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

# use our own Euler integrator
#Vz  = myodeint(derivative,vz0,tarray)     # derivative is our function, Vz is an array of velocities vs time 
    
# or use a canned integrator
Vz  = odeint(derivative,vz0,tarray)[:,0]       # derivative is our function, this is using scipy canned integration
                                               # NOTE, odeint actually returns the results in the 1st col of an array

# 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]

# 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 an aggregate coefficent of {:6.3f}".format(c))
mang = plt.get_current_fig_manager()                                    # this is somewhat archane, but Matplotlib now requires you to
                                                                        # use the window manager to set the figure tilte 
mang.set_window_title ("Plots for falling bodies, using an aggregate coefficent of {:6.3f}".format(c))


# Make a 2x2 array of plots, and 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 in secs")  
ax1.set_ylabel("velocity m/s")    

# 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 Distance for a falling body")
ax2.set_xlabel("distance (meters)")  
ax2.set_ylabel("velocity m/s",color='blue') 

#  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",color='green') 

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

# illustrating setting specific plot elements
ax4 = fig.add_subplot(2,2,4)
ax4.grid(True)
line1=ax4.plot(tarray[1:],np.diff(Vz)/delt, 'g-')  
ax4.set_title("Acceleration vs Time for a falling body")
ax4.set_xlabel("time in secs")  
ax4.set_ylabel("Accel m/s2")
plt.setp(line1,marker='o',markerfacecolor='r',mec='k')
           

plt.show()