diff --git a/examples/metadetect/config.yml b/examples/metadetect/config.yml index 64b34625..ee741f61 100755 --- a/examples/metadetect/config.yml +++ b/examples/metadetect/config.yml @@ -359,8 +359,8 @@ PZRealizationsPlotLens: photoz_realizations_plot: lens_photoz_realizations_plot TXShearCalibration: - add_fiducial_distance: True - copy_redshift: True + add_fiducial_distance: true + copy_redshift: true redshift_name: '00/redshift_true' TXTwoPointSCIA: diff --git a/examples/metadetect/pipeline.yml b/examples/metadetect/pipeline.yml index 8f2800ab..5cf705b0 100644 --- a/examples/metadetect/pipeline.yml +++ b/examples/metadetect/pipeline.yml @@ -120,6 +120,8 @@ stages: - name: HOSFSB # Disabled since not yet synchronised with new Treecorr MPI - name: TXTwoPointSCIA # Self-calibration intrinsic alignments of galaxies + - name: TXDeltaSigma + - name: TXDeltaSigmaPlots # Disabling these as they takes too long for a quick test. # The latter two need NaMaster diff --git a/txpipe/__init__.py b/txpipe/__init__.py index 39dbd2a3..2166f0ad 100644 --- a/txpipe/__init__.py +++ b/txpipe/__init__.py @@ -41,7 +41,7 @@ from .simulation import TXLogNormalGlass from .magnification import TXSSIMagnification from .covariance_nmt import TXFourierNamasterCovariance, TXRealNamasterCovariance - +from .delta_sigma import TXDeltaSigma # We no longer import all the extensions automatically here to avoid # some dependency problems when running under the LSST environment on NERSC. # You can still import them explicitly in your pipeline scripts by doing: diff --git a/txpipe/data_types/types.py b/txpipe/data_types/types.py index adf0787b..2414dfbd 100755 --- a/txpipe/data_types/types.py +++ b/txpipe/data_types/types.py @@ -468,6 +468,27 @@ def close(self): class FiducialCosmology(YamlFile): + + def to_astropy(self): + from astropy.cosmology import w0waCDM + + with open(self.path, "r") as fp: + params = yaml.load(fp, Loader=yaml.Loader) + + astropy_params = dict( + H0=params["H0"], + Om0=params["Omega_m"], + Ode0=1 - params["Omega_k"] - params["Omega_m"], + w0=params["w0"], + wa=params["wa"], + Neff=params["Neff"], + m_nu=params["sum_nu_masses"] + + ) + + cosmo = w0waCDM(**astropy_params) + + return cosmo # TODO replace when CCL has more complete serialization tools. def to_ccl(self, **kwargs): import pyccl as ccl diff --git a/txpipe/delta_sigma.py b/txpipe/delta_sigma.py new file mode 100644 index 00000000..9fce791e --- /dev/null +++ b/txpipe/delta_sigma.py @@ -0,0 +1,302 @@ +from .twopoint import TXTwoPoint, TREECORR_CONFIG, SHEAR_POS +from .base_stage import PipelineStage +from .data_types import SACCFile, ShearCatalog, HDFFile, QPNOfZFile, FiducialCosmology, TextFile, PNGFile +import numpy as np +from ceci.config import StageParameter +import os + + +class TXDeltaSigma(TXTwoPoint): + """Compute Delta-Sigma, the excess surface density around lenses. + + This is a subclass of TXTwoPoint that re-weights the lenses by the + inverse critical surface density based on the source redshift distribution. + + The strategy used here is to pre-compute sigma_crit^-1(z_lens) splines + for each source bin, and then use Treecorr's "w_eval" feature to apply + the appropriate weight to each lens based on its redshift when loading + the lens catalog. + """ + + name = "TXDeltaSigma" + + inputs = [ + ("binned_shear_catalog", ShearCatalog), + ("binned_lens_catalog", HDFFile), + ("binned_random_catalog", HDFFile), + ("shear_photoz_stack", QPNOfZFile), + ("lens_photoz_stack", QPNOfZFile), + ("tracer_metadata", HDFFile), + ("patch_centers", TextFile), + ("fiducial_cosmology", FiducialCosmology), + ] + outputs = [("delta_sigma", SACCFile)] + + config_options = TREECORR_CONFIG | { + "source_bins": StageParameter(list, [-1], msg="List of source bins to use (-1 means all)"), + "lens_bins": StageParameter(list, [-1], msg="List of lens bins to use (-1 means all)"), + "var_method": StageParameter(str, "jackknife", msg="Method for computing variance (jackknife, sample, etc.)"), + "use_randoms": StageParameter(bool, True, msg="Whether to use random catalogs"), + "low_mem": StageParameter(bool, False, msg="Whether to use low memory mode"), + "patch_dir": StageParameter(str, "./cache/delta-sigma-patches", msg="Directory for storing patch files"), + "chunk_rows": StageParameter(int, 100_000, msg="Number of rows to process in each chunk when making patches"), + "share_patch_files": StageParameter(bool, False, msg="Whether to share patch files across processes"), + "metric": StageParameter(str, "Euclidean", msg="Distance metric to use (Euclidean, Arc, etc.)"), + "gaussian_sims_factor": StageParameter( + list, + default=[1.0], + msg="Factor by which to decrease lens density to account for increased density contrast.", + ), + # "use_subsampled_randoms": StageParameter(bool, False, msg="Use subsampled randoms file for RR calculation"), + "delta_z": StageParameter(float, 0.001, msg="Z bin width for sigma_crit spline computation"), + } + + + def select_calculations(self, source_list, lens_list): + calcs = [] + # Fill in these config items that are not relevant for Delta Sigma + # but are needed by the parent class. + self.config["do_shear_pos"] = True + self.config["do_shear_shear"] = False + self.config["do_pos_pos"] = False + self.config["use_subsampled_randoms"] = False + + # For Delta Sigma everything is just shear-position + k = SHEAR_POS + for i in source_list: + for j in lens_list: + calcs.append((i, j, k)) + + if self.rank == 0: + print(f"Running {len(calcs)} calculations: {calcs}") + + return calcs + + + def read_nbin(self): + source_list, lens_list = super().read_nbin() + # This is a convenient place to also compute the sigma_crit_inverse splines + # that we need later + self.compute_sigma_crit_inverse_splines(source_list) + return source_list, lens_list + + + def compute_sigma_crit_inverse_spline(self, z, nz): + import scipy.interpolate + import scipy.integrate + + # We need a fiducial cosmology to convert redshifts to distances. + # This is why I don't like Delta Sigma. + with self.open_input("fiducial_cosmology", wrapper=True) as f: + cosmo = f.to_ccl() + + # Another option would be to use the QP object directly and use + # its .pdf() method. That's probably better than this. + zs_spline = scipy.interpolate.InterpolatedUnivariateSpline(z, nz, ext="zeros") + zmax = z.max() + 0.1 + dz = self.config["delta_z"] + + # We integrate this over source redshift for each lens redshift + # It's just a call to CCL weighted by the n(z) + def integrand(zl, zs): + w = np.where(zs > zl) + a_l = 1 / (1 + zl) + a_s = 1 / (1 + zs) + out = np.zeros_like(zs) + # https://ccl.readthedocs.io/en/latest/api/pyccl.cosmology.html#pyccl.cosmology.Cosmology.sigma_critical + sigma_crit = cosmo.sigma_critical(a_lens=a_l, a_source=a_s[w]) + n_z = zs_spline(zs[w]) + out[w] = n_z / sigma_crit + return out + + # Set up the grids needed by the integrals + zl_grid = np.arange(0, zmax, dz) + zs_grid = np.arange(0, zmax, dz) + sigma_crit_inv = np.zeros_like(zl_grid) + + # Do the integrals, just trapezoidal rule + for i, zli in enumerate(zl_grid[1:-1]): + integrand_vals = integrand(zli, zs_grid) + sigma_crit_inv[i] = scipy.integrate.trapezoid(integrand_vals, zs_grid) + + # This gives us the mean sigma_crit_inverse for each lens z + spline = scipy.interpolate.UnivariateSpline(zl_grid, sigma_crit_inv) + return spline + + + def compute_sigma_crit_inverse_splines(self, source_list): + """ + For each source bin, compute and store the spline + for sigma_crit_inverse as a function of lens redshift. + """ + from .utils import blank + splines = {} + with self.open_input("shear_photoz_stack", wrapper=True) as f: + for i in source_list: + if i == "all": + z, Nz = f.get_2d_n_of_z() + else: + z, Nz = f.get_bin_n_of_z(i) + if self.rank == 0: + print("Computing sigma_crit_inverse spline for source bin", i) + splines[i] = self.compute_sigma_crit_inverse_spline(z, Nz) + blank.sigma_crit_inverse_splines = splines + + def get_shear_catalog(self, i): + """ + This is a total hack. We want the get_lens_catalog to re-scale + the weights based on which source bin is being used. So we store + the source bin index in self.active_shear_catalog here, so that + get_lens_catalog can access it. + + This only works because the parent class calls get_shear_catalog + before get_lens_catalog. + + If that does change it should be obvious because self.active_shear_catalog + will be undefined. + """ + self.active_shear_catalog = i + return super().get_shear_catalog(i) + + + def get_lens_catalog(self, i): + import treecorr + from .utils import blank + treecorr.Catalog.eval_modules.append("txpipe.utils.blank as splines") + + # To get the scaling into Treecorr without re-creating multiple + # copies of all the files we have to pass the redshift into + # another module that treecorr can access and use the "w_eval" + # parameter to tell it to do it. + with self.open_input("binned_lens_catalog") as f: + redshift = f[f"/lens/bin_{i}/z"][:] + blank.redshift = redshift + + # Load the catalog. Note that we are not loading the weight + # from the file, because we want to scale it by 1/sigma_crit + # here. + cat = treecorr.Catalog( + self.get_input("binned_lens_catalog"), + ext=f"/lens/bin_{i}", + ra_col="ra", + dec_col="dec", + w_col="weight", + ra_units="degree", + dec_units="degree", + # Apply a function that rescales the weights + w_eval=f"weight / splines.sigma_crit_inverse_splines[{self.active_shear_catalog}](splines.redshift)", + patch_centers=self.get_input("patch_centers"), + save_patch_dir=self.get_patch_dir("binned_lens_catalog", i), + ) + + return cat + + def write_output(self, source_list, lens_list, meta, results): + import sacc + with self.open_input("fiducial_cosmology", wrapper=True) as f: + cosmo = f.to_ccl() + + # Get the source n(z) which are just passed into the output file + S = sacc.Sacc() + with self.open_input("shear_photoz_stack", wrapper=True) as f: + for i in source_list: + z, Nz = f.get_bin_n_of_z(i) + S.add_tracer("NZ", f"source_{i}", z, Nz) + + # Get the lens n(z), from which we also need the mean z + # so that we can turn angles into physical radii. + with self.open_input("lens_photoz_stack", wrapper=True) as f: + # For both source and lens + qp_object = f.read_ensemble() + + # we skip the "all" bin for now as we are not running it anyway + lens_mean_z = qp_object.mean()[:-1] + for i in lens_list: + z, Nz = f.get_bin_n_of_z(i) + S.add_tracer("NZ", f"lens_{i}", z, Nz) + + # Loop through calculated results and add to SACC + for d in results: + tracer1 = f"source_{d.i}" + tracer2 = f"lens_{d.j}" + + xi = d.object.xi + theta = np.exp(d.object.meanlogr) + # Convert theta to a radius in physical units at the lens redshift + a_lens = 1 / (1 + np.mean(lens_mean_z[d.j])) + r_mpc = cosmo.angular_diameter_distance(a_lens) * np.radians(theta / 60) + + weight = d.object.weight + err = np.sqrt(d.object.varxi) + n = len(xi) + for i in range(n): + S.add_data_point( + "galaxy_shearDensity_deltasigma", + (tracer1, tracer2), + xi[i], + theta=theta[i], + r_mpc=r_mpc[i], + error=err[i], + weight=weight[i], + ) + + # add a few bits of handy metadata + meta["nbin_source"] = len(source_list) + meta["nbin_lens"] = len(lens_list) + self.write_metadata(S, meta) + S.save_fits(self.get_output("delta_sigma"), overwrite=True) + + +class TXDeltaSigmaPlots(PipelineStage): + """Make plots of Delta Sigma results. + + """ + name = "TXDeltaSigmaPlots" + inputs = [ + ("delta_sigma", SACCFile), + ("fiducial_cosmology", FiducialCosmology), + + ] + outputs = [ + ("delta_sigma_plot", PNGFile), + ("delta_sigma_r_plot", PNGFile), +] + config_options = {} + + def run(self): + import sacc + import matplotlib.pyplot as plt + + sacc_data = sacc.Sacc.load_fits(self.get_input("delta_sigma")) + + # Plot in theta coordinates + nbin_source = sacc_data.metadata['nbin_source'] + nbin_lens = sacc_data.metadata['nbin_lens'] + with self.open_output("delta_sigma_plot", wrapper=True, figsize=(5*nbin_lens, 4*nbin_source)) as fig: + axes = fig.file.subplots(nbin_source, nbin_lens, squeeze=False) + for s in range(nbin_source): + for l in range(nbin_lens): + axes[s, l].set_title(f"Source {s}, Lens {l}") + axes[s, l].set_xlabel("Radius [arcmin]") + axes[s, l].set_ylabel(r"$\Delta \Sigma [M_\odot / pc^2]") + axes[s, l].grid() + x = sacc_data.get_tag("theta", tracers=(f"source_{s}", f"lens_{l}")) + y = sacc_data.get_mean(tracers=(f"source_{s}", f"lens_{l}")) + axes[s, l].plot(x, y) + plt.subplots_adjust(hspace=0.3, wspace=0.3) + + # Plot in r coordinates + nbin_source = sacc_data.metadata['nbin_source'] + nbin_lens = sacc_data.metadata['nbin_lens'] + with self.open_output("delta_sigma_r_plot", wrapper=True, figsize=(5*nbin_lens, 4*nbin_source)) as fig: + axes = fig.file.subplots(nbin_source, nbin_lens, squeeze=False) + for s in range(nbin_source): + for l in range(nbin_lens): + axes[s, l].set_title(f"Source {s}, Lens {l}") + axes[s, l].set_xlabel("Radius [Mpc]") + axes[s, l].set_ylabel(r"$R \cdot \Delta \Sigma [M_\odot / pc^2]$") + axes[s, l].grid() + x = sacc_data.get_tag("r_mpc", tracers=(f"source_{s}", f"lens_{l}")) + y = sacc_data.get_mean(tracers=(f"source_{s}", f"lens_{l}")) + axes[s, l].plot(x, y * np.array(x)) + plt.subplots_adjust(hspace=0.3, wspace=0.3) \ No newline at end of file diff --git a/txpipe/lens_selector.py b/txpipe/lens_selector.py index 0c264ec9..c6bd0bc6 100755 --- a/txpipe/lens_selector.py +++ b/txpipe/lens_selector.py @@ -608,7 +608,7 @@ def run(self): extra_cols = [c for c in self.config["extra_cols"] if c] # Regular columns. - cols = ["ra", "dec", "weight", "comoving_distance"] + cols = ["ra", "dec", "weight", "comoving_distance", "z"] # Object we use to make the separate lens bins catalog cat_output = self.open_output(self.get_binned_lens_name(), parallel=True) @@ -683,6 +683,7 @@ def add_redshifts(self, iterator): d = np.zeros(len(a)) d[z > 0] = pyccl.comoving_radial_distance(cosmo, a[z > 0]) data["comoving_distance"] = d + data["z"] = z yield s, e, data diff --git a/txpipe/twopoint.py b/txpipe/twopoint.py index 7202fc95..2d5b0b8a 100755 --- a/txpipe/twopoint.py +++ b/txpipe/twopoint.py @@ -81,7 +81,7 @@ class TXTwoPoint(PipelineStage): "use_randoms": StageParameter(bool, True, msg="Whether to use random catalogs"), "low_mem": StageParameter(bool, False, msg="Whether to use low memory mode"), "patch_dir": StageParameter(str, "./cache/patches", msg="Directory for storing patch files"), - "chunk_rows": StageParameter(int, 100_000, msg="Number of rows to process in each chunk"), + "chunk_rows": StageParameter(int, 100_000, msg="Number of rows to process in each chunk when making patches"), "share_patch_files": StageParameter(bool, False, msg="Whether to share patch files across processes"), "metric": StageParameter(str, "Euclidean", msg="Distance metric to use (Euclidean, Arc, etc.)"), "gaussian_sims_factor": StageParameter( @@ -423,7 +423,7 @@ def write_output(self, source_list, lens_list, meta, results): if self.config["do_shear_pos"] == True: S2.add_tracer("NZ", f"source_{i}", z, Nz) else: - sys.exit("Requesting a measurement that requires source galaxies but no source_list provided") + raise ValueError("Requesting a measurement that requires source galaxies but no source_list provided") if self.config["do_pos_pos"] or self.config["do_shear_pos"]: if lens_list: @@ -435,7 +435,7 @@ def write_output(self, source_list, lens_list, meta, results): if self.config["do_shear_pos"] == True: S2.add_tracer("NZ", f"lens_{i}", z, Nz) else: - sys.exit("Requesting a measurement that requires lens galaxies but no lens_list provided") + raise ValueError("Requesting a measurement that requires lens galaxies but no lens_list provided") # Now build up the collection of data points, adding them all to # the sacc data one by one. diff --git a/txpipe/utils/blank.py b/txpipe/utils/blank.py new file mode 100644 index 00000000..b3281116 --- /dev/null +++ b/txpipe/utils/blank.py @@ -0,0 +1,3 @@ +# This file is intentionally left blank. +# As part of a workaround in TreeCorr we use it as a namespace to +# store things needed to modify the weights in the DeltaSigma calculation. \ No newline at end of file