Plotting the same function multiple times on the same set of axes but with different parameters [Python]

Question

I am currently trying to plot a function which describes linear perturbation growth in cosmology for different world models. I would like to be able to have all the curves on the same set of axes, but am struggling with setting it up.

This is the function

$D(z) = \frac{5 \Omega_M H_0^2}{2} \int_z^{\infty} \frac{1+z}{H^3(z)}dz $

where $H(z) = H_0 [\Omega_m(1+z)^3+\Omega_{\Lambda}]$

My aim is to plot this function D, with respect to z, but have multiple plots with varying density parameters ($\Omega$).

I have managed two solutions but both aren't working perfectly, the first is very inefficient (adding new functions for each set of parameters):

z = np.arange(0.0,10,0.1)

MOm = 1.0
MOv = 0.0

COm = 0.3
COv = 0.7

H0 = 70

def Mf(z):
    A = (5/2)*MOm*(H0**2)
    H = H0 * np.sqrt( MOm*((1+z)**3) + MOv )
    return A * ((1+z)/(H**3))

def MF(z):
    res = np.zeros_like(z)
    for i,val in enumerate(z):
        y,err = integrate.quad(Mf,val,np.inf)
        res[i] = y
    return res

def MD(z):
    return (H0 * np.sqrt( MOm*((1+z)**3) + MOv )) * MF(z)

def Cf(z):
    A = (5/2)*COm*(H0**2)
    H = H0 * np.sqrt( COm*((1+z)**3) + COv )
    return A * ((1+z)/(H**3))

def CF(z):
    res = np.zeros_like(z)
    for i,val in enumerate(z):
        y,err = integrate.quad(Cf,val,np.inf)
        res[i] = y
    return res

def CD(z):
    return (H0 * np.sqrt( COm*((1+z)**3) + COv )) * CF(z) 


plt.plot(z,MD(z),label="Matter Dominated")
plt.plot(z,CD(z),label="Current Epoch")

So I tried to make it simpler with a for loop, but have been unable to work out how to add labels to each plot inside the loop:

Om = (1.0,0.3)
Ov = (0.0,0.7)

for param1,param2 in zip(Om,Ov):
    def f(z):
        A = (5/2)*param1*(H0**2)
        H = H0 * np.sqrt( param1*((1+z)**3) + param2 )
        return A * ((1+z)/(H**3))

    def F(z):
        res = np.zeros_like(z)
        for i,val in enumerate(z):
            y,err = integrate.quad(f,val,np.inf)
            res[i] = y
        return res

    def D(z):
        return (H0 * np.sqrt( param1*((1+z)**3) + param2 )) * F(z)

    plt.plot(z,D(z))

Could someone please help explain an efficient method of doing so? Or how to add labels to plots on the fly with a for loop. Any help would be greatly appreciated.

Answer 1

The following code gives the plot in less than a sec on my pc.

from scipy.integrate import quad

def H(z, H0, MOm, MOv):
    return H0 * np.sqrt( MOm*((1+z)**3) + MOv )

# fixed values
COm = 0.3
COv = 0.7
H0 = 70

z_range = np.linspace(0, 1, 100)
Om_range = (1.0,0.3)
Ov_range = (0.0,0.7)

for (MOm, MOv) in zip(Om_range, Ov_range):
    d_range = [(5/2)*MOm*(H0**2) * 
               quad(lambda x: (1 + x)/H(z, H0, MOm, MOv)**3, z, np.inf)[0] for z in z_range]
    plt.plot(z_range, d_range, 
             label = '$\\Omega_m = ' + str(MOm)+ ', \\Omega_\\Lambda = ' + str(MOv)+'$')
plt.legend(loc = 0)
plt.xlabel('$z$')
plt.ylabel('$D(z)$')