granttremblay/make_stars.py

## make_stars.py
def make_stars(kintable, fovimage, redshift, name, snthresh=100, velocity_thresh=None, project_by_median=True, project_by_redshift=False, cropbox=None, zoom=False, zoombox=None, wcs=True, nancolor='white', colorbar_scalefactor=0.047, vel_vmin=-500, vel_vmax=500, disp_vmin=0, disp_vmax=300, save=True, file_save_directory="./"):

    table = fits.getdata(kintable)
    #hdr = fits.getheader(fovimage)

    x = table['x_cor']
    y = table['y_cor']

    # Get the 2D dimensions into which you'll paint this data
    fovdata = fits.getdata(fovimage)
    dim = fovdata.shape

    # Make empty maps of NaNs
    velmap = np.full((dim[0], dim[1]), np.nan)
    dispmap = np.full((dim[0], dim[1]), np.nan)

    # Threshold & Crop
    if velocity_thresh is not None:
        mask = (table['vel_fit'] / table['vel_fit_err'] > snthresh) & (table['vel_fit'] < np.nanmedian(table['vel_fit'] + velocity_thresh)) & (table['vel_fit'] > np.nanmedian(table['vel_fit'] - velocity_thresh))
    else:
        mask = table['vel_fit'] / table['vel_fit_err'] > snthresh

    # YES, y,x, in that order. I know it's confusing.
    velmap[y[mask], x[mask]] = table['vel_fit'][mask]
    dispmap[y[mask], x[mask]] = table['disp_fit'][mask]

    if cropbox is not None:
        x1, x2, y1, y2 = cropbox

        # YES, you should be confused by the below line.
        # This is *inverted* for the actual mask, but NOT for the zoom.
        # So what the user would naturally expect to be x1 is actually y1
        keep = (x > x1) & (x < x2) & (y > y1) & (y < y2) # YES
        crop = np.logical_not(keep)

        velmap[x[crop], y[crop]] = np.nan
        dispmap[x[crop], y[crop]] = np.nan


    if project_by_redshift is True:
        if project_by_median is True:
            project_by_median = False
            print("project_by_redshift=True which overrides project_by_median=True")
        velmap = velmap - redshift * const.c.to(u.km / u.s).value
    if project_by_median is True:
        redshift = np.nanmedian(velmap) / const.c.to(u.km / u.s).value
        velmap = velmap - np.nanmedian(velmap)


    # Create the Figures

    if wcs is True:
        wcs = WCS(fits.getheader(fovimage, 1))

        velfig = plt.figure(1, figsize=(10,10))
        velax = velfig.add_subplot(111, projection=wcs)

        dispfig = plt.figure(2, figsize=(10,10))
        dispax = dispfig.add_subplot(111, projection=wcs)

        velax.coords[0].set_axislabel('Right Ascension')
        velax.coords[1].set_axislabel('Declination')

        dispax.coords[0].set_axislabel('Right Ascension')
        dispax.coords[1].set_axislabel('Declination')

    elif wcs is False:
        velfig = plt.figure(1, figsize=(10,10))
        velax = velfig.add_subplot(111)

        dispfig = plt.figure(2, figsize=(10,10))
        dispax = dispfig.add_subplot(111)

        velax.set_xlabel("X")
        velax.set_ylabel("Y")

        dispax.set_xlabel("X")
        dispax.set_ylabel("Y")

    velax.grid(False)
    dispax.grid(False)

    cmap_vel = cm.RdBu_r
    cmap_vel.set_bad(nancolor)

    cmap_disp = cm.plasma
    cmap_disp.set_bad(nancolor)

    if project_by_median is True or project_by_redshift is True:
        velframe = velax.imshow(velmap, origin='lower', vmin=vel_vmin, vmax=vel_vmax, cmap=cmap_vel, interpolation='nearest')
        velcbar = velfig.colorbar(velframe, fraction=colorbar_scalefactor, pad=0.01)
        velcbar.set_label(r"Stellar Velocity (km s$^{{-1}}$) relative to z = {}".format(round(redshift, 4)))
    else:
        velframe = velax.imshow(velmap, origin='lower', cmap=cmap_vel, interpolation='nearest')
        velcbar = velfig.colorbar(velframe, fraction=colorbar_scalefactor, pad=0.01)
        velcbar.set_label(r"Stellar Velocity (km s$^{-1}$)")
        vmin=None
        vmax=None

    dispframe = dispax.imshow(dispmap, origin='lower', vmin=disp_vmin, vmax=disp_vmax, cmap=cmap_disp, interpolation='nearest')
    dispcbar = dispfig.colorbar(dispframe, fraction=colorbar_scalefactor, pad=0.01)
    dispcbar.set_label(r"Stellar Velocity Dispersion (km s$^{-1}$)")

    if zoom is True:
        if zoombox is not None:
            x1, x2, y1, y2 = zoombox
            print("Zooming to {}".format(zoombox))
        elif zoombox is None and cropbox is not None:
            print("Using cropbox as zoombox")
        elif zoombox is None and cropbox is None:
            raise Exception("Zoom is TRUE but you don't have a crop or zoombox. You must specify at least one!")
        velax.set_xlim(x1, x2)
        velax.set_ylim(y1, y2)

        dispax.set_xlim(x1, x2)
        dispax.set_ylim(y1, y2)

    velax.grid(False)
    dispax.grid(False)

    # Save Everything

    if save is True:
        # Check that the file save directory exists. If not, create it.
        if not os.path.exists(file_save_directory):
            os.makedirs(file_save_directory)
            print("Creating file save directory: {}".format(file_save_directory))
        else:
            print("Found file save directory: {}".format(file_save_directory))

        # Save the PDF figure
        velfig_pdf_file = "{}_stellar_velocity.pdf".format(name.replace(" ", ""))
        velfig.savefig(file_save_directory + velfig_pdf_file, dpi=300, bbox_inches='tight')
        print("Saved Stellar Velocity Figure to {}".format(velfig_pdf_file))

        dispfig_pdf_file = "{}_stellar_dispersion.pdf".format(name.replace(" ", ""))
        dispfig.savefig(file_save_directory + dispfig_pdf_file, dpi=300, bbox_inches='tight')
        print("Saved Stellar Velocity Dispersion Figure to {}".format(dispfig_pdf_file))

        # Save the FITS file
        vel_fits_file = "{}_stellar_velocity.fits".format(name.replace(" ", ""))
        hdr = WCS(fits.getheader(fovimage, 1)).to_header()
        hdu = fits.PrimaryHDU(velmap, header=hdr)
        hdulist = fits.HDUList([hdu])
        hdulist.writeto(file_save_directory + vel_fits_file, overwrite=True, output_verify='silentfix')
        print("Saved Stellar Velocity Map FITS image to {}".format(vel_fits_file))


        # Save the FITS figure, along with a WCS

        disp_fits_file = "{}_stellar_dispersion.fits".format(name.replace(" ", ""))

        disp_fits_file = "{}_stellar_dispersion.fits".format(name.replace(" ", ""))
        hdu = fits.PrimaryHDU(dispmap, header=hdr)
        hdulist = fits.HDUList([hdu])
        hdulist.writeto(file_save_directory + disp_fits_file, overwrite=True, output_verify='silentfix')
        print("Saved Stellar Velocity Dispersion Map FITS image to {}".format(disp_fits_file))

        # Make Kinemetry Table

        kinemetry_table_filename = "{}_kinemetry_table.dat".format(name.replace(" ", ""))

        bin_number =np.arange(1, len(table["x_cor"][mask]) + 1)
        kinemetry_table = Table([bin_number,
                                table["x_cor"][mask],
                                table["y_cor"][mask],
                                table["vel_fit"][mask],
                                table["vel_fit_err"][mask],
                                table["disp_fit"][mask],
                                table["disp_fit_err"][mask]],
                                names=["#", "XBIN", "YBIN", "VEL", "ER_VEL", "SIG", "ER_SIG"])
        print("Saved Kinemetry Table to {}".format(kinemetry_table_filename))
        ascii.write(kinemetry_table, file_save_directory + kinemetry_table_filename, overwrite=True)


## make_stars2.py
make_stars(he0351_kintable,
                      he0351_fovimage,
                      he0351_redshift,
                      name = "HE 0351",
                      snthresh=50,
                      velocity_thresh=300, # in km/s
                      cropbox=(20, 180, 20, 180), # (x1, x2, y1, y2),
                      zoom=False, # It will zoom on the cropbox
                      wcs=True,
                      vel_vmin=-300,
                      vel_vmax=300,
                      disp_vmin=0,
                      disp_vmax=300,
                      project_by_median=True,
                      project_by_redshift=False,
                      save=True,
                      file_save_directory=file_save_directory)
	def make_stars(kintable, fovimage, redshift, name, snthresh=100, velocity_thresh=None, project_by_median=True, project_by_redshift=False, cropbox=None, zoom=False, zoombox=None, wcs=True, nancolor='white', colorbar_scalefactor=0.047, vel_vmin=-500, vel_vmax=500, disp_vmin=0, disp_vmax=300, save=True, file_save_directory="./"):

	table = fits.getdata(kintable)
	#hdr = fits.getheader(fovimage)

	x = table['x_cor']
	y = table['y_cor']

	# Get the 2D dimensions into which you'll paint this data
	fovdata = fits.getdata(fovimage)
	dim = fovdata.shape

	# Make empty maps of NaNs
	velmap = np.full((dim[0], dim[1]), np.nan)
	dispmap = np.full((dim[0], dim[1]), np.nan)

	# Threshold & Crop
	if velocity_thresh is not None:
	mask = (table['vel_fit'] / table['vel_fit_err'] > snthresh) & (table['vel_fit'] < np.nanmedian(table['vel_fit'] + velocity_thresh)) & (table['vel_fit'] > np.nanmedian(table['vel_fit'] - velocity_thresh))
	else:
	mask = table['vel_fit'] / table['vel_fit_err'] > snthresh

	# YES, y,x, in that order. I know it's confusing.
	velmap[y[mask], x[mask]] = table['vel_fit'][mask]
	dispmap[y[mask], x[mask]] = table['disp_fit'][mask]

	if cropbox is not None:
	x1, x2, y1, y2 = cropbox

	# YES, you should be confused by the below line.
	# This is inverted for the actual mask, but NOT for the zoom.
	# So what the user would naturally expect to be x1 is actually y1
	keep = (x > x1) & (x < x2) & (y > y1) & (y < y2) # YES
	crop = np.logical_not(keep)

	velmap[x[crop], y[crop]] = np.nan
	dispmap[x[crop], y[crop]] = np.nan


	if project_by_redshift is True:
	if project_by_median is True:
	project_by_median = False
	print("project_by_redshift=True which overrides project_by_median=True")
	velmap = velmap - redshift * const.c.to(u.km / u.s).value
	if project_by_median is True:
	redshift = np.nanmedian(velmap) / const.c.to(u.km / u.s).value
	velmap = velmap - np.nanmedian(velmap)


	# Create the Figures

	if wcs is True:
	wcs = WCS(fits.getheader(fovimage, 1))

	velfig = plt.figure(1, figsize=(10,10))
	velax = velfig.add_subplot(111, projection=wcs)

	dispfig = plt.figure(2, figsize=(10,10))
	dispax = dispfig.add_subplot(111, projection=wcs)

	velax.coords[0].set_axislabel('Right Ascension')
	velax.coords[1].set_axislabel('Declination')

	dispax.coords[0].set_axislabel('Right Ascension')
	dispax.coords[1].set_axislabel('Declination')

	elif wcs is False:
	velfig = plt.figure(1, figsize=(10,10))
	velax = velfig.add_subplot(111)

	dispfig = plt.figure(2, figsize=(10,10))
	dispax = dispfig.add_subplot(111)

	velax.set_xlabel("X")
	velax.set_ylabel("Y")

	dispax.set_xlabel("X")
	dispax.set_ylabel("Y")

	velax.grid(False)
	dispax.grid(False)

	cmap_vel = cm.RdBu_r
	cmap_vel.set_bad(nancolor)

	cmap_disp = cm.plasma
	cmap_disp.set_bad(nancolor)

	if project_by_median is True or project_by_redshift is True:
	velframe = velax.imshow(velmap, origin='lower', vmin=vel_vmin, vmax=vel_vmax, cmap=cmap_vel, interpolation='nearest')
	velcbar = velfig.colorbar(velframe, fraction=colorbar_scalefactor, pad=0.01)
	velcbar.set_label(r"Stellar Velocity (km s$^{{-1}}$) relative to z = {}".format(round(redshift, 4)))
	else:
	velframe = velax.imshow(velmap, origin='lower', cmap=cmap_vel, interpolation='nearest')
	velcbar = velfig.colorbar(velframe, fraction=colorbar_scalefactor, pad=0.01)
	velcbar.set_label(r"Stellar Velocity (km s$^{-1}$)")
	vmin=None
	vmax=None

	dispframe = dispax.imshow(dispmap, origin='lower', vmin=disp_vmin, vmax=disp_vmax, cmap=cmap_disp, interpolation='nearest')
	dispcbar = dispfig.colorbar(dispframe, fraction=colorbar_scalefactor, pad=0.01)
	dispcbar.set_label(r"Stellar Velocity Dispersion (km s$^{-1}$)")

	if zoom is True:
	if zoombox is not None:
	x1, x2, y1, y2 = zoombox
	print("Zooming to {}".format(zoombox))
	elif zoombox is None and cropbox is not None:
	print("Using cropbox as zoombox")
	elif zoombox is None and cropbox is None:
	raise Exception("Zoom is TRUE but you don't have a crop or zoombox. You must specify at least one!")
	velax.set_xlim(x1, x2)
	velax.set_ylim(y1, y2)

	dispax.set_xlim(x1, x2)
	dispax.set_ylim(y1, y2)

	velax.grid(False)
	dispax.grid(False)

	# Save Everything

	if save is True:
	# Check that the file save directory exists. If not, create it.
	if not os.path.exists(file_save_directory):
	os.makedirs(file_save_directory)
	print("Creating file save directory: {}".format(file_save_directory))
	else:
	print("Found file save directory: {}".format(file_save_directory))

	# Save the PDF figure
	velfig_pdf_file = "{}_stellar_velocity.pdf".format(name.replace(" ", ""))
	velfig.savefig(file_save_directory + velfig_pdf_file, dpi=300, bbox_inches='tight')
	print("Saved Stellar Velocity Figure to {}".format(velfig_pdf_file))

	dispfig_pdf_file = "{}_stellar_dispersion.pdf".format(name.replace(" ", ""))
	dispfig.savefig(file_save_directory + dispfig_pdf_file, dpi=300, bbox_inches='tight')
	print("Saved Stellar Velocity Dispersion Figure to {}".format(dispfig_pdf_file))

	# Save the FITS file
	vel_fits_file = "{}_stellar_velocity.fits".format(name.replace(" ", ""))
	hdr = WCS(fits.getheader(fovimage, 1)).to_header()
	hdu = fits.PrimaryHDU(velmap, header=hdr)
	hdulist = fits.HDUList([hdu])
	hdulist.writeto(file_save_directory + vel_fits_file, overwrite=True, output_verify='silentfix')
	print("Saved Stellar Velocity Map FITS image to {}".format(vel_fits_file))


	# Save the FITS figure, along with a WCS

	disp_fits_file = "{}_stellar_dispersion.fits".format(name.replace(" ", ""))

	disp_fits_file = "{}_stellar_dispersion.fits".format(name.replace(" ", ""))
	hdu = fits.PrimaryHDU(dispmap, header=hdr)
	hdulist = fits.HDUList([hdu])
	hdulist.writeto(file_save_directory + disp_fits_file, overwrite=True, output_verify='silentfix')
	print("Saved Stellar Velocity Dispersion Map FITS image to {}".format(disp_fits_file))

	# Make Kinemetry Table

	kinemetry_table_filename = "{}_kinemetry_table.dat".format(name.replace(" ", ""))

	bin_number =np.arange(1, len(table["x_cor"][mask]) + 1)
	kinemetry_table = Table([bin_number,
	table["x_cor"][mask],
	table["y_cor"][mask],
	table["vel_fit"][mask],
	table["vel_fit_err"][mask],
	table["disp_fit"][mask],
	table["disp_fit_err"][mask]],
	names=["#", "XBIN", "YBIN", "VEL", "ER_VEL", "SIG", "ER_SIG"])
	print("Saved Kinemetry Table to {}".format(kinemetry_table_filename))
	ascii.write(kinemetry_table, file_save_directory + kinemetry_table_filename, overwrite=True)
	make_stars(he0351_kintable,
	he0351_fovimage,
	he0351_redshift,
	name = "HE 0351",
	snthresh=50,
	velocity_thresh=300, # in km/s
	cropbox=(20, 180, 20, 180), # (x1, x2, y1, y2),
	zoom=False, # It will zoom on the cropbox
	wcs=True,
	vel_vmin=-300,
	vel_vmax=300,
	disp_vmin=0,
	disp_vmax=300,
	project_by_median=True,
	project_by_redshift=False,
	save=True,
	file_save_directory=file_save_directory)