Merge pull request #3456 from jzuhone/fix_join_fits_images

fix combining FITS images

yt/visualization/fits_image.py

@@ -735,12 +735,12 @@ def from_images(cls, image_list):
  w = first_image.wcs
  img_shape = first_image.shape
  data = []
- for is_first, _is_last, fid in mark_ends(image_list):
+ for fid in image_list:
  assert_same_wcs(w, fid.wcs)
  if img_shape != fid.shape:
  raise RuntimeError("Images do not have the same shape!")
  for hdu in fid.hdulist:
- if is_first:
+ if len(data) == 0:
  data.append(_astropy.pyfits.PrimaryHDU(hdu.data, header=hdu.header))
  else:
  data.append(_astropy.pyfits.ImageHDU(hdu.data, header=hdu.header))

yt/visualization/tests/test_fits_image.py

@@ -102,6 +102,13 @@ def test_fits_image():
  assert_equal(combined_fid.magnetic_unit, dens_img.magnetic_unit)
  assert_equal(combined_fid.current_time, dens_img.current_time)
 
+ # Make sure that we can combine FITSImageData instances with more
+ # than one image each
+ combined_fid2 = FITSImageData.from_images([combined_fid, combined_fid])
+ # Writing the FITS file ensures that we have assembled the images
+ # together correctly
+ combined_fid2.writeto("combined.fits", overwrite=True)
+
  cut = ds.cutting([0.1, 0.2, -0.9], [0.5, 0.42, 0.6])
  cut_frb = cut.to_frb((0.5, "unitary"), 128)