DMVisualization.py

import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
import pickle
import yt
import ipywidgets as widgets
from ast import literal_eval as make_tuple

SIGOsToInvestigate = [] # This list is generated by you clicking the buttons of interesting SIGOs in the output

def on_button_clicked(b):
    madeTuple = make_tuple(b.description)
    if madeTuple not in SIGOsToInvestigate:
        SIGOsToInvestigate.append(madeTuple)
        
def rgb_to_hex(rgb):
    return '#%02x%02x%02x' %rgb


print('Started')

filename = "D:/Star_Movie_768/"
res = '14Mpc'
vel = 'Sig2'
s_res = res
s_vel = vel

fb = 0.6



SIGOsToPlot = range(len(filteredSIGOsWithStars)) 
# filteredSIGOsWithStars is a 3D list: data takes the form of snap number, object idx from shrinker, grouped by 
# SIGO across snaps. You can change the format of this as need be, just change the two for loops below to fit your use.
# The important thing is that snap is the number of the snap data, SOI is the idx of the SIGO in shrinker

for SIGOnum in SIGOsToPlot:
    
    for snap,SOI in filteredSIGOsWithStars[SIGOnum]:
        
        button = widgets.Button(description = str((snap,SOI)))
        button.on_click(on_button_clicked)
        display(button)

        
        with open(filename + 'shrinker'+res+'_'+vel+'_'+str(snap)+'.dat','rb') as f:
            shrunken = pickle.load(f, encoding='latin1')
        centerCodeUnits = shrunken['cm'][SOI]
        

        snapnum = snap


        ds = yt.load(filename + "snap_"+str(snapnum).zfill(3)+".hdf5")
        redshift = ds.current_redshift

        with open(filename + 'SIGOidx'+res+'_'+vel+'_'+str(snapnum)+'_fb'+str(fb)+'.dat','rb') as f:
            SIGOidx = pickle.load(f, encoding='latin1')

        

        with open(filename + "HaloData"+s_res+"_"+s_vel+"_"+str(snapnum)+".dat", 'rb') as f:
            HaloData = pickle.load(f, encoding='latin1')
        # HaloData is a helper file. Make a dictionary with 2 keys: HaloData['cmHalo'][halonum] and HaloData['R200dm'][halonum] all in code units. Should be easy to make from the DM FOF. Just used for showing the size and location of halos, optional really.
    
        
        width = 6. * (.71 * (redshift + 1.)) # Sets the size of the box you're visualizing
        widthwithoutbuffer = 6.
        Around_SIGO = ds.region(centerCodeUnits, np.array(centerCodeUnits) - width/4., np.array(centerCodeUnits) + width/4.)

        #Generate the DM Density field
        ad = Around_SIGO



        fig, axs = plt.subplots(1, 2)
        plotwidth = 0.49
        fig.set_size_inches([16,6])
        (ax1, ax2) = axs
        # This part makes the gas density plot, left
        p = yt.ProjectionPlot(ds, "z", ("gas", "density"), center=centerCodeUnits, width=(widthwithoutbuffer/2,"kpc"),data_source=Around_SIGO) 
        p.set_cmap(field=("gas", "density"), cmap="magma")
        center=np.array(centerCodeUnits)  / (.71 * (redshift + 1.))
        print(center)
    #     if len(starLocations) > 0:
    #         for star in starLocations:
    #             p.annotate_marker(star, coord_system="data", plot_args={"color": 'black', "s": 60})

        plot = p.plots[('gas', 'density')]
        plot.figure = fig
        plot.axes = ax1
        p._setup_plots()

        ax = ax1
        print(ax.get_xlim())
        print(ax.get_ylim())


        mpl.rc('font', size=26) #controls default text size  
        p.set_zlim(("gas", "density"), zmin='min',zmax= .4)


        rect = (0., 0., 1., 1.)

        pos = ax1.get_position()
        print(pos)
        pos.x0 = 0.
        pos.x1 = plotwidth       # for example 0.2, choose your value
        pos.y0 = .5 - plotwidth / 2.
        pos.y1 = .5 + plotwidth / 2.
        ax1.set_position(pos)

        new_ax = fig.add_axes(ax1.get_position(), frameon=False)

        new_ax.set_xlim([0,1])
        new_ax.set_ylim([0,1])

        new_ax.set_xlim(ax.get_xlim())
        new_ax.set_ylim(ax.get_ylim())



        new_ax.patch.set_alpha(0.)
        new_ax.set_alpha(0.)
        new_ax.get_xaxis().set_ticks([])
        new_ax.get_yaxis().set_ticks([])
        
        res = '14Mpc'
        vel = 'Sig2'

        fb = 0.6





       
        rotation = shrunken['rotation']
        radii = shrunken['radii']  / (.71 * (redshift + 1.))

        listOfSIGOs = []
        for i in SIGOidx:
            listOfSIGOs.append(i)

        viewXMin = center[0] - widthwithoutbuffer / 4.
        viewXMax = center[0] + widthwithoutbuffer / 4.
        viewYMin = center[1] - widthwithoutbuffer / 4.
        viewYMax = center[1] + widthwithoutbuffer / 4.
        viewZMin = center[2] - widthwithoutbuffer / 4.
        viewZMax = center[2] + widthwithoutbuffer / 4.

        view = [0,0,0]
        Rx = np.array([[1,0,0],[0,np.cos(view[1]),-np.sin(view[1])],[0,np.sin(view[1]),np.cos(view[1])]])
        Ry = np.array([[np.cos(view[0]),0,np.sin(view[0])],[0,1,0],[-np.sin(view[0]),0,np.cos(view[0])]])
        Rz = np.array([[np.cos(view[2]),-np.sin(view[2]),0],[np.sin(view[2]),np.cos(view[2]),0],[0,0,1]])
        viewR = np.matmul(np.matmul(Rx,Ry),Rz)

        

        rotation = shrunken['rotation']
        radii = shrunken['radii']  / (.71 * (redshift + 1.))


        gasLen = len(shrunken['cm'])
        batch = 0


        if SOI >= 0:
            listOfSIGOs = [SOI]
            for SIGO in listOfSIGOs:
                # This visualizes the SIGO you want to look at in yellow
                color = 'yellow'
                
                cmSIGO = (shrunken['cm'][SIGO])  / (.71 * (redshift + 1.))
                if cmSIGO[0] < viewXMin or cmSIGO[0] > viewXMax or cmSIGO[1] < viewYMin or cmSIGO[1] > viewYMax or cmSIGO[2] < viewZMin or cmSIGO[2] > viewZMax:
                    print('SIGO number: ', SIGO)
                    print(cmSIGO)
                    print(center)
                    print(radii[SIGO])
                    continue

                cmSIGO = (shrunken['cm'][SIGO] - centerCodeUnits)  / (.71 * (redshift + 1.)) #- centerCodeUnits
                matrix = np.array([[1/radii[SIGO][0]**2,0.,0.],[0.,1/radii[SIGO][1]**2,0.],[0.,0.,1/radii[SIGO][2]**2]])
                matrix = np.matmul(rotation[SIGO].T,np.matmul(matrix,rotation[SIGO]))
                matrix = np.matmul(viewR.T,np.matmul(matrix,viewR))
                matrix = np.array([[matrix[0][0],matrix[0][1],matrix[0][2],0],[matrix[1][0],matrix[1][1],matrix[1][2],0],[matrix[2][0],matrix[2][1],matrix[2][2],0],[0,0,0,-1]])
                projection = np.array([[1.,0.,0.,0.],[0.,1.,0.,0.],[0.,0.,0.,1.]])
                outline = np.linalg.inv(np.matmul(np.matmul(projection,np.linalg.inv(matrix)),projection.T))
                x = -np.linspace(-1.01*radii[SIGO][2],1.01*radii[SIGO][2],5000)
                y = np.linspace(-1.01*radii[SIGO][2],1.01*radii[SIGO][2],5000)
                X,Y = np.meshgrid(x,y)
                eqn = outline[0][0] * (X)**2 + (outline[0][1] + outline[1][0]) * (X)*(Y) + outline[1][1] * (Y )**2
                Z=1
                new_ax.contour((X+cmSIGO[0]),(Y+cmSIGO[1]),eqn,[Z],colors=color,linewidths=2.)


        for gasObject in range(len(shrunken['cm'])):
            # This visualizes all other gas objects, colored by gas fraction. Red = dm-rich, green = gas-rich
            if gasObject == SOI:
                continue
            gasFrac = shrunken['gasFrac'][gasObject]
            r,g,b,a = plt.cm.PiYG(gasFrac)
            color = rgb_to_hex((int(256*r),int(256*g),int(256*b)))

            cmSIGO = (shrunken['cm'][gasObject])  / (.71 * (redshift + 1.))
            if cmSIGO[0] < viewXMin or cmSIGO[0] > viewXMax or cmSIGO[1] < viewYMin or cmSIGO[1] > viewYMax or cmSIGO[2] < viewZMin or cmSIGO[2] > viewZMax:
                continue

            cmSIGO = (shrunken['cm'][gasObject] - centerCodeUnits)  / (.71 * (redshift + 1.))
            matrix = np.array([[1/radii[gasObject][0]**2,0.,0.],[0.,1/radii[gasObject][1]**2,0.],[0.,0.,1/radii[gasObject][2]**2]])
            matrix = np.matmul(rotation[gasObject].T,np.matmul(matrix,rotation[gasObject]))
            matrix = np.matmul(viewR.T,np.matmul(matrix,viewR))
            matrix = np.array([[matrix[0][0],matrix[0][1],matrix[0][2],0],[matrix[1][0],matrix[1][1],matrix[1][2],0],[matrix[2][0],matrix[2][1],matrix[2][2],0],[0,0,0,-1]])
            projection = np.array([[1.,0.,0.,0.],[0.,1.,0.,0.],[0.,0.,0.,1.]])
            outline = np.linalg.inv(np.matmul(np.matmul(projection,np.linalg.inv(matrix)),projection.T))
            x = -np.linspace(-1.01*radii[gasObject][2],1.01*radii[gasObject][2],5000)
            y = np.linspace(-1.01*radii[gasObject][2],1.01*radii[gasObject][2],5000)
            X,Y = np.meshgrid(x,y)
            eqn = outline[0][0] * (X)**2 + (outline[0][1] + outline[1][0]) * (X)*(Y) + outline[1][1] * (Y )**2
            Z=1
            new_ax.contour((X+cmSIGO[0]),(Y+cmSIGO[1]),eqn,[Z],colors=color,linewidths=2.)

        for halo in range(len(HaloData['cmHalo'])):
            
            #This visualizes the DM halos
            color='white'
            cmHalo = HaloData['cmHalo'][halo] / (.71 * (redshift + 1.))
            R200Halo = HaloData['R200dm'][halo] / (.71 * (redshift + 1.))

            if cmHalo[0] < viewXMin - R200Halo  or cmHalo[0] > viewXMax + R200Halo or cmHalo[1] < viewYMin - R200Halo or cmHalo[1] > viewYMax + R200Halo or cmHalo[2] < viewZMin - R200Halo - 2. or cmHalo[2] > viewZMax + R200Halo or R200Halo < .33:#- widthwithoutbuffer / 8.
                continue
            print(halo)
            cmHalo = (HaloData['cmHalo'][halo] - centerCodeUnits) / (.71 * (redshift + 1.)) #- centerCodeUnits

            matrix = np.array([[1/R200Halo**2,0.,0.],[0.,1/R200Halo**2,0.],[0.,0.,1/R200Halo**2]])
            matrix = np.array([[matrix[0][0],matrix[0][1],matrix[0][2],0],[matrix[1][0],matrix[1][1],matrix[1][2],0],[matrix[2][0],matrix[2][1],matrix[2][2],0],[0,0,0,-1]])
            projection = np.array([[1.,0.,0.,0.],[0.,1.,0.,0.],[0.,0.,0.,1.]])
            outline = np.linalg.inv(np.matmul(np.matmul(projection,np.linalg.inv(matrix)),projection.T))
            x = -np.linspace(-1.01*R200Halo,1.01*R200Halo,5000)
            y = np.linspace(-1.01*R200Halo,1.01*R200Halo,5000)
            X,Y = np.meshgrid(x,y)
            eqn = outline[0][0] * (X)**2 + (outline[0][1] + outline[1][0]) * (X)*(Y) + outline[1][1] * (Y )**2
            Z=1

            new_ax.contour((X+cmHalo[0]),(Y+cmHalo[1]),eqn,[Z],colors=color,linewidths=2.)

        #-------------------------------------------------------------------------------------------------------------------

        width = 6 * (.71 * (redshift + 1.))
        widthwithoutbuffer = 3
        Around_SIGO2 = ds.region(centerCodeUnits, np.array(centerCodeUnits) - width/4., np.array(centerCodeUnits) + width/4.)

        #Generate the DM Density field
        ad = Around_SIGO
        pt = "PartType1"
        fields = ["particle_mass"] + [f"particle_position_{ax}" for ax in "xyz"]
        data = {field: ad[pt, field] for field in fields}
        ds_dm = yt.load_particles(data, data_source=ad)
        # Generate the missing SPH fields
        ds_dm.add_sph_fields()
        # This makes the DM field. Most of this code is repeated from above
        p2 = yt.ProjectionPlot(ds_dm, "z", ("io", "density"), center=centerCodeUnits, width=(widthwithoutbuffer,"kpc")) 
        p2.set_cmap(field=("io", "density"), cmap="B-W LINEAR")
        p2.set_zlim(("io", "density"), zmin='min',zmax= .4)
        center=np.array(centerCodeUnits)  / (.71 * (redshift + 1.))
        print(center)

    #     if len(starLocations) > 0:
    #         for star in starLocations:
    #             p2.annotate_marker(star, coord_system="data", plot_args={"color": 'black', "s": 60})

        plot = p2.plots[('io', 'density')]
        plot.figure = fig
        plot.axes = ax2
        p2._setup_plots()
        ax = ax2
        print(ax.get_xlim())
        print(ax.get_ylim())

        pos = ax2.get_position()
        print(pos)
        pos.x0 = 1. - plotwidth       # for example 0.2, choose your value
        pos.x1 = 1.
        pos.y0 = .5 - plotwidth / 2.
        pos.y1 = .5 + plotwidth / 2.
        ax2.set_position(pos)      # for example 0.2, choose your value


        rect = (0., 0., 1., 1.)
        new_ax2 = fig.add_axes(ax2.get_position(), frameon=False)

        new_ax2.set_xlim([0,1])
        new_ax2.set_ylim([0,1])

        new_ax2.set_xlim(ax.get_xlim())
        new_ax2.set_ylim(ax.get_ylim())



        new_ax2.patch.set_alpha(0.)
        new_ax2.set_alpha(0.)
        new_ax2.get_xaxis().set_ticks([])
        new_ax2.get_yaxis().set_ticks([])
        
        for gasObject in range(len(shrunken['cm'])):
            
            if gasObject == SOI:
                continue
            
            gasFrac = shrunken['gasFrac'][gasObject]
            r,g,b,a = plt.cm.PiYG(gasFrac)
            color = rgb_to_hex((int(256*r),int(256*g),int(256*b)))

            cmSIGO = (shrunken['cm'][gasObject])  / (.71 * (redshift + 1.))
            if cmSIGO[0] < viewXMin or cmSIGO[0] > viewXMax or cmSIGO[1] < viewYMin or cmSIGO[1] > viewYMax or cmSIGO[2] < viewZMin or cmSIGO[2] > viewZMax:
                continue

            cmSIGO = (shrunken['cm'][gasObject] - centerCodeUnits)  / (.71 * (redshift + 1.))
            matrix = np.array([[1/radii[gasObject][0]**2,0.,0.],[0.,1/radii[gasObject][1]**2,0.],[0.,0.,1/radii[gasObject][2]**2]])
            matrix = np.matmul(rotation[gasObject].T,np.matmul(matrix,rotation[gasObject]))
            matrix = np.matmul(viewR.T,np.matmul(matrix,viewR))
            matrix = np.array([[matrix[0][0],matrix[0][1],matrix[0][2],0],[matrix[1][0],matrix[1][1],matrix[1][2],0],[matrix[2][0],matrix[2][1],matrix[2][2],0],[0,0,0,-1]])
            projection = np.array([[1.,0.,0.,0.],[0.,1.,0.,0.],[0.,0.,0.,1.]])
            outline = np.linalg.inv(np.matmul(np.matmul(projection,np.linalg.inv(matrix)),projection.T))
            x = -np.linspace(-1.01*radii[gasObject][2],1.01*radii[gasObject][2],5000)
            y = np.linspace(-1.01*radii[gasObject][2],1.01*radii[gasObject][2],5000)
            X,Y = np.meshgrid(x,y)
            eqn = outline[0][0] * (X)**2 + (outline[0][1] + outline[1][0]) * (X)*(Y) + outline[1][1] * (Y )**2
            Z=1
            new_ax.contour((X+cmSIGO[0]),(Y+cmSIGO[1]),eqn,[Z],colors=color,linewidths=2.)


        if SOI >= 0:
            listOfSIGOs = [SOI]
            for SIGO in listOfSIGOs:

                color = 'yellow'

                
                cmSIGO = (shrunken['cm'][SIGO])  / (.71 * (redshift + 1.))
                if cmSIGO[0] < viewXMin or cmSIGO[0] > viewXMax or cmSIGO[1] < viewYMin or cmSIGO[1] > viewYMax or cmSIGO[2] < viewZMin or cmSIGO[2] > viewZMax:
                    continue

                cmSIGO = (shrunken['cm'][SIGO] - centerCodeUnits)  / (.71 * (redshift + 1.)) #- centerCodeUnits
                matrix = np.array([[1/radii[SIGO][0]**2,0.,0.],[0.,1/radii[SIGO][1]**2,0.],[0.,0.,1/radii[SIGO][2]**2]])
                matrix = np.matmul(rotation[SIGO].T,np.matmul(matrix,rotation[SIGO]))
                matrix = np.matmul(viewR.T,np.matmul(matrix,viewR))
                matrix = np.array([[matrix[0][0],matrix[0][1],matrix[0][2],0],[matrix[1][0],matrix[1][1],matrix[1][2],0],[matrix[2][0],matrix[2][1],matrix[2][2],0],[0,0,0,-1]])
                projection = np.array([[1.,0.,0.,0.],[0.,1.,0.,0.],[0.,0.,0.,1.]])
                outline = np.linalg.inv(np.matmul(np.matmul(projection,np.linalg.inv(matrix)),projection.T))
                x = -np.linspace(-1.01*radii[SIGO][2],1.01*radii[SIGO][2],5000)
                y = np.linspace(-1.01*radii[SIGO][2],1.01*radii[SIGO][2],5000)
                X,Y = np.meshgrid(x,y)
                eqn = outline[0][0] * (X)**2 + (outline[0][1] + outline[1][0]) * (X)*(Y) + outline[1][1] * (Y )**2
                Z=1
                
                new_ax2.contour((X+cmSIGO[0]),(Y+cmSIGO[1]),eqn,[Z],colors=color,linewidths=2.)



        for halo in range(len(HaloData['cmHalo'])):
            color='white'
            cmHalo = HaloData['cmHalo'][halo] / (.71 * (redshift + 1.))
            R200Halo = HaloData['R200dm'][halo] / (.71 * (redshift + 1.))

            if cmHalo[0] < viewXMin - R200Halo  or cmHalo[0] > viewXMax + R200Halo or cmHalo[1] < viewYMin - R200Halo or cmHalo[1] > viewYMax + R200Halo or cmHalo[2] < viewZMin - R200Halo - 2. or cmHalo[2] > viewZMax + R200Halo or R200Halo < .33:#- widthwithoutbuffer / 8.
                continue
            print(halo)
            cmHalo = (HaloData['cmHalo'][halo] - centerCodeUnits) / (.71 * (redshift + 1.)) #- centerCodeUnits

            matrix = np.array([[1/R200Halo**2,0.,0.],[0.,1/R200Halo**2,0.],[0.,0.,1/R200Halo**2]])
            matrix = np.array([[matrix[0][0],matrix[0][1],matrix[0][2],0],[matrix[1][0],matrix[1][1],matrix[1][2],0],[matrix[2][0],matrix[2][1],matrix[2][2],0],[0,0,0,-1]])
            projection = np.array([[1.,0.,0.,0.],[0.,1.,0.,0.],[0.,0.,0.,1.]])
            outline = np.linalg.inv(np.matmul(np.matmul(projection,np.linalg.inv(matrix)),projection.T))
            x = -np.linspace(-1.01*R200Halo,1.01*R200Halo,5000)
            y = np.linspace(-1.01*R200Halo,1.01*R200Halo,5000)
            X,Y = np.meshgrid(x,y)
            eqn = outline[0][0] * (X)**2 + (outline[0][1] + outline[1][0]) * (X)*(Y) + outline[1][1] * (Y )**2
            Z=1
            

            new_ax2.contour((X+cmHalo[0]),(Y+cmHalo[1]),eqn,[Z],colors=color,linewidths=2.)



        plt.gcf().savefig('D:/Star_Movie_768/SIGO_Movies_Paper_4/SIGO_'+str(SIGOnum) + '/snap_' + str(snapnum)+'_SOI_'+str(SOI)+'.png',bbox_inches='tight')