[WIP] Experimentally turn off precounting AMR datasets

yt/data_objects/construction_data_containers.py

@@ -1105,6 +1105,7 @@ def _fill_fields(self, fields):
             refine_by = [refine_by, refine_by, refine_by]
         refine_by = np.array(refine_by, dtype="i8")
         for chunk in parallel_objects(self._data_source.chunks(fields, "io")):
+            chunk[fields[0]]
             input_fields = [chunk[field] for field in fields]
             # NOTE: This usage of "refine_by" is actually *okay*, because it's
             # being used with respect to iref, which is *already* scaled!
@@ -1504,7 +1505,6 @@ def _fill_fields(self, fields):
             domain_dims = domain_dims.astype("int64")
             tot = ls.current_dims.prod()
             for chunk in ls.data_source.chunks(fields, "io"):
-                chunk[fields[0]]
                 input_fields = [chunk[field] for field in fields]
                 tot -= fill_region(
                     input_fields,

yt/data_objects/data_containers.py

@@ -297,7 +297,7 @@ def _generate_spatial_fluid(self, field, ngz):
             raise YTSpatialFieldUnitError(field)
         units = finfo.units
         try:
-            rv = self.ds.arr(np.zeros(self.ires.size, dtype="float64"), units)
+            rv = np.zeros(self.ires.size, dtype="float64")
             accumulate = False
         except YTNonIndexedDataContainer:
             # In this case, we'll generate many tiny arrays of unknown size and
@@ -312,7 +312,7 @@ def _generate_spatial_fluid(self, field, ngz):
                 for _chunk in self.chunks([], "spatial", ngz=0, preload_fields=deps):
                     o = self._current_chunk.objs[0]
                     if accumulate:
-                        rv = self.ds.arr(np.empty(o.ires.size, dtype="float64"), units)
+                        rv = np.empty(o.ires.size, dtype="float64")
                         outputs.append(rv)
                         ind = 0  # Does this work with mesh?
                     with o._activate_cache():
@@ -327,16 +327,15 @@ def _generate_spatial_fluid(self, field, ngz):
                     gz.field_parameters = self.field_parameters
                     wogz = gz._base_grid
                     if accumulate:
-                        rv = self.ds.arr(
-                            np.empty(wogz.ires.size, dtype="float64"), units
-                        )
+                        rv = np.empty(wogz.ires.size, dtype="float64")
                         outputs.append(rv)
                     ind += wogz.select(
                         self.selector,
                         source=gz[field][ngz:-ngz, ngz:-ngz, ngz:-ngz],
                         dest=rv,
                         offset=ind,
                     )
+        rv = self.ds.arr(rv, units)
         if accumulate:
             rv = uconcatenate(outputs)
         return rv

yt/data_objects/index_subobjects/grid_patch.py

@@ -388,7 +388,11 @@ def _get_selector_mask(self, selector):
         if self._cache_mask and hash(selector) == self._last_selector_id:
             mask = self._last_mask
         else:
-            mask, count = selector.fill_mask_regular_grid(self)
+            if selector is None or getattr(selector, "is_all_data", False):
+                mask = self.child_mask.copy()
+                count = mask.sum()
+            else:
+                mask, count = selector.fill_mask_regular_grid(self)
             if self._cache_mask:
                 self._last_mask = mask
             self._last_selector_id = hash(selector)
@@ -406,10 +410,14 @@ def select(self, selector, source, dest, offset):
         dim = np.squeeze(self.ds.dimensionality)
         nodal_flag = source.shape[:dim] - self.ActiveDimensions[:dim]
         if sum(nodal_flag) == 0:
+            dest.resize(offset + count, refcheck=False)
+            dest.shape = (offset + count,)
             dest[offset : offset + count] = source[mask]
         else:
             slices = get_nodal_slices(source.shape, nodal_flag, dim)
             for i, sl in enumerate(slices):
+                dest.resize((offset + count, i), refcheck=False)
+                dest.shape = (offset + count, i)
                 dest[offset : offset + count, i] = source[tuple(sl)][np.squeeze(mask)]
         return count
 

yt/data_objects/selection_objects/data_selection_objects.py

@@ -4,6 +4,7 @@
 import uuid
 from collections import defaultdict
 from contextlib import contextmanager
+from functools import cached_property
 
 import numpy as np
 from more_itertools import always_iterable
@@ -70,7 +71,10 @@ def __init__(self, ds, field_parameters, data_source=None):
                     % (data_source._dimensionality, self._dimensionality)
                 )
             self.field_parameters.update(data_source.field_parameters)
-        self.quantities = DerivedQuantityCollection(self)
+
+    @cached_property
+    def quantities(self):
+        return DerivedQuantityCollection(self)
 
     @property
     def selector(self):

yt/frontends/stream/io.py

@@ -30,8 +30,10 @@ def _read_fluid_selection(self, chunks, selector, fields, size):
         if any((ftype not in self.ds.fluid_types for ftype, fname in fields)):
             raise NotImplementedError
         rv = {}
+        if size is None:
+            size = 0
         for field in fields:
-            rv[field] = self.ds.arr(np.empty(size, dtype="float64"))
+            rv[field] = np.empty(size, dtype="float64")
 
         ng = sum(len(c.objs) for c in chunks)
         mylog.debug(

yt/geometry/grid_geometry_handler.py

@@ -344,10 +344,6 @@ def _identify_base_chunk(self, dobj):
         # These next two lines, when uncommented, turn "on" the fast index.
         # if dobj._type_name != "grid":
         #    fast_index = self._get_grid_tree()
-        if getattr(dobj, "size", None) is None:
-            dobj.size = self._count_selection(dobj, fast_index=fast_index)
-        if getattr(dobj, "shape", None) is None:
-            dobj.shape = (dobj.size,)
         dobj._current_chunk = list(
             self._chunk_all(dobj, cache=False, fast_index=fast_index)
         )[0]
@@ -361,6 +357,10 @@ def _count_selection(self, dobj, grids=None, fast_index=None):
         return count
 
     def _chunk_all(self, dobj, cache=True, fast_index=None):
+        if getattr(dobj, "size", None) is None:
+            dobj.size = self._count_selection(dobj, fast_index=fast_index)
+        if getattr(dobj, "shape", None) is None:
+            dobj.shape = (dobj.size,)
         gobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info)
         fast_index = fast_index or getattr(dobj._current_chunk, "_fast_index", None)
         yield YTDataChunk(dobj, "all", gobjs, dobj.size, cache, fast_index=fast_index)
@@ -440,7 +440,8 @@ def _chunk_io(
                     dobj,
                     "io",
                     grids,
-                    self._count_selection(dobj, grids),
+                    # self._count_selection(dobj, grids),
+                    None,
                     cache=cache,
                     fast_index=fast_index,
                 )

yt/utilities/io_handler.py

@@ -109,10 +109,10 @@ def _read_fluid_selection(
             nodal_flag = finfo.nodal_flag
             if np.any(nodal_flag):
                 num_nodes = 2 ** sum(nodal_flag)
-                rv[field] = np.empty((size, num_nodes), dtype="=f8")
+                rv[field] = np.empty((size or 0, num_nodes), dtype="=f8")
                 nodal_fields.append(field)
             else:
-                rv[field] = np.empty(size, dtype="=f8")
+                rv[field] = np.empty((size or 0,), dtype="=f8")
         ind = {field: 0 for field in fields}
         for field, obj, data in self.io_iter(chunks, fields):
             if data is None: