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)