MESA Profile File Aggregation Script (log_L, log_Teff)

extract_log_L_and_log_Teff.py

import os.path
import sys

import numpy as np
import pandas as pd

def main():
    """
    $ python extract_log_L_and_log_Teff.py LOGS/ log_L_and_log_Teff.csv
    """
    logs_directory = sys.argv[1]
    output_file = sys.argv[2]

    num_profiles = get_num_profiles(logs_directory)

    data = aggregate_profile_files(logs_directory, num_profiles)

    final_data = process_data(data)

    final_data.to_csv(output_file, index=False)

def get_num_profiles(logs_directory):
    """
    Read the "profiles.index" file in the LOGS directory to find out how many
    profile files MESA created.
    """
    profiles_file_name = "profiles.index"
    profiles_file = os.path.join(logs_directory, profiles_file_name)

    with open(profiles_file) as f:
        first_line = f.readline()

        # Pull out the total number of profile files from the first line of the
        # file by doing some funny string manipulation
        num_profiles_str = first_line.split(" models.")[0]
        num_profiles = int(num_profiles_str)

        return num_profiles

def aggregate_profile_files(logs_directory, num_profiles):
    """
    Reads in all of the profile files in the specified LOGS directory and
    aggregates them into a single Pandas DataFrame.
    """
    all_data = []

    for p_num in range(1, num_profiles + 1):
        profile_file_name = "profile{}.data".format(p_num)
        profile_file = os.path.join(logs_directory, profile_file_name)

        attributes = read_profile_attributes(profile_file)

        attributes["profile"] = p_num

        all_data.append(attributes)

    return pd.concat(all_data)

def read_profile_attributes(profile_file):
    """
    Reads in the attributes in the profile file at the given filepath into a
    Pandas DataFrame.
    """
    attributes_start_row = 1

    attributes = pd.read_fwf(profile_file, skiprows=attributes_start_row, nrows=1)

    return attributes

def process_data(data):
    """
    Filters down the data to just the desired columns and also log scales some
    of the relevant columns.
    """
    # Pull out the columns that we want to keep or log scale
    luminosity = data["photosphere_L"]
    teff = data["Teff"]
    profiles = data["profile"]

    # Log scale the luminosity and effective temperature
    log_l = np.log10(luminosity)
    log_teff = np.log10(teff)

    # Create a new data frame to store the desired columns
    final_data = pd.DataFrame()

    # Copy over the columns that we want to keep
    final_data["profile"] = profiles
    final_data["log_L"] = log_l
    final_data["log_Teff"] = log_teff

    return final_data

if __name__ == "__main__":
    main()