Fix unbound particle masking in haloquantity lookup and support 2D halo quantities (#233)

The unbound particle check in get_haloquantity_daskwrap compared against -1,
but get_hidx marks unbound particles with index_unbound (int64 max by default).
This caused unbound particles to be incorrectly looked up. Also adds support
for 2D halo quantities (e.g. positions/velocities) via drop_axis/new_axis.

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: cbyrohl

src/scida/customs/arepo/dataset.py

@@ -1109,10 +1109,12 @@ def get_hidx_daskwrap(gidx, halocelloffsets, index_unbound=None):
     )
 
 
-def get_haloquantity_daskwrap(gidx, halocelloffsets, valarr):
-    hidx = get_hidx_daskwrap(gidx, halocelloffsets)
+def get_haloquantity_daskwrap(
+    gidx, halocelloffsets, valarr, index_unbound=9223372036854775807
+):
+    hidx = get_hidx_daskwrap(gidx, halocelloffsets, index_unbound=index_unbound)
     dmax = np.iinfo(hidx.dtype).max
-    mask = ~((hidx == -1) | (hidx == dmax))
+    mask = ~((hidx == index_unbound) | (hidx == dmax))
     result = -1.0 * np.ones(hidx.shape, dtype=valarr.dtype)
     result[mask] = valarr[hidx[mask]]
     return result
@@ -1134,13 +1136,18 @@ def compute_haloquantity(gidx, halocelloffsets, hvals, *args):
     units = None
     if hasattr(hvals, "units"):
         units = hvals.units
+        hvals = hvals.magnitude
+    kw = dict()
+    if len(hvals.shape) == 2:
+        kw = dict(drop_axis=0, new_axis=0)
     res = map_blocks(
         get_haloquantity_daskwrap,
         gidx,
         halocelloffsets,
         hvals,
-        meta=np.array((), dtype=hvals.dtype),
+        dtype=hvals.dtype,
         output_units=units,
+        **kw,
     )
     return res