Skip to content

Commit

Permalink
Test halo accessor
Browse files Browse the repository at this point in the history
  • Loading branch information
cphyc committed Sep 6, 2023
1 parent 2c8d1d5 commit 9384474
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 2 deletions.
4 changes: 2 additions & 2 deletions yt/frontends/rockstar/data_structures.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import glob
import os
from functools import cached_property
from typing import Optional
from typing import List, Optional

import numpy as np

Expand All @@ -21,7 +21,7 @@ class RockstarBinaryFile(HaloCatalogFile):
_position_offset: int
_member_offset: int
_Npart: np.array
_ids_halos: list[int]
_ids_halos: List[int]
_file_size: int

def __init__(self, ds, io, filename, file_id, range):
Expand Down
20 changes: 20 additions & 0 deletions yt/frontends/rockstar/tests/test_outputs.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,3 +38,23 @@ def test_particle_selection():
ds = data_dir_load(r1)
psc = ParticleSelectionComparison(ds)
psc.run_defaults()


@requires_file(r1)
def test_halo_loading():
ds = data_dir_load(r1)

for halo_id, Npart in zip(
ds.r["halos", "particle_identifier"],
ds.r["halos", "num_p"],
):
halo = ds.halo("halos", halo_id)
assert halo is not None

# Try accessing properties
halo.position
halo.velocity
halo.mass

# Make sure we can access the member particles
assert len(halo.member_ids) == Npart

0 comments on commit 9384474

Please sign in to comment.