Fast seek through file

yt/frontends/ramses/data_structures.py

@@ -375,21 +375,24 @@ 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(
-            selector, self.domain_id, cell_count
-        )
-
         # Initializing data container
         for field in fields:
             data[field] = np.zeros(cell_count, "float64")
 
+        if cell_count == 0:
+            return data
+
+        level_inds, cell_inds, file_inds = self.oct_handler.file_index_octs(
+            selector, self.domain_id, cell_count
+        )
+
         cpu_list = [self.domain_id - 1]
         fill_hydro(
             fd,
             file_handler.offset,
             file_handler.level_count,
             cpu_list,
-            levels,
+            level_inds,
             cell_inds,
             file_inds,
             ndim,
@@ -416,38 +419,42 @@ def _fill_with_ghostzones(
             selector.count_octs(self.oct_handler, self.domain_id) * self.nz**ndim
         )
 
+        # Initializing data container
+        for field in fields:
+            tr[field] = np.zeros(cell_count, "float64")
+
+        if cell_count == 0:
+            return tr
+
         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
             )
             self._ghost_zone_cache = gz_cache
 
-        # Initializing data container
-        for field in fields:
-            tr[field] = np.zeros(cell_count, "float64")
         cpu_list = list(range(ncpu))
         fill_hydro(
             fd,
             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,
+            domain_inds=domain_inds,
         )
         return tr
 

yt/frontends/ramses/io_utils.pyx

@@ -18,6 +18,11 @@ cdef int INT32_SIZE = sizeof(np.int32_t)
 cdef int INT64_SIZE = sizeof(np.int64_t)
 cdef int DOUBLE_SIZE = sizeof(np.float64_t)
 
+
+cdef inline int skip_len(int Nskip, int record_len) noexcept nogil:
+    return Nskip * (record_len * DOUBLE_SIZE + INT64_SIZE)
+
+
 @cython.cpow(True)
 @cython.boundscheck(False)
 @cython.wraparound(False)
@@ -29,10 +34,13 @@ def read_amr(FortranFile f, dict headers,
 
     cdef INT64_t ncpu, nboundary, max_level, nlevelmax, ncpu_and_bound
     cdef DOUBLE_t nx, ny, nz
-    cdef INT64_t ilevel, icpu, n, ndim, skip_len
-    cdef INT32_t ng, buffer_size
+    cdef INT64_t ilevel, icpu, n, ndim, jump_len
+    cdef INT32_t ng
     cdef np.ndarray[np.int32_t, ndim=2] numbl
     cdef np.ndarray[np.float64_t, ndim=2] pos
+    cdef int i
+
+    cdef np.float64_t[::1, :] pos_view
 
     ndim = headers['ndim']
     numbl = headers['numbl']
@@ -43,17 +51,21 @@ def read_amr(FortranFile f, dict headers,
 
     ncpu_and_bound = nboundary + ncpu
 
-    pos = np.empty((0, 3), dtype=np.float64)
-    buffer_size = 0
+    # Allocate more memory if required
+    pos = np.empty((max(numbl.max(), ngridbound.max()), 3), dtype="d", order="F")
+    pos_view = pos
+
     # Compute number of fields to skip. This should be 31 in 3 dimensions
-    skip_len = (1          # father index
+    jump_len = (1          # father index
                 + 2*ndim   # neighbor index
                 + 2**ndim  # son index
                 + 2**ndim  # cpu map
                 + 2**ndim  # refinement map
     )
     # Initialize values
     max_level = 0
+    cdef int record_len
+
     for ilevel in range(nlevelmax):
         for icpu in range(ncpu_and_bound):
             if icpu < ncpu:
@@ -63,21 +75,25 @@ def read_amr(FortranFile f, dict headers,
 
             if ng == 0:
                 continue
-            # Skip grid index, 'next' and 'prev' arrays (they are used
-            # to build the linked list in RAMSES)
-            f.skip(3)
-
-            # Allocate more memory if required
-            if ng > buffer_size:
-                pos = np.empty((ng, 3), dtype="d")
-                buffer_size = ng
-
-            pos[:ng, 0] = f.read_vector("d") - nx
-            pos[:ng, 1] = f.read_vector("d") - ny
-            pos[:ng, 2] = f.read_vector("d") - nz
+            # Skip grid index, 'next' and 'prev' arrays and possibly
+            # records from previous iterations
+            record_len = (ng * INT32_SIZE + INT64_SIZE)
+            f.seek(record_len * 3, 1)
+
+            f.read_vector_inplace("d", <void*> &pos_view[0, 0])
+            f.read_vector_inplace("d", <void*> &pos_view[0, 1])
+            f.read_vector_inplace("d", <void*> &pos_view[0, 2])
+
+            for i in range(ng):
+                pos_view[i, 0] -= nx
+            for i in range(ng):
+                pos_view[i, 1] -= ny
+            for i in range(ng):
+                pos_view[i, 2] -= nz
 
             # Skip father, neighbor, son, cpu map and refinement map
-            f.skip(skip_len)
+            f.seek(record_len * jump_len, 1)
+
             # Note that we're adding *grids*, not individual cells.
             if ilevel >= min_level:
                 n = oct_handler.add(icpu + 1, ilevel - min_level, pos[:ng, :],
@@ -87,10 +103,6 @@ def read_amr(FortranFile f, dict headers,
 
     return max_level
 
-
-cdef inline int skip_len(int Nskip, int record_len) noexcept nogil:
-    return Nskip * (record_len * DOUBLE_SIZE + INT64_SIZE)
-
 @cython.cpow(True)
 @cython.boundscheck(False)
 @cython.wraparound(False)
@@ -151,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
@@ -174,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
@@ -190,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]
@@ -230,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: