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 to_fits_data(self,
fields=None,
other_keys=None,
length_unit=None,
**kwargs):
r"""Export the fields in this FixedResolutionBuffer instance
to a FITSImageData instance.
This will export a set of FITS images of either the fields specified
or all the fields already in the object.
Parameters
----------
fields : list of strings
These fields will be pixelized and output. If "None", the keys of
the FRB will be used.
other_keys : dictionary, optional
A set of header keys and values to write into the FITS header.
length_unit : string, optional
the length units that the coordinates are written in. The default
is to use the default length unit of the dataset.
"""
from yt.visualization.fits_image import FITSImageData
if length_unit is None:
length_unit = self.ds.length_unit
if "units" in kwargs:
issue_deprecation_warning("The 'units' keyword argument has been "
"replaced by the 'length_unit' keyword "
"argument and the former has been "
"deprecated. Setting 'length_unit' "
"to 'units'.")
length_unit = kwargs.pop("units")
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 "
"to_fits_data or access a field from the FixedResolutionBuffer "
"object.")
fid = FITSImageData(self, fields=fields, length_unit=length_unit)
if other_keys is not None:
for k, v in other_keys.items():
fid.update_all_headers(k, v)
return fid
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 create_spectral_slabs(filename, slab_centers, slab_width, **kwargs):
r"""
Given a dictionary of spectral slab centers and a width in
spectral units, extract data from a spectral cube at these slab
centers and return a `FITSDataset` instance containing the different
slabs as separate yt fields. Useful for extracting individual
lines from a spectral cube and separating them out as different fields.
Requires the SpectralCube (https://spectral-cube.readthedocs.io/en/latest/)
library.
All keyword arguments will be passed on to the `FITSDataset` constructor.
Parameters
----------
filename : string
The spectral cube FITS file to extract the data from.
slab_centers : dict of (float, string) tuples or YTQuantities
The centers of the slabs, where the keys are the names
of the new fields and the values are (float, string) tuples or
YTQuantities, specifying a value for each center and its unit.
slab_width : YTQuantity or (float, string) tuple
The width of the slab along the spectral axis.
Examples
--------
>>> slab_centers = {'13CN': (218.03117, 'GHz'),
... 'CH3CH2CHO': (218.284256, 'GHz'),
... 'CH3NH2': (218.40956, 'GHz')}
>>> slab_width = (0.05, "GHz")
>>> ds = create_spectral_slabs("intensity_cube.fits",
... slab_centers, slab_width,
... nan_mask=0.0)
"""
from spectral_cube import SpectralCube
from yt.visualization.fits_image import FITSImageData
from yt.frontends.fits.api import FITSDataset
cube = SpectralCube.read(filename)
if not isinstance(slab_width, YTQuantity):
slab_width = YTQuantity(slab_width[0], slab_width[1])
slab_data = {}
field_units = cube.header.get("bunit", "dimensionless")
for k, v in slab_centers.items():
if not isinstance(v, YTQuantity):
slab_center = YTQuantity(v[0], v[1])
else:
slab_center = v
mylog.info("Adding slab field %s at %g %s" %
(k, slab_center.v, slab_center.units))
slab_lo = (slab_center - 0.5 * slab_width).to_astropy()
slab_hi = (slab_center + 0.5 * slab_width).to_astropy()
subcube = cube.spectral_slab(slab_lo, slab_hi)
slab_data[k] = YTArray(subcube.filled_data[:, :, :], field_units)
width = subcube.header["naxis3"] * cube.header["cdelt3"]
w = subcube.wcs.copy()
w.wcs.crpix[-1] = 0.5
w.wcs.crval[-1] = -0.5 * width
fid = FITSImageData(slab_data, wcs=w)
for hdu in fid:
hdu.header.pop("RESTFREQ", None)
hdu.header.pop("RESTFRQ", None)
ds = FITSDataset(fid, **kwargs)
return ds
def test_fits_image():
curdir = os.getcwd()
tmpdir = tempfile.mkdtemp()
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 = prj_frb.to_fits_data(
fields=[("gas", "density"), ("gas", "temperature")], length_unit="cm"
)
fits_prj = FITSProjection(
ds,
"z",
[ds.fields.gas.density, "temperature"],
image_res=128,
width=(0.5, "unitary"),
)
assert_equal(fid1["density"].data, fits_prj["density"].data)
assert_equal(fid1["temperature"].data, fits_prj["temperature"].data)
fid1.writeto("fid1.fits", overwrite=True)
new_fid1 = FITSImageData.from_file("fid1.fits")
assert_equal(fid1["density"].data, new_fid1["density"].data)
assert_equal(fid1["temperature"].data, new_fid1["temperature"].data)
assert_equal(fid1.length_unit, new_fid1.length_unit)
assert_equal(fid1.time_unit, new_fid1.time_unit)
assert_equal(fid1.mass_unit, new_fid1.mass_unit)
assert_equal(fid1.velocity_unit, new_fid1.velocity_unit)
assert_equal(fid1.magnetic_unit, new_fid1.magnetic_unit)
assert_equal(fid1.current_time, new_fid1.current_time)
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.0
assert dw_cm[1].v == 50.0
slc = ds.slice(2, 0.5)
slc_frb = slc.to_frb((0.5, "unitary"), 128)
fid2 = slc_frb.to_fits_data(
fields=[("gas", "density"), ("gas", "temperature")], length_unit="cm"
)
fits_slc = FITSSlice(
ds,
"z",
[("gas", "density"), ("gas", "temperature")],
image_res=128,
width=(0.5, "unitary"),
)
assert_equal(fid2["density"].data, fits_slc["density"].data)
assert_equal(fid2["temperature"].data, fits_slc["temperature"].data)
dens_img = fid2.pop("density")
temp_img = fid2.pop("temperature")
combined_fid = FITSImageData.from_images([dens_img, temp_img])
assert_equal(combined_fid.length_unit, dens_img.length_unit)
assert_equal(combined_fid.time_unit, dens_img.time_unit)
assert_equal(combined_fid.mass_unit, dens_img.mass_unit)
assert_equal(combined_fid.velocity_unit, dens_img.velocity_unit)
assert_equal(combined_fid.magnetic_unit, dens_img.magnetic_unit)
assert_equal(combined_fid.current_time, dens_img.current_time)
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 = cut_frb.to_fits_data(
fields=[("gas", "density"), ds.fields.gas.temperature], length_unit="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["density"].data, fits_cut["density"].data)
assert_equal(fid3["temperature"].data, fits_cut["temperature"].data)
fid3.create_sky_wcs([30.0, 45.0], (1.0, "arcsec/kpc"))
fid3.writeto("fid3.fits", overwrite=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=(0.5, "unitary"),
)
assert_equal(fid4["density"].data, fits_oap["density"].data)
fid4.create_sky_wcs([30.0, 45.0], (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 = cvg.to_fits_data(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"
fid6 = FITSImageData.from_images(fid5)
fid5.change_image_name("density", "mass_per_volume")
assert fid5["mass_per_volume"].name == "mass_per_volume"
assert fid5["mass_per_volume"].header["BTYPE"] == "mass_per_volume"
assert "mass_per_volume" in fid5.fields
assert "mass_per_volume" in fid5.field_units
assert "density" not in fid5.fields
assert "density" not in fid5.field_units
assert "density" in fid6.fields
assert_equal(fid6["density"].data, fid5["mass_per_volume"].data)
fid7 = FITSImageData.from_images(fid4)
fid7.convolve("density", (3.0, "cm"))
sigma = 3.0 / fid7.wcs.wcs.cdelt[0]
kernel = _astropy.conv.Gaussian2DKernel(x_stddev=sigma)
data_conv = _astropy.conv.convolve(fid4["density"].data.d, kernel)
assert_allclose(data_conv, fid7["density"].data.d)
# We need to manually close all the file descriptors so
# that windows can delete the folder that contains them.
ds2.close()
for fid in (fid1, fid2, fid3, fid4, fid5, fid6, fid7, new_fid1, new_fid3):
fid.close()
os.chdir(curdir)
shutil.rmtree(tmpdir)
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)