DSEP Portable Execution Script

This Python script is designed to facilitate the portable execution of DSEP (Dartmouth Stellar Evolution Program) models by automatically updating paths for data tables and executables based on the user's environment. It includes functionalities for setting up the environment, configuring and linking required files, and running the DSEP executable with options for calibration, debugging, and output management.

#!/usr/bin/env python3
"""
Author: Thomas M. Boudreaux
Date  : November 10 2020
Description: Python script to run DSEP portably. This script will automatically
             update paths that are relevant so you don't have to manually do
             that.
"""
 
import os
import argparse
import shutil
import f90nml
import pandas as pd
import numpy as np
 
def config_env(branch, default):
    """
    Function: config_env
    Description: Configure enviroment for DSEP to run. Sets paths to data
                 tables nessisarity for DSEP to run. These paths are then
                 key-valued to the relevant unit number DSEP expects.
    Inputs:
                 - branch: git branch of executable to run.
                 - default: if true use the despX executable in the dsepX root
                            root path
    Outputs:
                 - fort: dictionary with key-value pairs for DSEP data table
                         absolute paths and their associate unit numbers.
                 - prog: absolute path to the DSEP executable.
    Pre-State:
                 - DSEPRoot Enviromental variable set.
    """
    local = os.environ['DSEPRoot']
    if not default:
        if os.path.exists(os.path.join(local, 'dsepX/build')):
            prog = os.path.join(local, f"dsepX/build/dsepX.{branch}")
        else:
            prog = os.path.join(local, f"dsepX/bin/dsepX.{branch}")
    else:
        prog = os.path.join(local, "dsepX/dsepX")
    zams = os.path.join(local, "DSEPmisc/zams")
    prems = os.path.join(local, "DSEPmisc/prems")
    opac = os.path.join(local, "DSEPmisc/opac")
    atm = os.path.join(local, "DSEPmisc/surfBC")
    opal = os.path.join(opac, "OPAL")
    tops = os.path.join(opac, "TOPS")
    phx = os.path.join(opac, "phoenix")
    bcphx = os.path.join(atm, "phoenix/GS98")
    bckur = os.path.join(atm, "atmk")
 
    fort = {
                12: os.path.join(prems, "m100.GS98"),
                15: os.path.join(opac, "FERMI.TAB"),
                38: os.path.join(bckur, "atmk1990p00.tab"),
                # 48: os.path.join(opal, "GS98hz"),
                48: os.path.join(tops, "TOPALGS98.dat"),
                95: os.path.join(bcphx, "z_p0d0.afe_p0d0.dat"),
                96: os.path.join(bcphx, "z_m0d5.afe_p0d0.dat"),
                97: os.path.join(bcphx, "z_m0d7.afe_p0d0.dat"),
                98: os.path.join(bcphx, "z_m1d0.afe_p0d0.dat"),
                99: os.path.join(bcphx, "z_m1d5.afe_p0d0.dat")
            }
 
    if not os.path.exists(prog):
        raise OSError(2,f"DSEP Executable {prog} not fond on disk, are you sure you have the right branch?")
 
    return fort, prog
 
 
def clear_fort_files(path):
    """
    Function: clear_fort_files
    Description: Remove all fortran unit files (fort.n) in a given path
    Inputs:
                 - path: path to remove fortran unit files from
    Outputs:
                 N/A
    Post-State:
                - All files glob(fort.*) at a depth of 1 in path are removed.
    """
    for file in os.listdir(path):
        if os.path.islink(file):
            if file.startswith('fort'):
                print(f'deleting: {file}')
                os.remove(file)
 
def symlink_file(src, dst):
    """
    Function: symlink_file
    Description: soft link src to dst. if src does not exist, create src.
    Inputs:
                 - src: path to source of symbolic link
                 - dst: path to destination of symbolic link
    Output:
                 N/A
    Post-State:
                 - dst is a symbolic soft link to src, and if src file did not
                   exist, it does now.
    """
    # if the src file does not exist already, create it (equivilent to touch)
    if not os.path.exists(src):
        with open(src, 'w') as f:
            pass
    os.symlink(src, dst)
 
def config_fort_files(fort, template, cnml, pnml, updatedParams=None):
    """
    Function: config_fort_files
    Description: Read in the control namelist file and update all releavant
                 paths in that file for the local system.
 
                 Clear all fortran unit files in the current working
                 directory. Sesolve physics and control namelist file
                 paths to absolute paths, and link these to the correct unit
                 numbers in the dictionary. Link the output files to the
                 correct unit numbers in the dictionary. Finally, symbolically
                 link all the key-value pairs in the dictioary to unit files
                 and actual files on disk.
    Inputs:
                 - fort: dictionary linking unit number to file path
                 - template: template path to where output files names with .ext
                            extension. .ext extension will be replaced with
                            relevant extensions.
                 - cnml: path to control namelist file
                 - pnml: path to physics namelist file
                 - updatedParams: dictionary of parameters to update in the
                                  control namelist file
    Outputs:
                 - cnml_filename: name of the updated control name list files
                                  with paths for the local system.
    Post-State:
                 - All Unit files nessisairy for DSEP to run are present in the
                   current working directory and linked to their requisite files
                   eleswhere on disk.
                 - A new file exists, with the same name as the control namelist
                   file but with _ul appened (for updated local). This file has
                   paths correct for the local system inside of it. This is
                   generated automatically every time, and this is actually the
                   control namelist file ultimately used by DSEP.
    """
 
    # Update the namelist file and generate the new namelist file
    nml=f90nml.read(cnml)
    nml['control']['fo95cobin'] = os.path.join(os.environ['DSEPRoot'],
                                               'DSEPmisc/opac/OPAL/OPAL_GS98.bin')
    units = [0, 1, 2, 35, 5, 7, 8, 9]
    for i, (_, n) in enumerate(zip(nml['control']['opecalex'], units)):
        nml['control']['opecalex'][i] = os.path.join(os.environ['DSEPRoot'],
                                                     f'DSEPmisc/opac/ferg04/gs98.{n}.tron')
    if updatedParams:
        for paramTitle in updatedParams:
            nml['control'][paramTitle] = updatedParams[paramTitle]
 
    cnml_filename = f"{cnml.strip('.nml')}_lu.nml"
 
    # f90nml cannot by default overwrite an existing file, so if the _ul file already
    #  exists, remove it so it may be rewritten.
    if os.path.exists(cnml_filename):
        os.remove(cnml_filename)
    nml.write(cnml_filename)
 
 
    # update the fort dictionary
    clear_fort_files('.')
    runFort = fort.copy()
    runFort[13] = os.path.abspath(pnml)
    runFort[14] = os.path.abspath(cnml_filename)
    runFort[19] = template.replace('.ext', '.track')
    runFort[20] = template.replace('.ext', '.short')
    runFort[37] = template.replace('.ext', '.iso')
    runFort[11] = template.replace('.ext', '.last')
    runFort[24] = template.replace('.ext', '.pmod')
    runFort[25] = template.replace('.ext', '.penv')
    runFort[26] = template.replace('.ext', '.patm')
 
    # Symbolically link all required files
    for fortKey in runFort:
        print(f'linking {runFort[fortKey]} -> fort.{fortKey}')
        symlink_file(runFort[fortKey], f'fort.{fortKey}')
 
    return cnml_filename
 
def run_dsep(fort, outDir, cnml, pnml, prog, updatedParams=None, debug=False, rr=False):
    """
    Function: run_dsep
    Description: Configure the system and Invoke the DSEP executable
    Inputs:
                 - fort: dictionary linking unit number to file path
                 - outDir: directory to where output files will be saved
                 - cnml: path to control namelist file
                 - pnml: path to physics namelist file
                 - prog: path to DSEP executable
                 - updatedParams: dictionary of parameters to update in the
                                  control namelist file
                 - debug: run dsep with gdb debugger
                 - rr: run with the rr reverse debuger (will overwrite gdb)
    Outputs:
                 N/A
    Post-State:
                 - DSEP runs and output files have been written to outDir
    """
    absOutDir = os.path.abspath(outDir)
 
    # template filename will have its extension replaced by all the standard dsep
    #  file extensions
    templateFilename = f"{absOutDir}/{fort[12].split('/')[-1]}.ext"
 
    cnml_filename = config_fort_files(fort, templateFilename, cnml, pnml, updatedParams=updatedParams)
 
    if debug and rr:
        debug = False
 
    print("Running DSEP")
    if debug:
        os.system(f"gdb {prog}")
    elif rr:
        os.system(f"rr record -n {prog} && rr replay")
    else:
        os.system(f"{prog}")
 
def load_iso(path):
    header = ["Age", "Log_T", "Log_g", "Log_L", "Log_R", "Y_core", "Z_core", "(Z/X)_surf", "L_H", "L_He", "M_He_core", "M_CO_core"]
    iso = pd.read_csv(path, delim_whitespace=True, names=header, comment='#')
    with open(path, 'r') as f:
        metadata_str = f.readline()
    metadata_str = metadata_str.strip('#').lstrip().rstrip().split(' ')[:3]
    metadata = dict()
    for i, data_elem in enumerate(metadata_str):
        keyvalue = data_elem.split('=')
        metadata[keyvalue[0]] = float(keyvalue[1])
    return iso, metadata
 
 
def solar_calilbrate(fort, outDir, cnml, pnml, prog, refx, refz, refml, trials=10):
    """
    Function: solar_calibrate
    Description: Run multiple instances of dsep to calibrate a solar model
    Inputs:
                     - fort: dictionary linking unit number to file path
                     - outDir: directory to where output files will be saved
                     - cnml: path to control namelist file
                     - pnml: path to physics namelist file
                     - prog: path to DSEP executable
    Outputs:
                     - Parameters for calibrated solar model
    Post-State:
                     - DSEP saves output of the calibrates solar model
    """
    run_dsep(fort, outDir, cnml, pnml, prog, {"CMIXLA":[refml, refml], "RSCLX":[refx], "RSCLZ": [refz]})
    iso, metadata = load_iso("fort.37")
    refR = iso.iloc[-1]['Log_R']
    refL = iso.iloc[-1]['Log_L']
    refZX = iso.iloc[-1]['(Z/X)_surf']
 
 
    for trial in range(trials):
        cxx = refx + np.random.uniform(-0.01, 0.01)
        cxml = refml
        cxz = refz
        run_dsep(fort, outDir, cnml, pnml, prog, {"CMIXLA":[cxml, cxml], "RSCLX":[cxx], "RSCLZ": [cxz]})
        iso, metadata = load_iso("fort.37")
        cxR = iso.iloc[-1]['Log_R']
        cxL = iso.iloc[-1]['Log_L']
        cxZX = iso.iloc[-1]['(Z/X)_surf']
 
        cmlx = refx
        cmlml = refml + np.random.uniform(-0.01, 0.01)
        cmlz = refml
        run_dsep(fort, outDir, cnml, pnml, prog, {"CMIXLA":[cmlml, cmlml], "RSCLX":[cmlx], "RSCLZ": [cmlz]})
        iso, metadata = load_iso("fort.37")
        cmlR = iso.iloc[-1]['Log_R']
        cmlL = iso.iloc[-1]['Log_L']
        cmlZX = iso.iloc[-1]['(Z/X)_surf']
 
        czx = refx
        czml = refml
        czz = refz = np.random.uniform(-0.01, 0.01)
        run_dsep(fort, outDir, cnml, pnml, prog, {"CMIXLA":[czml, czml], "RSCLX":[czx], "RSCLZ": [czz]})
        iso, metadata = load_iso("fort.37")
        czR = iso.iloc[-1]['Log_R']
        czL = iso.iloc[-1]['Log_L']
        czZX = iso.iloc[-1]['(Z/X)_surf']
 
        dx = refx - cxx
        dldx = (refL - cxL)/dx
        drdx = (refR - cxR)/dx
        dzxdx = (refZX - cxZX)/dx
 
        da = refml - cmlml
        dlda = (refL - cmlL)/da
        drda = (refR - cmlR)/da
        dzxda = (refZX - cmlZX)/da
 
        dz = refz - czz
        dldz = (refL - czL)/dz
        drdz = (refR - czR)/dz
        dzxdz = (refZX - czZX)/dz
 
        da = (refL*drdx/dldx - refR)/(drda - dlda*drdx/dldx)
        dx = - (refL + dlda*da)/dldx
 
        newml = refml + da
        newx = refx + dx
 
        refml = newml
        refx = newx
        print(f"Current Best Guess is : X: {refx}, Ml: {refml}")
 
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run DSEP and automatically update paths for local system")
    parser.add_argument('cnamelist', type=str, help="control namelist file")
    parser.add_argument('pnamelist', type=str, help="physics namelist file")
    parser.add_argument("-b", "--branch", type=str, help="git branch executable to run")
    parser.add_argument("-o", "--output", type=str, help="Directory to save DSEP result too", default="output")
    parser.add_argument("-c", type=float, nargs=2, help="calibrate with [refX, refMl]")
    parser.add_argument("-g", action="store_true", help="run with gdb debugger", default=False)
    parser.add_argument("-rr", action="store_true", help="run with r reverse debugger", default=False)
 
    args = parser.parse_args()
    if args.branch:
        default = False
    else:
        default = True
    fort, prog = config_env(args.branch, default)
 
    # if the directory that the user wants to save do already exists, overwrite it
    if os.path.exists(args.output):
        shutil.rmtree(args.output)
    os.mkdir(args.output)
 
    if not args.c:
        run_dsep(fort, args.output, args.cnamelist, args.pnamelist, prog, debug=args.g, rr=args.rr)
    else:
        solar_calilbrate(fort, args.output, args.cnamelist, args.pnamelist, prog, args.c[0], args.c[1])