tuningfork.py

import numpy as np
import pynbody
import math
import pickle
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib import cm
import scipy
from scipy.stats import binned_statistic_2d
import astropy
import skimage
from astropy.convolution import convolve, Tophat2DKernel

def rescale(array2D_inpixels, npixels, minx, maxx, miny, maxy):
#x goes from minx to maxx (0 is -25 kpc, 1000 is 25 kpc)
#y goes from maxy to miny (0 corresponds to 25 kpc, 1000 corresponds to -25kpc!)
    print("min and max x: "+str(minx)+", "+str(maxx))
    print("min and max y: "+str(miny)+", "+str(maxy))
    print("number of pixels: "+str(npixels))
    array2D = np.copy(array2D_inpixels).astype(float) #makes a copy because otherwise it changes the original value in the dictionary! which might be fine but also might have unintended consequences...
    xscale = float(npixels / (maxx - minx))
    yscale = float(npixels / (maxy - miny))
    
    hshift = 0.5*(1/xscale)
    vshift = 0.5*(1/yscale)
    
    array2D[:,0] = vshift+(array2D[:,0] / float(-yscale)) + maxy #center of y coordinate is (pixel / -scale) + maxy value + vertical shift
    array2D[:,1] = hshift+(array2D[:,1] / float(xscale)) - maxx
    return array2D 
        
class TuningFork:
    def __init__(self, sim_path, H2image_type, SFRimage_type, image_width, pixel_width_kpc, sphkernel='CubicSpline'):
        '''
        Here we initialize the tuning fork object, which will have the initial attributes: simulation This will be an object with the (initial) attributes sim_path, H2image_type, SFRimage_type, image_width_kpc, pixel_width_kpc. Other attributes will be set via methods. 
            
        '''

        ##### initial attributes ##########
        self.sim_path = sim_path
        self.H2image_type = H2image_type #string, either 'sph' or 'simple'
        self.SFRimage_type = SFRimage_type #string, either 'sph' or 'simple'
        self.image_width = image_width #image width in kpc units
        self.pixel_kpc = pixel_width_kpc #you will vary this. used with 'simple' image type
        self.Npixels = math.floor(self.image_width / self.pixel_kpc) #gives total number of pixels
        self.sphkernel = sphkernel #default cubic spline. BUT COULD CHANGE IN THE FUTURE GIVEN VDISP STUFF
        ######## attributes set by create_image() methods ########
        self.H2image = None #set by create_H2image()
        self.SFRimage = None #set by create_SFRimage()
        self.H2smooth_type = None #set by create_H2image()
        self.SFRsmooth_type = None #set by create_SFRimage()
        self.H2_lsmooth = None #set by create_H2image()
        self.SFR_lsmooth = None #set by create_SFRimage()

        ######## attributes set by get_peaks() ########
        self.minimum_H2distance = None #in pixel units, set by get_peaks()
        self.minimum_SFRdistance = None #in pixel units, set by get_peaks()
        self.minimum_H2threshold = None #physical units; set by get_peaks()
        self.minimum_SFRthreshold = None #physical units; set by get_peaks()
        self.H2peak_coords = None  #set by get_peaks()
        self.SFRpeak_coords = None
        self.H2peak_coords_physical = None
        self.SFRpeak_coords_physical = None

        ######## attributes set by calculate_fork() ########
        self.Lkpc = None #set by calculate_fork()
        self.depletion_gas = None #set by calculate_fork()
        self.depletion_sfr = None #set by calculate_fork()
        self.depletion_total = None #set by calculate_fork()

    def create_H2image(self, makeplot=True, H2smooth_type=None, H2_lsmooth=None, vmin=None, vmax=None):
        sim_path = self.sim_path
        half_width = self.image_width / 2.
        self.H2smooth_type = H2smooth_type
        self.H2_lsmooth = H2_lsmooth
        
        if self.H2image_type == "sph":
            pynbody.config['sph']['smooth-particles'] = 200
            pynbody.config['sph']['kernel'] = self.sphkernel
            s = pynbody.load(sim_path)
            s.physical_units()
            s.g['rhoH2'] = s.g['H2']*s.g['rho']
            s.g['rhoH2'].units = s.g['rho'].units

            imgwidth = str(self.image_width)+' kpc'
            H2image = pynbody.plot.sph.image(s.g, qty='rhoH2', units='Msol pc**-2', width=imgwidth, cmap='plasma', return_array=True, qtytitle=r'$\Sigma_{H_2}$ [M$_\odot$ pc$^{-2}$]', resolution=self.Npixels, vmin=vmin, vmax=vmax)                   
            self.H2image = H2image

        elif self.H2image_type == "simple":
            s = pynbody.load(sim_path)
            s.physical_units()

            pos_cut = (np.abs(s.g['pos'].in_units('kpc')[:,0] <= half_width) & (np.abs(s.g['pos'].in_units('kpc')[:,1]) <= half_width))
            Nbins = math.floor(2*half_width / self.pixel_kpc)
            bins = np.linspace(-1*half_width, half_width, Nbins+1)
            total_H2mass_perbin, xedge, yedge, binnum = binned_statistic_2d(s.g['pos'].in_units('kpc')[:,0][pos_cut], s.g['pos'].in_units('kpc')[:,1][pos_cut], s.g['mass'].in_units('Msol')[pos_cut]*s.g['H2'][pos_cut], statistic='sum', bins=bins)
            
            H2image = total_H2mass_perbin / ((self.pixel_kpc*1000.)**2.)
            if self.H2smooth_type == 'gaussian':
                H2_smooth = scipy.ndimage.gaussian_filter(H2image, self.H2_lsmooth)
                H2image = H2_smooth
            elif self.H2smooth_type == 'tophat':
                tophat2Dkernel = Tophat2DKernel(self.H2_lsmooth)
                H2_smooth = convolve(H2image, tophat2Dkernel)
                H2image = H2_smooth

            
            self.H2image = H2image
            if makeplot == True:
                if vmin is None:
                    vmin=1e-3
                else:
                    vmin=vmin
                if vmax is None:
                    vmax=10**math.ceil(np.log10(max(H2image.flatten())))
                else:
                    vmax=vmax
                vmax_pow = math.ceil(np.log10(max(H2image.flatten())))
                
                H2norm = mpl.colors.LogNorm(vmin=1e-2, vmax=10**vmax_pow)
                H2sm = plt.cm.ScalarMappable(cmap='plasma', norm=H2norm)

                #get rid of zero values bc it makes imshow show weird background noise
                zero_mask = (H2image == 0.)
                H2image[zero_mask] = 10**(np.log10(vmin) - 2) #2 orders of magnitude below the minimum value, but not zero!
                fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(4,3), facecolor='white', layout='tight')
                ax.imshow(H2image, norm=H2norm, cmap='plasma', extent=(-1*half_width, half_width, -1*half_width, half_width))
                ax.set_title(r'H$_2$ surface density map')
                ax.set_xlabel('x (kpc)')
                ax.set_ylabel('y (kpc)')
                H2cbar = fig.colorbar(H2sm, ax=ax, fraction=0.046, pad=0.04)
                H2cbar.set_label(r'$\Sigma_{H_2}$ (M$_\odot$ pc$^{-2}$)')
        return H2image

    def create_SFRimage(self, age_range=(3,10), vmin=None, vmax=None, makeplot=True, SFRsmooth_type=None, SFR_lsmooth=None): ###CAREFUL! specify age_range in Myrs. 
        sim_path = self.sim_path
        half_width = self.image_width/2.
        self.SFRsmooth_type = SFRsmooth_type
        self.SFR_lsmooth = SFR_lsmooth
        
        if self.SFRimage_type == "sph":
            age_norm = age_range[1] - age_range[0]
            @pynbody.derived_array
            def SFR(sim):
                return (1/age_norm)*sim.s['mass'].in_units('Msol')
                
            pynbody.config['sph']['smooth-particles'] = 200
            pynbody.config['sph']['kernel'] = self.sphkernel
            s = pynbody.load(sim_path)
            s.physical_units()
            
            s.s['SFR'].units = None
            s.s['SFR'].units = pynbody.units.Unit("Msol pc**-3") / pynbody.units.Unit("yr")

            ages = s.properties['time'].in_units('Myr') - s.s['tform'].in_units('Myr')
            age_cuts = (ages > age_range[0]) & (ages < age_range[1])

            
            imgwidth = str(self.image_width)+' kpc'
            
        
            SFRimage = pynbody.plot.image(s.s[age_cuts], qty='SFR', qtytitle=r"SFR (M$_\odot$ yr$^{-1}$ pc$^{-2}$)", units="Msol yr**-1 pc**-2", width=imgwidth, cmap='plasma', return_array=True, resolution=self.Npixels, vmin=vmin, vmax=vmax) 
            self.SFRimage = SFRimage
            
        elif self.SFRimage_type == "simple":
            
            #define half the image size, number of bins to use, and the bin edges (to be used later)
            
            Nbins = math.floor(self.image_width / self.pixel_kpc)
            print("number of bins: "+str(Nbins))
            bins = np.linspace(-1*half_width, half_width, Nbins+1)
            
            s = pynbody.load(sim_path)
            s.physical_units()
            #get stellar positions and masses
            stellar_positions = s.s['pos'].in_units('kpc')
            stellar_masses = s.s['mass'].in_units('Msol')
            
            #get the time of the snapshot and then the ages of the stars
            sim_age = s.properties['time'].in_units('Myr')
            stellar_ages = sim_age - s.s['tform'].in_units('Myr')
            
            #create a mask that takes only stars between the specified range
            age_cut = (stellar_ages <= age_range[1]) & (stellar_ages >= age_range[0])

            #apply mask to the positions and masses
            stellar_positions = stellar_positions[age_cut]
            stellar_masses = stellar_masses[age_cut]

            #create a msk that takes only positions within the specified image size
            pos_cut = (np.abs(stellar_positions[:,0]) <= half_width) & (np.abs(stellar_positions[:,1]) <= half_width)

            #apply mask to positions and masses
            stellar_positions = stellar_positions[pos_cut]
            stellar_masses = stellar_masses[pos_cut]

            #use scipy binned_statistic_2d to find the sum of masses in eaach bin. 
            total_Mstar_perbin, x_edge, y_edge, binnum = binned_statistic_2d(stellar_positions[:,0], stellar_positions[:,1], stellar_masses, statistic='sum', bins=bins, expand_binnumbers=True)
            total_SFR_perbin = (total_Mstar_perbin / ((age_range[1] - age_range[0])*1e6)) / (self.pixel_kpc*self.pixel_kpc)
                  
            SFRimage = total_SFR_perbin
            if self.SFRsmooth_type == 'gaussian':
                SFR_smooth = scipy.ndimage.gaussian_filter(SFRimage, self.SFR_lsmooth)
                SFRimage = SFR_smooth
            elif self.SFRsmooth_type == 'tophat':
                tophat2Dkernel = Tophat2DKernel(self.SFR_lsmooth)
                SFR_smooth = convolve(SFRimage, tophat2Dkernel)
                SFRimage = SFR_smooth
            self.SFRimage = SFRimage
            
            if makeplot == True:
                if vmin is None:
                    vmin=1e-3
                else:
                    vmin=vmin
                if vmax is None:
                    vmax=10**math.ceil(np.log10(max(SFRimage.flatten())))
                else:
                    vmax=vmax
                SFRnorm = mpl.colors.LogNorm(vmin=vmin, vmax=vmax)
                SFRsm = plt.cm.ScalarMappable(cmap='plasma', norm=SFRnorm)

                #get rid of zero values because it makes imshow plot things weirdly
                zero_mask = (SFRimage == 0.)
                SFRimage[zero_mask] = 10**(np.log10(vmin) - 2) #2 orders of magnitude below the minimum value, but not zero!
                
                fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(4,3), facecolor='white', layout='tight')
                ax.imshow(SFRimage, norm=SFRnorm, cmap='plasma', extent=(-1*half_width, half_width, -1*half_width, half_width))
                ax.set_title(r'SFR surface density map')
                ax.set_xlabel('x (kpc)')
                ax.set_ylabel('y (kpc)')
                SFRcbar = fig.colorbar(SFRsm, ax=ax, fraction=0.046, pad=0.04)
                SFRcbar.set_label(r'$\dot{{\Sigma}}_{{SFR}}$ (M$_{{\star, ({mn:2.0f}-{mx:2.0f} Myr}}$ / $\Delta t$  kpc$^{{-2}}$)'.format(mn=age_range[0], mx=age_range[1]))
                                            
        return SFRimage
        
        
    def create_circular_mask(self, center, radius):
        ######### center is a list of points in (Y,X) given by the get_peaks method
        ######### radius is L in pixels
        Y, X = np.ogrid[:self.Npixels, :self.Npixels] #make a grid the size of the image
        dist_from_center = np.sqrt((X - center[1])**2 + (Y - center[0])**2) #distance of each pixel from the point specified via center
        mask = (dist_from_center <= radius/2) #divide by two because L is aperture width (not radius)
        return mask
        

        
    def get_peaks(self, **kwargs):
        # minimum distance is in pixel units -- minimum distance specified between various peaks. THIS IS SOMETHING YOU WILL VARY
        if 'min_H2dist' in kwargs:
            self.minimum_H2distance = int(kwargs['min_H2dist'])
        if 'H2_thresh' in kwargs:
            self.minimum_H2threshold = kwargs['H2_thresh']
        if 'min_SFRdist' in kwargs:
            self.minimum_SFRdistance = int(kwargs['min_SFRdist'])
        if 'SFR_thresh' in kwargs:
            self.minimum_SFRthreshold = kwargs['SFR_thresh']
                
        if self.H2image_type == 'sph':
            H2image = self.H2image[::-1,:] #rotate image when it's sph so resulting coordinates are valid
        else:
            H2image = self.H2image

        if self.SFRimage_type == 'sph':
            SFRimage = self.SFRimage[::-1,:]
        else:
            SFRimage = self.SFRimage
        
        H2peaks = skimage.feature.peak_local_max(H2image, min_distance=self.minimum_H2distance, threshold_abs=self.minimum_H2threshold)
        print("minimum H2 threshold: "+str(self.minimum_H2threshold))
        SFRpeaks = skimage.feature.peak_local_max(SFRimage, min_distance=self.minimum_SFRdistance, threshold_abs=self.minimum_SFRthreshold)
        self.H2peak_coords = H2peaks
        self.SFRpeak_coords = SFRpeaks
        
        H2peaks_rescale = rescale(array2D_inpixels=H2peaks, npixels=self.Npixels, minx=-0.5*self.image_width, maxx=0.5*self.image_width, miny=-0.5*self.image_width, maxy=0.5*self.image_width)
        
        SFRpeaks_rescale = rescale(SFRpeaks, self.Npixels, -0.5*self.image_width, 0.5*self.image_width, -0.5*self.image_width, 0.5*self.image_width)
        
        self.H2peak_coords_physical = H2peaks_rescale
        self.SFRpeak_coords_physical = SFRpeaks_rescale
        
        return None
    
    def plot_peaks(self, fname, **kwargs):
        #use after you run get_peaks() and create_[]image()
        if 'min_H2dist' in kwargs:
            self.minimum_H2distance = int(kwargs['min_H2dist'])
        if 'H2_thresh' in kwargs:
            self.minimum_H2threshold = kwargs['min_H2thresh']
        if 'min_SFRdist' in kwargs:
            self.minimum_SFRdistance = int(kwargs['min_SFRdist'])
        if 'SFR_thresh' in kwargs:
            self.minimum_SFRthreshold = kwargs['SFR_thresh']
            
        if self.H2image_type == 'sph':
            H2image = self.H2image[::-1,:] #if the image type is sph, you gotta transform the image. basically, reverse the rows (or y axes) so that the origin is the bottom left corner of the plot.  
        
        else:
            H2image = self.H2image
            
        if self.SFRimage_type == 'sph':
            SFRimage = self.SFRimage[::-1,:]
            
        else:
            SFRimage = self.SFRimage

        if (self.H2peak_coords_physical is None) or (self.SFRpeak_coords_physical is None):
            self.get_peaks(min_H2dist=minH2dist, H2_threshold=H2_threshold, min_SFRdist=min_SFRdist, SFR_threshold=SFR_threshold)
        
        H2norm = mpl.colors.LogNorm(vmin=1e-2, vmax=10**math.ceil(np.log10(max(H2image.flatten()))))
        SFRnorm = mpl.colors.LogNorm(vmin=1e-3, vmax=10**math.ceil(np.log10(max(SFRimage.flatten()))))

        H2sm = plt.cm.ScalarMappable(cmap='plasma', norm=H2norm)
        SFRsm = plt.cm.ScalarMappable(cmap='plasma', norm=SFRnorm)
        
        fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(7,3), facecolor='white', layout='tight')
        half_width = self.image_width / 2.
        ax[0].imshow(H2image, norm=H2norm, cmap='plasma', extent=(-1*half_width, half_width, -1*half_width, half_width))
        ax[0].set_title(r'H$_2$ surface density map')
        ax[0].set_xlabel('x (kpc)')
        ax[0].set_ylabel('y (kpc)')
        ax[0].plot(self.H2peak_coords_physical[:,1], self.H2peak_coords_physical[:,0], 'ro', markersize=3.)

        ax[1].imshow(SFRimage, norm=SFRnorm, cmap='plasma', extent=(-1*half_width, half_width, -1*half_width, half_width))
        ax[1].set_title(r'SFR surface density map')
        ax[1].set_xlabel('x (kpc)')
        ax[1].set_ylabel('y (kpc)')
        ax[1].plot(self.SFRpeak_coords_physical[:,1], self.SFRpeak_coords_physical[:,0], 'ro', markersize=3.)

        fig.savefig(fname, dpi=300)
        
    def calculate_fork(self, minL, maxL, nL):
        H2image = self.H2image
        SFRimage = self.SFRimage
        print('got images of the objects...')
        Lkpc = np.logspace(np.log10(minL), np.log10(maxL), nL)
        Lpixel = self.Npixels*Lkpc / self.image_width #L in pixel size

        if self.H2peak_coords is None:
            print('no peaks for object! Please call TuningFork.get_peaks() :)')
            raise Exception("no peaks for object! Please call TuningFork.get_peaks() on the H2 image:)")
            # print('enter minimum H2 distance (in pixel units): ')
            # min_H2dist = input()
            # print('enter minimum H2 threshold: ')
            # min_H2thresh = input()
            # print('enter minimum SFR distance (in pixel units): ')
            # min_SFRdist = input()
            # print('enter minimum SFR threshold: ')
            # min_SFRthresh = input()
            # self.get_peaks(min_H2dist, min_H2thresh, min_SFRdist, min_SFRthresh)
            # print('got the peaks of the objects...')
        else:
            print('peaks already defined, grabbing...')
            H2peaks = self.H2peak_coords
            SFRpeaks = self.SFRpeak_coords
            print('got peaks')
        pixel_area_pc = (self.pixel_kpc*1000)**2. #pixel width in kpc, converted to pc, then squared
        
        depletion_gas = []
        depletion_sfr = []
        for l in range(len(Lpixel)):
            L = Lpixel[l]
            H2vals_perL_H2peak = []
            SFRvals_perL_H2peak = []
            for i in range(len(H2peaks)):
                print('on '+str(i)+' out of '+str(len(H2peaks))+' peaks')
                H2_cent = H2peaks[i] #in (Y, X) format
                image_mask = self.create_circular_mask(H2_cent, L)
                
                H2val_L = H2image[image_mask].sum()*pixel_area_pc / (len(H2image[image_mask].flatten())*pixel_area_pc)
                SFRval_L = SFRimage[image_mask].sum()*pixel_area_pc / (len(SFRimage[image_mask].flatten())*pixel_area_pc)

                H2vals_perL_H2peak.append(H2val_L)
                SFRvals_perL_H2peak.append(SFRval_L)

            H2_avg_H2peak = np.average(H2vals_perL_H2peak)
            SFR_avg_H2peak = np.average(SFRvals_perL_H2peak)
            depletion_gas.append(H2_avg_H2peak / SFR_avg_H2peak)

            H2vals_perL_SFRpeak = []
            SFRvals_perL_SFRpeak = []
            
            for j in range(len(SFRpeaks)):
                print('on '+str(j)+' out of '+str(len(SFRpeaks))+' peaks')
                SFR_cent = SFRpeaks[j]
                image_mask = self.create_circular_mask(SFR_cent, L)

                H2val_L = H2image[image_mask].sum()*pixel_area_pc / (len(H2image[image_mask].flatten())*pixel_area_pc)
                SFRval_L = SFRimage[image_mask].sum()*pixel_area_pc / (len(SFRimage[image_mask].flatten())*pixel_area_pc)

                H2vals_perL_SFRpeak.append(H2val_L)
                SFRvals_perL_SFRpeak.append(SFRval_L)

            H2_avg_SFRpeak = np.average(H2vals_perL_SFRpeak)
            SFR_avg_SFRpeak = np.average(SFRvals_perL_SFRpeak)
            depletion_sfr.append(H2_avg_SFRpeak / SFR_avg_SFRpeak)
            pct_done = 100*(l / len(Lpixel))
            print(str(pct_done)+'% done with tuning fork calculation')
        H2_total = H2image.sum()*pixel_area_pc / (len(H2image.flatten())*pixel_area_pc)
        SFR_total = SFRimage.sum()*pixel_area_pc / (len(SFRimage.flatten())*pixel_area_pc)
        depletion_total = H2_total / SFR_total

        self.Lkpc = Lkpc
        self.depletion_gas = depletion_gas
        self.depletion_sfr = depletion_sfr
        self.depletion_total = depletion_total
        
        # self.depletion_gas = depletion_gas
        # self.depletion_sfr = depletion_sfr
        # self.depletion_total = depletion_total

    def tuningfork_plot(self, figtitle, fname):
        depletion_norm = self.depletion_total
        depletion_gas = self.depletion_gas
        depletion_sfr = self.depletion_sfr
        Lkpc = self.Lkpc

        fig, ax = plt.subplots(figsize=(4,3))
        ax.plot(Lkpc, np.repeat(depletion_norm/depletion_norm, len(Lkpc)), label='total')
        ax.plot(Lkpc, depletion_gas/depletion_norm, label='gas peaks')
        ax.plot(Lkpc, depletion_sfr/depletion_norm, label='sfr peaks')

        ax.set_xscale('log')
        ax.set_yscale('log')
        ax.set_xlabel('L (kpc)')
        ax.set_ylabel(r'$\tau (L) / \tau_{dep}$')

        ax.legend()
        ax.set_title(figtitle)

        fig.savefig(fname, dpi=300)