ENH: Enable exporting GAMER and FLASH datasets to octrees

yt/frontends/stream/data_structures.py

@@ -836,7 +836,7 @@ def _setup_data_io(self):
 
     def _initialize_oct_handler(self):
         header = dict(
-            dims=[1, 1, 1],
+            dims=self.ds.domain_dimensions // self.ds.num_zones,
             left_edge=self.ds.domain_left_edge,
             right_edge=self.ds.domain_right_edge,
             octree=self.ds.octree_mask,

yt/geometry/oct_container.pyx

@@ -171,7 +171,7 @@ cdef class OctreeContainer:
                 pos[2] = self.DLE[2] + dds[2]/2.0
                 for k in range(self.nn[2]):
                     if self.root_mesh[i][j][k] == NULL:
-                        raise RuntimeError
+                        raise KeyError(i,j,k)
                     visitor.pos[0] = i
                     visitor.pos[1] = j
                     visitor.pos[2] = k
@@ -185,6 +185,16 @@ cdef class OctreeContainer:
     cdef np.int64_t get_domain_offset(self, int domain_id):
         return 0
 
+    def check_root_mesh(self):
+        cdef count = 0
+        for i in range(self.nn[0]):
+            for j in range(self.nn[1]):
+                for k in range(self.nn[2]):
+                    if self.root_mesh[i][j][k] == NULL:
+                        print("Missing ", i, j, k)
+                        count += 1
+        print("Missing total of %s out of %s" % (count, self.nn[0] * self.nn[1] * self.nn[2]))
+
     cdef int get_root(self, int ind[3], Oct **o) nogil:
         cdef int i
         for i in range(3):
@@ -567,7 +577,11 @@ cdef class OctreeContainer:
     def add(self, int curdom, int curlevel,
             np.ndarray[np.float64_t, ndim=2] pos,
             int skip_boundary = 1,
-            int count_boundary = 0):
+            int count_boundary = 0,
+            np.ndarray[np.uint64_t, ndim=1, cast=True] levels = None
+            ):
+        # In this function, if we specify curlevel = -1, then we query the
+        # (optional) levels array for the oct level.
         cdef int no, p, i
         cdef int ind[3]
         cdef int nb = 0
@@ -582,6 +596,12 @@ cdef class OctreeContainer:
         cdef int in_boundary = 0
         # How do we bootstrap ourselves?
         for p in range(no):
+            # We allow specifying curlevel = -1 to query from the levels array
+            # instead.
+            if curlevel == -1:
+                this_level = levels[p]
+            else:
+                this_level = curlevel
             #for every oct we're trying to add find the
             #floating point unitary position on this level
             in_boundary = 0
@@ -599,7 +619,7 @@ cdef class OctreeContainer:
             if cur == NULL: raise RuntimeError
             # Now we find the location we want
             # Note that RAMSES I think 1-findiceses levels, but we don't.
-            for _ in range(curlevel):
+            for _ in range(this_level):
                 # At every level, find the cell this oct
                 # lives inside
                 for i in range(3):

yt/loaders.py

@@ -930,6 +930,7 @@ def load_octree(
     default_species_fields=None,
     *,
     parameters=None,
+    domain_dimensions=None,
 ):
     r"""Load an octree mask into yt.
 
@@ -980,6 +981,9 @@ def load_octree(
     parameters: dictionary, optional
         Optional dictionary used to populate the dataset parameters, useful
         for storing dataset metadata.
+    domain_dimensions : array_like, optional
+        This is the domain dimensions of the root *mesh*, which can be used to
+        specify (indirectly) the number of root oct nodes.
 
     Example
     -------
@@ -1015,7 +1019,11 @@ def load_octree(
     # for compatibility
     if over_refine_factor is not None:
         nz = 1 << over_refine_factor
-    domain_dimensions = np.array([nz, nz, nz])
+    if domain_dimensions is None:
+        # We assume that if it isn't specified, it defaults to the number of
+        # zones (i.e., a single root oct.)
+        domain_dimensions = [nz, nz, nz]
+    domain_dimensions = np.array(domain_dimensions)
     nprocs = 1
     if bbox is None:
         bbox = np.array([[0.0, 1.0], [0.0, 1.0], [0.0, 1.0]], "float64")