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.
"""


[docs]def grids(): print("\nGrids:") for item in config.grids: print("\t" + str(item)) print("\n")
[docs]def get_model(grid_name, teff_new, tinner_new, tau_new): interpolate_dusty(grid_name, float(teff_new), float(tinner_new), float(tau_new))
[docs]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)
[docs]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()
[docs]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()
[docs]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)
[docs]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
[docs]def least2(data, model_l2): return np.nansum(np.square(model_l2 - data))
[docs]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]])
[docs]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)
[docs]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()
[docs]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()