csiborgtools/notebooks/flow/selection.ipynb

Selection fitting¶

In [ ]:

import numpy as np
import matplotlib.pyplot as plt
from tqdm import trange
from h5py import File
from jax.random import PRNGKey
from numpyro.infer import MCMC, NUTS, init_to_median
from astropy.cosmology import FlatLambdaCDM 
from corner import corner

import csiborgtools

%matplotlib inline
%load_ext autoreload
%autoreload 2

Om0 = 0.3
H0 = 100
cosmo = FlatLambdaCDM(H0=H0, Om0=Om0)

Fit parameters of the toy selection model¶

Choose either CF4 TFR or SFI.

In [ ]:

# with File("/mnt/extraspace/rstiskalek/catalogs/PV_compilation.hdf5", 'r') as f:
#     grp = f["SFI_gals"]
# #     # print(grp.keys())
#     mag = grp["mag"][...]


# with File("/mnt/extraspace/rstiskalek/catalogs/PV/CF4/CF4_TF-distances.hdf5", 'r') as f:
    # mag = f["w1"][...]
# mag = mag[mag > 3]

model = csiborgtools.flow.ToyMagnitudeSelection()

In [ ]:

nuts_kernel = NUTS(model, init_strategy=init_to_median(num_samples=5000))
mcmc = MCMC(nuts_kernel, num_warmup=15_000, num_samples=15_000)
mcmc.run(PRNGKey(42), extra_fields=("potential_energy",), mag=mag)
samples = mcmc.get_samples()

mcmc.print_summary()

In [ ]:

keys = ["alpha", "a", "m1", "m2"]
data = np.vstack([samples[key] for key in keys]).T
labels = [r"$\alpha$", r"$a$", r"$m_1$", r"$m_2$"]

fig = corner(data, labels=labels, show_titles=True, smooth=True)
# fig.savefig("../../plots/selection_corner_CF4.png", dpi=450)

fig.show()

In [ ]:

for key in keys:
    print(f"{key}: {np.mean(samples[key]):.3f} +/- {np.std(samples[key]):.3f}")

In [ ]:

mrange = np.linspace(mag.min(), mag.max(), 1000)
nsamples = len(samples["m1"])

indx = np.random.choice(nsamples, 500)

y = [model.log_observed_pdf(mrange, samples["alpha"][i], samples["m1"][i], samples["m2"][i], samples["a"][i]) for i in indx]
y = np.asarray(y)
y = 10**y

In [ ]:

plt.figure()
plt.hist(mag, bins="auto", density=True, histtype="step", color="blue",
         label="Data", zorder=1)

for i in range(100):
    plt.plot(mrange, y[i], color="black", alpha=0.25, lw=0.25)

plt.xlabel(r"$m$")
plt.ylabel(r"$p(m)$")
plt.tight_layout()

plt.savefig("../../plots/CF4_selection.png", dpi=450)

plt.show()

Hubble¶

$p(m) \propto 10^{0.6 m}$ ?

In [ ]:

from scipy.integrate import quad
from scipy.interpolate import interp1d

zmin=0.00001
zmax=5
z_range = np.linspace(zmin, zmax, 100000)
r_range = cosmo.comoving_distance(z_range).value
distmod_range = cosmo.distmod(z_range).value
r2mu = interp1d(r_range, distmod_range, kind="cubic")


def schechter_LF(M, M0=-20.83, alpha=-1):
    return 10**(0.4 * (M0 - M) * (alpha + 1)) * np.exp(-10**(0.4 * (M0 - M)))


def sample_schechter_LF(M0=-20.83, alpha=-1, Mfaint=-16, Mbright=-30, npoints=1):
    norm = quad(schechter_LF, Mbright, Mfaint, args=(M0, alpha))[0]

    samples = np.full(npoints, np.nan)
    for i in trange(npoints):
        while np.isnan(samples[i]):
            M = np.random.uniform(Mbright, Mfaint)
            if np.random.uniform(0, 1) < schechter_LF(M, M0, alpha) / norm:
                samples[i] = M

    return samples


def sample_radial_distance(rmax, npoints):
    return rmax * np.random.rand(npoints)**(1/3)


# z = np.linspace(0.001, 0.15, 100000)
# r = cosmo.comoving_distance(z).value
# mu = cosmo.distmod(z).value
# 
# 
# drdmu = np.gradient(r, mu)

In [ ]:

rmax = 300
npoints = 5000

r_150 = sample_radial_distance(100, npoints)
r_300 = sample_radial_distance(300, npoints)
r_1000 = sample_radial_distance(5000, npoints)

mu_150 = r2mu(r_150)
mu_300 = r2mu(r_300)
mu_1000 = r2mu(r_1000)

In [ ]:

def p_hubble(m, a, b):
    norm = np.log10(- 5 / np.log(1000) * (10**(3 / 5 * a) - 10**(3 / 5 * b)))
    return 10**(0.6 * m - norm)

In [ ]:

M_LF = sample_schechter_LF(npoints=npoints)

M_LF2 = sample_schechter_LF(npoints=npoints, M0=-20.83, alpha=-1.5)

In [ ]:

plt.figure()
M = -20.3

# m = mu + M
# x = np.linspace(11, m.max(), 1000)
# plt.plot(x, p_hubble(x, m.min(), m.max()) * 5.5, color="black")

# plt.hist(m, bins="auto", density=True, histtype="step", color="blue",)


cols = ["red", "green", "blue"]
rmax = [150, 300, 1000]
# for i, mu in enumerate([mu_150, mu_300, mu_1000]):
for i, mu in enumerate([mu_150, mu_300, mu_1000]):
    plt.hist(mu + M_LF, bins="auto", density=True,
             histtype="step", color=cols[i], label=rmax[i])

    plt.hist(mu + M_LF2, bins="auto", density=True,
             histtype="step", color=cols[i], label=rmax[i], ls="--")


plt.hist(mag, bins="auto", density=True, histtype="step", color="black", label="Data")

plt.yscale("log")
# plt.axvline(r2mu(rmax) + M, c="red")
plt.legend()

plt.show()

In [ ]:

In [ ]:

In [ ]:

In [ ]:

In [ ]:

M = sample_schechter_LF(npoints=10000)

In [ ]:

plt.figure()
plt.hist(x, bins="auto", density=True, histtype="step", color="blue",)
# plt.yscale("log")
plt.show()

In [ ]:

In [ ]:

yeuclidean = 10**(0.6 * mu)
ycomoving = r**2 * drdmu



k = np.argmin(np.abs(mu - 35)) 

yeuclidean /= yeuclidean[k]
ycomoving /= ycomoving[k]



plt.figure()
plt.plot(z, yeuclidean, label="Euclidean")
plt.plot(z, ycomoving, label="Comoving")

# plt.yscale('log')
plt.xlabel(r"$z$")
plt.ylabel(r"$p(\mu)$")

plt.legend()
plt.tight_layout()
plt.savefig("../../plots/pmu_comoving_vs_euclidean.png")
plt.show()

In [ ]:

In [ ]:

from scipy.interpolate import interp1d
from scipy.integrate import quad
from scipy.stats import norm

z = np.linspace(0.001, 0.1, 100000)
r = cosmo.comoving_distance(z).value
mu = cosmo.distmod(z).value


drdmu = np.gradient(r, mu)



mu2drdmu = interp1d(mu, drdmu, kind='cubic')
mu2r = interp1d(mu, r, kind='cubic')



def schechter_LF(M):
    M0 = -20.83
    alpha = -1
    return 10**(0.4 * (M0 - M) * (alpha + 1)) * np.exp(-10**(0.4 * (M0 - M)))
    



def phi(M):
    # return 1
    # return schechter_LF(M)# * norm.pdf(M, loc=-22, scale=1)
    loc = -22
    std = 0.1

    return norm.pdf(M, loc=loc, scale=std)

    # if -22 < M < -21:
    #     return 1
    # else:
    #     return 0

In [ ]:

xrange = np.linspace(-24, -18, 1000)

plt.figure()
plt.plot(xrange, schechter_LF(xrange))
# plt.yscale("log")
plt.show()

In [ ]:

mu_min = mu.min()
mu_max = mu.max()


m = 12


m_range = np.linspace(10, 16, 100)
y = np.full_like(m_range, np.nan)
for i in trange(len(m_range)):
    m = m_range[i]
    # y[i] = quad(lambda x: mu2drdmu(x) * mu2r(x)**2 * phi(m - x), mu_min, mu_max)[0]
    y[i] = quad(lambda x: 10**(0.6 * x) * phi(m - x), mu_min, mu_max)[0]



y_hubble = 10**(0.6 * m_range)
ycomoving = r**2 * drdmu


k = np.argmin(np.abs(m_range - 12))

y_hubble /= y_hubble[k]
y /= y[k]

In [ ]:

mu_max - 18

In [ ]:

plt.figure()
plt.plot(m_range, y, label="Numerical")
plt.plot(m_range, y_hubble, label="Hubble")
# plt.plot(mu, ycomoving, label="Comoving")

plt.xlabel(r"$m$")
plt.ylabel(r"$p(m)$")
plt.legend()

# plt.yscale("log")
plt.tight_layout()
# plt.xlim(10, 14)

plt.savefig("../../plots/pm.png", dpi=450)
plt.show()

In [ ]:

Simple simulation¶

In [ ]:

npoints = 10000
rmax = 30000

# pos = np.random.uniform(-boxsize, boxsize, (npoints, 3))


r = rmax * np.random.rand(npoints)**(1/3)

mu = 5 * np.log10(r) + 25

# M = np.ones(npoints) * -22
# M = np.random.normal(-22, 100, npoints)
M = np.random.uniform(-24, -18, npoints)


m = mu + M

In [ ]:

def f(m, a, b):
    norm = np.log10(- 5 / np.log(1000) * (10**(3 / 5 * a) - 10**(3 / 5 * b)))
    return 10**(0.6 * m - norm)

In [ ]:

plt.figure()
plt.hist(m, bins="auto", density=True, histtype="step")
m_range = np.linspace(m.min(), m.max(), 100)
# plt.plot(m_range, f(m_range, m.min(), m.max()))
# plt.yscale("log")
plt.show()

In [ ]: