示例#1
0
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)
示例#2
0
文件: ppv_cube.py 项目: victorgabr/yt
    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)
示例#3
0
    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)
示例#4
0
    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 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)
示例#6
0
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
示例#7
0
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)
示例#8
0
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)