import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import emcee
import apogee_tools as ap
import corner
import argparse
import os


# =============================================================
# MCMC prior and likelihood functions
# =============================================================

def lnlike(theta, lsf, tell_sp):

	"""
	Log-likelihood, computed from chi-squared
	"""

	# If theta is entered as a list, make it into a dictionary
	theta_keys = [key for key in ap.init.keys()]
	if type(theta) == np.ndarray:
		theta = dict(zip(theta_keys, theta))

	# Choose the appropriate Spectrum class to read the data
	if ap.data['instrument'] == 'APOGEE':
		data = ap.Apogee(id=ap.data['ID'], type=ap.data["dtype"], visit=ap.data['visit'])
	else:
		print('No Spectrum class to read data for instrument', ap.data['instrument'])

	chisq = ap.returnModelFit(data, theta, lsf=lsf, telluric=tell_sp)

	print('\n chisq', chisq, '\n')

	return -0.5 * chisq


def lnprior(theta):

	"""
	Specifies a flat prior
	"""

	# keys = theta.keys()
	theta_keys = [key for key in ap.init.keys()]
	if type(theta) == np.ndarray:
		theta = dict(zip(theta_keys, theta))
	keys = theta.keys()

	for k in keys:
		if (ap.prior[k][0] < theta[k] < ap.prior[k][1]):
			pass
		else:
			return -np.inf
			break
	return 0.0


def lnprob(theta, lsf, tell_sp):

	lnp = lnprior(theta)
	if not np.isfinite(lnp):
	    return -np.inf

	return lnp + lnlike(theta, lsf, tell_sp)


#########################################################################################

"""
Conventions:
param - full list of parameters
theta - only parameters that are being sampled
"""

if __name__ == "__main__":

	if 'config.yaml' not in os.listdir():
		print('\nError: config.yaml not found in the current working directory. \
			Using default file found inside apogee_tools.\n')

	# =============================================================
	# Command line input
	# =============================================================

	parser = argparse.ArgumentParser(description='Specify plotting directory.')
	parser.add_argument("plot", action="store", type=str)
	args = parser.parse_args()


	# =============================================================
	# Testing...
	# =============================================================

	if 'make_model' in args.plot:

		init_param, step_param, init_theta, step_theta, fiber, tell_sp, lsf = ap.initialize()
		mdl = ap.makeModel(params=init_param, lsf=lsf, telluric=tell_sp, plot=True)

	if 'test_fit' in args.plot:

		data = ap.Apogee(id=ap.data['ID'], type=ap.data["dtype"], visit=ap.data['visit'])
		chi_sq = ap.returnModelFit(data, init_param, lsf=lsf, plot=True)

		print('chi^2', chi_sq)

	if 'test_telluric' in args.plot:

		mdl = ap.getModel(params=[3200, 5.0, 0.0], grid='BTSETTL', xrange=[15200,16940])
		tell_sp = ap.applyTelluric(mdl, ap.getTelluric(), cut_rng=[min(mdl.wave), max(mdl.wave)])

		tell_sp.plot()


	# =============================================================
	# Run MCMC!
	# =============================================================

	if (ap.out["mcmc_sampler"] == True) or ('mcmc' in args.plot):

		init_param, step_param, init_theta, step_theta, fiber, tell_sp, lsf = ap.initialize()

		theta_keys = list(init_theta.keys())
		theta_vals = list(init_theta.values())

		ndim = len(init_theta)
		nsteps = ap.mcmc["nsteps"]
		nwalkers = ap.mcmc["nwalkers"]

		pos = [list(init_theta.values()) + 1e-4*np.random.randn(ndim) for i in range(nwalkers)]

		sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=(lsf, tell_sp))
		sampler.run_mcmc(pos, nsteps)

		np.save('sampler_chain', sampler.chain[:, :, :])

		samples = sampler.chain[:, :, :].reshape((-1, ndim))

		np.save('samples', samples)


	# =============================================================
	# Output corner/walker plots
	# =============================================================

	lbl = ['Teff', 'logg', '[Fe/H]', 'rv', 'vsini', r'$\alpha$']

	if ap.out["corner"] == True:

		try:
			fig = corner.corner(samples, labels=lbl, truths=theta_vals)
			fig.savefig("triangle.png")

		except:
			print('Traingle plot failed.')


	if 'corner' in args.plot:

		try:
			samples = np.load('samples.npy')
			fig = corner.corner(samples, labels=lbl, truths=theta_vals)
			fig.savefig("triangle.png")

		except:
			print('Traingle plot failed.')


	if (ap.out["walkers"] == True) or ('walkers' in args.plot):

		try:
			sampler_chain = np.load('sampler_chain.npy')

			ndim = len(sampler.chain.T)

			fig, ax = plt.subplots(ndim, sharex=True, figsize=[8,12])
			for i in range(ndim):
				ax[i].plot(sampler_chain.T[i], '-k', alpha=0.2);
				ax[i].set_ylabel(str(lbl[i]))
				if i == ndim:
					ax[i].set_xlabel(step)
			plt.tight_layout()
			plt.savefig('Walkers.png', dpi=300, bbox_inches='tight')
			plt.show()
			plt.close()

		except:
			print('Corner plot failed.')


23cf7b50ca22b47fed0f0bd510b9d2cdf57118a4
<mask token>


def lnlike(theta, lsf, tell_sp):
    """
	Log-likelihood, computed from chi-squared
	"""
    theta_keys = [key for key in ap.init.keys()]
    if type(theta) == np.ndarray:
        theta = dict(zip(theta_keys, theta))
    if ap.data['instrument'] == 'APOGEE':
        data = ap.Apogee(id=ap.data['ID'], type=ap.data['dtype'], visit=ap.
            data['visit'])
    else:
        print('No Spectrum class to read data for instrument', ap.data[
            'instrument'])
    chisq = ap.returnModelFit(data, theta, lsf=lsf, telluric=tell_sp)
    print('\n chisq', chisq, '\n')
    return -0.5 * chisq


def lnprior(theta):
    """
	Specifies a flat prior
	"""
    theta_keys = [key for key in ap.init.keys()]
    if type(theta) == np.ndarray:
        theta = dict(zip(theta_keys, theta))
    keys = theta.keys()
    for k in keys:
        if ap.prior[k][0] < theta[k] < ap.prior[k][1]:
            pass
        else:
            return -np.inf
            break
    return 0.0


<mask token>

<mask token>


def lnlike(theta, lsf, tell_sp):
    """
	Log-likelihood, computed from chi-squared
	"""
    theta_keys = [key for key in ap.init.keys()]
    if type(theta) == np.ndarray:
        theta = dict(zip(theta_keys, theta))
    if ap.data['instrument'] == 'APOGEE':
        data = ap.Apogee(id=ap.data['ID'], type=ap.data['dtype'], visit=ap.
            data['visit'])
    else:
        print('No Spectrum class to read data for instrument', ap.data[
            'instrument'])
    chisq = ap.returnModelFit(data, theta, lsf=lsf, telluric=tell_sp)
    print('\n chisq', chisq, '\n')
    return -0.5 * chisq


def lnprior(theta):
    """
	Specifies a flat prior
	"""
    theta_keys = [key for key in ap.init.keys()]
    if type(theta) == np.ndarray:
        theta = dict(zip(theta_keys, theta))
    keys = theta.keys()
    for k in keys:
        if ap.prior[k][0] < theta[k] < ap.prior[k][1]:
            pass
        else:
            return -np.inf
            break
    return 0.0


def lnprob(theta, lsf, tell_sp):
    lnp = lnprior(theta)
    if not np.isfinite(lnp):
        return -np.inf
    return lnp + lnlike(theta, lsf, tell_sp)


<mask token>

<mask token>


def lnlike(theta, lsf, tell_sp):
    """
	Log-likelihood, computed from chi-squared
	"""
    theta_keys = [key for key in ap.init.keys()]
    if type(theta) == np.ndarray:
        theta = dict(zip(theta_keys, theta))
    if ap.data['instrument'] == 'APOGEE':
        data = ap.Apogee(id=ap.data['ID'], type=ap.data['dtype'], visit=ap.
            data['visit'])
    else:
        print('No Spectrum class to read data for instrument', ap.data[
            'instrument'])
    chisq = ap.returnModelFit(data, theta, lsf=lsf, telluric=tell_sp)
    print('\n chisq', chisq, '\n')
    return -0.5 * chisq


def lnprior(theta):
    """
	Specifies a flat prior
	"""
    theta_keys = [key for key in ap.init.keys()]
    if type(theta) == np.ndarray:
        theta = dict(zip(theta_keys, theta))
    keys = theta.keys()
    for k in keys:
        if ap.prior[k][0] < theta[k] < ap.prior[k][1]:
            pass
        else:
            return -np.inf
            break
    return 0.0


def lnprob(theta, lsf, tell_sp):
    lnp = lnprior(theta)
    if not np.isfinite(lnp):
        return -np.inf
    return lnp + lnlike(theta, lsf, tell_sp)


<mask token>
if __name__ == '__main__':
    if 'config.yaml' not in os.listdir():
        print(
            """
Error: config.yaml not found in the current working directory. 			Using default file found inside apogee_tools.
"""
            )
    parser = argparse.ArgumentParser(description='Specify plotting directory.')
    parser.add_argument('plot', action='store', type=str)
    args = parser.parse_args()
    if 'make_model' in args.plot:
        (init_param, step_param, init_theta, step_theta, fiber, tell_sp, lsf
            ) = ap.initialize()
        mdl = ap.makeModel(params=init_param, lsf=lsf, telluric=tell_sp,
            plot=True)
    if 'test_fit' in args.plot:
        data = ap.Apogee(id=ap.data['ID'], type=ap.data['dtype'], visit=ap.
            data['visit'])
        chi_sq = ap.returnModelFit(data, init_param, lsf=lsf, plot=True)
        print('chi^2', chi_sq)
    if 'test_telluric' in args.plot:
        mdl = ap.getModel(params=[3200, 5.0, 0.0], grid='BTSETTL', xrange=[
            15200, 16940])
        tell_sp = ap.applyTelluric(mdl, ap.getTelluric(), cut_rng=[min(mdl.
            wave), max(mdl.wave)])
        tell_sp.plot()
    if ap.out['mcmc_sampler'] == True or 'mcmc' in args.plot:
        (init_param, step_param, init_theta, step_theta, fiber, tell_sp, lsf
            ) = ap.initialize()
        theta_keys = list(init_theta.keys())
        theta_vals = list(init_theta.values())
        ndim = len(init_theta)
        nsteps = ap.mcmc['nsteps']
        nwalkers = ap.mcmc['nwalkers']
        pos = [(list(init_theta.values()) + 0.0001 * np.random.randn(ndim)) for
            i in range(nwalkers)]
        sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=(lsf,
            tell_sp))
        sampler.run_mcmc(pos, nsteps)
        np.save('sampler_chain', sampler.chain[:, :, :])
        samples = sampler.chain[:, :, :].reshape((-1, ndim))
        np.save('samples', samples)
    lbl = ['Teff', 'logg', '[Fe/H]', 'rv', 'vsini', '$\\alpha$']
    if ap.out['corner'] == True:
        try:
            fig = corner.corner(samples, labels=lbl, truths=theta_vals)
            fig.savefig('triangle.png')
        except:
            print('Traingle plot failed.')
    if 'corner' in args.plot:
        try:
            samples = np.load('samples.npy')
            fig = corner.corner(samples, labels=lbl, truths=theta_vals)
            fig.savefig('triangle.png')
        except:
            print('Traingle plot failed.')
    if ap.out['walkers'] == True or 'walkers' in args.plot:
        try:
            sampler_chain = np.load('sampler_chain.npy')
            ndim = len(sampler.chain.T)
            fig, ax = plt.subplots(ndim, sharex=True, figsize=[8, 12])
            for i in range(ndim):
                ax[i].plot(sampler_chain.T[i], '-k', alpha=0.2)
                ax[i].set_ylabel(str(lbl[i]))
                if i == ndim:
                    ax[i].set_xlabel(step)
            plt.tight_layout()
            plt.savefig('Walkers.png', dpi=300, bbox_inches='tight')
            plt.show()
            plt.close()
        except:
            print('Corner plot failed.')

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import emcee
import apogee_tools as ap
import corner
import argparse
import os


def lnlike(theta, lsf, tell_sp):
    """
	Log-likelihood, computed from chi-squared
	"""
    theta_keys = [key for key in ap.init.keys()]
    if type(theta) == np.ndarray:
        theta = dict(zip(theta_keys, theta))
    if ap.data['instrument'] == 'APOGEE':
        data = ap.Apogee(id=ap.data['ID'], type=ap.data['dtype'], visit=ap.
            data['visit'])
    else:
        print('No Spectrum class to read data for instrument', ap.data[
            'instrument'])
    chisq = ap.returnModelFit(data, theta, lsf=lsf, telluric=tell_sp)
    print('\n chisq', chisq, '\n')
    return -0.5 * chisq


def lnprior(theta):
    """
	Specifies a flat prior
	"""
    theta_keys = [key for key in ap.init.keys()]
    if type(theta) == np.ndarray:
        theta = dict(zip(theta_keys, theta))
    keys = theta.keys()
    for k in keys:
        if ap.prior[k][0] < theta[k] < ap.prior[k][1]:
            pass
        else:
            return -np.inf
            break
    return 0.0


def lnprob(theta, lsf, tell_sp):
    lnp = lnprior(theta)
    if not np.isfinite(lnp):
        return -np.inf
    return lnp + lnlike(theta, lsf, tell_sp)


<mask token>
if __name__ == '__main__':
    if 'config.yaml' not in os.listdir():
        print(
            """
Error: config.yaml not found in the current working directory. 			Using default file found inside apogee_tools.
"""
            )
    parser = argparse.ArgumentParser(description='Specify plotting directory.')
    parser.add_argument('plot', action='store', type=str)
    args = parser.parse_args()
    if 'make_model' in args.plot:
        (init_param, step_param, init_theta, step_theta, fiber, tell_sp, lsf
            ) = ap.initialize()
        mdl = ap.makeModel(params=init_param, lsf=lsf, telluric=tell_sp,
            plot=True)
    if 'test_fit' in args.plot:
        data = ap.Apogee(id=ap.data['ID'], type=ap.data['dtype'], visit=ap.
            data['visit'])
        chi_sq = ap.returnModelFit(data, init_param, lsf=lsf, plot=True)
        print('chi^2', chi_sq)
    if 'test_telluric' in args.plot:
        mdl = ap.getModel(params=[3200, 5.0, 0.0], grid='BTSETTL', xrange=[
            15200, 16940])
        tell_sp = ap.applyTelluric(mdl, ap.getTelluric(), cut_rng=[min(mdl.
            wave), max(mdl.wave)])
        tell_sp.plot()
    if ap.out['mcmc_sampler'] == True or 'mcmc' in args.plot:
        (init_param, step_param, init_theta, step_theta, fiber, tell_sp, lsf
            ) = ap.initialize()
        theta_keys = list(init_theta.keys())
        theta_vals = list(init_theta.values())
        ndim = len(init_theta)
        nsteps = ap.mcmc['nsteps']
        nwalkers = ap.mcmc['nwalkers']
        pos = [(list(init_theta.values()) + 0.0001 * np.random.randn(ndim)) for
            i in range(nwalkers)]
        sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=(lsf,
            tell_sp))
        sampler.run_mcmc(pos, nsteps)
        np.save('sampler_chain', sampler.chain[:, :, :])
        samples = sampler.chain[:, :, :].reshape((-1, ndim))
        np.save('samples', samples)
    lbl = ['Teff', 'logg', '[Fe/H]', 'rv', 'vsini', '$\\alpha$']
    if ap.out['corner'] == True:
        try:
            fig = corner.corner(samples, labels=lbl, truths=theta_vals)
            fig.savefig('triangle.png')
        except:
            print('Traingle plot failed.')
    if 'corner' in args.plot:
        try:
            samples = np.load('samples.npy')
            fig = corner.corner(samples, labels=lbl, truths=theta_vals)
            fig.savefig('triangle.png')
        except:
            print('Traingle plot failed.')
    if ap.out['walkers'] == True or 'walkers' in args.plot:
        try:
            sampler_chain = np.load('sampler_chain.npy')
            ndim = len(sampler.chain.T)
            fig, ax = plt.subplots(ndim, sharex=True, figsize=[8, 12])
            for i in range(ndim):
                ax[i].plot(sampler_chain.T[i], '-k', alpha=0.2)
                ax[i].set_ylabel(str(lbl[i]))
                if i == ndim:
                    ax[i].set_xlabel(step)
            plt.tight_layout()
            plt.savefig('Walkers.png', dpi=300, bbox_inches='tight')
            plt.show()
            plt.close()
        except:
            print('Corner plot failed.')

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import emcee
import apogee_tools as ap
import corner
import argparse
import os


# =============================================================
# MCMC prior and likelihood functions
# =============================================================

def lnlike(theta, lsf, tell_sp):

	"""
	Log-likelihood, computed from chi-squared
	"""

	# If theta is entered as a list, make it into a dictionary
	theta_keys = [key for key in ap.init.keys()]
	if type(theta) == np.ndarray:
		theta = dict(zip(theta_keys, theta))

	# Choose the appropriate Spectrum class to read the data
	if ap.data['instrument'] == 'APOGEE':
		data = ap.Apogee(id=ap.data['ID'], type=ap.data["dtype"], visit=ap.data['visit'])
	else:
		print('No Spectrum class to read data for instrument', ap.data['instrument'])

	chisq = ap.returnModelFit(data, theta, lsf=lsf, telluric=tell_sp)

	print('\n chisq', chisq, '\n')

	return -0.5 * chisq


def lnprior(theta):

	"""
	Specifies a flat prior
	"""

	# keys = theta.keys()
	theta_keys = [key for key in ap.init.keys()]
	if type(theta) == np.ndarray:
		theta = dict(zip(theta_keys, theta))
	keys = theta.keys()

	for k in keys:
		if (ap.prior[k][0] < theta[k] < ap.prior[k][1]):
			pass
		else:
			return -np.inf
			break
	return 0.0


def lnprob(theta, lsf, tell_sp):

	lnp = lnprior(theta)
	if not np.isfinite(lnp):
	    return -np.inf

	return lnp + lnlike(theta, lsf, tell_sp)


#########################################################################################

"""
Conventions:
param - full list of parameters
theta - only parameters that are being sampled
"""

if __name__ == "__main__":

	if 'config.yaml' not in os.listdir():
		print('\nError: config.yaml not found in the current working directory. \
			Using default file found inside apogee_tools.\n')

	# =============================================================
	# Command line input
	# =============================================================

	parser = argparse.ArgumentParser(description='Specify plotting directory.')
	parser.add_argument("plot", action="store", type=str)
	args = parser.parse_args()


	# =============================================================
	# Testing...
	# =============================================================

	if 'make_model' in args.plot:

		init_param, step_param, init_theta, step_theta, fiber, tell_sp, lsf = ap.initialize()
		mdl = ap.makeModel(params=init_param, lsf=lsf, telluric=tell_sp, plot=True)

	if 'test_fit' in args.plot:

		data = ap.Apogee(id=ap.data['ID'], type=ap.data["dtype"], visit=ap.data['visit'])
		chi_sq = ap.returnModelFit(data, init_param, lsf=lsf, plot=True)

		print('chi^2', chi_sq)

	if 'test_telluric' in args.plot:

		mdl = ap.getModel(params=[3200, 5.0, 0.0], grid='BTSETTL', xrange=[15200,16940])
		tell_sp = ap.applyTelluric(mdl, ap.getTelluric(), cut_rng=[min(mdl.wave), max(mdl.wave)])

		tell_sp.plot()


	# =============================================================
	# Run MCMC!
	# =============================================================

	if (ap.out["mcmc_sampler"] == True) or ('mcmc' in args.plot):

		init_param, step_param, init_theta, step_theta, fiber, tell_sp, lsf = ap.initialize()

		theta_keys = list(init_theta.keys())
		theta_vals = list(init_theta.values())

		ndim = len(init_theta)
		nsteps = ap.mcmc["nsteps"]
		nwalkers = ap.mcmc["nwalkers"]

		pos = [list(init_theta.values()) + 1e-4*np.random.randn(ndim) for i in range(nwalkers)]

		sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=(lsf, tell_sp))
		sampler.run_mcmc(pos, nsteps)

		np.save('sampler_chain', sampler.chain[:, :, :])

		samples = sampler.chain[:, :, :].reshape((-1, ndim))

		np.save('samples', samples)


	# =============================================================
	# Output corner/walker plots
	# =============================================================

	lbl = ['Teff', 'logg', '[Fe/H]', 'rv', 'vsini', r'$\alpha$']

	if ap.out["corner"] == True:

		try:
			fig = corner.corner(samples, labels=lbl, truths=theta_vals)
			fig.savefig("triangle.png")

		except:
			print('Traingle plot failed.')


	if 'corner' in args.plot:

		try:
			samples = np.load('samples.npy')
			fig = corner.corner(samples, labels=lbl, truths=theta_vals)
			fig.savefig("triangle.png")

		except:
			print('Traingle plot failed.')


	if (ap.out["walkers"] == True) or ('walkers' in args.plot):

		try:
			sampler_chain = np.load('sampler_chain.npy')

			ndim = len(sampler.chain.T)

			fig, ax = plt.subplots(ndim, sharex=True, figsize=[8,12])
			for i in range(ndim):
				ax[i].plot(sampler_chain.T[i], '-k', alpha=0.2);
				ax[i].set_ylabel(str(lbl[i]))
				if i == ndim:
					ax[i].set_xlabel(step)
			plt.tight_layout()
			plt.savefig('Walkers.png', dpi=300, bbox_inches='tight')
			plt.show()
			plt.close()

		except:
			print('Corner plot failed.')

