Merge remote-tracking branch 'origin/main' into openpmd_update

doc/source/examining/loading_data.rst

@@ -665,8 +665,8 @@ volume rendering, and many of the analysis modules will not work.
 
 .. _loading-pluto-data:
 
-Pluto Data
-----------
+Pluto Data (AMR)
+----------------
 
 Support for Pluto AMR data is provided through the Chombo frontend, which
 is currently maintained by Andrew Myers. Pluto output files that don't use
@@ -691,6 +691,18 @@ To load it, you can navigate into that directory and do:
 The ``pluto.ini`` file must also be present alongside the HDF5 file.
 By default, all of the Pluto fields will be in code units.
 
+
+.. _loading-idefix-data:
+
+Idefix, Pluto VTK and Pluto XDMF Data
+-------------------------------------
+
+Support for Idefix ``.dmp``, ``.vtk`` data is provided through the ``yt_idefix``
+extension.
+It also supports monogrid ``.vtk`` and ``.h5`` data from Pluto.
+See `the PyPI page <https://pypi.org/project/yt-idefix/>`_ for details.
+
+
 .. _loading-enzo-data:
 
 Enzo Data

pyproject.toml

@@ -113,6 +113,7 @@ gdf = ["yt[HDF5]"]
 gizmo = ["yt[HDF5]"]
 halo-catalog = ["yt[HDF5]"]
 http-stream = ["requests>=2.20.0"]
+idefix = ["yt_idefix[HDF5]>=2.3.0"] # externally packaged frontend
 moab = ["yt[HDF5]"]
 nc4-cm1 = ["yt[netCDF4]"]
 open-pmd = ["yt[openpmd_api]"]
@@ -169,6 +170,7 @@ full = [
     "yt[gizmo]",
     "yt[halo_catalog]",
     "yt[http_stream]",
+    "yt[idefix]",
     "yt[moab]",
     "yt[nc4_cm1]",
     "yt[open_pmd]",

yt/fields/particle_fields.py

@@ -302,12 +302,7 @@ def particle_vectors(field, data):
 
 
 def get_angular_momentum_components(ptype, data, spos, svel):
-    if data.has_field_parameter("normal"):
-        normal = data.get_field_parameter("normal")
-    else:
-        normal = data.ds.arr(
-            [0.0, 0.0, 1.0], "code_length"
-        )  # default to simulation axis
+    normal = data.ds.arr([0.0, 0.0, 1.0], "code_length")  # default to simulation axis
     pos = data.ds.arr([data[ptype, spos % ax] for ax in "xyz"]).T
     vel = data.ds.arr([data[ptype, f"relative_{svel % ax}"] for ax in "xyz"]).T
     return pos, vel, normal

yt/frontends/ramses/fields.py

@@ -137,7 +137,7 @@ class RAMSESFieldInfo(FieldInfoContainer):
         ("particle_angular_momentum_y", (ang_mom_units, [], None)),
         ("particle_angular_momentum_z", (ang_mom_units, [], None)),
         ("particle_formation_time", ("code_time", [], None)),
-        ("particle_accretion_Rate", ("code_mass/code_time", [], None)),
+        ("particle_accretion_rate", ("code_mass/code_time", [], None)),
         ("particle_delta_mass", ("code_mass", [], None)),
         ("particle_rho_gas", (rho_units, [], None)),
         ("particle_cs**2", (vel_units, [], None)),

yt/frontends/ramses/hilbert.py

@@ -7,6 +7,7 @@
     bbox_intersects,
     fully_contains,
 )
+from yt.utilities.lib.geometry_utils import get_hilbert_indices
 
 # State diagram to compute the hilbert curve
 _STATE_DIAGRAM = np.array(
@@ -48,63 +49,22 @@
 
 
 def hilbert3d(
-    X: "np.ndarray[Any, np.dtype[np.float64]]", bit_length: int
+    ijk: "np.ndarray[Any, np.dtype[np.int64]]", bit_length: int
 ) -> "np.ndarray[Any, np.dtype[np.float64]]":
     """Compute the order using Hilbert indexing.
 
     Arguments
     ---------
-    X : (N, ndim) float array
+    ijk : (N, ndim) integer array
       The positions
     bit_length : integer
       The bit_length for the indexing.
     """
-    X = np.atleast_2d(X)
-
-    x_bit_mask, y_bit_mask, z_bit_mask = (
-        np.zeros(bit_length, dtype=bool) for _ in range(3)
-    )
-    i_bit_mask = np.zeros(3 * bit_length, dtype=bool)
-
-    npoint = X.shape[0]
-    order = np.zeros(npoint)
-
-    # Convert positions to binary
-    for ip in range(npoint):
-        for i in range(bit_length):
-            mask = 0b01 << i
-            x_bit_mask[i] = X[ip, 0] & mask
-            y_bit_mask[i] = X[ip, 1] & mask
-            z_bit_mask[i] = X[ip, 2] & mask
-
-        for i in range(bit_length):
-            # Interleave bits
-            i_bit_mask[3 * i + 2] = x_bit_mask[i]
-            i_bit_mask[3 * i + 1] = y_bit_mask[i]
-            i_bit_mask[3 * i] = z_bit_mask[i]
-
-        # Build Hilbert ordering using state diagram
-        cstate = 0
-        for i in range(bit_length - 1, -1, -1):
-            sdigit = (
-                4 * i_bit_mask[3 * i + 2]
-                + 2 * i_bit_mask[3 * i + 1]
-                + 1 * i_bit_mask[3 * i]
-            )
-            nstate = _STATE_DIAGRAM[sdigit, 0, cstate]
-            hdigit = _STATE_DIAGRAM[sdigit, 1, cstate]
-
-            i_bit_mask[3 * i + 2] = hdigit & 0b100
-            i_bit_mask[3 * i + 1] = hdigit & 0b010
-            i_bit_mask[3 * i] = hdigit & 0b001
-
-            cstate = nstate
-
-        # Compute ordering
-        for i in range(3 * bit_length):
-            order[ip] = order[ip] + i_bit_mask[i] * 2**i
-
-    return order
+    ijk = np.atleast_2d(ijk)
+    # A note here: there is a freedom in the way hilbert indices are
+    # being computed (should it be xyz or yzx or zxy etc.)
+    # and the yt convention is not the same as the RAMSES one.
+    return get_hilbert_indices(bit_length, ijk[:, [1, 2, 0]].astype(np.int64))
 
 
 def get_intersecting_cpus(
@@ -114,6 +74,7 @@ def get_intersecting_cpus(
     dx: float = 1.0,
     dx_cond: Optional[float] = None,
     factor: float = 4.0,
+    bound_keys: Optional["np.ndarray[Any, np.dtype[np.float64]]"] = None,
 ) -> set[int]:
     """
     Find the subset of CPUs that intersect the bbox in a recursive fashion.
@@ -123,15 +84,20 @@ def get_intersecting_cpus(
     if dx_cond is None:
         bbox = region.get_bbox()
         dx_cond = float((bbox[1] - bbox[0]).min().to("code_length"))
+    if bound_keys is None:
+        ncpu = ds.parameters["ncpu"]
+        bound_keys = np.empty(ncpu + 1, dtype="float64")
+        bound_keys[:ncpu] = [ds.hilbert_indices[icpu + 1][0] for icpu in range(ncpu)]
+        bound_keys[ncpu] = ds.hilbert_indices[ncpu][1]
 
     # If the current dx is smaller than the smallest size of the bbox
     if dx < dx_cond / factor:
         # Finish recursion
-        return get_cpu_list_cuboid(ds, np.asarray([LE, LE + dx]))
+        return get_cpu_list_cuboid(ds, np.asarray([LE, LE + dx]), bound_keys)
 
     # If the current cell is fully within the selected region, stop recursion
     if fully_contains(region.selector, LE, dx):
-        return get_cpu_list_cuboid(ds, np.asarray([LE, LE + dx]))
+        return get_cpu_list_cuboid(ds, np.asarray([LE, LE + dx]), bound_keys)
 
     dx /= 2
 
@@ -144,12 +110,18 @@ def get_intersecting_cpus(
 
                 if bbox_intersects(region.selector, LE_new, dx):
                     ret.update(
-                        get_intersecting_cpus(ds, region, LE_new, dx, dx_cond, factor)
+                        get_intersecting_cpus(
+                            ds, region, LE_new, dx, dx_cond, factor, bound_keys
+                        )
                     )
     return ret
 
 
-def get_cpu_list_cuboid(ds, X: "np.ndarray[Any, np.dtype[np.float64]]") -> set[int]:
+def get_cpu_list_cuboid(
+    ds,
+    X: "np.ndarray[Any, np.dtype[np.float64]]",
+    bound_keys: "np.ndarray[Any, np.dtype[np.float64]]",
+) -> set[int]:
     """
     Return the list of the CPU intersecting with the cuboid containing the positions.
     Note that it will be 0-indexed.
@@ -166,7 +138,6 @@ def get_cpu_list_cuboid(ds, X: "np.ndarray[Any, np.dtype[np.float64]]") -> set[i
         raise NotImplementedError("This function is only implemented in 3D.")
 
     levelmax = ds.parameters["levelmax"]
-    ncpu = ds.parameters["ncpu"]
     ndim = ds.parameters["ndim"]
 
     xmin, ymin, zmin = X.min(axis=0)
@@ -193,7 +164,7 @@ def get_cpu_list_cuboid(ds, X: "np.ndarray[Any, np.dtype[np.float64]]") -> set[i
     if bit_length > 0:
         ndom = 8
 
-    ijkdom = idom, jdom, kdom = np.empty((3, 8), dtype="int64")
+    ijkdom = idom, jdom, kdom = np.empty((3, 8), dtype=np.int64)
 
     idom[0], idom[1] = imin, imax
     idom[2], idom[3] = imin, imax
@@ -221,12 +192,8 @@ def get_cpu_list_cuboid(ds, X: "np.ndarray[Any, np.dtype[np.float64]]") -> set[i
         bounding_min[i] = omin * dkey
         bounding_max[i] = (omin + 1) * dkey
 
-    bound_key = np.empty(ncpu + 1, dtype="float64")
-    for icpu in range(1, ncpu + 1):
-        bound_key[icpu - 1], bound_key[icpu] = ds.hilbert_indices[icpu]
-
-    cpu_min = np.searchsorted(bound_key, bounding_min, side="right") - 1
-    cpu_max = np.searchsorted(bound_key, bounding_max, side="right")
+    cpu_min = np.searchsorted(bound_keys, bounding_min, side="right") - 1
+    cpu_max = np.searchsorted(bound_keys, bounding_max, side="right")
 
     cpu_read: set[int] = set()
 

yt/frontends/ramses/io.py

@@ -388,7 +388,7 @@ def _read_part_csv_file_descriptor(fname: Union[str, "os.PathLike[str]"]):
         "ly": "particle_angular_momentum_y",
         "lz": "particle_angular_momentum_z",
         "tform": "particle_formation_time",
-        "acc_Rate": "particle_accretion_Rate",
+        "acc_rate": "particle_accretion_rate",
         "del_mass": "particle_delta_mass",
         "rho_gas": "particle_rho_gas",
         "cs**2": "particle_sound_speed",

yt/frontends/ramses/tests/test_hilbert.py

@@ -42,8 +42,14 @@ def test_get_cpu_list():
     )
     outputs = ([0, 15], [0, 15], [0, 1, 15], [0, 13, 14, 15], [0])
 
+    ncpu = ds.parameters["ncpu"]
+    bound_keys = np.array(
+        [ds.hilbert_indices[icpu][0] for icpu in range(1, ncpu + 1)]
+        + [ds.hilbert_indices[ds.parameters["ncpu"]][1]],
+        dtype="float64",
+    )
     for i, o in zip(inputs, outputs):
         bbox = i
-        ls = list(get_cpu_list_cuboid(ds, bbox))
+        ls = list(get_cpu_list_cuboid(ds, bbox, bound_keys=bound_keys))
         assert len(ls) > 0
         assert all(np.array(o) == np.array(ls))

yt/frontends/tipsy/io.py

@@ -178,14 +178,11 @@ def _read_particle_data_file(self, data_file, ptf, selector=None):
                         )
                     )
                 else:
-                    par = self.ds.parameters
-                    nlines = 1 + par["nsph"] + par["ndark"] + par["nstar"]
                     aux_fh[afield].seek(0)
                     sh = aux_fields_offsets[afield][ptype]
-                    sf = nlines - count - sh
                     if tp[ptype] > 0:
                         aux = np.genfromtxt(
-                            aux_fh[afield], skip_header=sh, skip_footer=sf
+                            aux_fh[afield], skip_header=sh, max_rows=count
                         )
                         if aux.ndim < 1:
                             aux = np.array([aux])

yt/utilities/lib/bitarray.pxd

@@ -14,7 +14,21 @@ cimport numpy as np
 
 cdef inline void ba_set_value(np.uint8_t *buf, np.uint64_t ind,
                               np.uint8_t val) noexcept nogil:
-    # This assumes 8 bit buffer
+    # This assumes 8 bit buffer.  If value is greater than zero (thus allowing
+    # us to use 1-255 as 'True') then we identify first the index in the buffer
+    # we are setting.  We do this by truncating the index by bit-shifting to
+    # the left three times, essentially dividing it by eight (and taking the
+    # floor.)
+    # The next step is to turn *on* what we're attempting to turn on, which
+    # means taking our index and truncating it to the first 3 bits (which we do
+    # with an & operation) and then turning on the correct bit.
+    #
+    # So if we're asking for index 33 in the bitarray, we would want the 4th
+    # uint8 element, then the 2nd bit (index 1).
+    #
+    # To turn it on, we logical *or* with that.  To turn it off, we logical
+    # *and* with the *inverse*, which will allow everything *but* that bit to
+    # stay on.
     if val > 0:
         buf[ind >> 3] |= (1 << (ind & 7))
     else:
@@ -25,13 +39,63 @@ cdef inline np.uint8_t ba_get_value(np.uint8_t *buf, np.uint64_t ind) noexcept n
     if rv == 0: return 0
     return 1
 
+cdef inline void ba_set_range(np.uint8_t *buf, np.uint64_t start_ind,
+                              np.uint64_t stop_ind, np.uint8_t val) nogil:
+    # Should this be inclusive of both end points?  I think it should not, to
+    # match slicing semantics.
+    #
+    # We need to figure out the first and last values, and then we just set the
+    # ones in-between to 255.
+    if stop_ind < start_ind: return
+    cdef np.uint64_t i
+    cdef np.uint8_t j, bitmask
+    cdef np.uint64_t buf_start = start_ind >> 3
+    cdef np.uint64_t buf_stop = stop_ind >> 3
+    cdef np.uint8_t start_j = start_ind & 7
+    cdef np.uint8_t stop_j = stop_ind & 7
+    if buf_start == buf_stop:
+        for j in range(start_j, stop_j):
+            ba_set_value(&buf[buf_start], j, val)
+        return
+    bitmask = 0
+    for j in range(start_j, 8):
+        bitmask |= (1 << j)
+    if val > 0:
+        buf[buf_start] |= bitmask
+    else:
+        buf[buf_start] &= ~bitmask
+    if val > 0:
+        bitmask = 255
+    else:
+        bitmask = 0
+    for i in range(buf_start + 1, buf_stop):
+        buf[i] = bitmask
+    bitmask = 0
+    for j in range(0, stop_j):
+        bitmask |= (1 << j)
+    if val > 0:
+        buf[buf_stop] |= bitmask
+    else:
+        buf[buf_stop] &= ~bitmask
+
+
+cdef inline np.uint8_t _num_set_bits( np.uint8_t b ):
+    # https://stackoverflow.com/questions/30688465/how-to-check-the-number-of-set-bits-in-an-8-bit-unsigned-char
+    b = b - ((b >> 1) & 0x55)
+    b = (b & 0x33) + ((b >> 2) & 0x33)
+    return (((b + (b >> 4)) & 0x0F) * 0x01)
+
 cdef class bitarray:
     cdef np.uint8_t *buf
     cdef np.uint64_t size
     cdef np.uint64_t buf_size # Not exactly the same
+    cdef np.uint8_t final_bitmask
     cdef public object ibuf
 
     cdef void _set_value(self, np.uint64_t ind, np.uint8_t val)
     cdef np.uint8_t _query_value(self, np.uint64_t ind)
-    #cdef void set_range(self, np.uint64_t ind, np.uint64_t count, int val)
-    #cdef int query_range(self, np.uint64_t ind, np.uint64_t count, int *val)
+    cdef void _set_range(self, np.uint64_t start, np.uint64_t stop, np.uint8_t val)
+    cdef np.uint64_t _count(self)
+    cdef bitarray _logical_and(self, bitarray other, bitarray result = *)
+    cdef bitarray _logical_or(self, bitarray other, bitarray result = *)
+    cdef bitarray _logical_xor(self, bitarray other, bitarray result = *)