import numpy as np # Numerical python
import pylab as plt # Plotting for python
# '#' comments that are interpretted as code
#
# Basics
#
#Scalars
x = 4 #define variable
print 'x =', x
print '2x =', 2*x
print 'x/2 =', x/2
print 'x squared =', x**2
#Arrays
a = np.array([1,2,3,4,5]) # Need np.array ([...] is just a 'list')
print 'a =', a
print '2a =', 2*a
print 'a squared =', a**2
print 'length of a =', len(a)
print 'first element of a =', a[0] #zero indexing: 0 to N-1
print 'mean of a =', np.mean(a), a.mean() #(two ways)
print 'variance of a =', np.var(a), a.var()
print 'explicit varance of a =', np.sum((a-a.mean())**2)/(len(a)-1)
#np.var does not use N-1
#
# Analysis Example
#
dx, dy = np.loadtxt('data_python_101.txt', unpack=True)
print dx, dy
Nx = len(dx)
Ny = len(dy)
print 'We have %d of dx data and %d of dy data' % (Nx, Ny) #format string
#Use %d for integers and %f and %e for floats (decimals)
#Plot a histogram
fs = 16 #fontsize for axis labels
plt.hist(dx, bins = 15, range = [-4,4], histtype='step', color='blue', label=r'$\Delta$x')
plt.hist(dy, bins = 15, range = [-4,4], histtype='step', color='red', label=r'$\Delta$y')
plt.xlim([-4,4])
plt.legend(loc = 'upper right')
plt.xlabel(r'Distance ($\mu$m)')
plt.ylabel('Step Distribution')
plt.savefig('step_histogram.png')
plt.clf() #clear the figure
#Calculate Results
dx_mean = dx.mean()
dy_mean = dy.mean()
dx_var = np.sum((dx-dx_mean)**2)/(Nx-1)
dy_var = np.sum((dy-dy_mean)**2)/(Ny-1)
print "dx result: mean = %e , var = %e" % (dx_mean, dx_var)
print "dy result: mean = %e , var = %e" % (dy_mean, dy_var)
#Save results
results = np.array([[dx_mean, dx_var], [dy_mean, dy_var]])
print 'results =', results
print 'results[0,1] =', results[0,1] #2d arrays: row then column indexing
np.savetxt('results.txt', results) #save results