mynet_new.py

#Modified version of networks_info.py in krome/
#to return reactants and products of every reaction listed
#in the given file
#Piyush Sharda (2022)
import glob
import numpy as np
import re
import sympy as sp

#folder name
folder = "networks/"

#parse format reactants and products
#returns indexes with positions
def parseFormat(arow):
	idxs = []
	for i in range(len(arow)):
		if(arow[i].decode('utf-8') in ["R","P"]): idxs.append(i)
	return idxs

def parseFormatReactants(arow):
	idxs = []
	for i in range(len(arow)):
		if(arow[i].decode('utf-8') in ["R"]): idxs.append(i)
	return idxs

def parseFormatProducts(arow):
        idxs = []
        for i in range(len(arow)):
                if(arow[i].decode('utf-8') in ["P"]): idxs.append(i)
        return idxs

def parseFormatTmin(arow):
	idxs = []
	for i in range(len(arow)):
		if(arow[i].decode('utf-8') in ["Tmin"]): idxs.append(i)
	return idxs

def parseFormatTmax(arow):
	idxs = []
	for i in range(len(arow)):
		if(arow[i].decode('utf-8') in ["Tmax"]): idxs.append(i)
	return idxs

def decode(l):
	if isinstance(l, list):
		return [decode(x) for x in l]
	else:
		return l.decode('utf-8')

#check if variable is an integer number
def isNumber(arg):
	try:
		int(arg)
		return True
	except:
		return False

#symbols used in the function rate_new.SympyChemRate.eval()
#T = sp.symbols('T', positive=True)
#redshift = sp.symbols('z', real=True)
#user_Tdust=sp.Symbol('user_Tdust')
#user_crate=sp.Symbol('user_crate')
#user_Av=sp.Symbol('user_Av')
#user_ionH=sp.Symbol('user_ionH')
#user_ionH2=sp.Symbol('user_ionH2')
#user_dissH2=sp.Symbol('user_dissH2')
#user_ionC=sp.Symbol('user_ionC')
#user_ionO=sp.Symbol('user_ionO')
#user_dissCO=sp.Symbol('user_dissCO')
#user_dust2gas_ratio=sp.Symbol('user_dust2gas_ratio')

#select from 3 or _wD to determine if Deuterium is included or not
fstr = '_wD' #str(input("Which primordial+metal network do you want to read in: '3' or '_wD'? "))
if fstr == '3':
    withD = 0
    print('Excluding deuterium reactions')
elif fstr == '_wD':
    withD = 1
    print('Including deuterium reactions')
else:
    raise ValueError('Unsupported primordial+metal network name!')

ffstr = 'my_react_flash_metals' + fstr

for ntw in sorted(glob.glob(folder+ffstr)):

	print("****************")
	print(ntw)
	reactionCount = 0
	reactants = []
	products = []
	Tmins = []
	Tmaxs = []
	Rates = []

	#loop on file
	for row in open(ntw,"rb"):
		srow = row.strip()
		#skip comments and blanks
		if(srow==""): continue
		if(srow.startswith(b"#")): continue
		if(srow.startswith(b"!")): continue
		arow = srow.split(b",")
		#get format
		if(srow.startswith(b"idx")):
			idxs = parseFormat(arow)
			idxs_r = parseFormatReactants(arow)
			idxs_p = parseFormatProducts(arow)
			idxs_Tmin = parseFormatTmin(arow)
			idxs_Tmax = parseFormatTmax(arow)
			#idxs_rate = parseFormatRate(arow)
			print(idxs, idxs_r, idxs_p, idxs_Tmin, idxs_Tmax)
		#skip lines not starting with integers (index)
		if(not(isNumber(arow[0]))): continue

		#get all the reactants
		good_idxs_r = []
		for j in range(len(idxs_r)):
			if np.array(arow)[idxs_r[j]].decode('utf-8') != '':
				good_idxs_r.append(idxs_r[j])
		reactants.extend([list(np.array(arow)[good_idxs_r])])

		#get all the products
		good_idxs_p = []
		for k in range(len(idxs_p)):
			if np.array(arow)[idxs_p[k]].decode('utf-8') != ""  and np.array(arow)[idxs_p[k]].decode('utf-8') != 'g':
				good_idxs_p.append(idxs_p[k])
		products.extend([list(np.array(arow)[good_idxs_p])])

		#get all the Tmins
		good_idxs_Tmin = []
		for k in range(len(idxs_Tmin)):
			if np.array(arow)[idxs_Tmin[k]].decode('utf-8') != "":
				good_idxs_Tmin.append(idxs_Tmin[k])
		Tmins.extend([list(np.array(arow)[good_idxs_Tmin])])

		#get all the Tmaxs
		good_idxs_Tmax = []
		for k in range(len(idxs_Tmax)):
			if np.array(arow)[idxs_Tmax[k]].decode('utf-8') != "":
				good_idxs_Tmax.append(idxs_Tmax[k])
		Tmaxs.extend([list(np.array(arow)[good_idxs_Tmax])])

		#get all the Rates
		#first get the index of the column that stores the rates
		ncols_before_rate = max(idxs_r + idxs_p + idxs_Tmin + idxs_Tmax, default=0) + 1
		#step below is necessary to join stuff in the rate after a comma
		combined_rate_field = b",".join(arow[ncols_before_rate:])
		Rates.append(combined_rate_field)

		reactionCount += 1


	#print output
	print("Number of reactions:", reactionCount)
	print()


reactants = decode(reactants)
products = decode(products)
Tmins = decode(Tmins)
Tmaxs = decode(Tmaxs)
Rates = decode(Rates)

#Transform all the rates to sympy expressions
#Read whatever alphanumeric characters the rates use and convert them to sympy symbols

#this is to tell sympy that the string exp should be replaced by sp.exp
local_dict = {'exp': sp.exp, 'log10': lambda x: sp.log(x, 10)}

tokensymbols = []
symRates = []
for i in range(0, len(Rates)):
	# Find all unique non-numeric tokens in the string
	tokens = set(re.findall(r'[a-zA-Z_]\w*', Rates[i]))
    
	# Define all found tokens as sympy symbols
	#for token in tokens:

		#if token is already an existing sympy symbol, then dont sympy it
	#	if not token in tokensymbols and token != 'exp' and token!='log10':
	#		globals()[token] = sp.symbols(token)
	#		tokensymbols.append(token)

	# Convert the string to a sympy expression
	try:
		expr = sp.sympify(Rates[i], locals=local_dict)
		symRates.append(expr)
	except:
		print('No sympify for: ', i, Rates[i])
		symRates.append(None)
		#remove symbols of expression that could not be sympified
		#tokensymbols = [item for item in tokensymbols if item not in tokens]


#below code transforms the list of reactants and products in pynucastro readable formats
import rate_new as rate_new

rr = []
pp = []
for i in range(len(reactants)):
	subrr = []
	for j in range(len(reactants[i])):
		subrr.append(rate_new.ChemSpecie(reactants[i][j]))
	rr.extend([subrr])

	subpp = []
	for j in range(len(products[i])):
		subpp.append(rate_new.ChemSpecie(products[i][j]))
	pp.extend([subpp])


#find all the unique species actually present in the network
#this assumes 
# Flatten the nested list rr and convert to a set to get unique elements
unique_species_rr = list(set(element for sublist in rr for element in sublist))
unique_species_pp = list(set(element for sublist in pp for element in sublist))
print("Number of unique reactant species: ", len(unique_species_rr))
print("Number of unique product species: ", len(unique_species_pp))
check_rrpp = list(set(unique_species_rr) - set(unique_species_pp))
if len(check_rrpp) !=0 :
	raise ValueError('List of reactants and products contain different unique species')


#use these unique elements to define your composition
fcomp = rate_new.ChemComposition(specie=unique_species_rr).sympy()


#set all Tmins/Tmaxs that are strtype 'NONE' to NoneType
#set all other Tmins/Tmaxs to float type
def properify_T(nested_list):
	for sublist in nested_list:
		for i, value in enumerate(sublist):
			if isinstance(value, str):
				if value == 'NONE':
					sublist[i] = None
				else:
					sublist[i] = float(value)

		#for reactions where Tmin and Tmax are not a part of the format
		if len(sublist) == 0:
			sublist.append(None)
	return nested_list

Tmins = properify_T(Tmins)
Tmaxs = properify_T(Tmaxs)


#now make a dictionary containing all the rates to pass to pynucastro ChemRateCollection
sym_rates = {}

for x in range(len(rr)):
	reactants = []
	products = []
	for y in range(len(rr[x])):
		reactants.append(rr[x][y].sym_name)
	for y in range(len(pp[x])):
		products.append(pp[x][y].sym_name)
	sym_rates['r{0}'.format(x)] = rate_new.SympyChemRate(reactants=reactants, products=products, rate_expr=symRates[x], 
		                                                 Tmins=Tmins[x][0], Tmaxs=Tmaxs[x][0])