<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Two Dimensional Galactic Orbits
set initial conditions (x0,y0) and (vx0,vy0) in the plane z=0
set integration time step
set number of integrations or a final integration stop time
define the potential and the forces as derivatives of the potential
Step6:  The Plummer potential for mass $M_p$ and core radius $r_c$ is given by
$$
\Phi = -  {  M_p  \over {   {(r_c^2 + r^2)}^{1/2} }  }    \tag{1a}
$$
and is also used to described softened gravity of a point mass (think of the case $r_c = 0$) for N-body calculations.
The force is the gradient of the potential
$$
f = -\nabla \Phi                   \tag{1b}
$$
We also want to record the total energy (kinetic and potential)
Step11:  Integrator
Here we write a step function to solve
$$
  {  d\boldsymbol{x} \over dt } = \boldsymbol{v}   \tag{2a}  
$$
and
$$
  {  d\boldsymbol{v} \over dt } = f(\boldsymbol{x})   \tag{2b}  
$$
in the discretization we write this as (in a simple first order Euler algorithm)
xnew = xold + dt * vold
        vnew = vold + dt * fold
Step13:  Helper functions
Step14:  Initial conditions
For 2D orbits we only specify the X coordinate and Y velocity. The remaining values of the 6 phase space coordinates are 0. Why is this?
Step15:  Integrate
The following cell takes the last (pos,vel) and takes n steps in time dt
The cell after this will plot the orbit. If you re-execute the stepper cell, it will append, and shows how the orbit "grows" (or not).
Step16:  Saving data
There are many good and less ideal ways to save data. In astronomy standard formats such has FITS and HDF5 are common.  For our work here we use a simple and fast native python method, called pickle. You can save whole objects, and reading them back in will ensure the whole object structure and hierarchy is preserved.
Step18:  Questions
If we are just doing two dimensional orbits, can't we just leave the Z off and speed up computations? What do you need to change to do this?
How would we look for the period orbit?
If want to squash the potential and make it slightly oval, what would the changes be. Here we would define an ellipsoidal radius on which the potential is constant
Step19:  Energy and Angular Momentum conservation?
Now plot the energy conservation for this method as we did for the hand crafted one; did "odeint" do better?  Do the same for angular momentum.
<ASSISTANT_TASK:>
Python Code:
%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import math
Explanation: Two Dimensional Galactic Orbits
set initial conditions (x0,y0) and (vx0,vy0) in the plane z=0
set integration time step
set number of integrations or a final integration stop time
define the potential and the forces as derivatives of the potential
End of explanation
def radius(x):
    length of a vector
    if len(x.shape) == 1:
        return math.sqrt(np.inner(x,x))
    # elif len(x.shape) == 2:
        
def potential(pos):
    potential, defined as a negative number
    r = radius(pos)
    y1 = 1+r*r
    return -1.0/math.sqrt(y1)
def angmomz(pos,vel):
    Angular momentum in Z
    return pos[0]*vel[1] - pos[1]*vel[0]
def energy(pos,vel):
    Kinetic and Potential energy
    return 0.5*np.inner(vel,vel) + potential(pos)
def force(pos):
    force/acceleration (in our units mass=1 scale-length=1)
    # note we might be able to use sympy
    r = radius(pos)
    y2 = 1.0/math.sqrt(1+r*r)
    return -pos*y2*y2*y2
Explanation: The Plummer potential for mass $M_p$ and core radius $r_c$ is given by
$$
\Phi = -  {  M_p  \over {   {(r_c^2 + r^2)}^{1/2} }  }    \tag{1a}
$$
and is also used to described softened gravity of a point mass (think of the case $r_c = 0$) for N-body calculations.
The force is the gradient of the potential
$$
f = -\nabla \Phi                   \tag{1b}
$$
We also want to record the total energy (kinetic and potential):
$$
E = { 1\over 2} v^2 + \Phi          \tag{1c}
$$
and angular momentum
$$
J = r \times v                      \tag{1d}
$$
although we will only be using the Z component of this vector since we are computing orbits restricted to the Z plane:
$$
J_z = x \times v_y - y \times v_x         \tag{1e}
$$
Q1: We use a smooth function for the potential with smooth derivatives. What would happen with the "galaxy" potential from the previous lecture.
Force Field
We need some helper functions to compute (1a)..(1e).  For simplicity we use $M_p = 1$ and $r_c=1$.
End of explanation
def step0(pos,vel, dt):
    step0: simple first order Euler
    old = pos
    pos = pos + dt*vel
    vel = vel + dt*force(old)
    return (pos,vel)
def step1(pos,vel, dt):
    step1: simple first order Euler - updating position first
    pos = pos + dt*vel
    vel = vel + dt*force(pos)
    return (pos,vel)
def step2(pos,vel, dt):
    step2: simple first order Euler - updating velocity first
    vel = vel + dt*force(pos)
    pos = pos + dt*vel
    return (pos,vel)
def step4(pos,vel,dt):
    step4:  Runge Kutta 4 
    # not implemented yet
    return None
Explanation: Integrator
Here we write a step function to solve
$$
  {  d\boldsymbol{x} \over dt } = \boldsymbol{v}   \tag{2a}  
$$
and
$$
  {  d\boldsymbol{v} \over dt } = f(\boldsymbol{x})   \tag{2b}  
$$
in the discretization we write this as (in a simple first order Euler algorithm)
xnew = xold + dt * vold
        vnew = vold + dt * fold
End of explanation
def show_stats(data):
    Show some stats of a numpy array
    m = data.mean()
    s = data.std()
    dmin = data.min()
    dmax = data.max()
    rmin = (dmin-m)/s
    rmax = (dmax-m)/s
    print("Mean/Std:",m,s)
    print("Min/Max:",dmin,dmax)
    print("Rmin/Rmax:",rmin,rmax)
    print("Goodness: ",s/m)
Explanation: Helper functions
End of explanation
x0 = 1.0                       # initial X coordinate
v0 = 0.1                   # initial Y launch velocity (0.5946 would be a circular orbit)
n  = 200                    # number of steps to take
dt = 0.1                   # integration time step
step = step1                   # pick an integration method
print(step.__doc__)
                                       # Derived variables for the remainder
t = 0.0                                # always start at t=0
pos = np.array([x0,  0.0, 0.0])        # keeps the current pos
vel = np.array([0.0,  v0, 0.0])        # and vel
e = energy(pos,vel)                                        
j = angmomz(pos,vel)
time = np.zeros(1)                     # time array (we'll append to this)
time[0] = t
phase = np.concatenate(([t,e,j],pos,vel)).reshape(1,9)     # watch this peculiar 
print("e0 =",e)
print("phase = ",phase)
#  at x0=1.0 this should be the correct speed for a circular orbit
print("v0_circular=",1/math.pow(2.0,0.75))
Explanation: Initial conditions
For 2D orbits we only specify the X coordinate and Y velocity. The remaining values of the 6 phase space coordinates are 0. Why is this?
End of explanation
%%time
for i in range(n):
    (pos,vel) = step(pos,vel,dt)
    t = t + dt
    e = energy(pos,vel)
    j = angmomz(pos,vel)
    #print(i,pos,vel)
    p = np.concatenate(([t,e,j],pos,vel)).reshape(1,9)
    phase = np.concatenate((phase, p),axis=0)
    time = np.append(time,t)
#print(phase)
plt.scatter(phase[:,3],phase[:,4],c=time)
plt.axis('equal')
plt.title("Orbit")
x = phase[:,3]
y = phase[:,4]
rad = np.sqrt(x*x+y*y)-1
plt.scatter(phase[:,0],rad)
plt.scatter(phase[:,0], phase[:,1])
plt.title("Conserving Energy?")
show_stats(phase[:,1])
plt.scatter(phase[:,0], phase[:,2])
plt.title("Conserving Angular Momentum?")
show_stats(phase[:,2])
Explanation: Integrate
The following cell takes the last (pos,vel) and takes n steps in time dt
The cell after this will plot the orbit. If you re-execute the stepper cell, it will append, and shows how the orbit "grows" (or not).
End of explanation
try:
    import cPickle as pickle
    print("using cPickle")
except:
    import pickle
    print("using pickle")
# write it
pickle.dump(phase,open("orbit1.p","wb"))
# read it again
phase2 = pickle.load(open("orbit1.p","rb"))
print(phase[0])
print(phase2[0])
Explanation: Saving data
There are many good and less ideal ways to save data. In astronomy standard formats such has FITS and HDF5 are common.  For our work here we use a simple and fast native python method, called pickle. You can save whole objects, and reading them back in will ensure the whole object structure and hierarchy is preserved.
End of explanation
from scipy.integrate import odeint
def ofunc(y,t):
     function to integrate 
        Note we are re-using the force() function from the first part of this notebook
    
    pos = y[0:3]
    vel = y[3:]
    return np.concatenate((vel,force(pos)))
n=200
phase0 = np.array([x0,0,0, 0,v0,0])    # initial conditions
times  = np.arange(0.0,(n+1)*dt,dt)    # requested times where we want a solution
# times  = np.linspace(0.0,n*dt,n+1)
print(ofunc(phase0,0.0))
%%time
orbit = odeint(ofunc, phase0, times)
plt.scatter(orbit[:,0],orbit[:,1],c=times)
plt.axis('equal')
#plt.scatter(phase[:,3],phase[:,4])
plt.title("Orbit")
# plot the old one again 
plt.scatter(phase[:,3],phase[:,4],c=time)
plt.axis('equal')
plt.title("Orbit")
# compare the last
p1 = phase[-1,3:]
p2 = orbit[-1,:]
#
print(phase[-1,0],p1)
print(times[-1],p2)
print(0.0,p1-p2)
Explanation: Questions
If we are just doing two dimensional orbits, can't we just leave the Z off and speed up computations? What do you need to change to do this?
How would we look for the period orbit?
If want to squash the potential and make it slightly oval, what would the changes be. Here we would define an ellipsoidal radius on which the potential is constant:
$$
r^2 = { x^2 \over a^2} + { y^2 \over b^2 }
$$
instead of the normal 
$$
r^2 = x^2 + y^2
$$
A 2009 IAS lecture by Tremaine is an excellent lecture for (symplectic) orbit integrators. See https://video.ias.edu/PiTP2009-Tremaine
Orbits using scipy
For many scientific applications there are canned routines made available by the community.  The scipy package is one such module. We will derive the same orbit integration using scipy.odeint
See e.g. 
https://docs.scipy.org/doc/scipy-0.18.1/reference/generated/scipy.integrate.odeint.html#scipy.integrate.odeint
or better
https://docs.scipy.org/doc/scipy-0.18.1/reference/tutorial/integrate.html
However, this function uses the usual ODE notation (cf. (2a) and (2b))
$$
{  d\boldsymbol{y} \over dt } = f(\boldsymbol{y},t)   \tag{3}
$$
End of explanation
# e = energy(pos,vel)
# j = angmomz(pos,vel)
et = np.zeros(len(times))
jt = np.zeros(len(times))
for i in range(len(times)):
    pos=orbit[i,:3]
    vel=orbit[i,3:]
    et[i] = energy(pos,vel)
    jt[i] = angmomz(pos,vel)
plt.plot(times,et)
plt.show()
plt.plot(times,jt)
plt.show()
show_stats(et)
show_stats(jt)
Explanation: Energy and Angular Momentum conservation?
Now plot the energy conservation for this method as we did for the hand crafted one; did "odeint" do better?  Do the same for angular momentum.
End of explanation
<END_TASK>