Sketch of a `Synfit.py`

synfit_sketch.py

import os
import numpy as np
import matplotlib.pyplot as plt
import s4


#Synplot path
SPATH = os.getenv('HOME') + '/synplot/synplot/'


def synfit_wrp(fits_file_path, teff, wstart, wend, params):
    """
    Run synfit for a specific Teff and find the best value for logg.

    Parameters
    ----------

    fits_file_path: str;
        Path to the FITS spectrum to be fitted.

    teff: float;
        Effective temperature of the to be fitted spectrum.

    wstart: float;
        Starting wavelength.

    wend: float;
        Ending wavelength.

    params: dic;
        A dictionary containing synplot inputs

    Returns
    -------

    loggm: float;
        Best logg found

    chi2m: float;
        Value of the chi-square of the fitting
    """

    #Load the FITS file and save in the SPATH directory
    spec_data = s4.io.load_spectrum(os.path.realpath(fits_file_path))
    fits_file = 'spec.dat'
    np.savetxt(SPATH + fits_file, spec_data)

    #Make a deep copy of the params dic
    params_copy = params.copy()

    #Add wstart and wend to params_copy
    params_copy['wstart'] = wstart
    params_copy['wend'] = wend

    #write the parameters on a single string
    params_string = ','.join([k+'='+str(v) for k,v in params_copy.iteritems()])

    #generate the synfit command
    synfit_cmd = 'idl -e "CD, \''+SPATH+"' & synfit, " + \
                 "tllim="+ str(teff) + ", tulim=" + str(teff) + ", " + \
                 "gstep=0.1, obs='"+ fits_file +"',"+ \
                 params_string + '"'

    #run synfit
    s4.io.run_command(synfit_cmd)

    # Delete FITS file from synplot folder
    os.remove(SPATH + fits_file)

    #obtain the best fit
    chi2 = np.loadtxt(SPATH + 'chisq.dat')

    _, loggm, _, _, chi2m = chi2[chi2[:,-1].argmin()]

    if loggm == 4.5:
        loggm = 4.499

    #plot both
    # The radial correction fo synplot is different from the S4.
    # On synplot, the synthetic spectrum is dislocated to the stellar frame
    # while on S4, the stellar spectrum is dislocated to the synthetic frame
    # i.e., the rest frame.
    # We have a problem with tha masks/windows.
    # They are obtained in the stellar frame, but we are plotting on the rest
    #frame. So we have to apply a radial velocity correction for them too.
    if 'mask' in params_copy:
        params_copy['mask'] = list(np.array(params_copy['mask'])* \
                              s4.spectools.rvcorr(params_copy['rv']))

    # Calculate the synthetic spectrum of the best fit
    best_fit = s4.synthesis.Synplot(teff, loggm, SPATH, **params_copy)
    best_fit.run()
    if 'scale' in params_copy:
        best_fit.apply_scale()

    # Plot figure
    fig = plt.figure()

    # Set radial velocity eo 0 if not defined
    if 'rv' not in params_copy:
        params_copy['rv'] = 0

    # Plot observation spectrum
    plt.plot(spec_data[:, 0] * s4.spectools.rvcorr(params_copy['rv']),
             spec_data[:, 1], 'b-', alpha=0.8, label = 'Observation')

    # Calculate two syntehtic spectra, each one with a lower and upper value
    # of logg than the best_fit
    logg_step = 0.2

    if loggm - logg_step >= 3.5:
        lower_logg = s4.synthesis.Synplot(teff, loggm-logg_step, SPATH,
                                          **params_copy)
        lower_logg.run()

        if 'scale' in params_copy:
            lower_logg.apply_scale()

        plt.plot(lower_logg.spectrum[: , 0], lower_logg.spectrum[:, 1],
                'm-', lw=2.5, label = '$\log g - 0.2$')

    if loggm + logg_step < 4.5:
        upper_logg = s4.synthesis.Synplot(teff, loggm+logg_step, SPATH,
                                          **params_copy)
        upper_logg.run()

        if 'scale' in params_copy:
            upper_logg.apply_scale()

        plt.plot(upper_logg.spectrum[:, 0], upper_logg.spectrum[:, 1],
                'g-', lw=2.5,  label = '$\log g + 0.2$')

    # Plot best fit
    plt.plot(best_fit.spectrum[:, 0], best_fit.spectrum[:, 1],
             'r-', lw=3.5, label = 'Best fit')

    if 'mask' in params_copy:
        s4.plottools.plot_windows(params_copy['mask'])
    plt.legend(loc='lower right')
    plt.xlim(wstart, wend)
    plt.ylim(0, 1.1)
    plt.xlabel(r'Wavelength $(\AA)$')
    plt.ylabel('Normalized Flux')
    plt.title(r'$T$={} K, $\log g$={:.2f}, $\chi^2$={:.4f}'.format(teff, loggm, chi2m))

    return loggm, chi2m


def synfit_loop(fits_path, wstart, wend, teff_range, params_dic):
    """
    Run synfit on a loop.

    Parameters
    ----------

    fits_path: str;
        Path to the FITS spectrum to be fitted.

    teff: float;
        Effective temperature of the to be fitted spectrum.

    wstart: float;
        Starting wavelength.

    wend: float;
        Ending wavelength.

    teff_range: list;
        List of teff values in which will run synfit.

    param_dics: dic;
        A dictionary containing synplot inputs

    Return
    ------

    fits: numpy.ndarray;
        N-dimension array with [teff, logg, chi2]
    """

    fits = []
    for temp in teff_range:
        loggm, chi2m = synfit_wrp(fits_path, temp, wstart, wend, params_dic)

        for i in [temp, loggm, chi2m]:
            fits.append(i)

    fits = np.array(fits).reshape(-1, 3)

    return fits