<jupyter_start><jupyter_text>PIP examplesIn this notebook we will show how to compute correlation functions with PIP and angular upweights, with major inspiration from https://arxiv.org/pdf/1703.02070.pdf.Look first at notebook basic_examples.ipynb to understand the *pycorr* API. One can also look at the notebook https://github.com/cosmodesi/pypower/blob/main/nb/pip_examples.ipynb for similar examples for the power spectrum.<jupyter_code>import os
import tempfile

import numpy as np
from scipy import spatial
from matplotlib import pyplot as plt

from pycorr import TwoPointCorrelationFunction, utils, setup_logging

# To activate logging
setup_logging()
# First, generate fake data with cartesian positions and weights, in a box of size 800 (Mpc/h) at 1000 (Mpc/h).

def generate_catalogs(size=50000, boxsize=(800,)*3, offset=(1000., 0, 0), n_individual_weights=1, seed=42):
    rng = np.random.RandomState(seed=seed)
    positions = [o + rng.uniform(-0.5, 0.5, size)*b for o, b in zip(offset, boxsize)]
    weights = [np.ones(size, dtype='f8') for i in range(n_individual_weights)]
    return positions, weights

parent_positions, parent_weights = generate_catalogs()
randoms_positions, randoms_weights = generate_catalogs(size=200000, seed=84)<jupyter_output><empty_output><jupyter_text>Three scenarii- uncorrelated: objects within close pairs are randomly removed. That an object in close pair does not receive a fiber does not impact on the fiber assignment for the other object: selection probabilities of close objects are uncorrelated. There are no zero probability pairs.- onepass: we randomly pick a close pair and randomly discard one of the two objects; we iterate the procedure until there are no more close pairs. This is the "maximum randomness" algorithm used in https://arxiv.org/pdf/1703.02070.pdf. Selection probabilites of close objects are correlated. There are zero probability pairs.- twopass: same as singlepass, with an extra pass on objects left without fiber in the first pass. All pairs are non-zero probability with multiple passes.- twopass-fsky: same as singlepass, with an extra pass on objects left without fiber in the first pass on some fraction of the sky. All pairs are non-zero probability in the region with multiple passes.<jupyter_code># First find all (unique) pairs within 0.08 deg.
def find_collision_pairs(positions, radius_in_degree=0.08):
    distance = sum(p**2 for p in positions)**0.5
    unit_positions = np.array([p/distance for p in positions]).T
    tree = spatial.cKDTree(unit_positions, leafsize=16, compact_nodes=True, copy_data=False, balanced_tree=True)
    radius = 2.*np.sin(np.radians(radius_in_degree) / 2.) # distance on the unit sphere
    pairs = tree.query_pairs(radius, p=2., eps=0, output_type='ndarray')
    return pairs

parent_size = len(parent_positions[0])
target_pairs = find_collision_pairs(parent_positions)
print('Found {:d}/{:d} targets in pairs.'.format(target_pairs.shape[0], parent_size))
def run_uncorrelated_fiber_assignment(size, target_pairs, prob=0.5, nrealizations=65, seed=42):
    # Objects within close pairs are randomly removed with probability prob.
    bits = []
    indices_close_pairs = np.unique(target_pairs)
    rng = np.random.RandomState(seed=seed)
    for ireal in range(nrealizations):
        #print(ireal, end=" ")
        bit = np.ones(size, dtype='?')
        bit[indices_close_pairs] = rng.uniform(0., 1., indices_close_pairs.size) < prob
        bits.append(bit)
    return bits


def mr_fiber_collisions(rng, size, pairs, groupids):
    # We randomly pick a close pair and randomly discard one of the two objects.
    # We iterate the procedure until there are no more close pairs.
    bit = np.ones(size, dtype='?')
    collision_groups_pairs = {}

    for pair in pairs:
        groupid = groupids[pair[0]]
        assert groupid == groupids[pair[1]]
        if groupid not in collision_groups_pairs:
            collision_groups_pairs[groupid] = []
        collision_groups_pairs[groupid].append(pair)
    
    while len(collision_groups_pairs):
        for igrp, grp in list(collision_groups_pairs.items()):
            index = grp[rng.randint(len(grp))][rng.randint(2)]
            bit[index] = 0
            grp = [pair for pair in grp if index not in pair]
            if not grp:
                del collision_groups_pairs[igrp]
            else:
                collision_groups_pairs[igrp] = grp

    return bit


def find_collision_groups(size, pairs):
    # Group pairs into collision groups
    collision_groups = []
    groupids = -np.ones(size, dtype='i4')
    for pair in pairs:
        groupid1, groupid2 = groupids[pair[0]], groupids[pair[1]]
        if groupid1 < 0 and groupid2 < 0:
            groupids[pair[0]] = groupids[pair[1]] = len(collision_groups)
            collision_groups.append(set(pair))
        elif groupid1 >= 0 and groupid2 < 0:
            groupids[pair[1]] = groupid1
            collision_groups[groupid1].add(pair[1])
        elif groupid2 >= 0 and groupid1 < 0:
            groupids[pair[0]] = groupid2
            collision_groups[groupid2].add(pair[0])
        elif groupid1 != groupid2:
            for idx in collision_groups[groupid2]:
                groupids[idx] = groupid1
            collision_groups[groupid1] |= collision_groups[groupid2]
            collision_groups[groupid2] = set() # to keep groupids the same
    return groupids


def _run_mr_fiber_assignment(ireal, size=None, pairs=None, groupids=None, npasses=1, indices_multiple_pass=None, seed=42):
    rng = np.random.RandomState(seed=seed+ireal)
    bit = np.zeros(size, dtype='?')
    n_pairs = pairs.shape[0] + 1
    for ipass in range(npasses):
        if ipass == 0:
            indices = np.arange(size)
        else:
            indices = indices_multiple_pass
        collided = np.flatnonzero(~bit[indices])
        _pairs = pairs[np.all(np.isin(pairs, collided), axis=-1)]
        _n_pairs = _pairs.shape[0]
        assert _n_pairs < n_pairs
        n_pairs = _n_pairs
        bit[indices] = mr_fiber_collisions(rng, size, _pairs, groupids)[indices]
    return bit


def run_mr_fiber_assignment(size, pairs, npasses=1, fsky_multiple_pass=1., nrealizations=65, seed=42):
    # Run mr_fiber_collisions with multiple passes.
    # The first pass is run full-sky, the other passes are run on fsky_multiple_pass of the sky.
    indices_multiple_pass = np.arange(int(fsky_multiple_pass * size + 0.5))
    bits = []
    groupids = find_collision_groups(size, pairs)
    
    from multiprocessing import Pool
    from functools import partial
    with Pool(4) as pool:
        bits = pool.map(partial(_run_mr_fiber_assignment, size=size, pairs=pairs, groupids=groupids,
                                npasses=npasses, indices_multiple_pass=indices_multiple_pass, seed=seed), range(nrealizations))
    return bits

# getting fiber assignment runs
nrealizations = 65
cases = {}

cases['uncorrelated'] = {'bits': run_uncorrelated_fiber_assignment(parent_size, target_pairs, nrealizations=nrealizations, prob=0.5, seed=42)}
cases['onepass'] = {'bits': run_mr_fiber_assignment(parent_size, target_pairs, nrealizations=nrealizations, npasses=1, fsky_multiple_pass=1., seed=42)}
cases['twopass'] = {'bits': run_mr_fiber_assignment(parent_size, target_pairs, nrealizations=nrealizations, npasses=2, fsky_multiple_pass=1., seed=42)}
cases['twopass-fsky'] = {'bits': run_mr_fiber_assignment(parent_size, target_pairs, nrealizations=nrealizations, npasses=2, fsky_multiple_pass=0.5, seed=42)}

for name, case in cases.items():
    this = case['bits'][0]
    fiberassign_bits = utils.pack_bitarrays(*case['bits'][1:], dtype=np.uint64)
    # restricting to the objects that received a fiber in the actual pass
    positions = [p[this] for p in parent_positions]
    weights = [w[this] for w in parent_weights]
    bitwise_weights = [w[this] for w in fiberassign_bits]
    case.update(positions=positions, weights=weights, bitwise_weights=bitwise_weights)
# Routine to run correlation function estimations with different weighting schemes
def run_case(case, mode, edges, nthreads=4):
    positions = case['positions']
    weights = case['weights']
    bitwise_weights = case['bitwise_weights']
    # Parent sample
    result_parent = TwoPointCorrelationFunction(mode, edges, data_positions1=parent_positions, data_weights1=parent_weights,
                                                randoms_positions1=randoms_positions, randoms_weights1=randoms_weights,
                                                estimator='landyszalay', engine='corrfunc', position_type='xyz', compute_sepsavg=False, nthreads=nthreads)

    # No weights
    result_fibered = TwoPointCorrelationFunction(mode, edges, data_positions1=positions, data_weights1=weights,
                                                 randoms_positions1=randoms_positions, randoms_weights1=randoms_weights,
                                                 estimator='landyszalay', engine='corrfunc', position_type='xyz', compute_sepsavg=False, nthreads=nthreads)
    # With iip weights (nrealizations already includes current realization)
    wiip = nrealizations/(1 + utils.popcount(*bitwise_weights))
    result_iip = TwoPointCorrelationFunction(mode, edges, data_positions1=positions, data_weights1=wiip*weights[0],
                                             randoms_positions1=randoms_positions, randoms_weights1=randoms_weights,
                                             estimator='landyszalay', engine='corrfunc', position_type='xyz', compute_sepsavg=False, nthreads=nthreads)
    # With pip weights
    result_pip = TwoPointCorrelationFunction(mode, edges, data_positions1=positions, data_weights1=bitwise_weights + weights,
                                             randoms_positions1=randoms_positions, randoms_weights1=randoms_weights,
                                             estimator='landyszalay', engine='corrfunc', position_type='xyz', compute_sepsavg=False, nthreads=nthreads)
    # With pip + angular weights
    # We provide the parent sample as randoms, and switch to "weight" estimator
    # First D1_parentD2_parent/D1D2_PIP
    wangedges = np.logspace(-3, 0., 31)
    wangD1D2 = TwoPointCorrelationFunction('theta', wangedges, data_positions1=positions, data_weights1=bitwise_weights + weights,
                                           randoms_positions1=parent_positions, randoms_weights1=parent_weights,
                                           estimator='weight', engine='corrfunc', position_type='xyz', compute_sepsavg=False, nthreads=nthreads)

    # First D1_parentR2/D1R2_IIP
    # Note that input bitwise weights are automatically turned into IIP when other weight is float
    wangD1R2 = TwoPointCorrelationFunction('theta', wangedges, data_positions1=positions, data_weights1=bitwise_weights + weights,
                                           data_positions2=randoms_positions, data_weights2=randoms_weights,
                                           randoms_positions1=parent_positions, randoms_weights1=parent_weights,
                                           randoms_positions2=randoms_positions, randoms_weights2=randoms_weights,
                                           estimator='weight', engine='corrfunc', position_type='xyz', compute_sepsavg=False, nthreads=nthreads)
    # angular upweights D1D2_twopoint_weights and D1R2_twopoint_weights
    result_ang = TwoPointCorrelationFunction(mode, edges, data_positions1=positions, data_weights1=wiip*weights[0],
                                                randoms_positions1=randoms_positions, randoms_weights1=randoms_weights,
                                                D1D2_twopoint_weights=wangD1D2, D1R2_twopoint_weights=wangD1R2,
                                                engine='corrfunc', position_type='xyz', compute_sepsavg=False, nthreads=nthreads)
    
    # PIP weights + angular upweights D1D2_twopoint_weights and D1R2_twopoint_weights
    result_pipang = TwoPointCorrelationFunction(mode, edges, data_positions1=positions, data_weights1=bitwise_weights + weights,
                                                randoms_positions1=randoms_positions, randoms_weights1=randoms_weights,
                                                D1D2_twopoint_weights=wangD1D2, D1R2_twopoint_weights=wangD1R2,
                                                engine='corrfunc', position_type='xyz', compute_sepsavg=False, nthreads=nthreads)
    
    return result_parent, result_fibered, result_iip, result_pip, result_ang, result_pipang
# Look at the angular correlation function
mode = 'theta'
edges = np.logspace(-2, 0, 31)

fig, lax = plt.subplots(ncols=2, nrows=2, sharex=False, sharey=False, figsize=(20, 10))
fig.subplots_adjust(hspace=0.4, wspace=0.3)
lax = lax.flatten()

# Deactivate logging as it is a bit wordy
setup_logging('warning')
for ax, name in zip(lax, cases):
    
    result_parent, result_fibered, result_iip, result_pip, result_ang, result_pipang = run_case(cases[name], mode, edges)
    ax.set_title(name)
    ax.plot(result_parent.sep, result_parent.corr, label='parent')
    ax.plot(result_fibered.sep, result_fibered.corr, label='fibered')
    ax.plot(result_iip.sep, result_iip.corr, label='iip')
    ax.plot(result_pip.sep, result_pip.corr, label='pip')
    ax.plot(result_ang.sep, result_ang.corr, label='ang')
    ax.plot(result_pipang.sep, result_pipang.corr, label='pip + ang')
    ax.set_xscale('log')
    ax.set_xlabel(r'$\theta$ [deg]')
    ax.set_ylabel(r'$w(\theta)$')
    ax.grid(True)
    ax.legend()

plt.show()<jupyter_output><empty_output><jupyter_text>The different curves- The "parent" curve corresponds to the autocorrelation of the initial target or "unfibered" sample. This curve is zero as input positions are uniformly random.- The "fibered" curve corresponds to the autocorrelation of the "fibered" sample, i.e. the sample of objects that received a fiber. This curve drops below 0.08 deg due to introduced fiber collisions.- The "iip" curve corresponds to the autocorrelation of the fibered sample with individual inverse weights defined by nrealizations/(1 + popcount(bitwise_weight)), which ignores correlation in selection probabilites of close objects ("popcount" counts the number of 1 bits, i.e. returns the number of fiber assignment realizations where this object was selected).- The "pip" curve corresponds to the autocorrelation of the fibered sample with pairwise inverse probability pairs defined for a *pair* of objects by nrealisations/(1 + popcount(bitwise_weight_i & bitwise_weight_j)) (here "popcount" counts the number of fiber assignment realizations where this *pair* was selected).- The "pip + ang" curve is similar as "pip", but with additional angular weights for D1D2 and D1R2 pairs, defined as: wangD1D2 = D1_parentD2_parent/D1D2_PIP and wangD1R2 = D1_parentR2/D1R2_IIP. The weights correct for zero-probability close pairs, assuming that non-zero probability pairs are representative of the zero-probability ones. The different plots (cases)- In the "uncorrelated" case, selection probabilities of close objects are uncorrelated. Hence the correlation function is fully corrected by IIP weights. PIP weights are statistically equivalent to IIP weights in this case. Angular weights are useless since there are no zero-probability pairs.- In the "onepass" case, selection probabilities of close objects are correlated. Hence PIP weights are required. Yet we see no weighting scheme can recover an unbiased estimate of the correlation function. This is because all pairs are zero-probability: what we want to avoid in a real survey.- In the "twopass" case, fiber assignment is run twice. Hence there are no zero-probability pairs. We see that PIP weights are required (IIP weights are not sufficient), because selection probabilities of close objects are correlated. Angular upweights bring no improvement, since all pairs can be selected with non zero-probability.- In the "twopass-fsky" case, fiber assignment is twice only one 50% of the footprint. There are zero-probability pairs, hence the angular upweights are required.<jupyter_code># Also note that normalization factors for PIP and IIP are different, due to finite number of fiber assignment realizations:
norm_iip = result_iip.D1D2.wnorm
norm_pip = result_pip.D1D2.wnorm
print('PIP v.s. IIP normalization is {} v.s. {} (ratio = {:.8f})'.format(norm_pip, norm_iip, norm_pip/norm_iip))
# Look at the "monopole" correlation function

mode = 's'
edges = np.logspace(-0.5, 1., 31)

fig, lax = plt.subplots(ncols=2, nrows=2, sharex=False, sharey=False, figsize=(20, 10))
fig.subplots_adjust(hspace=0.4, wspace=0.3)
lax = lax.flatten()

for ax, name in zip(lax, cases):
    
    result_parent, result_fibered, result_iip, result_pip, result_ang, result_pipang = run_case(cases[name], mode, edges)
    ax.set_title(name)
    ax.plot(result_parent.sep, result_parent.corr, label='parent')
    ax.plot(result_fibered.sep, result_fibered.corr, label='fibered')
    ax.plot(result_iip.sep, result_iip.corr, label='iip')
    ax.plot(result_pip.sep, result_pip.corr, label='pip')
    ax.plot(result_ang.sep, result_ang.corr, label='ang')
    ax.plot(result_pipang.sep, result_pipang.corr, label='pip + ang')
    ax.set_xscale('log')
    ax.set_xlabel(r'$s$ [Mpc/h]')
    ax.set_ylabel(r'$\xi(s)$')
    ax.grid(True)
    ax.legend(loc=1)

plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09047564338235294, 'fraction_numerical': 0.01711856617647059, 'mean_word_length': 3.7025931928687195, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 15, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 40, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '385452', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 5405, 'n_tokens_neox': 5129, 'n_words': 1441, 'path': 'nb/pip_examples.ipynb', 'repo_name': 'cosmodesi/pycorr', 'size': '17408'}