def export_rgba(
image,
fn,
h5=True,
fits=False,
):
"""
This function accepts an *image*, of shape (N,M,4) corresponding to r,g,b,a,
and saves to *fn*. If *h5* is True, then it will save in hdf5 format. If
*fits* is True, it will save in fits format.
"""
if (not h5 and not fits) or (h5 and fits):
raise ValueError("Choose either HDF5 or FITS format!")
if h5:
f = h5py.File('%s.h5' % fn, "w")
f.create_dataset("R", data=image[:, :, 0])
f.create_dataset("G", data=image[:, :, 1])
f.create_dataset("B", data=image[:, :, 2])
f.create_dataset("A", data=image[:, :, 3])
f.close()
if fits:
from yt.visualization.fits_image import FITSImageData
data = {}
data["r"] = image[:, :, 0]
data["g"] = image[:, :, 1]
data["b"] = image[:, :, 2]
data["a"] = image[:, :, 3]
fib = FITSImageData(data)
fib.writeto('%s.fits' % fn, clobber=True)
def write_fits(self,
filename,
clobber=False,
length_unit=None,
sky_scale=None,
sky_center=None):
r""" Write the PPVCube to a FITS file.
Parameters
----------
filename : string
The name of the file to write to.
clobber : boolean, optional
Whether to overwrite a file with the same name that already
exists. Default False.
length_unit : string, optional
The units to convert the coordinates to in the file.
sky_scale : tuple, optional
Conversion between an angle unit and a length unit, if sky
coordinates are desired, e.g. (1.0, "arcsec/kpc")
sky_center : tuple, optional
The (RA, Dec) coordinate in degrees of the central pixel. Must
be specified with *sky_scale*.
Examples
--------
>>> cube.write_fits("my_cube.fits", clobber=False,
... sky_scale=(1.0,"arcsec/kpc"), sky_center=(30.,45.))
"""
vunit = fits_info[self.axis_type][0]
vtype = fits_info[self.axis_type][1]
v_center = 0.5 * (self.vbins[0] + self.vbins[-1]).in_units(vunit).value
if length_unit is None:
units = str(self.ds.get_smallest_appropriate_unit(self.width))
else:
units = length_unit
units = sanitize_fits_unit(units)
dx = self.width.in_units(units).v / self.nx
dy = self.width.in_units(units).v / self.ny
dv = self.dv.in_units(vunit).v
w = _astropy.pywcs.WCS(naxis=3)
w.wcs.crpix = [
0.5 * (self.nx + 1), 0.5 * (self.ny + 1), 0.5 * (self.nv + 1)
]
w.wcs.cdelt = [dx, dy, dv]
w.wcs.crval = [0.0, 0.0, v_center]
w.wcs.cunit = [units, units, vunit]
w.wcs.ctype = ["LINEAR", "LINEAR", vtype]
fib = FITSImageData(self.data.transpose(), fields=self.field, wcs=w)
fib.update_all_headers("bunit", re.sub('()', '', str(self.proj_units)))
fib.update_all_headers("btype", self.field)
if sky_scale is not None and sky_center is not None:
fib.create_sky_wcs(sky_center, sky_scale)
fib.writeto(filename, clobber=clobber)
def write_fits(self,
filename,
sky_scale=None,
sky_center=None,
overwrite=True,
**kwargs):
r""" Export images to a FITS file. Writes the SZ distortion in all
specified frequencies as well as the mass-weighted temperature and the
optical depth. Distance units are in kpc, unless *sky_center*
and *scale* are specified.
Parameters
----------
filename : string
The name of the FITS file to be written.
sky_scale : tuple
Conversion between an angle unit and a length unit, if sky
coordinates are desired, e.g. (1.0, "arcsec/kpc")
sky_center : tuple, optional
The (RA, Dec) coordinate in degrees of the central pixel. Must
be specified with *sky_scale*.
overwrite : boolean, optional
If the file already exists, do we overwrite?
Additional keyword arguments are passed to
:meth:`~astropy.io.fits.HDUList.writeto`.
Examples
--------
>>> # This example just writes out a FITS file with kpc coords
>>> szprj.write_fits("SZbullet.fits", overwrite=False)
>>> # This example uses sky coords
>>> sky_scale = (1., "arcsec/kpc") # One arcsec per kpc
>>> sky_center = (30., 45., "deg")
>>> szprj.write_fits("SZbullet.fits", sky_center=sky_center, sky_scale=sky_scale)
"""
from yt.visualization.fits_image import FITSImageData
dx = self.dx.in_units("kpc")
dy = dx
w = _astropy.pywcs.WCS(naxis=2)
w.wcs.crpix = [0.5 * (self.nx + 1)] * 2
w.wcs.cdelt = [dx.v, dy.v]
w.wcs.crval = [0.0, 0.0]
w.wcs.cunit = ["kpc"] * 2
w.wcs.ctype = ["LINEAR"] * 2
fib = FITSImageData(self.data, fields=self.data.keys(), wcs=w)
if sky_scale is not None and sky_center is not None:
fib.create_sky_wcs(sky_center, sky_scale)
fib.writeto(filename, overwrite=overwrite, **kwargs)
def export_fits(self,
filename,
fields=None,
overwrite=False,
other_keys=None,
units="cm",
**kwargs):
r"""Export a set of pixelized fields to a FITS file.
This will export a set of FITS images of either the fields specified
or all the fields already in the object.
Parameters
----------
filename : string
The name of the FITS file to be written.
fields : list of strings
These fields will be pixelized and output. If "None", the keys of the
FRB will be used.
overwrite : boolean
If the file exists, this governs whether we will overwrite.
other_keys : dictionary, optional
A set of header keys and values to write into the FITS header.
units : string, optional
the length units that the coordinates are written in, default 'cm'.
"""
from yt.visualization.fits_image import FITSImageData
if fields is None:
fields = list(self.data.keys())
else:
fields = ensure_list(fields)
if len(fields) == 0:
raise RuntimeError(
"No fields to export. Either pass a field or list of fields to "
"export_fits or access a field from the fixed resolution buffer "
"object.")
fib = FITSImageData(self, fields=fields, units=units)
if other_keys is not None:
for k, v in other_keys.items():
fib.update_all_headers(k, v)
fib.writeto(filename, overwrite=overwrite, **kwargs)
def test_fits_image():
tmpdir = tempfile.mkdtemp()
curdir = os.getcwd()
os.chdir(tmpdir)
fields = ("density", "temperature")
units = (
'g/cm**3',
'K',
)
ds = fake_random_ds(64,
fields=fields,
units=units,
nprocs=16,
length_unit=100.0)
prj = ds.proj("density", 2)
prj_frb = prj.to_frb((0.5, "unitary"), 128)
fid1 = FITSImageData(prj_frb,
fields=["density", "temperature"],
units="cm")
fits_prj = FITSProjection(ds,
"z", ["density", "temperature"],
image_res=128,
width=(0.5, "unitary"))
assert_equal(fid1.get_data("density"), fits_prj.get_data("density"))
assert_equal(fid1.get_data("temperature"),
fits_prj.get_data("temperature"))
fid1.writeto("fid1.fits", clobber=True)
new_fid1 = FITSImageData.from_file("fid1.fits")
assert_equal(fid1.get_data("density"), new_fid1.get_data("density"))
assert_equal(fid1.get_data("temperature"),
new_fid1.get_data("temperature"))
ds2 = load("fid1.fits")
ds2.index
assert ("fits", "density") in ds2.field_list
assert ("fits", "temperature") in ds2.field_list
dw_cm = ds2.domain_width.in_units("cm")
assert dw_cm[0].v == 50.
assert dw_cm[1].v == 50.
slc = ds.slice(2, 0.5)
slc_frb = slc.to_frb((0.5, "unitary"), 128)
fid2 = FITSImageData(slc_frb,
fields=["density", "temperature"],
units="cm")
fits_slc = FITSSlice(ds,
"z", ["density", "temperature"],
image_res=128,
width=(0.5, "unitary"))
assert_equal(fid2.get_data("density"), fits_slc.get_data("density"))
assert_equal(fid2.get_data("temperature"),
fits_slc.get_data("temperature"))
dens_img = fid2.pop("density")
temp_img = fid2.pop("temperature")
# This already has some assertions in it, so we don't need to do anything
# with it other than just make one
FITSImageData.from_images([dens_img, temp_img])
cut = ds.cutting([0.1, 0.2, -0.9], [0.5, 0.42, 0.6])
cut_frb = cut.to_frb((0.5, "unitary"), 128)
fid3 = FITSImageData(cut_frb,
fields=["density", "temperature"],
units="cm")
fits_cut = FITSOffAxisSlice(ds, [0.1, 0.2, -0.9],
["density", "temperature"],
image_res=128,
center=[0.5, 0.42, 0.6],
width=(0.5, "unitary"))
assert_equal(fid3.get_data("density"), fits_cut.get_data("density"))
assert_equal(fid3.get_data("temperature"),
fits_cut.get_data("temperature"))
fid3.create_sky_wcs([30., 45.], (1.0, "arcsec/kpc"))
fid3.writeto("fid3.fits", clobber=True)
new_fid3 = FITSImageData.from_file("fid3.fits")
assert_same_wcs(fid3.wcs, new_fid3.wcs)
assert new_fid3.wcs.wcs.cunit[0] == "deg"
assert new_fid3.wcs.wcs.cunit[1] == "deg"
assert new_fid3.wcs.wcs.ctype[0] == "RA---TAN"
assert new_fid3.wcs.wcs.ctype[1] == "DEC--TAN"
buf = off_axis_projection(ds, ds.domain_center, [0.1, 0.2, -0.9], 0.5, 128,
"density").swapaxes(0, 1)
fid4 = FITSImageData(buf, fields="density", width=100.0)
fits_oap = FITSOffAxisProjection(ds, [0.1, 0.2, -0.9],
"density",
width=(0.5, "unitary"),
image_res=128,
depth_res=128,
depth=(0.5, "unitary"))
assert_equal(fid4.get_data("density"), fits_oap.get_data("density"))
fid4.create_sky_wcs([30., 45.], (1.0, "arcsec/kpc"), replace_old_wcs=False)
assert fid4.wcs.wcs.cunit[0] == "cm"
assert fid4.wcs.wcs.cunit[1] == "cm"
assert fid4.wcs.wcs.ctype[0] == "linear"
assert fid4.wcs.wcs.ctype[1] == "linear"
assert fid4.wcsa.wcs.cunit[0] == "deg"
assert fid4.wcsa.wcs.cunit[1] == "deg"
assert fid4.wcsa.wcs.ctype[0] == "RA---TAN"
assert fid4.wcsa.wcs.ctype[1] == "DEC--TAN"
cvg = ds.covering_grid(ds.index.max_level, [0.25, 0.25, 0.25],
[32, 32, 32],
fields=["density", "temperature"])
fid5 = FITSImageData(cvg, fields=["density", "temperature"])
assert fid5.dimensionality == 3
fid5.update_header("density", "time", 0.1)
fid5.update_header("all", "units", "cgs")
assert fid5["density"].header["time"] == 0.1
assert fid5["temperature"].header["units"] == "cgs"
assert fid5["density"].header["units"] == "cgs"
os.chdir(curdir)
shutil.rmtree(tmpdir)