Beispiel #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)
Beispiel #2
0
    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)
Beispiel #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)
    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)
Beispiel #5
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)