desk.sed_fit — Dusty-Evolved-Star-Kit 1.4.7 documentation

# Source code for desk.sed_fit

    from __future__ import absolute_import

    import csv
    import glob
    import copy
    import importlib
    import math
    import os
    import pdb
    import time
    from fnmatch import fnmatch
    from multiprocessing import Pool, cpu_count
    from multiprocessing import Process, Value, cpu_count

    import astropy.units as u
    import matplotlib.pyplot as plt
    import numpy as np
    from astropy.io import ascii
    from astropy.table import Table, Column
    from desk import config, plotting_seds, get_remote_models, parameter_ranges
    from desk.parameter_ranges import create_par
    from desk.plotting_seds import create_fig, single_fig
    from desk.interpolate_dusty import interpolate_dusty
    from matplotlib import rc
    from scipy import interpolate

    importlib.reload(config)

    """
    Steve Goldman
    Space Telescope Science Institute
    Nov 16, 2018
    sgoldman@stsci.edu

    This package takes a photometry or a spectrum (or spectra) in csv format and fits it/them with a grid of models.
    The grids used here are converted from output from the DUSTY code (Elitzur & Ivezic 2001, MNRAS, 327, 403) using
    the other script dusty_to_grid.py. The code interpolates and trims a version of the data and calculates the least
    squares value for each grid in the model and the data. The DUSTY outputs are then scaled and returned in files:
    fitting_results.csv and fitting_plotting_output.csv (for plotting the results). An example plotting script has also
    been provided.
    """

    def grids():
        print("\nGrids:")
        for item in config.grids:
            print("\t" + str(item))
        print("\n")

    def get_model(grid_name, teff_new, tinner_new, tau_new):
        interpolate_dusty(grid_name, float(teff_new), float(tinner_new), float(tau_new))

    def check_models(model_grid, full_path):
        global csv_file
        global fits_file
        csv_file = full_path + "models/" + model_grid + "_outputs.csv"
        fits_file = full_path + "models/" + model_grid + "_models.fits"
        if os.path.isfile(csv_file) and os.path.isfile(fits_file):
            print("\nYou already have the grid!\n")
            # print("Great job")

        else:
            user_proceed = input("Models not found locally, download the models [y]/n?: ")
            if user_proceed == "y" or user_proceed == "":
                get_remote_models.get_models(model_grid)
            elif user_proceed == "n":
                raise ValueError("Please make another model selection")
            elif user_proceed != "y" and user_proceed != "n":
                raise ValueError("Invalid selection")
        return (csv_file, fits_file)

    def make_output_files_dusty():
        with open("fitting_results.csv", "w") as f:
            f.write("source,L,vexp_predicted,teff,tinner,odep,mdot\n")
            f.close()
        with open("fitting_plotting_outputs.csv", "w") as f:
            f.write("target_name,data_file,norm,index,grid_name,teff,tinner,odep\n")
            f.close()

    def make_output_files_grams():
        with open("fitting_results.csv", "w") as f:
            f.write("source,L,rin,teff,tinner,odep,mdot\n")
            f.close()
        with open("fitting_plotting_outputs.csv", "w") as f:
            f.write("target_name,data_file,norm,index,grid_name,teff,tinner,odep\n")
            f.close()

    def get_data(filename):
        """
        :param filename: filename of input data. Should be csv with Column 0: wavelength in um and Col 1: flux in Jy
        :return: two arrays of wavelength (x) and flux (y) in unit specified in config.py
        """
        global log_average_flux_wm2
        table = ascii.read(filename, delimiter=",")
        table.sort(table.colnames[0])
        x = np.array(table.columns[0])
        y = np.array(table.columns[1])
        index = np.where(
            (x > config.fitting["wavelength_min"]) & (x < config.fitting["wavelength_max"])
        )
        x = x[index]
        y = y[index]
        y = y * u.Jy
        y = y.to(u.W / (u.m * u.m), equivalencies=u.spectral_density(x * u.um))
        log_average_flux_wm2 = np.log10(np.median(y).value)
        return x, np.array(y)

    def find_closest(target_wave, model_wave):
        """
        :param target_wave: Target wavelength in um
        :param model_wave: model wavelength in um
        :return: a 1-D array of the closest data wavelength values, to the model wavelength values
        """
        idx = np.searchsorted(model_wave[0], target_wave[0])
        idx = np.clip(idx, 1, len(model_wave[0]) - 1)
        left = model_wave[0][idx - 1]
        right = model_wave[0][idx]
        idx -= target_wave[0] - left < right - target_wave[0]
        closest_data_flux = model_wave[1][idx]
        return closest_data_flux

    def least2(data, model_l2):
        return np.nansum(np.square(model_l2 - data))

    def trim(data, model_trim):
        """
        :param data: input data in tuple of x and y arrays
        :param model_trim: input model
        :return: trims motimedel to wavelength range of data
        """
        indexes = np.where(
            np.logical_and(
                model_trim[0] >= np.min(data[0]), model_trim[0] <= np.max(data[0])
            )
        )
        return np.vstack([model_trim[0][indexes], model_trim[1][indexes]])

    def fit_norm(data, norm_model):
        """
        :param data: input data in tuple of x and y arrays
        :param norm_model: closest wavelength values to data in 1-D array (from trim)
        :return: trimmed model in 2 column np.array
        """
        global trials
        stats = []
        # normalization range
        # trials = np.linspace(config.fitting['min_norm'], config.fitting['max_norm'], config.fitting['ntrials'])
        trials = np.logspace(
            log_average_flux_wm2 - 2, log_average_flux_wm2 + 2, number_of_tries
        )
        for t in trials:
            stat = least2(data[1], norm_model * t)
            stats.append(stat)
        return stats

    # for each target, fit spectra with given models (.fits file)
    def sed_fitting(target):
        stat_values = []
        raw_data = get_data(target)  # gets target data
        # model_x = grid_dusty[0][0][:]  # gets model wavelengths
        for model in np.array(grid_dusty):
            trimmed_model = trim(
                raw_data, model
            )  # gets fluxes for corresponding wavelengths of data and models
            matched_model = find_closest(raw_data, trimmed_model)
            stat_values.append(fit_norm(raw_data, matched_model))
        stat_array = np.vstack(stat_values)
        argmin = np.argmin(stat_array)
        model_index = argmin // stat_array.shape[1]
        trial_index = argmin % stat_array.shape[1]
        target_name = (target.split("/")[-1][:15]).replace("IRAS-", "IRAS ")

        def dusty_fit():
            # calculates luminosity and scales outputs
            luminosity = int(
                np.power(10.0, distance_norm - math.log10(trials[trial_index]) * -1)
            )
            scaled_vexp = (
                float(grid_outputs[model_index]["vexp"]) * (luminosity / 10000) ** 0.25
            )
            scaled_mdot = (
                grid_outputs[model_index]["mdot"]
                * ((luminosity / 10000) ** 0.75)
                * (config.target["assumed_gas_to_dust_ratio"] / 200) ** 0.5
            )

            teff = int(grid_outputs[model_index]["teff"])
            tinner = int(grid_outputs[model_index]["tinner"])
            odep = grid_outputs[model_index]["odep"]

            # creates output file
            latex_array = [
                target_name,
                luminosity,
                np.round(scaled_vexp, 1),
                teff,
                tinner,
                odep,
                "%.3E" % float(scaled_mdot),
            ]

            plotting_array = [
                target_name,
                target,
                trials[trial_index],
                model_index,
                model_grid,
                teff,
                tinner,
                odep,
            ]

            # printed output
            if config.output["printed_output"] == "True":
                print()
                print()
                print(
                    (
                        "             Target: "
                        + target_name
                        + "        "
                        + str(counter.value + 1)
                        + "/"
                        + str(number_of_targets)
                    )
                )
                print("-------------------------------------------------")
                print(("Luminosity\t\t\t|\t" + str(round(luminosity))))
                print(
                    (
                        "Optical depth\t\t\t|\t"
                        + str(round(grid_outputs[model_index]["odep"], 3))
                    )
                )
                print(("Expansion velocity (scaled)\t|\t" + str(round(scaled_vexp, 2))))
                print(("Gas mass loss (scaled)\t\t|\t" + str("%.2E" % float(scaled_mdot))))
                print("-------------------------------------------------")
            with open("fitting_results.csv", "a") as f:
                writer = csv.writer(f, delimiter=",", lineterminator="\n")
                writer.writerow(np.array(latex_array))
                f.close()

            with open("fitting_plotting_outputs.csv", "a") as f:
                writer = csv.writer(f, delimiter=",", lineterminator="\n")
                writer.writerow(np.array(plotting_array))
                f.close()
            counter.value += 1

        def grams_fit():
            # pdb.set_trace()
            luminosity = grid_outputs[model_index]["lum"] * ((distance_value / 50) ** 2)
            teff = grid_outputs[model_index]["teff"]
            tinner = grid_outputs[model_index]["tinner"]
            odep = grid_outputs[model_index]["odep"]
            mdot = grid_outputs[model_index]["mdot"] * (distance_value / 50)
            rin = grid_outputs[model_index]["rin"] * (distance_value / 50)
            # creates output file
            latex_array = [
                target_name,
                luminosity,
                rin,
                teff,
                tinner,
                odep,
                "%.3E" % float(mdot),
            ]

            plotting_array = [
                target_name,
                target,
                trials[trial_index],
                model_index,
                model_grid,
                teff,
                odep,
            ]
            if config.output["printed_output"] == "True":
                print()
                print()
                print(
                    (
                        "             Target: "
                        + target_name
                        + "        "
                        + str(counter.value + 1)
                        + "/"
                        + str(number_of_targets)
                    )
                )
                print("-------------------------------------------------")
                print(("Luminosity\t\t\t|\t" + str(round(luminosity))))
                print(
                    (
                        "Optical depth\t\t\t|\t"
                        + str(round(grid_outputs[model_index]["odep"], 3))
                    )
                )
                print(("Inner Radius\t\t\t|\t" + str(rin)))
                print(("Dust production rate \t\t|\t" + str("%.2E" % float(mdot))))
                print("-------------------------------------------------")
            with open("fitting_results.csv", "a") as f:
                writer = csv.writer(f, delimiter=",", lineterminator="\n")
                writer.writerow(np.array(latex_array))
                f.close()

            with open("fitting_plotting_outputs.csv", "a") as f:
                writer = csv.writer(f, delimiter=",", lineterminator="\n")
                writer.writerow(np.array(plotting_array))
                f.close()
            counter.value += 1

        if fnmatch(model_grid, "grams*"):
            grams_fit()
        else:
            dusty_fit()

    def fit(
        source="default",
        distance=config.target["distance_in_kpc"],
        grid=config.fitting["model_grid"],
    ):
        """
        :param source: Name of target in array of strings (or one string)
        :param distance: distance to source(s) in kiloparsecs
        :param grid: Name of model grid
        :return:
        """
        # set variables
        global model_grid
        global counter
        global grid_dusty
        global grid_outputs
        global distance_norm
        global distance_value
        global full_path
        global files
        global number_of_tries
        global number_of_targets
        start = time.time()
        counter = Value("i", 0)
        number_of_tries = 200
        # normalization calculation
        # solar constant = 1379 W
        # distance to sun in kpc 4.8483E-9
        distance_value = float(copy.copy(distance))
        distance_norm = math.log10(((float(distance) / 4.8482e-9) ** 2) / 1379)
        full_path = str(__file__.replace("sed_fit.py", ""))

        # User input for models
        if grid == "carbon":
            model_grid = "Zubko-Crich-bb"
        elif grid == "oxygen":
            model_grid = "Oss-Orich-bb"
        else:
            if grid in config.grids:
                model_grid = grid
            else:
                raise ValueError(
                    "\n\nUnknown grid. Please make another model selection.\n\n To see options use: desk grids\n"
                )

        # check if models exist
        check_models(model_grid, full_path)

        # gets models
        grid_dusty = Table.read(fits_file)
        grid_outputs = Table.read(csv_file)

        # Model grid equal lengths check
        if len(grid_dusty) == len(grid_outputs):
            pass
        else:
            raise ValueError(
                "Model grid input error: mismatch in model spectra and model output"
            )

        # creates correct results files
        if fnmatch(model_grid, "grams*"):
            make_output_files_grams()
        else:
            make_output_files_dusty()

        # SED FITTING ###############################

        if source == "default":
            source = full_path + "put_target_data_here/"
        if fnmatch(source, "*.csv"):
            number_of_targets = 1
            sed_fitting(source)
        elif os.path.isdir(source):
            source_dir = (source + "/").replace("//", "/")
            if glob.glob(source_dir + "/*.csv"):
                files = os.listdir(source)
                files = glob.glob(source + "/" + "*.csv")
                number_of_targets = len(files)
                with Pool(processes=cpu_count() - 1) as pool:
                    pool.map(sed_fitting, [target_string for target_string in files])
            else:
                raise ValueError(
                    "\n\n\nERROR: No .csv files in that directory. Please make another selection.\n\n"
                )
        else:
            raise ValueError(
                "\n\n\nError: Not a .csv file. Please make another selection.\n\n"
            )

        # creating figures
        if config.output["create_figure"] == "yes":
            print("\n. . . Creating SED figure . . . . . . . . . . . .")
            create_fig()  # runs plotting script
        else:
            print(
                "No figure created. To automatically generate a figure change the "
                + '"create_figure" variable in the config.py script to "yes".'
            )

        if not os.path.isfile("parameter_ranges_" + config.fitting["model_grid"] + ".png"):
            print(". . . Creating parameter range figure . . . . . .")
            create_par()

        end = time.time()
        print()
        print("Time: " + str("%.2f" % float((end - start) / 60)) + " minutes")

    if __name__ == "__main__":
        fit()