<jupyter_start><jupyter_text>e2e examplesIn this notebook we will show how to compute a lognormal, apply and assess quality of reconstruction.<jupyter_code>import os
import tempfile

import numpy as np
from matplotlib import pyplot as plt

# For mockfactory installation, see https://github.com/adematti/mockfactory
from mockfactory import EulerianLinearMock, LagrangianLinearMock, utils, setup_logging
# For cosmoprimo installation, see https://cosmoprimo.readthedocs.io/en/latest/user/building.html
from cosmoprimo.fiducial import DESI

from pyrecon import MultiGridReconstruction, IterativeFFTReconstruction, PlaneParallelFFTReconstruction
# For pypower installation, see https://github.com/adematti/pypower
from pypower import CatalogFFTPower
from pyrecon.metrics import MeshFFTCorrelator, MeshFFTPropagator, CatalogMesh

# Set up logging
setup_logging()
z = 1.
# Load DESI fiducial cosmology
cosmo = DESI()
power = cosmo.get_fourier().pk_interpolator().to_1d(z=z)
f = cosmo.sigma8_z(z=z,of='theta_cb')/cosmo.sigma8_z(z=z,of='delta_cb') # growth rate<jupyter_output><empty_output><jupyter_text>Lognormal mock in box<jupyter_code>bias, nbar, nmesh, boxsize, boxcenter, los = 2.0, 1e-3, 200, 1000., 0., 'x'

mock = LagrangianLinearMock(power, nmesh=nmesh, boxsize=boxsize, boxcenter=boxcenter, seed=42, unitary_amplitude=False)
# This is Lagrangian bias, Eulerian bias - 1
mock.set_real_delta_field(bias=bias-1)
mesh_real = mock.mesh_delta_r + 1.
mock.set_analytic_selection_function(nbar=nbar)
mock.poisson_sample(seed=43)
mock.set_rsd(f=f, los=los)
data = mock.to_catalog()
offset = data.boxcenter - data.boxsize/2.
data['Position'] = (data['Position'] - offset) % data.boxsize + offset

from mockfactory.make_survey import RandomBoxCatalog
randoms = RandomBoxCatalog(nbar=nbar, boxsize=boxsize, boxcenter=boxcenter, seed=44)<jupyter_output>[000002.89] [0/1] 01-18 16:45  LagrangianLinearMock      INFO     White noise generated
[000004.74] [0/1] 01-18 16:45  LagrangianLinearMock      INFO     Poisson sampling done, total number of objects is 1000661.
[000005.73] [0/1] 01-18 16:45  LagrangianLinearMock      INFO     Catalog produced. Assigning in cell shift.<jupyter_text>Run reconstruction<jupyter_code>#ReconstructionAlgorithm = MultiGridReconstruction # choose your algorithm
ReconstructionAlgorithm = IterativeFFTReconstruction
#ReconstructionAlgorithm = PlaneParallelFFTReconstruction
recon = ReconstructionAlgorithm(f=f, bias=bias, los=los, nmesh=nmesh, boxsize=boxsize, boxcenter=boxcenter, fft_engine='fftw', fft_plan='estimate')
recon.assign_data(data.gget('Position'))
recon.set_density_contrast()
# Run reconstruction
recon.run()

data['Position_rec'] = data['Position'] - recon.read_shifts(data['Position'], field='disp+rsd')
randoms['Position_rec'] = randoms['Position'] - recon.read_shifts(randoms['Position'], field='disp')
offset = data.boxcenter - data.boxsize/2.
for catalog in [data, randoms]:
    catalog['Position_rec'] = (catalog['Position_rec'] - offset) % catalog.boxsize + offset<jupyter_output>[000007.95] [0/1] 01-18 16:45  IterativeFFTReconstruction INFO     Using mesh RealMesh(dtype=float64, nmesh=[200 200 200], boxsize=[1000. 1000. 1000.], boxcenter=[0. 0. 0.]).
[000008.64] [0/1] 01-18 16:45  IterativeFFTReconstruction INFO     Running iteration 0.
[000009.15] [0/1] 01-18 16:45  IterativeFFTReconstruction INFO     Running iteration 1.
[000009.63] [0/1] 01-18 16:45  IterativeFFTReconstruction INFO     Running iteration 2.<jupyter_text>Estimate power spectrum multipoles<jupyter_code>kedges = np.arange(0.01, 0.4, 0.005)

poles = CatalogFFTPower(data_positions1=data['Position'], boxsize=boxsize, boxcenter=boxcenter, nmesh=nmesh, resampler='cic',
                        interlacing=2, ells=(0, 2, 4), los=los, edges=kedges, position_type='pos').poles

poles_recon = CatalogFFTPower(data_positions1=data['Position_rec'], shifted_positions1=randoms['Position_rec'],
                            boxsize=boxsize, boxcenter=boxcenter, nmesh=nmesh, resampler='cic', interlacing=2, ells=(0, 2, 4), los=los, edges=kedges, position_type='pos').poles

ax = plt.gca()
for ill, ell in enumerate(poles.ells):
    ax.plot(poles.k, poles.k*poles(ell=ell, complex=False), color='C{:d}'.format(ill), linestyle='--', label='pre-recon' if ill == 0 else None)
    ax.plot(poles_recon.k, poles_recon.k*poles_recon(ell=ell, complex=False), color='C{:d}'.format(ill), linestyle='-', label='post-recon' if ill == 0 else None)
ax.legend()
ax.grid(True)
ax.set_xlabel('$k$ [$h/\mathrm{Mpc}$]')
ax.set_ylabel(r'$k P(k)$ [$(\mathrm{Mpc}/h)^{2}$]')
plt.show()<jupyter_output>[000011.85] [0/1] 01-18 16:45  CatalogFFTPower           INFO     Painting catalog 1 to mesh CatalogMesh(nmesh=[200 200 200], boxsize=[1000. 1000. 1000.], boxcenter=[0. 0. 0.], dtype=float64).
[000012.26] [0/1] 01-18 16:45  CatalogMesh               INFO     Running interlacing at order 2.
[000013.40] [0/1] 01-18 16:45  CatalogFFTPower           INFO     Done painting catalog 1 to mesh.
[000013.40] [0/1] 01-18 16:45  CatalogFFTPower           INFO     Using 77 k-bins between 0.010 and 0.395.
[000013.40] [0/1] 01-18 16:45  CatalogFFTPower           INFO     Running power spectrum estimation.
[000014.79] [0/1] 01-18 16:45  CatalogFFTPower           INFO     Power spectrum computed in elapsed time 1.39 s.
[000014.91] [0/1] 01-18 16:45  CatalogFFTPower           INFO     Painting catalog 1 to mesh CatalogMesh(nmesh=[200 200 200], boxsize=[1000. 1000. 1000.], boxcenter=[0. 0. 0.], dtype=float64).
[000015.79] [0/1] 01-18 16:45  CatalogMesh               INFO     Running interlacing at order [...]<jupyter_text>Compute correlator to assess reconstruction<jupyter_code>mesh_recon = CatalogMesh(data['Position_rec'], shifted_positions=randoms['Position_rec'],
                         boxsize=boxsize, boxcenter=boxcenter, nmesh=nmesh, resampler='cic',
                         interlacing=2, position_type='pos')

muedges = np.linspace(-1., 1., 5)
correlator = MeshFFTCorrelator(mesh_recon, mesh_real, edges=(kedges, muedges), los=los)
propagator = correlator.to_propagator(growth=bias)
transfer = correlator.to_transfer(growth=bias)
fig, lax = plt.subplots(nrows=1, ncols=3, figsize=(14,4))
fig.subplots_adjust(wspace=0.3)
lax = lax.flatten()
for imu, mu in enumerate(correlator.muavg[3:]):
    k = correlator.k[:,imu]
    lax[0].plot(k, correlator(k=k, mu=mu), label=r'$\mu = {:.2f}$'.format(mu))
    lax[1].plot(k, transfer(k=k, mu=mu), label=r'$\mu = {:.2f}$'.format(mu))
    lax[2].plot(k, propagator(k=k, mu=mu), label=r'$\mu = {:.2f}$'.format(mu))
for ax in lax:
    ax.legend()
    ax.grid(True)
    ax.set_xlabel('$k$ [$\mathrm{Mpc}/h$]')
lax[0].set_ylabel(r'$r(k) = P_{\mathrm{rec},\mathrm{init}}/\sqrt{P_{\mathrm{rec}}P_{\mathrm{init}}}$')
lax[1].set_ylabel(r'$t(k) = \sqrt{P_{\mathrm{rec}}/P_{\mathrm{init}}}$')
lax[2].set_ylabel(r'$g(k) = P_{\mathrm{rec},\mathrm{init}}/P_{\mathrm{init}}$')
plt.show()
# You can save the result (similarly for MeshFFTCorrelator and MeshFFTTransfer)
with tempfile.TemporaryDirectory() as tmp_dir:
    fn = os.path.join(tmp_dir, 'tmp.npy')
    propagator.save(fn)
    # And reload the result
    propagator = MeshFFTPropagator.load(fn)
    print(propagator.edges[0])<jupyter_output>[000028.30] [0/1] 01-18 16:45  MeshFFTPropagator         INFO     Saving /tmp/tmpfw8ilxs4/tmp.npy.
[000028.32] [0/1] 01-18 16:45  MeshFFTPropagator         INFO     Loading /tmp/tmpfw8ilxs4/tmp.npy.
[0.01  0.015 0.02  0.025 0.03  0.035 0.04  0.045 0.05  0.055 0.06  0.065
 0.07  0.075 0.08  0.085 0.09  0.095 0.1   0.105 0.11  0.115 0.12  0.125
 0.13  0.135 0.14  0.145 0.15  0.155 0.16  0.165 0.17  0.175 0.18  0.185
 0.19  0.195 0.2   0.205 0.21  0.215 0.22  0.225 0.23  0.235 0.24  0.245
 0.25  0.255 0.26  0.265 0.27  0.275 0.28  0.285 0.29  0.295 0.3   0.305
 0.31  0.315 0.32  0.325 0.33  0.335 0.34  0.345 0.35  0.355 0.36  0.365
 0.37  0.375 0.38  0.385 0.39  0.395]
{'fraction_non_alphanumeric': 0.14584124572730725, 'fraction_numerical': 0.09596151411571086, 'mean_word_length': 4.425824175824176, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '584181', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 3476, 'n_tokens_neox': 2946, 'n_words': 640, 'path': 'nb/e2e_examples.ipynb', 'repo_name': 'seshnadathur/pyrecon', 'size': '7899'}