evolutiontracker.py

"""
Program to track gas objects between snapshots.
Returns: pickled file with digraph containing important data on the gas objects on nodes, and matching objects connected by edges.
Indexed as (snapshot number, gas object index)

Added to Github 10/12/23
"""

import networkx as nx
from ipywidgets import IntProgress
from IPython.display import display
import time
from __future__ import print_function, division
from sys import argv
import numpy as np
import readsubfHDF5
import snapHDF5 
from annikaEllipsoid import *
try:
   import cPickle as pickle
except:
   import pickle



def dx_wrap(dx,box):
	#wraps to account for period boundary conditions. This mutates the original entry
	idx = dx > +box/2.0
	dx[idx] -= box
	idx = dx < -box/2.0
	dx[idx] += box 
	return dx
def dist2(dx,dy,dz,box):
	#Calculates distance taking into account periodic boundary conditions
	return dx_wrap(dx,box)**2 + dx_wrap(dy,box)**2 + dx_wrap(dz,box)**2

filename = "D:/Star_Movie_768/"
s_res = '14Mpc'
s_vel = 'Sig2'
box = 1775
snapkey = [0,10,20,30,40,50,60,70,80,90,100,110,120,130,140,150]

with open(filename+"IDDictionary"+s_res+"_"+s_vel+"_fb0.6.dat",'rb') as f:
    idDict = pickle.load(f) # This contains important data from shrinker about objects across snapshots, generated by IDDictionary.py

gasObjectIndices = [] # This is a good reference for all of our gas objects' indices ordered by snapshot index, for later
for snapnum2 in snapkey:
    gasObjectIndex = []
    #Should be run with a snap number input
	
    snapnum = int(snapnum2)

    #File paths
    filename = "D:/Star_Movie_768/"
    filename2 = filename +  "GasOnly_FOF" #Used for readsubfHDF5
    
    cat = readsubfHDF5.subfind_catalog(filename2, snapnum)

    halo100_indices= np.where((cat.GroupLenType[:,0]+cat.GroupLenType[:,4]) >100)[0]
    gasObjectIndex = halo100_indices
    gasObjectIndices.append(gasObjectIndex)

G = nx.DiGraph() # This graph contains all gas objects as nodes, with data about them stored as dictionaries on the nodes. Edges represent the same object (or a merger) across different timesteps, directed by time.
    
excludedList = []
for key in idDict.keys():
    # This loop adds all objects at all snapshots to the graph as nodes.
    snap, ind = key
    if idDict[key]['baryonFrac'] > -1:
        G.add_node(key)
        for key2 in idDict[key].keys():
            G.nodes[key][key2] = idDict[key][key2]
    else:
        excludedList.append(key)
    

nodeList = list(G.nodes)

cms = [] # List of coordinates of centers of mass of the objects
for snap in range(len(snapkey)):
    cmsAtSnap = []
    for gasObject in gasObjectIndices[snap]:
        node = (snapkey[snap],gasObject)
        cmsAtSnap.append(G.nodes[node]['cm'])
    cms.append(cmsAtSnap)

# Now we iterate through the nodes to find matching nodes. We only want the nearest matching nodes by time to create the graph, so there is a double for loop over the snapshots

matchCountAll = 0 # Track the matches. About 2/3 of objects should find a match -- if it is significantly less or greater, something is going wrong.
max_count = len(nodeList)
f = IntProgress(min=0, max=max_count) # instantiate the progress bar
display(f) # display the bar
for node in nodeList:
    f.value += 1
    snap, idx = node
    gasIDSet1 = set(G.nodes[node]['gasIDs'])
    starIDSet1 = set(G.nodes[node]['starIDs'])
    dmIDSet1 = set(G.nodes[node]['dmIDs'])
    cm = G.nodes[node]['cm']
    
    matchList = []
    weightList = [] # List of fraction of shared IDs
    if snap < maxSnap:
        snapIter = snapkey[snapkey.index(snap) + 1]
    else:
        snapIter = maxSnap + 1
    
    while len(matchList) == 0 and snapIter <= maxSnap:
        # We'll loop through this for each object until we find a match or run out of later snapshots (i.e., the object dissolved)
        snapIndex = snapkey.index(snapIter)
        gasObjectList = np.array(gasObjectIndices[snapIndex])
        numToFilter = int(np.sqrt(len(gasObjectList))) # We're only going to search the closest sqrt(N) of N objects to see if they match to save heavily on time.
        # This reduces the number of matches found by about 1%, which is a good tradeoff for a sqrt(N) decrease in time spent. Consider searching everything if the node you're looking for a match for is a SIGO.
        cms = np.array([G.nodes[(snapIter,gasObject)]['cm'] for gasObject in gasObjectList])
        distances = [dist2(cms[:,0] - cm[0],cms[:,1] - cm[1],cms[:,2] - cm[2],box)]
        filterSmallestDists = sorted(range(len(distances[0])), key=lambda sub: distances[0][sub])[:numToFilter]
        for gasObject in gasObjectList[filterSmallestDists]:
            node2 = (snapIter, gasObject)
            if node2 not in excludedList:
                gasIDSet2 = set(G.nodes[node2]['gasIDs'])
                starIDSet2 = set(G.nodes[node2]['starIDs'])
                dmIDSet2 = set(G.nodes[node2]['dmIDs'])
                objectsMatch = False

                gasIntersectionID = gasIDSet1.intersection(gasIDSet2)
                starIntersectionID = starIDSet1.intersection(starIDSet2)
                dmIntersectionID = dmIDSet1.intersection(dmIDSet2)

                # We'll accept an object as a match if any of the following hold: There are at least 50 gas, dm, or star particles in either the earlier or later object, and of those particles at least a third are present in the other object
                # We require 50 particles for statistical reasons, and a third matching gives good results qualitatively for matches.
              
                if len(gasIDSet1) > 50 and len(gasIntersectionID) > len(gasIDSet1) / 3.:
                    objectsMatch = True

                elif len(gasIDSet2) > 50 and len(gasIntersectionID) > len(gasIDSet2) / 3.:
                    objectsMatch = True

                elif len(starIDSet1) > 50 and len(starIntersectionID) > len(starIDSet1) / 3.:
                    objectsMatch = True

                elif len(starIDSet2) > 50 and len(starIntersectionID) > len(starIDSet2) / 3.:
                    objectsMatch = True

                elif len(dmIDSet1) > 50 and len(dmIntersectionID) > len(dmIDSet1) / 3.:
                    objectsMatch = True

                elif len(dmIDSet2) > 50 and len(dmIntersectionID) > len(dmIDSet2) / 3.:
                    objectsMatch = True

                if objectsMatch:
                    matchList.append(node2)
                    weightList.append((len(dmIntersectionID) + len(starIntersectionID) + len(gasIntersectionID)) / (len(dmIDSet1) + len(starIDSet1) + len(gasIDSet1)))
        if snapIter < snapkey[len(snapkey) - 1]:
            snapIter  = snapkey[snapkey.index(snapIter) + 1]
        else:
            snapIter += 1
    
    matchCount = 0
    for matchNode in matchList:
        G.add_edges_from([(node, matchNode)], weight= weightList[matchCount])
        matchCount += 1
    matchCountAll += matchCount
print('Total matches: ', matchCountAll)

for snap in snapkey:
    with open(filename+'shrinker'+s_res+'_'+s_vel+'_'+str(snap)+'.dat','rb') as f:
        shrunken = pickle.load(f)
    snapIndex = snapkey.index(snap)
    gasObjectList = gasObjectIndices[snapIndex]
    for gasObject in gasObjectList:
        node = (snap, gasObject)
        if node not in excludedList and node in G.nodes():
            mDM = max(shrunken['mDM'][gasObject] * 1e10 / 0.71,0)
            mGas = max(shrunken['mGas'][gasObject] * 1e10 / 0.71,0)
            mStar = max(shrunken['mStar'][gasObject] * 1e10 / 0.71,0)
            mTot = mDM + mGas + mStar
            G.nodes[node]['logMTot'] = np.log10(mTot) # Just want to add some total mass data to our nodes before we save it, for plots later
              
nx.write_gpickle(G, filename+'GasObjectHistoryGraph.dat')