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()