Merge pull request yt-project#1395 from MatthewTurk/dsr_frb

Adding FRB generation to ds.r

doc/source/analyzing/objects.rst

@@ -199,6 +199,21 @@ This accepts arguments the same way:::
 
 .. _available-objects:
 
+Making Image Buffers
+^^^^^^^^^^^^^^^^^^^^
+
+Using the slicing syntax above for choosing a slice, if you also provide an
+imaginary step value you can obtain a
+:class:`~yt.visualization.api.FixedResolutionBuffer` of the chosen resolution.
+
+For instance, to obtain a 1024 by 1024 buffer covering the entire
+domain but centered at 0.5 in code units, you can do:::
+
+   frb = ds.r[0.5, ::1024j, ::1024j]
+
+This `frb` object then can be queried like a normal fixed resolution buffer,
+and it will return arrays of shape (1024, 1024).
+
 Available Objects
 -----------------
 
@@ -650,7 +665,7 @@ Here are some examples:
    cutout = sp1 - sp3
    sp4 = sp1 ^ sp2
    sp5 = sp1 & sp2
-   
+
 
 Note that the ``+`` operation and the ``|`` operation are identical.  For when
 multiple objects are to be combined in an intersection or a union, there are

yt/data_objects/data_containers.py

@@ -495,16 +495,16 @@ def to_dataframe(self, fields = None):
     def save_as_dataset(self, filename=None, fields=None):
         r"""Export a data object to a reloadable yt dataset.
 
-        This function will take a data object and output a dataset 
-        containing either the fields presently existing or fields 
+        This function will take a data object and output a dataset
+        containing either the fields presently existing or fields
         given in the ``fields`` list.  The resulting dataset can be
         reloaded as a yt dataset.
 
         Parameters
         ----------
         filename : str, optional
-            The name of the file to be written.  If None, the name 
-            will be a combination of the original dataset and the type 
+            The name of the file to be written.  If None, the name
+            will be a combination of the original dataset and the type
             of data container.
         fields : list of string or tuple field names, optional
             If this is supplied, it is the list of fields to be saved to
@@ -1628,6 +1628,8 @@ def to_frb(self, width, resolution, center=None, height=None,
         elif iterable(height):
             h, u = height
             height = self.ds.quan(h, input_units = u)
+        elif not isinstance(height, YTArray):
+            height = self.ds.quan(height, 'code_length')
         if not iterable(resolution):
             resolution = (resolution, resolution)
         from yt.visualization.fixed_resolution import FixedResolutionBuffer
@@ -1857,7 +1859,7 @@ def calculate_isocontour_flux(self, field, value,
 
     def _calculate_flux_in_grid(self, grid, mask, field, value,
                     field_x, field_y, field_z, fluxing_field = None):
-        
+
         vc_fields = [field, field_x, field_y, field_z]
         if fluxing_field is not None:
             vc_fields.append(fluxing_field)

yt/data_objects/region_expression.py

@@ -38,25 +38,31 @@ def __getitem__(self, item):
         if isinstance(item, tuple) and isinstance(item[1], string_types):
             return self.all_data[item]
         if isinstance(item, slice):
+            # This is for the case where we give a slice as an index; one
+            # possible use case of this would be where we supply something
+            # like ds.r[::256j] .  This would be expanded, implicitly into
+            # ds.r[::256j, ::256j, ::256j].  Other cases would be if we do
+            # ds.r[0.1:0.9] where it will be expanded along three dimensions.
             item = (item, item, item)
         if len(item) != self.ds.dimensionality:
             # Not the right specification, and we don't want to do anything
-            # implicitly.
+            # implicitly.  Note that this happens *after* the implicit expansion
+            # of a single slice.
             raise YTDimensionalityError(len(item), self.ds.dimensionality)
         if self.ds.dimensionality != 3:
             # We'll pass on this for the time being.
             raise RuntimeError
 
         # OK, now we need to look at our slices.  How many are a specific
         # coordinate?
-        
+
         if not all(isinstance(v, slice) for v in item):
             return self._create_slice(item)
         else:
             if all(s.start is s.stop is s.step is None for s in item):
                 return self.all_data
             return self._create_region(item)
-            
+
     def _spec_to_value(self, input_tuple):
         if not isinstance(input_tuple, tuple):
             # We now assume that it's in code_length
@@ -66,6 +72,9 @@ def _spec_to_value(self, input_tuple):
         return value
 
     def _create_slice(self, slice_tuple):
+        # This is somewhat more complex because we want to allow for slicing
+        # in one dimension but also *not* using the entire domain; for instance
+        # this means we allow something like ds.r[0.5, 0.1:0.4, 0.1:0.4].
         axis = None
         new_slice = []
         for ax, v in enumerate(slice_tuple):
@@ -79,6 +88,23 @@ def _create_slice(self, slice_tuple):
         # This new slice doesn't need to be a tuple
         source = self._create_region(new_slice)
         sl = self.ds.slice(axis, coord, data_source = source)
+        # Now, there's the possibility that what we're also seeing here
+        # includes some steps, which would be for getting back a fixed
+        # resolution buffer.  We check for that by checking if we have
+        # exactly two imaginary steps.
+        xax = self.ds.coordinates.x_axis[axis]
+        yax = self.ds.coordinates.y_axis[axis]
+        if getattr(new_slice[xax].step, "imag", 0.0) != 0.0 and \
+           getattr(new_slice[yax].step, "imag", 0.0) != 0.0:
+            # We now need to convert to a fixed res buffer.
+            # We'll do this by getting the x/y axes, and then using that.
+            width = source.right_edge[xax] - source.left_edge[xax]
+            height = source.right_edge[yax] - source.left_edge[yax]
+            # Make a resolution tuple with
+            resolution = (int(new_slice[xax].step.imag),
+                          int(new_slice[yax].step.imag))
+            sl = sl.to_frb(width = width, resolution = resolution,
+                           height = height)
         return sl
 
     def _slice_to_edges(self, ax, val):

yt/data_objects/tests/test_dataset_access.py

@@ -66,6 +66,16 @@ def test_accessing_all_data():
     assert_equal(dd["density"]*2.0, ds.r["density"])
     assert_equal(dd["gas", "density"]*2.0, ds.r["gas", "density"])
 
+def test_slice_from_r():
+    ds = fake_amr_ds(fields = ["density"])
+    sl1 = ds.r[0.5, :, :]
+    sl2 = ds.slice("x", 0.5)
+    assert_equal(sl1["density"], sl2["density"])
+
+    frb1 = sl1.to_frb(width = 1.0, height = 1.0, resolution = (1024, 512))
+    frb2 = ds.r[0.5, ::1024j, ::512j]
+    assert_equal(frb1["density"], frb2["density"])
+
 def test_particle_counts():
     ds = fake_random_ds(16, particles=100)
     assert ds.particle_type_counts == {'io': 100}