Mask out levels that we won't read

yt/frontends/ramses/data_structures.py

@@ -375,7 +375,7 @@ def _fill_no_ghostzones(self, fd, fields, selector, file_handler):
         data = {}
         cell_count = selector.count_oct_cells(self.oct_handler, self.domain_id)
 
-        levels, cell_inds, file_inds = self.oct_handler.file_index_octs(
+        level_inds, cell_inds, file_inds = self.oct_handler.file_index_octs(
             selector, self.domain_id, cell_count
         )
 
@@ -389,7 +389,7 @@ def _fill_no_ghostzones(self, fd, fields, selector, file_handler):
             file_handler.offset,
             file_handler.level_count,
             cpu_list,
-            levels,
+            level_inds,
             cell_inds,
             file_inds,
             ndim,
@@ -418,13 +418,13 @@ def _fill_with_ghostzones(
 
         gz_cache = getattr(self, "_ghost_zone_cache", None)
         if gz_cache:
-            levels, cell_inds, file_inds, domains = gz_cache
+            level_inds, cell_inds, file_inds, domain_inds = gz_cache
         else:
             gz_cache = (
-                levels,
+                level_inds,
                 cell_inds,
                 file_inds,
-                domains,
+                domain_inds,
             ) = self.oct_handler.file_index_octs_with_ghost_zones(
                 selector, self.domain_id, cell_count, self._num_ghost_zones
             )
@@ -439,15 +439,15 @@ def _fill_with_ghostzones(
             file_handler.offset,
             file_handler.level_count,
             cpu_list,
-            levels,
+            level_inds,
             cell_inds,
             file_inds,
             ndim,
             all_fields,
             fields,
             tr,
             oct_handler,
-            domains=domains,
+            domains=domain_inds,
         )
         return tr
 

yt/frontends/ramses/io_utils.pyx

@@ -163,13 +163,13 @@ def fill_hydro(FortranFile f,
                np.ndarray[np.int64_t, ndim=2] offsets,
                np.ndarray[np.int64_t, ndim=2] level_count,
                list cpu_enumerator,
-               np.ndarray[np.uint8_t, ndim=1] levels,
+               np.ndarray[np.uint8_t, ndim=1] level_inds,
                np.ndarray[np.uint8_t, ndim=1] cell_inds,
                np.ndarray[np.int64_t, ndim=1] file_inds,
                INT64_t ndim, list all_fields, list fields,
                dict tr,
                RAMSESOctreeContainer oct_handler,
-               np.ndarray[np.int32_t, ndim=1] domains=np.array([], dtype='int32')):
+               np.ndarray[np.int32_t, ndim=1] domain_inds=np.array([], dtype='int32')):
     cdef INT64_t offset
     cdef dict tmp
     cdef str field
@@ -186,8 +186,9 @@ def fill_hydro(FortranFile f,
     cdef np.int64_t[::1] cpu_list = np.asarray(cpu_enumerator, dtype=np.int64)
 
     cdef np.int64_t[::1] jumps = np.zeros(nfields_selected + 1, dtype=np.int64)
-    cdef int jump_len
+    cdef int jump_len, Ncells
     cdef np.ndarray[np.float64_t, ndim=3] buffer
+    cdef np.uint8_t[::1] mask_level = np.zeros(nlevels, dtype=np.uint8)
 
     jump_len = 0
     j = 0
@@ -202,8 +203,18 @@ def fill_hydro(FortranFile f,
     cdef int first_field_index = jumps[0]
 
     buffer = np.empty((level_count.max(), twotondim, nfields_selected), dtype="float64", order='F')
+
+    Ncells = len(level_inds)
+    for i in range(Ncells):
+        mask_level[level_inds[i]] |= 1
+
+    # print("Reading hydro fields, unique levels=%s", np.asarray(levels_to_read))
+    # selected_levels = range(levels.min(), levels.max()+1)
+    # selected_cpus = range(cpu_list.min(), cpu_list.max()+1)
     # Loop over levels
     for ilevel in range(nlevels):
+        if mask_level[ilevel] == 0:
+            continue
         # Loop over cpu domains
         for ii in range(ncpu_selected):
             icpu = cpu_list[ii]
@@ -242,7 +253,7 @@ def fill_hydro(FortranFile f,
 
             if ncpu_selected > 1:
                 oct_handler.fill_level_with_domain(
-                    ilevel, levels, cell_inds, file_inds, domains, tr, tmp, domain=icpu+1)
+                    ilevel, level_inds, cell_inds, file_inds, domain_inds, tr, tmp, domain=icpu+1)
             else:
                 oct_handler.fill_level(
-                    ilevel, levels, cell_inds, file_inds, tr, tmp)
+                    ilevel, level_inds, cell_inds, file_inds, tr, tmp)

yt/geometry/oct_container.pxd

@@ -85,20 +85,20 @@ cdef class OctreeContainer:
     cpdef void fill_level(
         self,
         const int level,
-        const np.uint8_t[:] levels,
-        const np.uint8_t[:] cell_inds,
-        const np.int64_t[:] file_inds,
+        const np.uint8_t[::1] level_inds,
+        const np.uint8_t[::1] cell_inds,
+        const np.int64_t[::1] file_inds,
         dict dest_fields,
         dict source_fields,
         np.int64_t offset = ?
     )
     cpdef int fill_level_with_domain(
         self,
         const int level,
-        const np.uint8_t[:] levels,
-        const np.uint8_t[:] cell_inds,
-        const np.int64_t[:] file_inds,
-        const np.int32_t[:] domains,
+        const np.uint8_t[::1] level_inds,
+        const np.uint8_t[::1] cell_inds,
+        const np.int64_t[::1] file_inds,
+        const np.int32_t[::1] domain_inds,
         dict dest_fields,
         dict source_fields,
         const np.int32_t domain,

yt/geometry/oct_container.pyx

@@ -745,15 +745,15 @@ cdef class OctreeContainer:
     cpdef void fill_level(
         self,
         const int level,
-        const np.uint8_t[:] levels,
-        const np.uint8_t[:] cell_inds,
-        const np.int64_t[:] file_inds,
+        const np.uint8_t[::1] levels,
+        const np.uint8_t[::1] cell_inds,
+        const np.int64_t[::1] file_inds,
         dict dest_fields,
         dict source_fields,
         np.int64_t offset = 0
     ):
-        cdef np.ndarray[np.float64_t, ndim=2] source
-        cdef np.ndarray[np.float64_t, ndim=1] dest
+        cdef np.float64_t[:, :] source
+        cdef np.float64_t[::1] dest
         cdef int i, lvl
 
         for key in dest_fields:
@@ -831,10 +831,10 @@ cdef class OctreeContainer:
     cpdef int fill_level_with_domain(
         self,
         const int level,
-        const np.uint8_t[:] levels,
-        const np.uint8_t[:] cell_inds,
-        const np.int64_t[:] file_inds,
-        const np.int32_t[:] domains,
+        const np.uint8_t[::1] level_inds,
+        const np.uint8_t[::1] cell_inds,
+        const np.int64_t[::1] file_inds,
+        const np.int32_t[::1] domain_inds,
         dict dest_fields,
         dict source_fields,
         const np.int32_t domain,
@@ -846,18 +846,18 @@ cdef class OctreeContainer:
         These buffer oct cells have a different domain than the local one and
         are usually not read, but one has to read them e.g. to compute ghost zones.
         """
-        cdef np.ndarray[np.float64_t, ndim=2] source
-        cdef np.ndarray[np.float64_t, ndim=1] dest
+        cdef np.float64_t[:, :] source
+        cdef np.float64_t[::1] dest
         cdef int i, count, lev
         cdef np.int32_t dom
 
         for key in dest_fields:
             dest = dest_fields[key]
             source = source_fields[key]
             count = 0
-            for i in range(levels.shape[0]):
-                lev = levels[i]
-                dom = domains[i]
+            for i in range(level_inds.shape[0]):
+                lev = level_inds[i]
+                dom = domain_inds[i]
                 if lev != level or dom != domain: continue
                 count += 1
                 if file_inds[i] < 0: