Update volume rendering code to work with more modern yt installations

.gitignore

@@ -54,6 +54,7 @@ docs/_build/
 docs/api
 
 *.fits
+*/tests/*.png
 
 # Other generated stuff
 */version.py

docs/yt_example.rst

@@ -1,6 +1,10 @@
 Visualizing spectral cubes with yt
 ==================================
 
+.. note::
+
+    The minimum yt version required to analyze spectral cubes is 3.5.0.
+
 Extracting yt objects
 ---------------------
 
@@ -69,7 +73,7 @@ produce a 3-d isocontour visualization using an object returned by
 
     import numpy as np
     from spectral_cube import SpectralCube
-    from yt.mods import ColorTransferFunction, write_bitmap
+    import yt
     import astropy.units as u
 
     # Read in spectral cube
@@ -85,31 +89,46 @@ produce a 3-d isocontour visualization using an object returned by
     vmin = 0.05
     vmax = 4.0
     dv = 0.02
-
+    
     # Set up color transfer function
-    transfer = ColorTransferFunction((vmin, vmax))
+    transfer = yt.ColorTransferFunction((vmin, vmax))
     transfer.add_layers(n_v, dv, colormap='RdBu_r')
-
-    # Set up the camera parameters
-
+
     # Derive the pixel coordinate of the desired center
-    # from the corresponding world coordinate
-    center = ytcube.world2yt([51.424522,
-                              30.723611,
-                              5205.18071])
-    direction = np.array([1.0, 0.0, 0.0])
-    width = 100.  # pixels
+    # from the corresponding world coordinate, this will
+    # be the position we are focused on
+    center = ytcube.world2yt([51.40,
+                              30.76,
+                              4653.75])
+    # We need to give this units so that yt knows how to 
+    # handle it. Pixel units are "code_length"
+    center = ds.arr(center, "code_length")
+
+    # The vector from the camera to the focus
+    direction = np.array([0.0, 0.0, 1.0])
+    # The resolution of the image
     size = 1024
-
-    camera = ds.camera(center, direction, width, size, transfer,
-                       fields=['flux'])
-
+
+    # Create the scene
+    sc = yt.create_scene(ds, field='flux')
+
+    # Set the transfer function
+    source = sc[0]
+    source.tfh.tf = transfer
+    source.tfh.bounds = (vmin, vmax)
+
+    # Now grab the camera and set the resolution, focus,
+    # and normal vector (camera angle)
+    cam = sc.camera
+    cam.set_resolution(size)
+    cam.set_focus(center)
+    cam.switch_orientation(normal_vector=direction)
+
     # Take a snapshot and save to a file
-    snapshot = camera.snapshot()
-    write_bitmap(snapshot, 'cube_rendering.png', transpose=True)
+    sc.save('rendering.png', sigma_clip=6)
 
 You can move the camera around; see the `yt camera docs
-<http://yt-project.org/docs/dev/reference/api/generated/yt.visualization.volume_rendering.camera.Camera.html>`_.
+<https://yt-project.org/doc/visualizing/volume_rendering.html#camera>`_.
 
 Movie Making
 ------------
@@ -128,9 +147,8 @@ Example:
 
 .. raw:: html
 
-   <iframe src="http://player.vimeo.com/video/104489207" width=500 height=281
-   frameborder=0 webkitallowfullscreen mozallowfullscreen allowfullscreen>
-   </iframe>
+   <iframe title="vimeo-player" src="https://player.vimeo.com/video/104489207" 
+   width="640" height="480" frameborder="0" allowfullscreen></iframe>
 
 SketchFab Isosurface Contours
 -----------------------------

setup.cfg

@@ -34,7 +34,7 @@ all =
     regions ; python_version<'3.8'
     reproject
     scipy
-    yt ; python_version<'3.8'
+    yt>=3.5.0 ; python_version<'3.8'
 
 [options.package_data]
 spectral_cube.tests = data/*

spectral_cube/tests/setup_package.py

@@ -2,5 +2,6 @@
 
 def get_package_data():
     return {
-        _ASTROPY_PACKAGE_NAME_ + '.tests': ['coveragerc', 'data/*.fits', 'data/*.hdr', 'data/*.lmv', 'data/*reg']
+        _ASTROPY_PACKAGE_NAME_ + '.tests': ['coveragerc', 'data/*.fits', 'data/*.hdr', 
+                                            'data/*.lmv', 'data/*reg', 'data/*.npz']
     }

spectral_cube/tests/test_spectral_cube.py

@@ -845,6 +845,46 @@ def test_yt_roundtrip_wcs(self):
         assert_allclose(ytc3.world2yt(world_coord3), yt_coord3.value)
         self.cube = self.ytc1 = self.ytc2 = self.ytc3 = None
 
+    def test_yt_vr(self):
+        filepath = os.path.abspath(os.path.dirname(__file__))
+
+        ds = self.ytc1.dataset
+
+        # Test one rendering
+
+        n_v = 10
+        vmin = 0.05
+        vmax = 4.0
+        dv = 0.02
+
+        transfer = yt.ColorTransferFunction((vmin, vmax))
+        transfer.add_layers(n_v, dv, colormap='RdBu_r')
+
+        direction = (0.0, 0.0, 1.0)
+        sc = yt.create_scene(ds, field='flux')
+        source = sc[0]
+        source.tfh.tf = transfer
+        source.tfh.bounds = (vmin, vmax)
+        cam = sc.camera
+        cam.switch_orientation(normal_vector=direction)
+        cam.set_resolution(32)
+        im1 = sc.render()
+        yt_vr1 = os.path.join(filepath, "data", "yt_vr1.npz")
+        with np.load(yt_vr1) as i1:
+            assert_array_equal(im1, i1['arr_0'])
+
+        # Test movie
+        im2 = self.ytc1.quick_render_movie('.', size=32, nframes=4,
+                                           camera_angle=(0.0, 0.0, 1.0), 
+                                           rot_vector=(0.0, 1.0, 0.0),
+                                           north_vector=(0.0, 1.0, 0.0),
+                                           run_ffmpeg=False)
+        im2 = np.array(im2)
+        yt_vr2 = os.path.join(filepath, "data", "yt_vr2.npz")
+        with np.load(yt_vr2) as i2:
+            assert_array_equal(im2, i2['arr_0'])
+        self.cube = self.ytc1 = self.ytc2 = self.ytc3 = None
+
 def test_read_write_rountrip(tmpdir, data_adv):
     cube = SpectralCube.read(data_adv)
     tmp_file = str(tmpdir.join('test.fits'))

spectral_cube/ytcube.py

@@ -11,6 +11,7 @@
 
 __all__ = ['ytCube']
 
+
 class ytCube(object):
     """ Light wrapper of a yt object with ability to translate yt<->wcs
     coordinates """
@@ -57,18 +58,19 @@ def yt2world(self, yt_coord, first_index=0):
         world_coord = self.wcs.wcs_pix2world([yt_coord], first_index)[0]
         return world_coord
 
-
     def quick_render_movie(self, outdir, size=256, nframes=30,
-                           camera_angle=(0,0,1), north_vector=(0,0,1),
-                           rot_vector=(1,0,0),
+                           camera_angle=(0,0,1),
+                           north_vector=(0,1,0),
+                           rot_vector=None,
+                           zoom=None,
                            colormap='doom',
                            cmap_range='auto',
                            transfer_function='auto',
                            start_index=0,
                            image_prefix="",
                            output_filename='out.mp4',
-                           log_scale=False,
-                           rescale=True):
+                           log_scale=False, run_ffmpeg=True,
+                           rescale=True, sigma_clip=None):
         """
         Create a movie rotating the cube 360 degrees from
         PP -> PV -> PP -> PV -> PP
@@ -86,87 +88,110 @@ def quick_render_movie(self, outdir, size=256, nframes=30,
         camera_angle: 3-tuple
             The initial angle of the camera
         north_vector: 3-tuple
-            The vector of 'north' in the data cube.  Default is coincident with
-            the spectral axis
+            The vector of 'north' in the data cube. Default is the "y" direction
         rot_vector: 3-tuple
-            The vector around which the camera will be rotated
+            The vector around which the camera will be rotated. Default: None,
+            which means it will be set to the north_vector value.
+        zoom : float
+            Change the width of the FOV of the camera. Default: None, which 
+            does no zooming. 
         colormap: str
             A valid colormap.  See `yt.show_colormaps`
+        cmap_range : 2-tuple or "auto"
+            If a 2-tuple of floats, this will be the (vmin, vmax) for the colorbar.
+            Otherise, "auto" sets the values from the data.
         transfer_function: 'auto' or `yt.visualization.volume_rendering.TransferFunction`
             Either 'auto' to use the colormap specified, or a valid
             TransferFunction instance
-        log_scale: bool
-            Should the colormap be log scaled?
-        rescale: bool
-            If True, the images will be rescaled to have a common 95th
-            percentile brightness, which can help reduce flickering from having
-            a single bright pixel in some projections
         start_index : int
             The number of the first image to save
         image_prefix : str
             A string to prepend to the image name for each image that is output
         output_filename : str
             The movie file name to output.  The suffix may affect the file type
             created.  Defaults to 'out.mp4'.  Will be placed in ``outdir``
-
-        Returns
-        -------
-
-
+        log_scale : bool
+            Should the colormap be log scaled?
+        run_ffmpeg : bool
+            If True, ffmpeg will be used to make a movie out of the images.
+            Default: True
+        rescale: bool
+            If True, the images will be rescaled to have a common 95th
+            percentile brightness, which can help reduce flickering from having
+            a single bright pixel in some projections
+        sigma_clip: float, optional
+            Image values greater than this number times the standard deviation
+            plus the mean of the image will be clipped before saving. Useful
+            for enhancing images as it gets rid of rare high pixel values.
+            Default: None
         """
         try:
             import yt
         except ImportError:
-            raise ImportError("yt could not be imported.  Cube renderings are not possible.")
-
-        scale = np.max(self.cube.shape)
+            raise ImportError("yt could not be imported. Cube renderings are not possible.")
+        else:
+            from packaging import version
+            if version.parse(yt.__version__) < version.parse("3.5.0"):
+                raise RuntimeError("Only yt versions >= 3.5.0 are supported. Please upgrade yt.")
 
         if not os.path.exists(outdir):
             os.makedirs(outdir)
         elif not os.path.isdir(outdir):
             raise OSError("Output directory {0} exists and is not a directory.".format(outdir))
 
+        if rot_vector is None:
+            rot_vector = north_vector
+
         if cmap_range == 'auto':
             upper = self.cube.max().value
             lower = self.cube.std().value * 3
             cmap_range = [lower,upper]
 
+        data_source = self.dataset.all_data()
+        sc = yt.create_scene(data_source, field='flux')
+        source = sc[0]
+
         if transfer_function == 'auto':
             tfh = self.auto_transfer_function(cmap_range, log=log_scale)
             tfh.tf.map_to_colormap(cmap_range[0], cmap_range[1], colormap=colormap)
-            tf = tfh.tf
+            source.tfh = tfh
         else:
             tf = transfer_function
+            source.tfh.tf = tf
+            source.tfh.bounds = cmap_range
 
-        center = self.dataset.domain_center
-        cam = self.dataset.h.camera(center, camera_angle, scale, size, tf,
-                                    north_vector=north_vector, fields='flux')
+        cam = sc.camera
+        cam.set_focus(data_source.get_field_parameter("center"))
+        cam.set_resolution(size)
+        cam.switch_view(normal_vector=camera_angle, north_vector=north_vector)
+        if zoom is not None:
+            cam.zoom(zoom)
 
-        im  = cam.snapshot()
+        im = sc.render()
         images = [im]
 
         pb = ProgressBar(nframes)
-        for ii,im in enumerate(cam.rotation(2 * np.pi, nframes,
-                                            rot_vector=rot_vector)):
+        for ii in cam.iter_rotate(2*np.pi, nframes, rot_vector=rot_vector):
+            im = sc.render()
             images.append(im)
-            im.write_png(os.path.join(outdir,"%s%04i.png" % (image_prefix,
-                                                             ii+start_index)),
-                         rescale=False)
+            outfile = os.path.join(outdir, "%s%04i.png" % (image_prefix,
+                                                           ii+start_index))
+            im.write_png(outfile, sigma_clip=sigma_clip, rescale=False)
             pb.update(ii+1)
         log.info("Rendering complete in {0}s".format(time.time() - pb._start_time))
 
         if rescale:
             _rescale_images(images, os.path.join(outdir, image_prefix))
 
-        pipe = _make_movie(outdir, prefix=image_prefix,
-                           filename=output_filename)
+        if run_ffmpeg:
+            pipe = _make_movie(outdir, prefix=image_prefix,
+                               filename=output_filename)
 
         return images
 
-    def auto_transfer_function(self, cmap_range, log=False, colormap='doom',
-                               **kwargs):
-
-        from yt.visualization.volume_rendering.transfer_function_helper import TransferFunctionHelper
+    def auto_transfer_function(self, cmap_range, log=False):
+        from yt.visualization.volume_rendering.transfer_function_helper import \
+            TransferFunctionHelper
         tfh = TransferFunctionHelper(self.dataset)
         tfh.set_field('flux')
         tfh.set_bounds(bounds=cmap_range)
@@ -262,6 +287,7 @@ def _rescale_images(images, prefix):
         image = image.rescale(cmax=cmax, amax=amax).swapaxes(0,1)
         image.write_png("%s%04i.png" % (prefix, i), rescale=False)
 
+
 def _make_movie(moviepath, prefix="", filename='out.mp4', overwrite=True):
     """
     Use ffmpeg to generate a movie from the image series