This Python script is designed for astronomers and astrophysicists working with stellar abundance data. It involves reading an abundance map, interpolating data across a specified grid dimension, and outputting statistics regarding element mass fractions and changes in target values versus new calculations. The script leverages NumPy for numerical operations, SciPy for interpolation, argparse for command-line argument parsing, and pickle for object serialization. Its primary function is to aid in the analysis of stellar compositions by converting abundance ratios to mass fractions and vice versa, ultimately saving the calculated statistics for further analysis.

from generateTOPStables import convert_XYZ_FeHAlphaHaHe, call_num_frac, parse_abundance_file, format_TOPS_string
import argparse
import os
from scipy.interpolate import RegularGridInterpolator, interp2d
import pickle
import numpy as np
 
def get_stats(interpF, gridDim, FeHb, AlphaHb, aHeb, numFracProgram, abunTable, inputDir, abundanceMap):
    if interpF:
        MetalAbunMap = pickle.load(open(interpF, 'rb'))
    else:
        # This run numfrac gridDim^3 times for a variety of compositions,
        #  at each it finds the X, Y, and Z values and uses these to build
        #  a 3D interpolated vector field function which can take X, Y, and Z, 
        #  and return [Fe/H], [alpha/H], and a(He) 
        A = np.empty(shape=(gridDim, gridDim, gridDim, 3))
        FeH = np.linspace(FeHb[0], FeHb[1], gridDim)
        AlphaH = np.linspace(AlphaHb[0], AlphaHb[1], gridDim)
        aHe = np.linspace(aHeb[0], aHeb[1], gridDim)
        if not abunTable:
            saveName = "{}2XYZ.p".format(numFracProgram.split('/')[-1])
            for i, feh in enumerate(FeH):
                for j, alphah in enumerate(AlphaH):
                    for k, ahe in enumerate(aHe):
                        # this callibrate.tmp file is to store the output of one
                        #  run, it will be overwritten each time and eventually
                        #  deleted to clean up
                        status = call_num_frac(numFracProgram,
                                "calibrate.tmp", feh, alphah, ahe, force=True)
                        _, oX, oY, oZ = parse_abundance_file("calibrate.tmp")
                        A[i, j, k] = np.array([oX, oY, oZ])
        else:
            saveName = f"{abunTable.split('/')[-1].split('.')[0]}2XYZ.p"
            FeH_tup = (FeH.min(), FeH.max(), FeH.shape[0])
            AlphaH_tup = (AlphaH.min(), AlphaH.max(), AlphaH.shape[0])
            aHe_tup = (aHe.min(), aHe.max(), aHe.shape[0])
 
            status = call_num_frac(numFracProgram, "calibrate.tmp",
                                   FeH_tup, AlphaH_tup, aHe_tup, force=True,
                                   abunTable=abunTable)
            _, oX, oY, oZ = parse_abundance_file("calibrate.tmp", big=True)
            A[:, :, :, 0] = oX
            A[:, :, :, 1] = oY
            A[:, :, :, 2] = oZ
 
        os.remove("calibrate.tmp")
        MetalAbunMap = RegularGridInterpolator((FeH, AlphaH, aHe), A)
 
        # force this to save to disk so we don't have to do this every time
        pickle.dump(MetalAbunMap,
                    open(os.path.join(inputDir, saveName), 'wb'))
    with open(abundanceMap, 'r') as f:
        contents = f.read().split('\n')
        pContents = [[float(y) for y in x.split(',')]
                          for x in contents
                          if x != '' and x.lstrip()[0] != '#']
    stats = np.zeros(shape=(len(pContents), 4))
    for i, composition in enumerate(pContents):
        results = convert_XYZ_FeHAlphaHaHe(MetalAbunMap,composition[0], composition[2])
        status = call_num_frac(numFracProgram, "ValidationOutput.tmp", results[0], 0.0, results[1], force=True, abunTable=abunTable)
        if int(status) == 0:
            massFraction, X, Y, Z = parse_abundance_file("ValidationOutput.tmp")
        if os.path.exists("ValidationOutput.tmp"):
            os.remove("ValidationOutput.tmp")
        stats[i] = np.array([composition[0], composition[2], X, Z])
    return stats
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Given an abundance map find the bounds to test over")
    parser.add_argument("numFracProgram", help="path to program to generate"
                        " Number and Mass fractions", type=str)
    parser.add_argument("abundanceMap", help="Path to abundace map csv",
                        type=str)
    parser.add_argument("FeHb", help="lower/upper bound for"
                        " FeH interpolation", type=float, nargs=2)
    parser.add_argument("AlphaHb", help="lower/upper bound for"
                        " AlphaH interpolation", type=float, nargs=2)
    parser.add_argument("aHeb", help="lower/upper bound for"
                        " aHe interpolation", type=float, nargs=2)
    parser.add_argument("--interpF", help="Pickeled Interpolation File to go"
                        " between [Fe/H], [alpha/H], a(H) and X, Y, Z",
                        type=str)
    parser.add_argument("--gridDim", type=int, help="Grid dimension to use"
                        " when generating the enhancmnet -> abundence table.",
                        default=50)
    parser.add_argument("--inputDir", type=str, help="Directory to save input"
                        " file to (such as interpolation function)",
                        default=".")
    parser.add_argument("--abunTable", type=str, help="Table to pull abundances"
                        " from. If set numFracProgram must point to executable"
                        " expedciting table input")
 
    args = parser.parse_args()
 
    stats = get_stats(args.interpF, args.gridDim, args.FeHb, args.AlphaHb, args.aHeb, args.numFracProgram, args.abunTable, args.inputDir, args.abundanceMap)
    np.save("stats.npy", stats)
 
    diff = stats[:, :2] - stats[:, 2:]
    fractionalChange = diff/stats[:, :2]
    for comp in stats:
        print(f"target X: {comp[0]}, target Z: {comp[1]}, new X: {comp[2]}, new Z: {comp[3]}")