mirror of
https://github.com/Findus23/halo_comparison.git
synced 2024-09-19 16:03:50 +02:00
91 lines
2.9 KiB
Python
91 lines
2.9 KiB
Python
from pathlib import Path
|
|
|
|
import h5py
|
|
import pandas as pd
|
|
from pandas import DataFrame
|
|
import matplotlib.pyplot as plt
|
|
from paths import base_dir
|
|
|
|
from readfiles import read_file
|
|
|
|
|
|
def all_children(df, id):
|
|
subhalos: pd.DataFrame = df.loc[df.parent_halo_id == id]
|
|
|
|
if len(subhalos) == 0:
|
|
yield id
|
|
|
|
for sh, value in subhalos.iterrows():
|
|
yield from all_children(df, sh)
|
|
|
|
|
|
def subhalo_particles_recursively(halo: int, waveform: str, Nres: int):
|
|
directory = base_dir / f"{waveform}_{Nres}_100"
|
|
|
|
group_catalog = h5py.File(directory / "vroutput.catalog_groups")
|
|
df = pd.DataFrame(
|
|
{
|
|
"group_size": group_catalog["Group_Size"],
|
|
"offset": group_catalog["Offset"],
|
|
"offset_unbound": group_catalog["Offset_unbound"],
|
|
"parent_halo_id": group_catalog["Parent_halo_ID"],
|
|
}
|
|
)
|
|
|
|
particle_catalog = h5py.File(directory / "vroutput.catalog_particles")
|
|
particle_ids = particle_catalog["Particle_IDs"][:]
|
|
|
|
particle_catalog_unbound = h5py.File(
|
|
directory / "vroutput.catalog_particles.unbound"
|
|
)
|
|
particle_ids_unbound = particle_catalog_unbound["Particle_IDs"][:]
|
|
|
|
all_halo_ids = [halo] + list(all_children(df, halo))
|
|
|
|
all_halo_offsets = [[df["offset"][j], df["offset"][j + 1]] for j in all_halo_ids]
|
|
all_halo_offsets_unbound = [
|
|
[df["offset_unbound"][j], df["offset_unbound"][j + 1]] for j in all_halo_ids
|
|
]
|
|
|
|
all_halo_particles = [
|
|
set(particle_ids[all_halo_offsets[j][0] : all_halo_offsets[j][1]])
|
|
for j in range(len(all_halo_offsets))
|
|
]
|
|
all_halo_particles_unbound = [
|
|
set(
|
|
particle_ids_unbound[
|
|
all_halo_offsets_unbound[j][0] : all_halo_offsets_unbound[j][1]
|
|
]
|
|
)
|
|
for j in range(len(all_halo_offsets_unbound))
|
|
]
|
|
|
|
return all_halo_ids, all_halo_particles, all_halo_particles_unbound
|
|
|
|
|
|
# #This could be used to make a 2-D plot, but that doesn't teach us much about the halos:
|
|
# particles, _ = read_file(directory)
|
|
# print(particles)
|
|
# print(particle_ids_in_halo)
|
|
# bound_particles = particles.loc[list(particle_ids_in_halo)]
|
|
# unbound_particles = particles.loc[list(particle_ids_unbound_in_halo)]
|
|
|
|
# plt.scatter(bound_particles['X'], bound_particles['Z'], label='bound', s=1)
|
|
# plt.scatter(unbound_particles['X'], unbound_particles['Z'], label='unbound', s=1)
|
|
# plt.legend()
|
|
# plt.show()
|
|
|
|
|
|
# #This was an initial test to see if reading in works as expected
|
|
# lower_offset = df['offset'][test_group]
|
|
# higher_offset = df['offset'][test_group + 1]
|
|
# particle_ids_in_halo = set(particle_ids[lower_offset:higher_offset])
|
|
|
|
# lower_offset_unbound = df['offset_unbound'][test_group]
|
|
# higher_offset_unbound = df['offset_unbound'][test_group + 1]
|
|
# particle_ids_unbound_in_halo = set(particle_ids_unbound[lower_offset_unbound:higher_offset_unbound])
|
|
|
|
# size = df['group_size'][test_group]
|
|
|
|
# assert size == len(particle_ids_in_halo) + len(particle_ids_unbound_in_halo)
|
|
|