Repository: yt-project/yt

Issue:
ENH: Allow YTCoveringGrid to accept a data source
Presently, we set up a data source for covering grids in `_setup_data_source`.  We should be able to provide this, thus enabling spheres to be regularized for instance.

The output would still be a box, but this would be a way to produce regularized data from within specific regions.

Alternately, if we are OK with over-generating the data, we could simply apply a `selector` mask to the covering grid.


Discussion:
As a note, we probably do *not* want to immediately make this available for smoothed covering grids, as those include interpolation that we are not accounting for.

I believe that this would be a reasonably straightforward change (but would also benefit from something like a "clone" operation on data objects so that we don't mess with them via modifying `max_level` etc).

Resolving patch:
diff --git a/yt/data_objects/construction_data_containers.py b/yt/data_objects/construction_data_containers.py
--- a/yt/data_objects/construction_data_containers.py
+++ b/yt/data_objects/construction_data_containers.py
@@ -627,6 +627,9 @@ class YTCoveringGrid(YTSelectionContainer3D):
         A list of fields that you'd like pre-generated for your object
     num_ghost_zones : integer, optional
         The number of padding ghost zones used when accessing fields.
+    data_source :
+        An existing data object to intersect with the covering grid. Grid points
+        outside the data_source will exist as empty values.
 
     Examples
     --------
@@ -656,12 +659,14 @@ def __init__(
         num_ghost_zones=0,
         use_pbar=True,
         field_parameters=None,
+        *,
+        data_source=None,
     ):
         if field_parameters is None:
             center = None
         else:
             center = field_parameters.get("center", None)
-        YTSelectionContainer3D.__init__(self, center, ds, field_parameters)
+        super().__init__(center, ds, field_parameters, data_source=data_source)
 
         self.level = level
         self.left_edge = self._sanitize_edge(left_edge)
@@ -815,7 +820,17 @@ def shape(self):
         return tuple(self.ActiveDimensions.tolist())
 
     def _setup_data_source(self):
-        self._data_source = self.ds.region(self.center, self.left_edge, self.right_edge)
+
+        reg = self.ds.region(self.center, self.left_edge, self.right_edge)
+        if self._data_source is None:
+            # note: https://github.com/yt-project/yt/pull/4063 implemented
+            # a data_source kwarg for YTCoveringGrid, but not YTArbitraryGrid
+            # so as of 4063, this will always be True for YTArbitraryGrid
+            # instances.
+            self._data_source = reg
+        else:
+            self._data_source = self.ds.intersection([self._data_source, reg])
+
         self._data_source.min_level = 0
         self._data_source.max_level = self.level
         # This triggers "special" behavior in the RegionSelector to ensure we


Test patch:
diff --git a/yt/data_objects/tests/test_covering_grid.py b/yt/data_objects/tests/test_covering_grid.py
--- a/yt/data_objects/tests/test_covering_grid.py
+++ b/yt/data_objects/tests/test_covering_grid.py
@@ -91,6 +91,58 @@ def test_covering_grid():
         sp.quantities.total_mass()
 
 
+@requires_file(cyl_2d)
+@requires_file(cyl_3d)
+def test_covering_grid_data_source():
+    # test the data_source kwarg (new with PR 4063)
+
+    for level in [0, 1, 2]:
+        for nprocs in [1, 2, 4, 8]:
+            ds = fake_random_ds(16, nprocs=nprocs)
+            dn = ds.refine_by**level
+            cg = ds.covering_grid(level, [0.0, 0.0, 0.0], dn * ds.domain_dimensions)
+
+            # check that the regularized sphere is centered where it should be,
+            # to within the tolerance of the grid. Use a center offsect from
+            # the domain center by a bit
+            center = ds.domain_center + np.min(ds.domain_width) * 0.15
+            sp = ds.sphere(center, (0.1, "code_length"))
+            dims = dn * ds.domain_dimensions
+            cg_sp = ds.covering_grid(level, [0.0, 0.0, 0.0], dims, data_source=sp)
+
+            # find the discrete center of the sphere by averaging the position
+            # along each dimension where values are non-zero
+            cg_mask = cg_sp["gas", "density"] != 0
+            discrete_c = ds.arr(
+                [
+                    np.mean(cg_sp[("index", dim)][cg_mask])
+                    for dim in ds.coordinates.axis_order
+                ]
+            )
+
+            # should be centered to within the tolerance of the grid at least
+            grid_tol = ds.arr(
+                [cg_sp[("index", dim)][0, 0, 0] for dim in ds.coordinates.axis_order]
+            )
+            assert np.all(np.abs(discrete_c - center) <= grid_tol)
+
+            # check that a region covering the whole domain matches no data_source
+            reg = ds.region(ds.domain_center, ds.domain_left_edge, ds.domain_right_edge)
+            cg_reg = ds.covering_grid(level, [0.0, 0.0, 0.0], dims, data_source=reg)
+            assert np.all(cg[("gas", "density")] == cg_reg[("gas", "density")])
+
+            # check that a box covering a subset of the domain is the right volume
+            right_edge = ds.domain_left_edge + ds.domain_width * 0.5
+            c = (right_edge + ds.domain_left_edge) / 2
+            reg = ds.region(c, ds.domain_left_edge, right_edge)
+            cg_reg = ds.covering_grid(level, [0.0, 0.0, 0.0], dims, data_source=reg)
+            box_vol = cg_reg[("index", "cell_volume")][
+                cg_reg["gas", "density"] != 0
+            ].sum()
+            actual_vol = np.prod(right_edge - ds.domain_left_edge)
+            assert box_vol == actual_vol
+
+
 @requires_module("xarray")
 def test_xarray_export():
     def _run_tests(cg):