Energy Calculations

import numpy as np
import matplotlib.pyplot as plt
from pylab import *
import os

H_0 = 70
h = 0.7
path = 'testhalos/'

#KINEMATIC FUNCTIONS
#returns the relative velocity, taking into account hubble expansion
def relative_velocity(xnew,ynew,znew,x,y,z,vx,vy,vz,system):
   vexpx = vx[system] + H_0 * (x[system]-xnew)
   vexpy = vy[system] + H_0 * (y[system]-ynew)
   vexpz = vz[system] + H_0 * (z[system]-znew)
   vrel = ((vexpx - vx[0])**2 + (vexpy - vy[0])**2 + (vexpz - vz[0])**2)**.5

   return vrel

#defines center of mass of system and returns components for new origin
def center_of_mass(rad,x,y,z,mass,system):
   totmass = np.sum(mass[0:system])
   compr = rad[0:system] * mass[0:system]
   compx = x[0:system] * mass[0:system]
   compy = y[0:system] * mass[0:system]
   compz = z[0:system] * mass[0:system]
   radnew = sum(compr) / totmass
   xnew = sum(compx) / totmass
   ynew = sum(compy) / totmass
   znew = sum(compz) / totmass


   return radnew, xnew, ynew, znew

#scales the mass for the radial velocity figure
def scale_mass(mass, system):
   maximum = np.log10(np.amax(mass))
   minimum = np.log10(np.amin(mass))
   ranges = maximum - minimum
   boundmass = 50**((np.log10(mass[:system]) - minimum) / ranges + 1)
   unboundmass = 50**((np.log10(mass[system:]) - minimum) / ranges + 1)

   #print "These are the bound masses:", boundmass
   return boundmass, unboundmass


#defines radial velocity
def radial_velocity(x,y,z,vx,vy,vz,system):
   vxpol = vx + H_0 * (x - x[0])
   vypol = vy + H_0 * (y - y[0])
   vzpol = vz + H_0 * (z - z[0])
   numerator = vxpol*(x-x[0]) + vypol*(y-y[0]) + vzpol*(z-z[0])
   denominator = ((x-x[0])**2 + (y-y[0])**2 + (z-z[0])**2)**.5

   vradial = np.zeros(numerator.shape[0])
   vradial[1:] = numerator[1:]/denominator[1:]

   return vradial

########################BEGINNING OF ANALYSIS##############################

def plot_velocities(halos, halofile, system):
   h = 0.7
   listno = halos[0,:]
   rad = halos[1,:] / h
   mass = halos[2,:]
   x = halos[3,:] / h
   y = halos[4,:] / h
   z = halos[5,:] / h
   vx = halos[6,:]
   vy = halos[7,:]
   vz = halos[8,:]

   #print "What does system equal?:", system
   vrad = radial_velocity(x, y, z, vx, vy, vz, system)

   fig = plt.figure()
   ax1 = fig.add_subplot(111)
   xax = np.linspace(0, rad[system+200], 100)
   yax = 70 * xax
   ax1.plot(xax, yax, color = 'c', linewidth = 2)

   bmass, ubmass = scale_mass(mass, system)

   target = ax1.scatter(0, vrad[0], s = bmass[0], c = 'black', alpha = 0.7)
   bound = ax1.scatter(rad[1:system], vrad[1:system], s = bmass[1:], c = 'green', alpha = 0.7)
   unbound = ax1.scatter(rad[system:system+200], vrad[system:system+200], s = ubmass[0:200], c = 'red', alpha = 0.7)
   for i in range(0,system+200):
      ax1.annotate('%d'%listno[i], xy = (rad[i], vrad[i]), xytext = (rad[i], vrad[i]), fontweight = 'bold')

   plt.xlim(xmin = 0)
   ax = plt.axes()
   ax.xaxis.grid()
   ax.yaxis.grid()
   ax1.set_xlabel('radius [Mpc]')
   ax1.set_ylabel('velocity [km/s]')
   plt.legend((target, bound, unbound), ('Target Halo', 'Bound Halos', 'Unbound Halos'), scatterpoints = 1, loc = 'upper left', labelspacing = 1)
   massy = np.array_str(mass[0])
   plt.title('Target Halo Mass = ' + massy + " M_sun")
######################need to change to given directory!
   plt.savefig(path+halofile+'velocity.png')
   #plt.show()
   plt.close()
   #print "Plotting done!"


#Plots T/V vs. radius until the point where the system becomes unbound
def plot_energy_ratio(Kinetic, Potential, halos, halofile):
   rad = halos[1,:] / h
   T = np.array(Kinetic)
   V = np.array(Potential)
   #print "Kinetic", T
   #print "Potential", V
   ratio = T / V

   plt.clf()
   ratioo, = plt.plot(rad[0:ratio.shape[0]], np.abs(ratio), 'darkmagenta')
   plt.xlabel('r [Mpc]')
   plt.ylabel('T/V')
   plt.legend([ratioo],["T/V"], loc='best')
   plt.grid()
   massy = np.array_str(halos[2,0])
   plt.title('Target Halo Mass = ' + massy + " M_sun")
   plt.savefig(path+halofile+'ratio.png')
   #plt.show()


#############################MAIN FUNCTION##################################
#Produces energy output of every cluster step until the system is unbound
def calculate_energy(halos, halofile):
   h= 0.7
   rad = halos[1,:]/h     #Mpc
   mass = halos[2,:]    #M_sun
   x = halos[3,:] / h   #Mpc
   y = halos[4,:] / h   #Mpc
   z = halos[5,:] / h   #Mpc
   vx = halos[6,:]      #km/s
   vy = halos[7,:]      #km/s
   vz = halos[8,:]      #km/s
   G = 4.302 * 10**-9  #Mpc M_s^-1 (km/s)^2
   H_0 = 70
   Potential = []
   Kinetic = []
   Energy = []
   E = 0
   system = 1

   #Calculates the energy in the system starting with the closest halo
   #to the target halo. If the value is negative, an additional
   #closest halo is added into the calculation. This repeats until
   #the energy values go positive.
   while (E <= 0):
      system += 1
      rnew, xnew, ynew, znew = center_of_mass(rad,x,y,z,mass,system)
      KE = 0
      PE = 0

      for i in range(0, system):
         #Kinetic Energy Calculation
         targmass = mass[i]
         vrel = relative_velocity(xnew,ynew,znew,x,y,z,vx,vy,vz,i)
         KE += .5 * targmass * vrel**2
         "KE:", KE
         #print "targmass:", targmass
         #print "vrel:", vrel

         #Potential Energy Calculation
         PE_comp = G * mass[i] * mass[0:i] / np.absolute((rad[i] - rad[0:i]))
         #print "Sums of PE", PE_comp
         PE -= np.sum(PE_comp)
         #print "PE:", PE

      E = KE + PE
      print "Potential: ", PE
      print "Kinetic: ", KE
      print "Total Energy: ", E
      print "System: ", system
      print ""

      Kinetic.append(KE)
      Potential.append(PE)
      Energy.append(E)

      #If system includes over a certain amount of halos, target is trashed
      if (system > 200):
         print "All halos bound"
         return -1, -1, -1, -1

   #print "system: ", system
   number = system

   #Plots the velocity distribution and the |T/V| energy ratio as a function
   #of radius
   plot_velocities(halos, halofile, system)
   plot_energy_ratio(Kinetic, Potential, halos, halofile)
   value = rad[number]
   totmass = np.sum(mass[0:number+1])
   print totmass
   halomass = mass[0]

   return number, value, halomass, totmass

#######################################################################


files = []

for root, dirnames, filenames in os.walk(path):
   for file in filenames:
      files.append(file)

print files

halofiles = [x for x in files if '.npy' in x]
trial = 0

#print halofiles
haloabundance = []
boundradius = []
halomass = []
totalmass = []

############all halos in directory
for halofile in halofiles:
   if '.png' in halofile:
      continue
   trial += 1
   print "##########################################################"
   print "file #:", trial
   print "##########################################################"
   halos = np.load(path+halofile)
   number, radius, hmass, totmass = calculate_energy(halos, halofile)
   if number == -1:
      continue
   halomass.append(hmass)
   haloabundance.append(number)
   boundradius.append(radius)
   totalmass.append(totmass)

statistics = np.array([halomass,haloabundance,boundradius, totalmass])
np.save("halostatistics.npy", statistics)

#############single halo
#user = raw_input("choose halo\n")
#halos = np.load(path+user)
#calculate_energy(halos, user)