Aquila-Consortium/csiborgtools
0
0
Fork 0
mirror of https://github.com/Richard-Sti/csiborgtools.git synced 2024-10-18 18:25:06 +02:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add SFI groups (#121)

Browse Source 
* Update LOS sampling

* Bring SD back to speed

* Update nb

* Update script

* Update requirements

* Update flow chi2 calc

* Update script

* Add support for SFI groups

* Update plot

* Update script

* Update script
This commit is contained in:
Richard Stiskalek 2024-03-26 15:11:04 +01:00 committed by GitHub
parent 27c1f9249b
commit 29b37172fa
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 253 additions and 135 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

57
csiborgtools/flow/flow_model.py
View File

@ -254,7 +254,7 @@ class DataLoader:
for key in f.keys():
arr[key] = f[key][:]
elif catalogue in ["LOSS", "Foundation", "SFI_gals", "2MTF",
"Pantheon+"]:
"Pantheon+", "SFI_gals_masked", "SFI_groups"]:
with File(catalogue_fpath, 'r') as f:
grp = f[catalogue]
@ -746,10 +746,13 @@ class SD_PV_validation_model(BaseFlowValidationModel):
self._r_xrange = r_xrange
# Get the various vmapped functions
self._f_ptilde_wo_bias = lambda mu, err: calculate_ptilde_wo_bias(r_xrange, mu, err, r2_xrange, True) # noqa
self._f_simps = lambda y: simps(y, dr) # noqa
self._f_zobs = lambda beta, Vr, vpec_rad: predict_zobs(r_xrange, beta, Vr, vpec_rad, Omega_m) # noqa
self._vmap_ptilde_wo_bias = vmap(lambda mu, err: calculate_ptilde_wo_bias(r_xrange, mu, err, r2_xrange, True)) # noqa
self._vmap_simps = vmap(lambda y: simps(y, dr))
self._vmap_zobs = vmap(lambda beta, Vr, vpec_rad: predict_zobs(r_xrange, beta, Vr, vpec_rad, Omega_m), in_axes=(None, 0, 0)) # noqa
self._vmap_ll_zobs = vmap(lambda zobs, zobs_pred, sigma_v: calculate_ll_zobs(zobs, zobs_pred, sigma_v), in_axes=(0, 0, None)) # noqa
# Distribution of external velocity components
self._Vext = dist.Uniform(-500, 500)
@ -757,12 +760,15 @@ class SD_PV_validation_model(BaseFlowValidationModel):
self._alpha = dist.LogNormal(*lognorm_mean_std_to_loc_scale(1.0, 0.5)) # noqa
self._beta = dist.Normal(1., 0.5)
# Distribution of velocity uncertainty sigma_v
self._sv = dist.LogNormal(*lognorm_mean_std_to_loc_scale(150, 100))
self._sigma_v = dist.LogNormal(*lognorm_mean_std_to_loc_scale(150, 100)) # noqa
self._h = dist.LogNormal(*lognorm_mean_std_to_loc_scale(1., 0.25))
self._Omega_m = Omega_m
def predict_zobs_single(self, **kwargs):
raise NotImplementedError("This method is not implemented yet.")
def __call__(self, sample_alpha=False):
def __call__(self, sample_alpha=True, sample_beta=True):
"""
The simple distance NumPyro PV validation model.
@ -771,29 +777,39 @@ class SD_PV_validation_model(BaseFlowValidationModel):
sample_alpha : bool, optional
Whether to sample the density bias parameter `alpha`, otherwise
it is fixed to 1.
sample_beta : bool, optional
Whether to sample the velocity bias parameter `beta`, otherwise
it is fixed to 1.
"""
Vx = numpyro.sample("Vext_x", self._Vext)
Vy = numpyro.sample("Vext_y", self._Vext)
Vz = numpyro.sample("Vext_z", self._Vext)
alpha = numpyro.sample("alpha", self._alpha) if sample_alpha else 1.0
beta = numpyro.sample("beta", self._beta)
sigma_v = numpyro.sample("sigma_v", self._sv)
beta = numpyro.sample("beta", self._beta) if sample_beta else 1.0
sigma_v = numpyro.sample("sigma_v", self._sigma_v)
h = numpyro.sample("h", self._h)
Vext_rad = project_Vext(Vx, Vy, Vz, self._RA, self._dec)
# Calculate p(r) and multiply it by the galaxy bias
ptilde = self._vmap_ptilde_wo_bias(self._r_hMpc, self._e2_rhMpc)
ptilde *= self._los_density**alpha
def scan_body(ll, i):
# Calculate p(r) and multiply it by the galaxy bias
ptilde = self._f_ptilde_wo_bias(
h * self._r_hMpc[i], h**2 * self._e2_rhMpc[i])
ptilde *= self._los_density[i]**alpha
# Normalization of p(r)
pnorm = self._vmap_simps(ptilde)
# Normalization of p(r)
pnorm = self._f_simps(ptilde)
# Calculate p(z_obs) and multiply it by p(r)
zobs_pred = self._vmap_zobs(beta, Vext_rad, self._los_velocity)
ptilde *= self._vmap_ll_zobs(self._z_obs, zobs_pred, sigma_v)
# Calculate p(z_obs) and multiply it by p(r)
zobs_pred = self._f_zobs(beta, Vext_rad[i], self._los_velocity[i])
ptilde *= calculate_ll_zobs(self._z_obs[i], zobs_pred, sigma_v)
ll = jnp.sum(jnp.log(self._vmap_simps(ptilde) / pnorm))
return ll + jnp.log(self._f_simps(ptilde) / pnorm), None
ll = 0.
ll, __ = scan(scan_body, ll, jnp.arange(self.ndata))
numpyro.factor("ll", ll)
@ -1258,7 +1274,7 @@ def get_model(loader, zcmb_max=None, verbose=True):
los_overdensity[mask], los_velocity[mask], RA[mask], dec[mask],
zCMB[mask], mB[mask], x1[mask], c[mask], e_mB[mask], e_x1[mask],
e_c[mask], loader.rdist, loader._Omega_m)
elif kind in ["SFI_gals", "2MTF"]:
elif kind in ["SFI_gals", "2MTF", "SFI_gals_masked"]:
keys = ["RA", "DEC", "z_CMB", "mag", "eta", "e_mag", "e_eta"]
RA, dec, zCMB, mag, eta, e_mag, e_eta = (loader.cat[k] for k in keys)
@ -1271,6 +1287,15 @@ def get_model(loader, zcmb_max=None, verbose=True):
los_overdensity[mask], los_velocity[mask], RA[mask], dec[mask],
zCMB[mask], mag[mask], eta[mask], e_mag[mask], e_eta[mask],
loader.rdist, loader._Omega_m)
elif kind == "SFI_groups":
keys = ["RA", "DEC", "z_CMB", "r_hMpc", "e_r_hMpc"]
RA, dec, zCMB, r_hMpc, e_r_hMpc = (loader.cat[k] for k in keys)
mask = (zCMB < zcmb_max)
model = SD_PV_validation_model(
los_overdensity[mask], los_velocity[mask], RA[mask], dec[mask],
zCMB[mask], r_hMpc[mask], e_r_hMpc[mask], loader.rdist,
loader._Omega_m)
else:
raise ValueError(f"Catalogue `{kind}` not recognized.")

287
notebooks/flow_calibration.ipynb
View File

File diff suppressed because one or more lines are too long

16
notebooks/flow_calibration.py
View File

@ -35,8 +35,10 @@ def read_samples(catalogue, simname, ksmooth, include_calibration=False,
if catalogue in ["LOSS", "Foundation", "Pantheon+"]:
alpha_cal, beta_cal, mag_cal, e_mu_intrinsic = [], [], [], []
elif catalogue in ["2MTF", "SFI_gals"]:
elif catalogue in ["2MTF", "SFI_gals", "SFI_gals_masked"]:
a, b, e_mu_intrinsic = [], [], []
elif catalogue == "SFI_groups":
h = []
else:
raise ValueError(f"Catalogue {catalogue} not recognized.")
@ -86,6 +88,8 @@ def read_samples(catalogue, simname, ksmooth, include_calibration=False,
a.append(f[f"sim_{nsim}/a"][:])
b.append(f[f"sim_{nsim}/b"][:])
e_mu_intrinsic.append(f[f"sim_{nsim}/e_mu_intrinsic"][:])
elif catalogue == "SFI_groups":
h.append(f[f"sim_{nsim}/h"][:])
else:
raise ValueError(f"Catalogue {catalogue} not recognized.")
@ -95,8 +99,10 @@ def read_samples(catalogue, simname, ksmooth, include_calibration=False,
if catalogue in ["LOSS", "Foundation", "Pantheon+"]:
alpha_cal, beta_cal, mag_cal, e_mu_intrinsic = np.hstack(alpha_cal), np.hstack(beta_cal), np.hstack(mag_cal), np.hstack(e_mu_intrinsic) # noqa
elif catalogue in ["2MTF", "SFI_gals"]:
elif catalogue in ["2MTF", "SFI_gals", "SFI_gals_masked"]:
a, b, e_mu_intrinsic = np.hstack(a), np.hstack(b), np.hstack(e_mu_intrinsic) # noqa
elif catalogue == "SFI_groups":
h = np.hstack(h)
else:
raise ValueError(f"Catalogue {catalogue} not recognized.")
@ -114,9 +120,12 @@ def read_samples(catalogue, simname, ksmooth, include_calibration=False,
if catalogue in ["LOSS", "Foundation", "Pantheon+"]:
data += [alpha_cal, beta_cal, mag_cal, e_mu_intrinsic]
names += ["alpha_cal", "beta_cal", "mag_cal", "e_mu_intrinsic"]
elif catalogue in ["2MTF", "SFI_gals"]:
elif catalogue in ["2MTF", "SFI_gals", "SFI_gals_masked"]:
data += [a, b, e_mu_intrinsic]
names += ["a", "b", "e_mu_intrinsic"]
elif catalogue == "SFI_groups":
data += [h]
names += ["h"]
else:
raise ValueError(f"Catalogue {catalogue} not recognized.")
@ -162,6 +171,7 @@ def names_to_latex(names, for_corner=False):
"beta": r"$\beta$",
"Vmag": r"$V_{\rm ext}$",
"sigma_v": r"$\sigma_v$",
"h": r"$h$",
}
labels = copy(names)

3
scripts/field_los.py
View File

@ -58,7 +58,8 @@ def get_los(catalogue_name, simname, comm):
if comm.Get_rank() == 0:
folder = "/mnt/extraspace/rstiskalek/catalogs"
if catalogue_name in ["LOSS", "Foundation", "SFI_gals", "2MTF",
if catalogue_name in ["LOSS", "Foundation", "SFI_gals",
"SFI_gals_masked", "SFI_groups", "2MTF",
"Pantheon+"]:
fpath = join(folder, "PV_compilation_Supranta2019.hdf5")
with File(fpath, 'r') as f:

6
scripts/field_los.sh
View File

@ -1,13 +1,13 @@
nthreads=3
nthreads=5
memory=40
on_login=${1}
on_login=0
queue="berg"
env="/mnt/users/rstiskalek/csiborgtools/venv_csiborg/bin/python"
file="field_los.py"
simname=${1}
catalogue=${2}
nsims="-1"
simname="csiborg2_main"
MAS="SPH"
grid=1024

13
scripts/flow_validation.py
View File

@ -52,7 +52,7 @@ def get_model(args, nsim_iterator, get_model_kwargs):
if args.catalogue == "A2":
fpath = join(folder, "A2.h5")
elif args.catalogue in ["LOSS", "Foundation", "Pantheon+", "SFI_gals",
"2MTF"]:
"2MTF", "SFI_groups", "SFI_gals_masked"]:
fpath = join(folder, "PV_compilation_Supranta2019.hdf5")
else:
raise ValueError(f"Unknown catalogue: `{args.catalogue}`.")
@ -107,10 +107,13 @@ def run_model(model, nsteps, nburn, nchains, nsim, dump_folder,
# Calculate the chi2
keys = list(thinned_samples.keys())
nsamples = len(thinned_samples[keys[0]])
zobs_mean, zobs_std = model.predict_zobs(thinned_samples)
nu = model.ndata - len(keys)
chi2 = [np.sum((zobs_mean[:, i] - model._z_obs)**2 / zobs_std[:, i]**2) / nu # noqa
for i in range(nsamples)]
try:
zobs_mean, zobs_std = model.predict_zobs(thinned_samples)
nu = model.ndata - len(keys)
chi2 = [np.sum((zobs_mean[:, i] - model._z_obs)**2 / zobs_std[:, i]**2) / nu # noqa
for i in range(nsamples)]
except NotImplementedError:
chi2 = [0. for _ in range(nsamples)]
gof = csiborgtools.numpyro_gof(model, mcmc, model_kwargs)

6
scripts/flow_validation.sh
View File

@ -1,5 +1,5 @@
memory=4
on_login=1
on_login=0
nthreads=${1}
ksmooth=${2}
@ -7,8 +7,8 @@ queue="berg"
env="/mnt/users/rstiskalek/csiborgtools/venv_csiborg/bin/python"
file="flow_validation.py"
catalogue="LOSS"
simname="Carrick2015"
catalogue="SFI_groups"
simname="csiborg2_main"
pythoncm="$env $file --catalogue $catalogue --simname $simname --ksmooth $ksmooth"