#!/usr/bin/env python
# coding: utf-8

# $\renewcommand{\vec}[1]{\mathbf{#1}}$
# 
# 
# <style>
# @import url(https://www.numfys.net/static/css/nbstyle.css);
# </style>
# <a href="https://www.numfys.net"><img class="logo" /></a>
# 
# # Roller Coaster
# 
# ### Examples - Mechanics
# <section class="post-meta">
# By Jonas Tjemsland, Eilif Sommer Øyre and Jon Andreas Støvneng
# </section>
# Last edited: April 17th 2018 
# ___
# 
# 
# ## Introduction
# 
# In a series of laboratory sessions in a physics course at NTNU, the students are studying the rolling motion of objects on a curved track. Due to the curvature, the velocity and acceleration will vary. The students uses a high speed camera to record the motion and compare the results with numerical simulations.
# 
# In this notebook, we will simulate a rolling ball on some arbitrary track in two dimensions by solving Newton's second law. The track is made by performing cubic spline interpolation on a set of points. The setup will be based on the setup used in the aforementioned laboratory sessions. For more information, see [1].

# First, we import neccesary library packages.

# In[1]:


import numpy as np
import scipy.interpolate as interp
import matplotlib.pyplot as plt
plt.style.use('bmh')    # Nicer looking plots


# 
# 

# ## Theory
# #### Equation of Motion
# 
# An object, such as a ball, is rolling on a curved track. The rotation axis passes thorugh the center of mass. Let $m$ be the mass, $r$ be the radius of the ball and $I_n$ the moment of inertia. The forces on the center of mass is described by Newton's second law
# 
# \begin{equation}
# \mathbf F = m\mathbf a.
# \end{equation}
# 
# 
# The forces on the ball can be decomposed into a component parallel to the track and a component normal to the track, as shown in figure 1. Let the local slope of the track be described by an angle $\theta(x)$. The gravitational force acts on the center of mass vertically downwards. This amounts to a component parallel to the track $mg\sin \theta(x)$ and a normal component $mg\cos \theta(x)$. The forces from the track on the ball have a normal component $N$ (normal force) and the parallel component $f$ (friction).
# 
# 
# 
# <img src="https://www.numfys.net/media/notebooks/images/rolling-forces.png">
# **Figure 1:** *A cylinder or sphere is rolling on curve $y(x)$. The forces acting on the object are indicated on the figure. At each point $x$, the slope of the curve is defined by an angle $\theta(x)$.* 
# 
# 
# 
# Due to the curvature of the track, the center of mass has a centripetal acceleration normal to the track. The total acceleration is thus given by
# 
# \begin{equation}
# \mathbf a = \dot v \mathbf{e}_\parallel + \frac{v^2}{R(x)} \mathbf{e}_\perp,
# \end{equation}
#     
# where $R(x)$ is the *radius of the curvature* and $v$ is the speed (along the track). Thus, we obtain the two equations
# 
# \begin{equation}
# mg \sin \theta(x)-f=m\dot v,
# \label{eq:parallel}
# \end{equation}
# \begin{equation}
# N-mg\cos\theta(x) = \frac{mv^2}{R(x)}.
# \label{eq:normal}
# \end{equation}
# 
# The first equation describes the motion parallel to the track, while the second equation yields, with $N\geq 0$, a condition for when the ball loses contact with the surface.
# 
# An expression for the friction is found by using Newton's second law of rotation,
# 
# \begin{equation}
# \tau = I_0\omega.
# \end{equation}
# 
# The friction $f$ is the only force that do not pass through the rotation axis, and is thus the only force that contribute to the total torque $\tau$. The friction acts on the object a distance $r$ from the rotation axis (and at an angle $\pi/2$), so that $fr = I_0 \omega$. The ball is assumed to roll without any gliding. By using the rolling condition $v=\omega r$, equation \eqref{eq:parallel} can be written as
# 
# \begin{equation}
# \dot v = \frac{g\sin (\theta(x))}{1 + I_0/mr^2}.
# \label{eq:vdot}
# \end{equation}
# 
# A complete discussion on rotational dynamics can be found in e.g. [chap. 9-10, 2] and [chap. 6, 3].

# #### Curvature
# 
# Consider the curve from $y(x)$ from $A$ to $B$ as shown in figure 2. The curve from $A$ to $B$, with an arclenth $\Delta s$, can be considered as a small circle sector with angle $\Delta \theta$. The circle has a centre at $C$ and a radius $R\approx R_A\approx R_B$. The radius of the curvature thus becomes $R=\Delta s/\Delta \theta$. The *curvature* is in turn defined as $\kappa= 1/R$.
# 
# 
# <img src="https://www.numfys.net/media/notebooks/images/rolling-curvature.png">
# **Figure 2:** *Sketch used to describe curvature and the radius of the curvature. (Taken from [1])* 
# 
# 
# Consider now a infinitesimal arclenth $\Delta s \to\text{d} s$. From the figure one can see that
# 
# \begin{equation}
# \text{d}y/\text{d} x = \tan\theta,
# \label{eq:theta}
# \end{equation}
# 
# which gives
# 
# $$\frac{\text{d} \theta}{\text{d}x} = \frac{\text{d}}{\text{d}x}\arctan\left(\frac{\text{d}y}{\text{d}x}\right)=\frac{1}{1+\left(\text d y/\text dx\right)^2}\frac{\text d^2 y}{\text dx^2}.$$
# 
# The differential $\text d s$ is given by
# 
# $$(\text{d}s)^2 = (\text d x)^2 + (\text d y)^2 \quad \Longrightarrow \quad \text d s=\text d x\sqrt{1 + \left(\text d y/dx\right)^2}.$$
# 
# The curvature can thus be written as
# 
# \begin{equation}
# \kappa =\frac{\text{d}\theta}{\text{d}s} = \frac{\text{d}^2 y/\text{d}x^2}{\left(1 +\left(\text d y/\text d x\right)^2\right)^{3/2}},
# \end{equation}
# 
# and the radius of the curvature becomes
# 
# \begin{equation}
# R(x) = \frac{\text{d}s}{\text{d}\theta} = \frac{\left(1 +\left(\text d y/\text d x\right)^2\right)^{3/2}}{\text{d}^2 y/\text{d}x^2}.
# \label{eq:R}
# \end{equation}
# 
# Note that $\kappa$ is always finite as long as $\text d y/\text d x$ is continuous, while $R\to \infty$ if  $\text d^2 y/\text d x^2\to 0$. 
# 
# Curvature is discussed in more detail in e.g. [chap. 11, 4] and [chap. 13, 5].

# ## Setup
# 
# We will be considering a solid sphere (ball), which has a moment of inertia $I_0= 2mr^2/5$. The moment of inertia is easily computed for such objects, but this will not be discussed here. See e.g. [chap. 9, 2] for a general discussion and [6] for a list of some typical values for $I_0$. Assume that the ball has a radius $r=1\,\text{cm}$ and is made of iron with a density $\rho=7850\,\text{kg/m}^3$ (density found in [7]). The mass of the ball is in this case $m= 32.9\,\text{g}$ and the moment of inertia is thereby $I_0=13.2\,\mathrm{g\,cm^2}$.

# In[2]:


# Properties of the rolling object
r = 0.01                 # m      (radius)
rho = 7850               # kg/m^3 (density)
g = 9.81                 # m/s^2  (gravitational acceleration)

m = (4/3)*np.pi*r**3*rho # kg     (mass)
c = 2/5
I0 = c*m*r**2            # kg m^2 (moment of inertia)


# <img src="https://www.numfys.net/media/notebooks/images/roller-coaster.png">
# **Figure 3:** *The actual setup used in the laboratory frame is shown to the left. A sketch is shown to the right. (Right image taken from [1])*
# 
# 
# The track is made by a plastic bar which is attached to a solid frame at $N=8$ different mounts, as shown in figure 3. Let $(x_i, y_i)$ be the position of the $i$th mount. The $x$ positions are fixed and uniformly distributed across the frame. We will assume that the frame has a length $L=1.4\,\text{m}$, such that the distance between the mounts is $l = 20\,\text{cm}$. The $y$ positions can be changed.

# In[3]:


# Properties of the frame
L = 1.4                   # m (length)
yi = [.5, .3, .25, .3,
     .35, .34, .25, .15]  # m (y-positions)

N = len(yi)               #   (# of mounts)
xi = np.linspace(0, L, N) # m (x-positions)


# The track can be described by a cubic spline that interpolates $(x_i, y_i)$. The cubic spline consists of a set of cubic polynomials that has continuous first and second derivaties at the interpolation points. See our notebook on [Cubic Splines](https://nbviewer.jupyter.org/urls/www.numfys.net/media/notebooks/cubic_splines.ipynb) for a general discussion. We will be using the function [`CubicSpline`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.CubicSpline.html#scipy.interpolate.CubicSpline) from `scipy.interpolate` to perform the interpolation. This creates a callable class, which we can treat as a function $y(x)$ that describes the track.

# In[4]:


# Callable class for the track curve
get_y = interp.CubicSpline(xi, yi, bc_type="natural")


# As we saw earlier, $\theta$ and $R$ depends on $\text dy/\text dx$ and $\text d^2y/\text dx^2$. The class `CubicSpline` has a function called `derivative` which returns a class `PPoly`, which essentially is the derivative of the spline. `PPoly` has also a function `derivative`. We can thus easily compute $\theta(x)$ and $R(x)$ by using equations \eqref{eq:theta} and \eqref{eq:R}.

# In[5]:


get_dydx = get_y.derivative()      # Callable derivative of track
get_d2ydx2 = get_dydx.derivative() # Callable double derivative of track

def get_theta(x):
    """ Returns the angle of the track. """
    return -np.arctan(get_dydx(x))
    
def get_R(x):
    """ Returns the radius of the curvature. """
    return (1 + (get_dydx(x))**2)**1.5/get_d2ydx2(x)

def get_curvature(x):
    """ Returns the curvature (1/R). """
    return get_d2ydx2(x)/(1 + (get_dydx(x))**2)**1.5


# Let's plot the track!

# In[6]:


x = np.linspace(xi[0], xi[-1], 200)

# Create figure
fig, axarr = plt.subplots(3, 1, sharex=True, figsize=(12, 9), dpi=600)
fig.subplots_adjust(hspace=0.02)

# Axes 1:
axarr[0].plot(x, get_y(x), 'C0', label=r"$y(x)$")
axarr[0].plot(xi, yi, 'C1o', label="Mounts")
axarr[0].set_ylabel(r"$y(x)$, [m]", size='15')
#axarr[0].set_aspect('equal')

# Axes 2: 
axarr[1].plot(x, get_theta(x), 'C0')
axarr[1].set_ylabel(r"$\theta(x)$, [rad]", size='15')

# Axes 2: 
axarr[2].plot(x, get_curvature(x), 'C0')
axarr[2].set_ylabel(r"$\kappa(x)$, [1/m]", size='15')

plt.show()


# We are assuming that there is no loss of mechanical energy. Thus, if the highest point of the track is $y(x=0)$, then the ball will fall of to the right. If this is not the case, the ball will roll back and forth. 

# ## Numerics
# 
# We start by assuming that the ball is always in contact with the track. Equation \eqref{eq:vdot} yields a coupled set of ordinary differential equations (ODE)
# 
# $$\frac{\text dv}{\text dt} = \frac{g\sin (\theta(t))}{1 + I_0/mr^2},$$
# 
# where $v$ is the momentary velocity along the track. The $x$ position is in turn given by
# 
# $$ \text dx = \text ds \cos(\theta)\: \Longrightarrow \: \frac{\text dx}{\text dt}  = v\cos(\theta).$$

# In[7]:


def get_vdot(theta):
    """ Returns the time derivative of the (total) velocity. """
    return g*np.sin(theta)/(1 + c)

def RHS(z):
    """ Evaluates the right hand side of the two coupled 
    ODEs given in the text.
    
    Parameters:
        z : array-like, len(2), float. [v, x]
            The first element is the velocity, the second is the x-position.
        
    Returns:
        array-like, len(2), float. [a, v]
        The first element is the time derivative of the velocity (acceleration),
        the second is the time derivative of the position (velocity).
    """
    # z = [x, v]
    # w = [vx, a]
    w = np.zeros(2)
    theta = get_theta(z[0])
    w[0] = z[1]*1/np.sqrt(1+np.tan(theta)**2)
    w[1] = get_vdot(theta)
    return w


# $v(t)$ and $x(t)$ can be found by applying a method for solving ODEs, such as an [Euler method](https://nbviewer.jupyter.org/urls/www.numfys.net/media/notebooks/eulers_method.ipynb), a [Runge-Kutta method](https://nbviewer.jupyter.org/urls/www.numfys.net/media/notebooks/runge_kutta_method.ipynb) or a more advanced [addative method](https://nbviewer.jupyter.org/urls/www.numfys.net/media/notebooks/adaptive_runge_kutta_methods.ipynb). We refer you to the respective modules or an example that solves an ODE at [numfys.net](https://www.numfys.net). In this notebook we will be using the 4th order Runge-Kutta method.

# In[8]:


def rk4step(f, y, h):
    """ Performs one step of the 4th order Runge-Kutta method.
    
    Parameters:
        f : Callable function with one input parameter.
            The right hand side of the ODE. Note that the
            RHS is in our case not a function of time.
        y : array-like, float. Current position.
        h : float. Time step.
    """
    s1 = f(y)
    s2 = f(y + h*s1/2.0)
    s3 = f(y + h*s2/2.0)
    s4 = f(y + h*s3)
    
    return y + h/6.0*(s1 + 2.0*s2 + 2.0*s3 + s4)


# We choose $x=0$ and $v=0$ as initial conditions, and a time-step $\Delta t = 10^{-3}\, \text{s}$. We have assumed that the ball do not loose any mechanical energy. One method of checking whether the numerical computation gives realistic results is to check if the total mechanical energy
# 
# \begin{equation}
# E = \frac{1}{2}mv^2 + mgh + \frac{1}{2}I_0 \omega^2 = \frac{1}{2}m(1 + c)v^2 + mgh,
# \end{equation}
# 
# is constant.

# In[9]:


dt = 1e-3 # s
tstop = 5 # s. If the ball has not reached the end within 5 seconds, we stop.
x0 = 0    # m. Initial x-position
v0 = 0    # m/s. Inital velocity

def get_K(v):
    """ Returns the kinetic energy. """
    return .5*m*(1 + c)*v**2

def get_U(h):
    """ Returns the potential energy. """
    return m*g*h


# Everything is now set to roll the ball down the track. 

# In[10]:


# Set initial values
x = [x0]        # x-position
v = [v0]        # velocity
h = get_y(x0)   # height
K = [get_K(v0)] # kinetic energy
U = [get_U(h)]  # potential energy

it = 0 # Iterator
itmax = tstop/dt # Maximum number og iterations
while x0 <= L and it < itmax:
    # Perform one step of the Runge-Kutta method
    [x0, v0] = rk4step(RHS, [x0, v0], dt)
    # Append different values to their arrays
    x = np.append(x, x0)
    v = np.append(v, v0)
    h = get_y(x0)
    K = np.append(K, get_K(v0))
    U = np.append(U, get_U(h))
    it += 1

print("Iterations: %i"%(it))
print("Final time: %.2f s"%(it*dt))
dE = (K[0] - K[-1] + U[0] - U[-1])/(K[0] + U[0])
print("Relative change in mechanical energy: %.2e"%(dE))


# The relative change in the mechanical energy was minimal, which means that the time step was more than small enough. The ball used $1.07\,\text{s}$ to reach the end of the track.
# 
# Let's plot the computed quantities!

# In[11]:


t = np.linspace(0, it*dt, it + 1)

# Create figure
fig, axarr = plt.subplots(3, 1, sharex=True, figsize=(10, 8), dpi=400)
fig.subplots_adjust(hspace=0.02)

# Axes 1:
axarr[0].plot(t, x)
axarr[0].set_ylabel(r"$x(t)$, [m]")

# Axes 2: 
axarr[1].plot(t, v)
axarr[1].set_ylabel(r"$v(t)$, [m/s]")

# Axes 2: 
axarr[2].plot(t, U, label="Potential energy, $U$")
axarr[2].plot(t, K, label="Kinetic energy, $K$")
axarr[2].plot(t, K + U, label="Mechanical energy, $E=U+K$")
axarr[2].set_ylabel(r"Energy, [J]")
axarr[2].set_xlabel(r'$t$, [s]')
axarr[2].legend()

plt.show()


# Equation \eqref{eq:normal} gives a condition for when the ball looses its contact with the track:
# 
# $$g\cos\theta\left(x\right) + v^2\kappa\left(x\right) \leq 0.$$
# 
# Let's check if the ball lost contact in the above calculation.

# In[12]:


theta = get_theta(x)
kappa = get_curvature(x)
N = g*np.cos(theta) + v**2*kappa
index = np.argmax(N < 0)
if index <= 0:
    print("The ball did not loose contact.")
else:
    print("The ball lost contact at x = %.2f m."%(x[index]))


# We can also check whether the ball begins to slide or not. The maximum static friction force is proportional to the normal force:
# 
# \begin{equation}
# f_\mathrm{max} = \mu_s N.
# \end{equation}
# 
# $\mu_s$ is called the static friction coefficient. The static friction coefficient between plastic and metal is $0.25<\mu_s<0.4$ [8]. The normal force can be calculated using equation \eqref{eq:normal}. If the gravitational pull along the track exceeds the maximum static friction, $mg\sin\theta>f_\mathrm{max}$, the ball begins to glide.

# In[13]:


mu = 0.25
f = mu*N
Fgx = m*g*np.sin(theta)
index = np.argmax(f < Fgx)
if index <= 1:
    print("The ball did not glide.")
else:
    print("The ball began to glide at x = %.2f m."%(x[index]))


# ## Animation

# In[14]:


y = get_y(x)

from matplotlib import animation
from IPython.display import HTML

plt.style.use('default')

# Set up the figure
fig = plt.figure(figsize=(8, 4), dpi=120)
ax = plt.axes(xlim=(xi[0], xi[-1]), 
              ylim=(np.min(yi), np.max(yi)))
ax.set_aspect('equal')

# Define the different elements in the animation
track, = ax.plot(x,y, "k")
obj = plt.Circle((x[0], y[0]), r, fc="grey") # Obj. 1
ax.add_patch(obj)
ax.set_xlabel(r"$x$, [m]")
ax.set_ylabel(r"$y$, [m]")
props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
textbox = ax.text(0.775, 0.95, '', transform=ax.transAxes, fontsize=12,
                             verticalalignment='top', bbox=props)

# Calculates the number of frames
FPS = 30
n = len(t)
framesNum = int(FPS*t[-1])

# Animation function. This is called sequentially.
def animate(j):
    i = j*int(n/framesNum)
    theta = get_theta(x[i])
    obj.center = (x[i] + r*np.sin(theta), y[i] + r*np.cos(theta))
    t_str = r'$t = %.2f$ ' % (t[i])
    x_str = r'$x = %.2f$ ' % (x[i])
    y_str = r'$y = %.2f$ ' % (y[i])
    textbox.set_text(t_str + "\n" + x_str + "\n" + y_str)

# Create animation
anim = animation.FuncAnimation(fig, animate, frames=framesNum, interval=1000/FPS)

plt.close(anim._fig)

# Display the animation
HTML(anim.to_html5_video())


# ## Further Work
# 
# - If the ball slips, will it reach the end faster or slower?
# - How can we consider slipping?

# ## Resources
# 
# [1] NTNU fysikklab: http://home.phys.ntnu.no/brukdef/undervisning/fyslab/. In particular, the assignment text is found at http://home.phys.ntnu.no/brukdef/undervisning/fyslab/files/oppgave2.pdf.
# 
# [2] H.D. Young and R.A. Freedman: *University Physics*. Pearson Education, 14th edition, 2016 
# 
# [3] J.R. Lien and G. Løvhøiden: *Generell fysikk for universiteter og høgskoler*. Bind 1, Universitetsforlaget, 2001  
# 
# [4] R. A. Adams, C. Essex: *Calculus a Complete Course*. Pearson, 8th edition, 2013  
# 
# [5] J. Hass, M.D. Weir og G.B. Thomas Jr.: *Thomas’ Calculus*. Pearson Education, 13th edition, 2014  
# 
# [6] Wikipedia: *List of moments of inertia*. https://en.wikipedia.org/wiki/List_of_moments_of_inertia. 2017 (aquired 31th November 2017)  
# 
# [7] The Engineering ToolBox: *Metals and Alloys - Densities*. https://www.engineeringtoolbox.com/metal-alloys-densities-d_50.html. (aquired 31th November 2017)  
# 
# [8] Coefficient of friction: http://www.tribology-abc.com/abc/cof.htm