From 9384474c729f5a7bb3c80c9f20414d1aefac1968 Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Tue, 5 Sep 2023 16:39:43 +0100 Subject: [PATCH] Test halo accessor --- yt/frontends/rockstar/data_structures.py | 4 ++-- yt/frontends/rockstar/tests/test_outputs.py | 20 ++++++++++++++++++++ 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/yt/frontends/rockstar/data_structures.py b/yt/frontends/rockstar/data_structures.py index 5eecd24630..8baa377694 100644 --- a/yt/frontends/rockstar/data_structures.py +++ b/yt/frontends/rockstar/data_structures.py @@ -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 @@ -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): diff --git a/yt/frontends/rockstar/tests/test_outputs.py b/yt/frontends/rockstar/tests/test_outputs.py index 920881ec1c..617357753b 100644 --- a/yt/frontends/rockstar/tests/test_outputs.py +++ b/yt/frontends/rockstar/tests/test_outputs.py @@ -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