diff --git a/massive_star/convert_large_to_small.py b/massive_star/convert_large_to_small.py new file mode 100644 index 0000000..1e66b95 --- /dev/null +++ b/massive_star/convert_large_to_small.py @@ -0,0 +1,91 @@ +import numpy as np +from pynucastro.networks.rate_collection import Composition +import argparse + +implemented_networks = { + "aprox19": ['h1', 'he3', 'he4', 'c12', 'n14', 'o16', 'ne20', 'mg24', + 'si28', 's32', 'ar36', 'ca40', 'ti44', 'cr48', 'fe52', + 'fe54', 'ni56', 'n'], +} +# the aprox19 nuclei list excludes p_nse. We'll deal with that split by later setting 'p' to 0 + + +def main(file_name, new_file_name, network): + + # READ MESA FILE + # this code adapted from py_mesa_reader. + bulk_data = np.genfromtxt(file_name, skip_header=5, names=True, ndmin=1, dtype=None) + bulk_names = bulk_data.dtype.names + with open(file_name) as f: + for i, line in enumerate(f): + if i == 1: + header_names = line.split() + elif i == 2: + header_data = [eval(datum) for datum in line.split()] + elif i > 2: + break + header_data = dict(zip(header_names, header_data)) + + + # SHRINK WITH PYNUCASTRO + try: + start, end = bulk_names.index('neut'), bulk_names.index('opacity')-1 + large_network = bulk_names[start:end+1] + except ValueError as e: + print("The given start or end species is not listed in the MESA file.") + raise e + + small_network = implemented_networks[network] + + network_col_names = [] + network_col_bulk = [] + c = Composition(large_network) + for row in bulk_data: + c.set_array(list(row)[start:end+1]) + new_c = c.bin_as(small_network) + dictionary = new_c.data + dictionary['p'] = 0 + + if not network_col_names: + network_col_names = dictionary.keys() + + network_col_bulk.append(list(dictionary.values())) + + network_col_bulk = np.array(network_col_bulk) + + + # OUTPUT TO FILE + with open(new_file_name, 'w') as f: + def writeline(array): + f.write(''.join([str(x).rjust(40) for x in array])+'\n') + def replace_cols(data, replacement, start, end): + return list(data)[:start]+list(replacement)+list(data)[end+1:] + + writeline(range(1,len(header_data)+1)) + writeline(header_data.keys()) + writeline(header_data.values()) + f.write('\n') + + col_names = replace_cols(bulk_names, network_col_names, start, end) + writeline(range(1,len(col_names)+1)) + writeline(col_names) + for i in range(bulk_data.size): + writeline(replace_cols(bulk_data[i], network_col_bulk[i], start, end)) + f.write('\n') + + + print(f'Binned network (size={end-start+1}) into network {network} and saved file in {new_file_name}.') + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Bin a large network into a smaller network specified by user. Input and output files use the MESA file format.") + parser.add_argument('network', choices=implemented_networks.keys(), help="The name of the smaller network to bin into.") + parser.add_argument('file_name') + parser.add_argument('new_file_name', nargs='?') + args = parser.parse_args() + if args.new_file_name is None: + file_stem, file_ending = args.file_name.rsplit('.',1) + new_file_name = f"{file_stem}.{args.network}.{file_ending}" + else: + new_file_name = args.new_file_name + main(args.file_name, new_file_name, args.network)