import math

import numpy as np

import matplotlib

import matplotlib.pyplot as plt

def cosmos_at_z(z, H0, w_r, w_m, w_v, n):

# Some constants

c = 299792.458 # speed of light in km/s

HtoGyr = 977.8 # coefficient for converting 1/H to Gyr

h = H0/100.0

# do a little calculation

# w_r = 4.165E-5/(h*h) # includes 3 massless neutrino species, T CMB = 2.72528

w_k = 1 - w_m - w_r - w_v # calculate Omega curvature (1 - Omega(total))

az = 1.0/(1.0 + z) # scale factor of a given z

# calculate age of the universe at redshift z

zage = 0.0

delta_a = az/n # step, delta_a

for i in range(n):

a = delta_a*(i + 0.5)

_a = 1.0/a

a_dot = math.sqrt(w_k + w_m*_a + w_r*_a*_a + w_v*a*a)

zage = zage + 1.0/a_dot

zage = delta_a*zage # age of the universe at redshift z in units of 1/H0

# calculate time from z to now, and comoving radial distance

dtt = 0.0 # time from z to now in units of 1/H0

dcmr = 0.0 # r, comoving radial distance in units of c/H0

delta_a = (1.0 - az)/n # step, delta_a

for i in range(n):

a = az+delta_a*(i + 0.5)

_a = 1.0/a

a_dot = math.sqrt(w_k + w_m*_a + w_r*_a*_a + w_v*a*a)

dtt = dtt + 1.0/a_dot

dcmr = dcmr + 1.0/(a*a_dot)

dtt = delta_a*dtt

dcmr = delta_a*dcmr # r = proper distance at the time observation, in units of c/H0

dpte = dcmr*az # proper distance at the time of emission, in units of c/H0

age = dtt + zage # age of the universe now in units of 1/H0

# calculate tangential comoving distance

ratio = 1.0

x = math.sqrt(abs(w_k))*dcmr # x = r/R0 or R0 = r/x

if x > 0.1:

if w_k > 0:

ratio = math.sinh(x)/x

else:

ratio = math.sin(x)/x

else:

y = x*x

if w_k < 0: y = -y

ratio = 1.0 + y/6.0 + y*y/120.0 # series expansion for sin(x) and sinh(x)

dcmt = ratio*dcmr # R0 * sin(r/R0) = r/x * sin(x)

da = az*dcmt
 dl = dcmt/az

# # human-readable

zage_Gyr = HtoGyr/H0*zage

age_Gyr = (HtoGyr/H0)*age

dcmr_Mpc = (c/H0)*dcmr

dtt_Gyr = (HtoGyr/H0)*dtt

da_mpc = (c/H0)*da

dl_mpc = (c/H0)*dl

print('For z, H0, Omega_M, Omega_Vacuum = ', z, H0, w_m, w_v)

print('The age at redshift z was = '+ str(zage_Gyr) + ' Gyr')

print('The comoving radial distance '+ str(dcmr_Mpc) + ' Mpc')

print('The light travel time was '+ str(dtt_Gyr) + ' Gyr')

print('It is now ' + str(age_Gyr) + ' Gyr since the Big Bang')

print('DA = '+ str(da_mpc), ' Mpc')

print('DL = ' + str(dl_mpc), ' Mpc')

return dcmr, dpte, da, dl, dtt, zage, age

def cosmos_at_z_array(z, H0, w_r, w_m, w_v, n):

dcmrlist, dptelist, dalist, dllist, dttlist, zagelist, agelist = [], [], [], [], [], [], []

for zi in z:

dcmr, dpte, da, dl, dtt, zage, age = cosmos_at_z(zi, H0, w_r, w_m, w_v, n)

dcmrlist.append(dcmr)

dptelist.append(dpte)

dalist.append(da)

dllist.append(dl)

dttlist.append(dtt)

zagelist.append(zage)

agelist.append(age)
 return dcmrlist, dptelist, dalist, dllist, dttlist, zagelist, agelist

def plot_proper_dist():

z = np.linspace(0.01, 1000, 1000)

benchmark = cosmos_at_z_array(z, 70, 0.0, 0.3, 0.7, 1000)

matter = cosmos_at_z_array(z, 70, 0.0, 1.0, 0.0, 1000)

vacuum = cosmos_at_z_array(z, 70, 0.0, 0.0, 1.0, 1000)

# proper distance at the time of observation

plt.plot(z, benchmark[0], 'r', label='Benchmark')

plt.plot(z, matter[0], 'g', label='matter-only')

plt.plot(z, vacuum[0], 'b', label=r'$\Lambda$-only')

plt.xscale('log')

plt.yscale('log')

plt.xlabel(r'$z$', fontsize=20)

plt.ylabel(r'$(H_{0}/c) d_{p}(t_{0})$', fontsize=20)

plt.show()

plt.savefig('fahmi2.png')

# proper distance at the time of emission

plt.plot(z, benchmark[1], 'r', label='Benchmark')

plt.plot(z, matter[1], 'g', label='matter-only')

plt.plot(z, vacuum[1], 'b', label=r'$\Lambda$-only')

plt.xscale('log')

plt.yscale('log')

plt.xlabel(r'$z$', fontsize=20)

plt.ylabel(r'$(H_{0}/c) d_{p}(t_{e})$', fontsize=20)

plt.show()

plt.savefig('fahmi3.png')
if __name__ == "__main__":

cosmos_at_z(1100, 67.74, 4.165E-5*10000/(67.74*67.74), 0.3089, 0.6911, 1000)

plot_proper_dist()

ref: Ned Wright (www.astro.ucla.edu/~wright)