[FEAT] Allow building octree with particles

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

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"

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,10 @@ 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))
+    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]])
 
@@ -58,3 +62,58 @@ def test_octree():
     rho1.sort()
     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):
+    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

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(