From f097c9757d9744460fb437f0e356ada4bc0c240f Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Tue, 21 Apr 2026 16:49:29 +0200 Subject: [PATCH 1/4] feat: Allow building octree with particles --- yt/frontends/stream/io.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/yt/frontends/stream/io.py b/yt/frontends/stream/io.py index 75611af018d..0994fb294bc 100644 --- a/yt/frontends/stream/io.py +++ b/yt/frontends/stream/io.py @@ -264,6 +264,29 @@ def _read_fluid_selection(self, chunks, selector, fields, size): subset.fill(field_vals, rv, selector, ind) return rv + def _read_particle_fields(self, chunks, ptf, selector): + chunks = list(chunks) + for chunk in chunks: + for subset in chunk.objs: + sf = self.fields[subset.domain_id - subset._domain_offset] + for ptype, field_list in sorted(ptf.items()): + if (ptype, "particle_position") in sf: + x, y, z = sf[ptype, "particle_position"].T + else: + x, y, z = ( + sf[ptype, f"particle_position_{ax}"] + for ax in self.ds.coordinates.axis_order + ) + + mask = selector.select_points(x, y, z, 0.0) + + if mask is None: + continue + + for field in field_list: + data = np.asarray(sf[ptype, field]) + yield (ptype, field), data[mask] + class IOHandlerStreamUnstructured(BaseIOHandler): _dataset_type = "stream_unstructured" From 808450393f024d2d2e4dccec9bdafeeee74cab5c Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Tue, 21 Apr 2026 17:36:16 +0200 Subject: [PATCH 2/4] test: test that particles can be accessed --- nose_ignores | 1 + .../stream/tests/test_stream_octree.py | 50 ++++++++++++++++++- 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/nose_ignores b/nose_ignores index 6db1089c71f..5c531bcae5f 100644 --- a/nose_ignores +++ b/nose_ignores @@ -36,6 +36,7 @@ --ignore-file=test_set_zlim\.py --ignore-file=test_stream_particles\.py --ignore-file=test_stream_stretched\.py +--ignore-file=test_stream_octree\.py --ignore-file=test_version\.py --ignore-file=test_gadget_pytest\.py --ignore-file=test_vr_orientation\.py diff --git a/yt/frontends/stream/tests/test_stream_octree.py b/yt/frontends/stream/tests/test_stream_octree.py index 97a0b069eb5..bffc3bf9f62 100644 --- a/yt/frontends/stream/tests/test_stream_octree.py +++ b/yt/frontends/stream/tests/test_stream_octree.py @@ -1,4 +1,5 @@ import numpy as np +import pytest from numpy.testing import assert_equal import yt @@ -37,7 +38,8 @@ def test_octree(): octree_mask = np.array(OCT_MASK_LIST, dtype=np.uint8) quantities = {} - quantities["gas", "density"] = np.random.random((22, 1)) + rng = np.random.default_rng(12345) + quantities["gas", "density"] = rng.random((22, 1)) bbox = np.array([[-10.0, 10.0], [-10.0, 10.0], [-10.0, 10.0]]) @@ -58,3 +60,49 @@ def test_octree(): rho1.sort() rho2.sort() assert_equal(rho1, rho2) + + +@pytest.mark.parametrize("build_pos", ["component_by_component", "combined"]) +def test_octree_particles(build_pos): + octree_mask = np.array(OCT_MASK_LIST, dtype=np.uint8) + + L = [-10, -10, -10] + R = [10, 10, 10] + + quantities = {} + rng = np.random.default_rng(12345) + + pos = rng.uniform(-10, 10, size=(100, 3)) + quantities["gas", "density"] = rng.random((22, 1)) + if build_pos == "component_by_component": + quantities["io", "particle_position_x"] = pos[:, 0] + quantities["io", "particle_position_y"] = pos[:, 1] + quantities["io", "particle_position_z"] = pos[:, 2] + else: + quantities["io", "particle_position"] = pos + quantities["io", "particle_mass"] = (np.linspace(0, 1, 100), "Msun") + + bbox = np.array([L, R]).T + + ds = yt.load_octree( + octree_mask=octree_mask, + data=quantities, + bbox=bbox, + num_zones=1, + partial_coverage=0, + ) + + # Ensure "io" is registered as a particle type + assert "io" in ds.particle_types + assert "io" in ds.particle_types_raw + + # Ensure we can access the particle positions + _pos = ds.r[:]["particle_position"] + _pos_x = ds.r[:]["particle_position_x"] + _pos_y = ds.r[:]["particle_position_y"] + _pos_z = ds.r[:]["particle_position_z"] + + # Make sure other fields are defined + mass = ds.r[:]["particle_mass"].to("Msun") + assert mass.min() == 0 + assert mass.max() == 1 From c9c4dd08f052e3072d530ad8d58037cbe39e4388 Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Tue, 21 Apr 2026 17:37:38 +0200 Subject: [PATCH 3/4] test: callback should keep units --- yt/frontends/stream/tests/test_stream_octree.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/yt/frontends/stream/tests/test_stream_octree.py b/yt/frontends/stream/tests/test_stream_octree.py index bffc3bf9f62..a7e70c0b037 100644 --- a/yt/frontends/stream/tests/test_stream_octree.py +++ b/yt/frontends/stream/tests/test_stream_octree.py @@ -40,6 +40,8 @@ def test_octree(): quantities = {} rng = np.random.default_rng(12345) quantities["gas", "density"] = rng.random((22, 1)) + quantities["gas", "dinos"] = (np.linspace(0, 1, 22).reshape(22, 1), "Msun") + quantities["gas", "turtle"] = (lambda: np.linspace(0, 1, 22).reshape(22, 1), "Msun") bbox = np.array([[-10.0, 10.0], [-10.0, 10.0], [-10.0, 10.0]]) @@ -61,6 +63,15 @@ def test_octree(): rho2.sort() assert_equal(rho1, rho2) + # Make sure units aren't stripped + assert str(ds.r[:]["dinos"].units) == "Msun" + assert ds.r[:]["dinos"].min().to("Msun") == 0 + assert ds.r[:]["dinos"].max().to("Msun") == 1 + + assert str(ds.r[:]["turtle"].units) == "Msun" + assert ds.r[:]["turtle"].min().to("Msun") == 0 + assert ds.r[:]["turtle"].max().to("Msun") == 1 + @pytest.mark.parametrize("build_pos", ["component_by_component", "combined"]) def test_octree_particles(build_pos): From d6c370ea86ea61e04f38a8d2e5d551f99d96ad06 Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Tue, 21 Apr 2026 17:54:25 +0200 Subject: [PATCH 4/4] doc: add narrative docstring --- yt/loaders.py | 44 ++++++++++++++++++++++++++++++++++++++------ 1 file changed, 38 insertions(+), 6 deletions(-) diff --git a/yt/loaders.py b/yt/loaders.py index 27423a5c7d9..f5ae42e55d1 100644 --- a/yt/loaders.py +++ b/yt/loaders.py @@ -1078,8 +1078,7 @@ def load_octree( the root oct is not refined, there will be only one entry for the root, so the size of the mask will be (n_octs - 1)*8 + 1. data : dict - A dictionary of 1D arrays. Note that these must of the size of the - number of "False" values in the ``octree_mask``. + A dictionary of 1D arrays. See note below for the format. bbox : array_like (xdim:zdim, LE:RE), optional Size of computational domain in units of length sim_time : float, optional @@ -1115,17 +1114,50 @@ def load_octree( Optional string used to assign a name to the dataset. Stream datasets will use this value in place of a filename (in image prefixing, etc.) + Note + ---- + Data is a dictionary with keys (type, field_name). ``type`` is typically + ``gas`` for gas fields and ``io`` for particle fields, but these are not + strictly enforced. + + - Field values are 2D arrays of shape (Ncell, 1), with Ncell the number of + ``False`` values in ``octree_mask``. + - Particle values are 1D arrays of shape (Npart,), with Npart the number of + particles + + Each entry can be: + 1. A raw numpy array. If it is a known field (e.g., density), it is assumed + to be in code units. If it is not a known field, it is assumed to be + dimensionless. + 2. A tuple (numpy array, unit string). As above, but the type is now + explicit. + 3. Instead of a numpy array, a parameter-less function that returns a + numpy array can be provided. This is useful for cases where the data + is generated on-the-fly, or to implement a basic form of lazy loading. + **Note that this is only supported for fields, not particles.** + Example ------- >>> import numpy as np - >>> oct_mask = np.zeros(33) # 5 refined values gives 7 * 4 + 5 octs to mask + >>> oct_mask = np.zeros(33) # 5 refined values gives 7 * 4 + 5 octs to mask ... oct_mask[[0, 5, 7, 16]] = 8 >>> octree_mask = np.array(oct_mask, dtype=np.uint8) >>> quantities = {} - >>> quantities["gas", "density"] = np.random.random((29, 1)) # num of false's - >>> # Quantities can also contain parameter-less callbacks - >>> quantities["gas", "temperature"] = lambda: np.random.random((29, 1)) * 1e4 + ... # Note: there are 29 leaf cells + ... quantities["gas", "density"] = np.random.random((29, 1)) + ... quantities["gas", "dinosaurs"] = np.random.random((29, 1)) + ... quantities["gas", "turtle_number_density"] = (np.random.random((29, 1)), "1/cm**3") + ... quantities["gas", "costly_field"] = lambda: np.random.random((29, 1)) + ... quantities["gas", "costly_field_with_units"] = (lambda: np.random.random((29, 1)), "K") + >>> # Load in 123 particles + ... quantities["io", "particle_position_x"] = np.random.random(123) + ... quantities["io", "particle_position_y"] = np.random.random(123) + ... quantities["io", "particle_position_z"] = np.random.random(123) + ... # This also works + ... quantities["io", "particle_position"] = np.random.random((123, 3)) + ... quantities["io", "particle_mass"] = (np.random.random(123), "g") + >>> bbox = np.array([[-10.0, 10.0], [-10.0, 10.0], [-10.0, 10.0]]) >>> ds = load_octree(