<jupyter_start><jupyter_text>I want to add the local density to the MDPL2 catalog I have. However I no longer have the full particle catalog and don't feel like spending the time to recover it, when the cached particles I have should be sufficient. So I will attempt to do it with the cached particles I have.<jupyter_code>import numpy as np
from pearce.mocks.kittens import MDPL2, TestBox
cat = TestBox(0,0)#MDPL2()
cat.load(1.0, particles =True)
cat._dow
#halo_table = cat.halocat.halo_table
def add_local_density(cat, radius=[10]):  # [1,5,10]

    from fast3tree import fast3tree

    if type(radius) == float:
        radius = np.array([radius])
    elif type(radius) == list:
        radius = np.array(radius)
    densities = np.zeros((cat.halocat.halo_table['halo_x'].shape[0], radius.shape[0]))

    mean_particle_density = cat._downsample_factor * (cat.npart / cat.Lbox) ** 3
    all_particles = np.c_[cat.halocat.ptcl_table[c] for c in ['x', 'y','z']
    
    with fast3tree(all_particles) as tree:
        for r_idx, r in enumerate(radius):
            print  'Calculating Densities for radius %d' % r
            # densities[:, r_idx] = densities[:, r_idx]/ (downsample_factor * 4 * np.pi / 3 * r ** 3)
            for idx, halo_pos in enumerate(
                    izip(reader.halo_table['halo_x'], reader.halo_table['halo_y'], reader.halo_table['halo_z'])):
                # print idx
                particle_idxs = tree.query_radius(halo_pos, r, periodic=True)
                densities[idx, r_idx] += len(particle_idxs)
                
                volume = (4 * np.pi / 3 * r ** 3)
                cat.halocat.halo_table['halo_local_density_%d' % (int(r))] = densities[:, r_idx] / (
                            volume * mean_particle_density)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11859688195991092, 'fraction_numerical': 0.013919821826280624, 'mean_word_length': 2.8152866242038215, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, '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': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '344554', 'license': 'MIT', 'n_tokens_mistral': 620, 'n_tokens_neox': 592, 'n_words': 158, 'path': 'notebooks/Test MDPL2 Local Density.ipynb', 'repo_name': 'mclaughlin6464/pearce', 'size': '1796'}