diff --git a/.gitignore b/.gitignore index f8dd5e260..be9f9e604 100644 --- a/.gitignore +++ b/.gitignore @@ -54,6 +54,7 @@ docs/_build/ docs/api *.fits +*/tests/*.png # Other generated stuff */version.py diff --git a/docs/yt_example.rst b/docs/yt_example.rst index f248db33a..9c3002d19 100644 --- a/docs/yt_example.rst +++ b/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 -`_. +`_. Movie Making ------------ @@ -128,9 +147,8 @@ Example: .. raw:: html - + SketchFab Isosurface Contours ----------------------------- diff --git a/setup.cfg b/setup.cfg index 1504165f8..d7c43743d 100644 --- a/setup.cfg +++ b/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/* diff --git a/spectral_cube/tests/data/yt_vr1.npz b/spectral_cube/tests/data/yt_vr1.npz new file mode 100644 index 000000000..441a5340c Binary files /dev/null and b/spectral_cube/tests/data/yt_vr1.npz differ diff --git a/spectral_cube/tests/data/yt_vr2.npz b/spectral_cube/tests/data/yt_vr2.npz new file mode 100644 index 000000000..66671fbf9 Binary files /dev/null and b/spectral_cube/tests/data/yt_vr2.npz differ diff --git a/spectral_cube/tests/setup_package.py b/spectral_cube/tests/setup_package.py index a5a5a412f..662dba479 100644 --- a/spectral_cube/tests/setup_package.py +++ b/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'] } diff --git a/spectral_cube/tests/test_spectral_cube.py b/spectral_cube/tests/test_spectral_cube.py index 3c56bc44b..d63b3f21c 100644 --- a/spectral_cube/tests/test_spectral_cube.py +++ b/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')) diff --git a/spectral_cube/ytcube.py b/spectral_cube/ytcube.py index 2102d0afe..d48fec0c3 100644 --- a/spectral_cube/ytcube.py +++ b/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,21 +88,21 @@ 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 @@ -108,65 +110,88 @@ def quick_render_movie(self, outdir, size=256, nframes=30, 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