text:
#!/usr/bin/env python

"""SPECFIT.PY - Generic stellar abundance determination software

"""

from __future__ import print_function

__authors__ = 'David Nidever <dnidever@montana.edu>'
__version__ = '20200711'  # yyyymmdd

import os
import shutil
import contextlib, io, sys
import numpy as np
import warnings
from astropy.io import fits
from astropy.table import Table
from dlnpyutils.minpack import curve_fit
from dlnpyutils.least_squares import least_squares
from scipy.interpolate import interp1d
from dlnpyutils import utils as dln, bindata, astro
import doppler
from doppler.spec1d import Spec1D
from doppler import (cannon,utils,reader)
import copy
import logging
import time
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.legend import Legend
import tempfile
from . import models
from synple import synple

# Ignore these warnings, it's a bug
warnings.filterwarnings("ignore", message="numpy.dtype size changed")
warnings.filterwarnings("ignore", message="numpy.ufunc size changed")

cspeed = 2.99792458e5  # speed of light in km/s

def synmodel(spec,params,alinefile=None,mlinefile=None,verbose=False,normalize=True):
    """
    Synthetic spectrum model.

    Parameters
    ----------
    spec : Spec1D object or str
         The observed Spec1D spectrum to match or the name of a spectrum file.
    params : dict
         Dictionary of initial values to use or parameters/elements to hold fixed.
    normalize : bool, optional
         Renormalize the model spectrum using the observed spectrum's continuum function.  The
         synthetic spectrum will already have been normalized using the "true" continuum.  This
         step is to simulate any systematic effects of the spectrum normalization algorithm that
         the observed spectrum undergoes.  Default is True.
    verbose : int, optional
         Verbosity level (0, 1, or 2).  The default is 0 and verbose=2 is for debugging.
    alinefile : str, optional
         The atomic linelist to use.  Default is None which means the default synple linelist is used.
    mlinefile : str, optional
         The molecular linelist to use.  Default is None which means the default synple linelist is used.

    Returns
    -------
    model : Spec1D object
        The synthetic spectrum.  The "true" continuum is in model.cont.

    Example
    -------

    .. code-block:: python

         model = synmodel(spec,params)
    """

    # Read in the spectrum
    if type(spec) is str:
        filename = spec
        spec = doppler.read(filename)
        if spec is None:
            print('Problem loading '+filename)
            return

    params = dict((key.upper(), value) for (key, value) in params.items())  # all CAPS

    # Initialize the fitter
    fitparams = ['TEFF']  # "dummy" fitting variable
    spfitter = SpecFitter(spec,params,fitparams=fitparams,verbose=(verbose>=2),
                          alinefile=alinefile,mlinefile=mlinefile)
    spfitter.norm = normalize  # normalize the synthetic spectrum
    model = spfitter.model(spec.wave.flatten(),params['TEFF'],retobj=True)
    model.instrument = 'Model'

    return model

class SpecFitter:
    def __init__ (self,spec,params,fitparams=None,norm=True,verbose=False,
                  alinefile=None,mlinefile=None):
        # Parameters
        self.params = params
        if fitparams is not None:
            self.fitparams = fitparams
        else:
            self.fitparams = list(params.keys())  # by default fit all parameters
        self.nsynfev = 0  # number of synthetic spectra made
        self.njac = 0     # number of times jacobian called
        # Save spectrum information
        self.spec = spec.copy()
        self.flux = spec.flux.flatten()
        self.err = spec.err.flatten()
        self.wave = spec.wave.flatten()
        self.lsf = spec.lsf.copy()
        self.lsf.wavevac = spec.wavevac  # need this later for synspec prep
        self.wavevac = spec.wavevac
        self.verbose = verbose
        self.norm = norm    # normalize
        self.continuum_func = spec.continuum_func
        self.alinefile = alinefile
        self.mlinefile = mlinefile
        # Convert vacuum to air wavelengths
        #  synspec uses air wavelengths
        if spec.wavevac is True:
            wave = astro.vactoair(spec.wave.copy().flatten()).reshape(spec.wave.shape)
        else:
            wave = spec.wave.copy()
        if wave.ndim==1:
            wave = np.atleast_2d(wave).T
        # Figure out the wavelength parameters
        npix = spec.npix
        norder = spec.norder
        xp = np.arange(npix//20)*20
        wr = np.zeros((spec.lsf.norder,2),np.float64)
        dw = np.zeros(spec.lsf.norder,np.float64)
        mindw = np.zeros(norder,np.float64)
        for o in range(spec.norder):
            dw[o] = np.median(dln.slope(wave[:,o]))
            wr[o,0] = np.min(wave[:,o])
            wr[o,1] = np.max(wave[:,o])
            fwhm = spec.lsf.fwhm(wave[xp,o],xtype='Wave',order=o)
            # FWHM is in units of lsf.xtype, convert to wavelength/angstroms, if necessary
            if spec.lsf.xtype.lower().find('pix')>-1:
                fwhm *= np.abs(dw[o])

            # need at least ~4 pixels per LSF FWHM across the spectrum
            #  using 3 affects the final profile shape
            mindw[o] = np.min(fwhm/4)
        self._dwair = np.min(mindw)    # IN AIR WAVELENGTHS!!
        self._w0air = np.min(wave)
        self._w1air = np.max(wave)
        # parameters to save
        self._all_pars = []
        self._all_model = []
        self._all_chisq = []
        self._jac_array = None

    @property
    def params(self):
        return self._params

    @params.setter
    def params(self,params):
        """ Dictionary, keys must be all CAPS."""
        self._params = dict((key.upper(), value) for (key, value) in params.items())  # all CAPS

    @property
    def fitparams(self):
        return self._fitparams

    @fitparams.setter
    def fitparams(self,fitparams):
        """ list, keys must be all CAPS."""
        self._fitparams = [v.upper() for v in fitparams]      # all CAPS

    def mkinputs(self,args):
        """ Make INPUTS dictionary."""
        # Create INPUTS with all arguments needed to make the spectrum
        inputs = self.params.copy()  # initialize with initial/fixed values
        for k in range(len(self.fitparams)):        # this overwrites the values for the fitted values
            inputs[self.fitparams[k]] = args[k]
        inputs['DW'] = self._dwair          # add in wavelength parameters
        inputs['W0'] = self._w0air
        inputs['W1'] = self._w1air
        return inputs

    def chisq(self,model):
        return np.sqrt( np.sum( (self.flux-model)**2/self.err**2 )/len(self.flux) )

    def model(self, xx, *args, retobj=False):
        """ Return a model spectrum flux with the given input arguments."""
        # The input arguments correspond to FITPARAMS
        # This corrects for air/vacuum wavelength differences
        if self.verbose:
            print(args)
        # The arguments correspond to the fitting parameters
        inputs = self.mkinputs(args)
        if self.verbose:
            print(inputs)
        # Create the synthetic spectrum
        synspec = model_spectrum(inputs,verbose=self.verbose,   # always returns air wavelengths
                                 alinefile=self.alinefile,mlinefile=self.mlinefile)
        self.nsynfev += 1
        # Convolve with the LSF and do air/vacuum wave conversion
        pspec = prepare_synthspec(synspec,self.lsf,norm=self.norm,
                                  continuum_func=self.continuum_func)
        # Save models/pars/chisq
        self._all_pars.append(list(args).copy())
        self._all_model.append(pspec.flux.flatten().copy())
        self._all_chisq.append(self.chisq(pspec.flux.flatten()))
        # Return flattened spectrum
        if retobj:
            return pspec
        else:
            return pspec.flux.flatten()

    def getstep(self,name,val,relstep=0.02):
        """ Calculate step for a parameter."""
        # It mainly deals with edge cases
        #if val != 0.0:
        #    step = relstep*val
        #else:
        #    if name=='RV':
        #        step = 1.0
        #    elif name=='VROT':
        #        step = 0.5
        #    elif name=='VMICRO':
        #        step = 0.5
        #    elif name.endswith('_H'):
        #        step = 0.02
        #    else:
        #        step = 0.02
        if name=='TEFF':
            step = 5.0
        elif name=='RV':
            step = 0.1
        elif name=='VROT':
            step = 0.5
        elif name=='VMICRO':
            step = 0.5
        elif name.endswith('_H'):
            step = 0.01
        else:
            step = 0.01
        return step

        return step

    def jac(self,x,*args):
        """ Compute the Jacobian matrix (an m-by-n matrix, where element (i, j)
        is the partial derivative of f[i] with respect to x[j]). """

        if hasattr(self,'logger') is False:
            logger = dln.basiclogger()
        else:
            logger = self.logger

        logger.info(args)

        if self.verbose:
            logger.info(' ')
            logger.info('##### Calculating Jacobian Matrix #####')
            logger.info(' ')

        # A new synthetic spectrum does not need to be generated RV, vmicro or vsini.
        # Some time can be saved by not remaking those.
        # Use a one-sided derivative.

        # Boundaries
        lbounds,ubounds = mkbounds(self.fitparams)

        relstep = 0.02
        npix = len(x)
        npar = len(args)

        # Get INPUTS dictionary and make keys all CAPS
        inputs = self.mkinputs(args)
        inputs = dict((key.upper(), value) for (key, value) in inputs.items())

        # Some important parameters
        w0 = inputs['W0']
        w1 = inputs['W1']
        dw = inputs['DW']
        rv = inputs.get('RV')
        vrot = inputs.get('VROT')
        vmicro = inputs.get('VMICRO')

        # Create synthetic spectrum at current values
        #  set vrot=vmicro=rv=0, will modify later if necessary
        if self.verbose:
            logger.info('--- Current values ---')
            logger.info(args)
        tinputs = inputs.copy()
        tinputs['VMICRO'] = 0
        tinputs['VROT'] = 0
        tinputs['RV'] = 0
        origspec = model_spectrum(tinputs,keepextend=True,  # always are wavelengths
                                  alinefile=self.alinefile,mlinefile=self.mlinefile)
        self.nsynfev += 1
        # Smooth and shift
        smorigspec = smoothshift_spectrum(origspec,vrot=vrot,vmicro=vmicro,rv=rv)
        # Trim to final wavelengths
        smorigspec = trim_spectrum(smorigspec,w0,w1)
        # Convolve with the LSF and do air/vacuum wave conversion
        pspec = prepare_synthspec(smorigspec,self.lsf,norm=self.norm,
                                  continuum_func=self.continuum_func)
        # Flatten the spectrum
        f0 = pspec.flux.flatten()
        # Save models/pars/chisq
        self._all_pars.append(list(args).copy())
        self._all_model.append(f0.copy())
        self._all_chisq.append(self.chisq(f0))
        chisq = np.sqrt( np.sum( (self.flux-f0)**2/self.err**2 )/len(self.flux) )
        self
        if self.verbose:
            logger.info('chisq = '+str(chisq))

        # MASK PIXELS!?

        # Initialize jacobian matrix
        jac = np.zeros((npix,npar),np.float64)

        # Loop over parameters
        for i in range(npar):
            pars = np.array(copy.deepcopy(args))
            step = self.getstep(self.fitparams[i],pars[i],relstep)
            # Check boundaries, if above upper boundary
            #   go the opposite way
            if pars[i]>ubounds[i]:
                step *= -1
            pars[i] += step
            tinputs = self.mkinputs(pars)

            if self.verbose:
                logger.info(' ')
                logger.info('--- '+str(i+1)+' '+self.fitparams[i]+' '+str(pars[i])+' ---')
                logger.info(pars)
            # VROT/VMICRO/RV, just shift/smooth original spectrum
            if self.fitparams[i]=='VROT' or self.fitparams[i]=='VMICRO' or self.fitparams[i]=='RV':
                tvrot = tinputs.get('VROT')
                tvmicro = tinputs.get('VMICRO')
                trv = tinputs.get('RV')
                #import pdb; pdb.set_trace()
                # Smooth and shift
                synspec = smoothshift_spectrum(origspec,vrot=tvrot,vmicro=tvmicro,rv=trv)
                # Trim to final wavelengths
                synspec = trim_spectrum(synspec,w0,w1)
            else:
                synspec = model_spectrum(tinputs,alinefile=self.alinefile,
                                         mlinefile=self.mlinefile)  # always returns air wavelengths
                self.nsynfev += 1

            # Convert to vacuum wavelengths if necessary
            if self.wavevac:
                synspec.wave = astro.airtovac(synspec.wave)
                synspec.wavevac = True
            # Convolve with the LSF and do air/vacuum wave conversion
            pspec = prepare_synthspec(synspec,self.lsf,norm=self.norm,
                                      continuum_func=self.continuum_func)
            # Flatten the spectrum
            f1 = pspec.flux.flatten()

            # Save models/pars/chisq
            self._all_pars.append(list(pars).copy())
            self._all_model.append(f1.copy())
            self._all_chisq.append(self.chisq(f1))

            if np.sum(~np.isfinite(f1))>0:
                print('some nans/infs')
                import pdb; pdb.set_trace()

            jac[:,i] = (f1-f0)/step

        if np.sum(~np.isfinite(jac))>0:
            print('some nans/infs')
            import pdb; pdb.set_trace()

        self._jac_array = jac.copy()   # keep a copy

        self.njac += 1

        return jac

def trim_spectrum(spec,w0,w1):
    """ Trim a synthetic spectrum to [w0,w1]."""
    # This assumes that the spectrum has a single order
    wv1, ind1 = dln.closest(spec.wave,w0)
    wv2, ind2 = dln.closest(spec.wave,w1)
    # Nothing to do
    if ind1==0 and ind2==(spec.npix-1):
        return spec
    outspec = spec.copy()
    outspec.flux = outspec.flux[ind1:ind2+1]
    outspec.wave = outspec.wave[ind1:ind2+1]
    if outspec.err is not None:
        outspec.err = outspec.err[ind1:ind2+1]
    if outspec.mask is not None:
        outspec.mask = outspec.mask[ind1:ind2+1]
    if hasattr(outspec,'cont'):
        if outspec.cont is not None:
            outspec.cont = outspec.cont[ind1:ind2+1]
    outspec.npix = len(outspec.flux)
    return outspec

def getabund(inputs,verbose=False):
    """ Grab the abundances out of the input file and return array of abundances."""

    # Create the input 99-element abundance array
    codedir = os.path.dirname(os.path.abspath(__file__))
    pertab = Table.read(codedir+'/data/periodic_table.txt',format='ascii')

    feh = inputs.get('FEH')
    if feh is None:
        feh = inputs.get('FE_H')
    if feh is None:
        raise ValueError('FE_H missing from inputs')

    # Read model atmosphere
    modelfile = inputs.get('modelfile')
    if modelfile is None:
        raise ValueError('modelfile missing from inputs')
    atmostype, teff, logg, vmicro2, mabu, nd, atmos = synple.read_model(modelfile,verbose=verbose)
    mlines = dln.readlines(modelfile)

    # solar abundances
    # first two are Teff and logg
    # last two are Hydrogen and Helium
    solar_abund = np.array([ 4750., 2.5,
                            -10.99, -10.66,  -9.34,  -3.61,  -4.21,
                            -3.35,  -7.48,  -4.11,  -5.80,  -4.44,
                            -5.59,  -4.53,  -6.63,  -4.92,  -6.54,
                            -5.64,  -7.01,  -5.70,  -8.89,  -7.09,
                            -8.11,  -6.40,  -6.61,  -4.54,  -7.05,
                            -5.82,  -7.85,  -7.48,  -9.00,  -8.39,
                            -9.74,  -8.70,  -9.50,  -8.79,  -9.52,
                            -9.17,  -9.83,  -9.46, -10.58, -10.16,
                           -20.00, -10.29, -11.13, -10.47, -11.10,
                           -10.33, -11.24, -10.00, -11.03,  -9.86,
                           -10.49,  -9.80, -10.96,  -9.86, -10.94,
                           -10.46, -11.32, -10.62, -20.00, -11.08,
                           -11.52, -10.97, -11.74, -10.94, -11.56,
                           -11.12, -11.94, -11.20, -11.94, -11.19,
                           -12.16, -11.19, -11.78, -10.64, -10.66,
                           -10.42, -11.12, -10.87, -11.14, -10.29,
                           -11.39, -20.00, -20.00, -20.00, -20.00,
                           -20.00, -20.00, -12.02, -20.00, -12.58,
                           -20.00, -20.00, -20.00, -20.00, -20.00,
                           -20.00, -20.00])

    # Deal with alpha abundances
    #  only add the individual alpha abundance if it's not already there
    #  sometimes we might fit a single alpha element but want to use
    #  ALPHA_H to set the rest of them
    if inputs.get('ALPHA_H') is not None:
        alpha = inputs['ALPHA_H']
        elem = ['O','MG','SI','S','CA','TI']
        for k in range(len(elem)):
            if inputs.get(elem[k]+'_H') is None:
                inputs[elem[k]+'_H'] = alpha

    # Scale global metallicity
    abu = solar_abund.copy()
    abu[2:] += feh
    # Now offset the elements with [X/Fe], [X/Fe]=[X/H]-[Fe/H]
    g, = np.where( (np.char.array(list(inputs.keys())).find('_H') != -1) &
                   (np.char.array(list(inputs.keys())) != 'FE_H') )
    if len(g)>0:
        ind1,ind2 = dln.match(np.char.array(list(inputs.keys()))[g],np.char.array(pertab['symbol']).upper()+'_H')
        for k in range(len(ind1)):
            key1 = np.char.array(list(inputs.keys()))[g[ind1[k]]]
            abu[ind2[k]] += float(inputs[key1]) - feh
            if verbose:
                print('%s %f' % (key1,float(inputs[key1])))
    # convert to linear
    abu[2:] = 10**abu[2:]
    # Divide by N(H)
    g, = np.where(np.char.array(mlines).find('ABUNDANCE SCALE') != -1)
    nhtot = np.float64(mlines[g[0]].split()[6])
    abu[2:] /= nhtot
    # use model values for H and He
    abu[0:2] = mabu[0:2]

    return abu

def synple_wrapper(inputs,verbose=False,tmpbase='/tmp',alinefile=None,mlinefile=None):
    """ This is a wrapper around synple to generate a new synthetic spectrum."""
    # Wavelengths are all AIR!!

    # inputs is a dictionary with all of the inputs
    # Teff, logg, [Fe/H], some [X/Fe], and the wavelength parameters (w0, w1, dw).

    # Make temporary directory for synple to work in
    curdir = os.path.abspath(os.curdir)
    tdir = os.path.abspath(tempfile.mkdtemp(prefix="syn",dir=tmpbase))
    os.chdir(tdir)

    # Linelists to use
    linelist = ['gfallx3_bpo.19','kmol3_0.01_30.20']   # default values
    if alinefile is not None:   # atomic linelist input
        linelist[0] = alinefile
    if mlinefile is not None:   # molecular linelist input
        linelist[1] = mlinefile

    if verbose:
        print('Using linelist: ',linelist)

    # Make key names all CAPS
    inputs = dict((key.upper(), value) for (key, value) in inputs.items())

    # Make the model atmosphere file
    teff = inputs['TEFF']
    logg = inputs['LOGG']
    metal = inputs['FE_H']

    tid,modelfile = tempfile.mkstemp(prefix="mod",dir=".")
    os.close(tid)  # close the open file
    # Limit values
    #  of course the logg/feh ranges vary with Teff
    mteff = dln.limit(teff,3500.0,60000.0)
    mlogg = dln.limit(logg,0.0,5.0)
    mmetal = dln.limit(metal,-2.5,0.5)
    model, header, tail = models.mkmodel(mteff,mlogg,mmetal,modelfile)
    inputs['modelfile'] = modelfile
    if os.path.exists(modelfile) is False or os.stat(modelfile).st_size==0:
        print('model atmosphere file does NOT exist')
        import pdb; pdb.set_trace()

    # Create the synspec synthetic spectrum
    w0 = inputs['W0']
    w1 = inputs['W1']
    dw = inputs['DW']
    vmicro = inputs.get('VMICRO')
    vrot = inputs.get('VROT')
    if vrot is None:
        vrot = 0.0
    # Get the abundances
    abu = getabund(inputs,verbose=verbose)

    wave,flux,cont = synple.syn(modelfile,(w0,w1),dw,vmicro=vmicro,vrot=vrot,
                                abu=list(abu),verbose=verbose,linelist=linelist)

    # Delete temporary files
    shutil.rmtree(tdir)
    os.chdir(curdir)

    return (wave,flux,cont)

def smoothshift_spectrum(inpspec,vmicro=None,vrot=None,rv=None):
    """ This smoothes the spectrum by Vrot+Vmicro and
        shifts it by RV."""

    #vmicro = inputs.get('VMICRO')
    #vrot = inputs.get('VROT')
    #rv = inputs.get('RV')

    # Nothing to do
    if vmicro is None and vrot is None and rv is None:
        return inpspec.copy()

    # Initialize output spectrum
    spec = inpspec.copy()

    # Some broadening
    if vmicro is not None or vrot is not None:
        flux = utils.broaden(spec.wave,spec.flux,vgauss=vmicro,vsini=vrot)
        spec.flux = flux

    ## Vrot/Vsini (km/s) and Vmicro (in km/s)
    #if vrot is not None or vmicro is not None:
    #    wave, flux = synple.call_rotin(wave, flux, vrot, fwhm, space, steprot, stepfwhm, clean=False, reuseinputfiles=True)

    # Doppler shift only (in km/s)
    if rv is not None:
        if rv != 0.0:
            shiftwave = spec.wave*(1+rv/cspeed)
            gd,ngd,bd,nbd = dln.where( (spec.wave >= np.min(shiftwave)) & (spec.wave <= np.max(shiftwave)), comp=True)
            # Doppler shift and interpolate onto wavelength array
            if hasattr(spec,'cont'):
                cont = synple.interp_spl(spec.wave[gd], shiftwave, spec.cont)
                spec.cont *= 0
                spec.cont[gd] = cont
                # interpolate the continuing to the missing pixels
                if nbd>0:
                    contmissing = dln.interp(spec.wave[gd],spec.cont[gd],spec.wave[bd],kind='linear',assume_sorted=False)
                    spec.cont[bd] = contmissing
            flux = synple.interp_spl(spec.wave[gd], shiftwave, spec.flux)
            spec.flux *= 0
            spec.flux[gd] = flux
            if nbd>0:
                # Fill in missing values with interpolated values
                if np.sum(np.isfinite(spec.flux[gd]))>0:
                    coef = dln.poly_fit(spec.wave[gd],spec.flux[gd],2)
                    fluxmissing = dln.poly(spec.wave[bd],coef)
                    spec.flux[bd] = fluxmissing
                # Mask these pixels
                if spec.mask is None:
                    spec.mask = np.zeros(len(spec.flux),bool)
                spec.mask[bd] = True

    return spec

def model_spectrum(inputs,verbose=False,keepextend=False,alinefile=None,mlinefile=None):
    """
    This creates a model spectrum given the inputs:
    RV, Teff, logg, vmicro, vsini, [Fe/H], [X/Fe], w0, w1, dw.
    This creates the new synthetic spectrum and then convolves with vmicro, vsini and
    shifts to velocity RV.

    The returned spectrum always uses AIR wavelengths!!!

    Parameters
    ----------
    inputs : dictionary
       Input parameters, stellar parameters, abundances.
    keepextend : bool, optional
       Keep the extensions on the ends.  Default is False.
    alinefile : str, optional
       Atomic linelist filename.  Default is None (use synple's default one).
    mlinefile : str, optional
       Molecular linelist filename.  Default is None (use synple's default one).
    verbose : bool, optional
       Verbose output.  Default is False.

    Returns
    -------
    synspec : Spec1D
       The synthetic spectrum as Spec1D object.

    """

    # Make key names all CAPS
    inputs = dict((key.upper(), value) for (key, value) in inputs.items())

    # Extend on the ends for RV/convolution purposes
    w0 = inputs['W0']
    w1 = inputs['W1']
    dw = inputs['DW']
    rv = inputs.get('RV')
    vrot = inputs.get('VROT')
    vmicro = inputs.get('VMICRO')
    inputsext = inputs.copy()
    if rv is not None or vrot is not None or vmicro is not None:
        numext = int(np.ceil(w1*(1.0+1500/cspeed)-w1))
        inputsext['W0'] = w0-numext*dw
        inputsext['W1'] = w1+numext*dw
        if verbose:
            print('Extending wavelength by '+str(numext)+' pixels on each end')

    # Create the synthetic spectrum
    #  set vrot=vmicro=0, will convolve later if necessary
    inputsext['VMICRO'] = 0
    inputsext['VROT'] = 0
    wave1,flux1,cont1 = synple_wrapper(inputsext,verbose=verbose,alinefile=alinefile,
                                       mlinefile=mlinefile)

    # Get final wavelength array
    wv1, ind1 = dln.closest(wave1,w0)
    wv2, ind2 = dln.closest(wave1,w1)
    synspec = Spec1D(flux1/cont1,err=flux1*0,wave=wave1,lsfpars=np.array(0.0))
    synspec.cont = cont1
    synspec.wavevac = False
    # Smooth and shift
    if rv is not None or vrot is not None or vmicro is not None:
        synspec = smoothshift_spectrum(synspec,vrot=vrot,vmicro=vmicro,rv=rv)
    # Trim to final wavelengths
    if keepextend is False:
        synspec = trim_spectrum(synspec,w0,w1)

    return synspec

def prepare_synthspec(synspec,lsf,norm=True,continuum_func=None):
    """ Prepare a synthetic spectrum to be compared to an observed spectrum."""
    # Convolve with LSF and do air<->vacuum wavelength conversion

    # Convert wavelength from air->vacuum or vice versa
    if synspec.wavevac != lsf.wavevac:
        # Air -> Vacuum
        if synspec.wavevac is False:
            synspec.wave = astro.airtovac(synspec.wave)
            synspec.wavevac = True
        # Vacuum -> Air
        else:
            synspec.dispersion = astro.vactoair(synspec.wave)
            synspec.wavevac = False

    # Initialize the output spectrum
    if lsf.wave.ndim==2:
        npix,norder = lsf.wave.shape
    else:
        npix = len(lsf.wave)
        norder = 1
    pspec = Spec1D(np.zeros((npix,norder),np.float32),err=np.zeros((npix,norder),np.float32),
                   wave=lsf.wave,lsfpars=lsf.pars,lsftype=lsf.lsftype,lsfxtype=lsf.xtype)
    pspec.cont = np.zeros((npix,norder),np.float32)
    if continuum_func is not None:
        pspec.continuum_func = continuum_func

    # Loop over orders
    if lsf.wave.ndim==1:
        wave = np.atleast_2d(lsf.wave.copy()).T
    else:
        wave = lsf.wave.copy()
    for o in range(lsf.norder):
        wobs = wave[:,o]
        dw = np.median(dln.slope(wobs))
        wv1,ind1 = dln.closest(synspec.wave,np.min(wobs)-2*np.abs(dw))
        wv2,ind2 = dln.closest(synspec.wave,np.max(wobs)+2*np.abs(dw))
        modelflux = synspec.flux[ind1:ind2+1]
        modelwave = synspec.wave[ind1:ind2+1]
        modelcont = synspec.cont[ind1:ind2+1]

        # Rebin, if necessary
        #  get LSF FWHM (A) for a handful of positions across the spectrum
        xp = np.arange(npix//20)*20
        fwhm = lsf.fwhm(wobs[xp],xtype='Wave',order=o)
        # FWHM is in units of lsf.xtype, convert to wavelength/angstroms, if necessary
        if lsf.xtype.lower().find('pix')>-1:
            fwhm *= np.abs(dw)
        #  convert FWHM (A) in number of model pixels at those positions
        dwmod = dln.slope(modelwave)
        dwmod = np.hstack((dwmod,dwmod[-1]))
        xpmod = dln.interp(modelwave,np.arange(len(modelwave)),wobs[xp],kind='cubic',assume_sorted=False,extrapolate=True)
        xpmod = np.round(xpmod).astype(int)
        fwhmpix = np.abs(fwhm/dwmod[xpmod])
        # need at least ~4 pixels per LSF FWHM across the spectrum
        #  using 3 affects the final profile shape
        nbin = np.round(np.min(fwhmpix)//4).astype(int)

        if np.min(fwhmpix) < 3.7:
            warnings.warn('Model has lower resolution than the observed spectrum. Only '+str(np.min(fwhmpix))+' model pixels per resolution element')
        if np.min(fwhmpix) < 2.8:
            raise Exception('Model has lower resolution than the observed spectrum. Only '+str(np.min(fwhmpix))+' model pixels per resolution element')

        if nbin>1:
            npix2 = np.round(len(synspec.flux) // nbin).astype(int)
            modelflux = dln.rebin(modelflux[0:npix2*nbin],npix2)
            modelwave = dln.rebin(modelwave[0:npix2*nbin],npix2)
            modelcont = dln.rebin(modelcont[0:npix2*nbin],npix2)

        # Convolve
        lsf2d = lsf.anyarray(modelwave,xtype='Wave',order=o,original=False)
        cflux = utils.convolve_sparse(modelflux,lsf2d)
        # Interpolate onto final wavelength array
        flux = synple.interp_spl(wobs, modelwave, cflux)
        cont = synple.interp_spl(wobs, modelwave, modelcont)
        pspec.flux[:,o] = flux
        pspec.cont[:,o] = cont
        pspec.normalized = True

    # Normalize
    if norm is True:
        newcont = pspec.continuum_func(pspec)
        pspec.flux /= newcont
        pspec.cont *= newcont

    return pspec

def mkbounds(params,paramlims=None):
    """ Make lower and upper boundaries for parameters """
    params = np.char.array(params).upper()
    if paramlims is not None:
        limkeys = np.char.array(list(paramlims.keys())).upper()
    n = len(params)
    lbounds = np.zeros(n,np.float64)
    ubounds = np.zeros(n,np.float64)
    # Teff
    g, = np.where(params=='TEFF')
    if len(g)>0:
        lbounds[g[0]] = 3500
        ubounds[g[0]] = 60000
    # logg
    g, = np.where(params=='LOGG')
    if len(g)>0:
        lbounds[g[0]] = 0
        ubounds[g[0]] = 5
    # fe_h
    g, = np.where(params=='FE_H')
    if len(g)>0:
        lbounds[g[0]] = -3
        ubounds[g[0]] = 1
    # Vmicro
    g, = np.where(params=='VMICRO')
    if len(g)>0:
        lbounds[g[0]] = 0
        ubounds[g[0]] = 5
    # Vsini/vrot
    g, = np.where(params=='VROT')
    if len(g)>0:
        lbounds[g[0]] = 0
        ubounds[g[0]] = 500
    # RV
    g, = np.where(params=='RV')
    if len(g)>0:
        lbounds[g[0]] = -1500
        ubounds[g[0]] = 1500
    # abundances
    g, = np.where( (params.find('_H') != -1) & (params != 'FE_H') )
    if len(g)>0:
        lbounds[g] = -3
        ubounds[g] = 10

    # Use input parameter limits
    if paramlims is not None:
        for i,f in enumerate(params):
            g, = np.where(limkeys==f)
            if len(g)>0:
                lbounds[i] = paramlims[limkeys[g[0]]][0]
                ubounds[i] = paramlims[limkeys[g[0]]][1]

    bounds = (lbounds,ubounds)
    return bounds

def mkdxlim(fitparams):
    """ Make array of parameter changes at which curve_fit should finish."""
    npar = len(fitparams)
    dx_lim = np.zeros(npar,float)
    for k in range(npar):
        if fitparams[k]=='TEFF':
            dx_lim[k] = 1.0
        elif fitparams[k]=='LOGG':
            dx_lim[k] = 0.005
        elif fitparams[k]=='VMICRO':
            dx_lim[k] = 0.1
        elif fitparams[k]=='VROT':
            dx_lim[k] = 0.1
        elif fitparams[k]=='RV':
            dx_lim[k] = 0.01
        elif fitparams[k].endswith('_H'):
            dx_lim[k] = 0.005
        else:
            dx_lim[k] = 0.01
    return dx_lim

def initpars(params,fitparams,bounds=None):
    """ Make initial set of parameters given PARAMS and
        FITPARAMS."""

    params = dict((key.upper(), value) for (key, value) in params.items()) # all CAPS
    fitparams = [v.upper() for v in fitparams]  # all CAPS

    npars = len(fitparams)
    pinit = np.zeros(npars,np.float64)
    # Loop over parameters
    for k in range(npars):
        ind, = np.where(np.char.array(list(params.keys()))==fitparams[k])
        # This parameter is in PARAMS
        if len(ind)>0:
            pinit[k] = params[fitparams[k]]
        # Not in PARAMS
        else:
            if fitparams[k]=='RV':
                pinit[k] = 0.0
            elif fitparams[k]=='VMICRO':
                pinit[k] = 2.0
            elif fitparams[k]=='VROT':
                pinit[k] = 0.0
            elif fitparams[k]=='TEFF':
                pinit[k] = 5000.0
            elif fitparams[k]=='LOGG':
                pinit[k] = 3.0
            elif fitparams[k].endswith('_H'):
                # Abundances, use FE_H if possible
                if 'FE_H' in params.keys():
                    pinit[k] = params['FE_H']
                else:
                    pinit[k] = 0.0
            else:
                pinit[k] = 0.0

    # Make sure inital parameters are within the boundary limits
    if bounds is not None:
        for k in range(npars):
            pinit[k] = dln.limit(pinit[k],bounds[0][k],bounds[1][k])

    return pinit

def specfigure(figfile,spec,fmodel,out,original=None,verbose=True,figsize=10):
    """ Make diagnostic figure."""
    #import matplotlib
    matplotlib.use('Agg')
    #import matplotlib.pyplot as plt
    if os.path.exists(figfile): os.remove(figfile)
    norder = spec.norder
    nlegcol = 2
    if original is not None: nlegcol=3
    # Single-order plot
    if norder==1:
        fig,ax = plt.subplots()
        fig.set_figheight(figsize*0.5)
        fig.set_figwidth(figsize)
        if original is not None:
            plt.plot(original.wave,original.flux,color='green',label='Original',linewidth=1)
        plt.plot(spec.wave,spec.flux,'b',label='Masked Data',linewidth=0.5)
        plt.plot(fmodel.wave,fmodel.flux,'r',label='Model',linewidth=0.5,alpha=0.8)
        leg = ax.legend(loc='upper left', frameon=True, framealpha=0.8, ncol=nlegcol)
        plt.xlabel('Wavelength (Angstroms)')
        plt.ylabel('Normalized Flux')
        xr = dln.minmax(spec.wave)
        yr = [np.min([spec.flux,fmodel.flux]), np.max([spec.flux,fmodel.flux])]
        if original is not None:
            yr = [np.min([original.flux,spec.flux,fmodel.flux]), np.max([spec.flux,fmodel.flux])]
        yr = [yr[0]-dln.valrange(yr)*0.15,yr[1]+dln.valrange(yr)*0.005]
        yr = [np.max([yr[0],-0.2]), np.min([yr[1],2.0])]
        plt.xlim(xr)
        plt.ylim(yr)
        snr = np.nanmedian(spec.flux/spec.err)
        plt.title(spec.filename)
        #ax.annotate(r'S/N=%5.1f   Teff=%5.1f$\pm$%5.1f  logg=%5.2f$\pm$%5.2f  [Fe/H]=%5.2f$\pm$%5.2f   Vrel=%5.2f$\pm$%5.2f   chisq=%5.2f' %
        #            (snr, out['TEFF'], out['tefferr'], out['LOGG'], out['loggerr'], out['FE_H'], out['feherr'], out['RV'], out['vrelerr'], out['chisq']),
        #            xy=(np.mean(xr), yr[0]+dln.valrange(yr)*0.05),ha='center')
    # Multi-order plot
    else:
        fig,ax = plt.subplots(norder)
        fig.set_figheight(figsize)
        fig.set_figwidth(figsize)
        for i in range(norder):
            if original is not None:
                ax[i].plot(original.wave[:,i],original.flux[:,i],color='green',label='Original',linewidth=1)
            ax[i].plot(spec.wave[:,i],spec.flux[:,i],'b',label='Masked Data',linewidth=0.5)
            ax[i].plot(fmodel.wave[:,i],fmodel.flux[:,i],'r',label='Model',linewidth=0.5,alpha=0.8)
            if i==0:
                leg = ax[i].legend(loc='upper left', frameon=True, framealpha=0.8, ncol=nlegcol)
            ax[i].set_xlabel('Wavelength (Angstroms)')
            ax[i].set_ylabel('Normalized Flux')
            xr = dln.minmax(spec.wave[:,i])
            yr = [np.min([spec.flux[:,i],fmodel.flux[:,i]]), np.max([spec.flux[:,i],fmodel.flux[:,i]])]
            if original is not None:
                yr = [np.min([original.flux[:,i],spec.flux[:,i],fmodel.flux[:,i]]), np.max([spec.flux[:,i],fmodel.flux[:,i]])]
            yr = [yr[0]-dln.valrange(yr)*0.05,yr[1]+dln.valrange(yr)*0.05]
            if i==0:
                yr = [yr[0]-dln.valrange(yr)*0.15,yr[1]+dln.valrange(yr)*0.05]
            yr = [np.max([yr[0],-0.2]), np.min([yr[1],2.0])]
            ax[i].set_xlim(xr)
            ax[i].set_ylim(yr)
            # legend
            if i==0:
                snr = np.nanmedian(spec.flux/spec.err)
                ax[i].set_title(spec.filename)
                #ax[i].annotate(r'S/N=%5.1f   Teff=%5.1f$\pm$%5.1f  logg=%5.2f$\pm$%5.2f  [Fe/H]=%5.2f$\pm$%5.2f   Vrel=%5.2f$\pm$%5.2f   chisq=%5.2f' %
                #               (snr,out['teff'],out['tefferr'],out['logg'],out['loggerr'],out['feh'],out['feherr'],out['vrel'],out['vrelerr'],out['chisq']),
                #               xy=(np.mean(xr), yr[0]+dln.valrange(yr)*0.05),ha='center')
    plt.savefig(figfile,bbox_inches='tight')
    plt.close(fig)
    if verbose is True: print('Figure saved to '+figfile)

def dopvrot_lsq(spec,models=None,initpar=None,verbose=False,logger=None):
    """
    Least Squares fitting with forward modeling of the spectrum.

    Parameters
    ----------
    spec : Spec1D object
         The observed spectrum to match.
    models : list of Cannon models, optional
         A list of Cannon models to use.  The default is to load all of the Cannon
         models in the data/ directory and use those.
    initpar : numpy array, optional
         Initial estimate for [teff, logg, feh, RV, vsini], optional.
    verbose : bool, optional
         Verbose output of the various steps.  This is False by default.

    Returns
    -------
    out : numpy structured array
         The output structured array of the final derived RVs, stellar parameters and errors.
    bmodel : Spec1D object
         The best-fitting Cannon model spectrum (as Spec1D object).

    Example
    -------

    .. code-block:: python

         out, bmodel = fit_lsq(spec)

    """

    if logger is None:
        logger = dln.basiclogger()

    # Load and prepare the Cannon models
    #-------------------------------------------
    if models is None:
        models = cannon.models.copy()
        models.prepare(spec)

    # Get initial estimates
    if initpar is None:
        initpar = np.array([6000.0, 2.5, -0.5, 0.0, 0.0])
    initpar = np.array(initpar).flatten()

    # Calculate the bounds
    lbounds = np.zeros(5,float)+1e5
    ubounds = np.zeros(5,float)-1e5
    for p in models:
        lbounds[0:3] = np.minimum(lbounds[0:3],np.min(p.ranges,axis=1))
        ubounds[0:3] = np.maximum(ubounds[0:3],np.max(p.ranges,axis=1))
    lbounds[3] = -1000
    ubounds[3] = 1000
    lbounds[4] = 0.0
    ubounds[4] = 500.0
    bounds = (lbounds, ubounds)

    # function to use with curve_fit
    def spec_interp_vsini(x,teff,logg,feh,rv,vsini):
        """ This returns the interpolated model for a given spectrum."""
        # The "models" and "spec" must already exist outside of this function
        m = models(teff=teff,logg=logg,feh=feh,rv=rv)
        if m is None:      # there was a problem
            return np.zeros(spec.flux.shape,float).flatten()+1e30
        # Broaden to vsini
        if spec.norder>1:
            smflux = spec.flux*0
            for k in range(spec.norder):
                smflux[:,k] = utils.broaden(m.wave[:,k],m.flux[:,k],vsini=vsini)
        else:
            smflux = utils.broaden(m.wave.flatten(),m.flux.flatten(),vsini=vsini)
        return smflux.flatten()

    def spec_interp_vsini_jac(x,*args):
        """ Compute the Jacobian matrix (an m-by-n matrix, where element (i, j)
        is the partial derivative of f[i] with respect to x[j]). """

        relstep = 0.02
        npix = len(x)
        npar = len(args)

        # Current values
        f0 = spec_interp_vsini(x,*args)

        # Initialize jacobian matrix
        jac = np.zeros((npix,npar),np.float64)

        # Loop over parameters
        for i in range(npar):
            pars = np.array(copy.deepcopy(args))
            step = relstep*pars[i]
            if step<=0.0:
                step = 0.02
            pars[i] += step
            f1 = spec_interp_vsini(x,*pars)

            jac[:,i] = (f1-f0)/step

        return jac

    # Use curve_fit
    lspars, lscov = curve_fit(spec_interp_vsini, spec.wave.flatten(), spec.flux.flatten(), sigma=spec.err.flatten(),
                              p0=initpar, bounds=bounds, jac=spec_interp_vsini_jac)
    # If it hits a boundary then the solution won't change much compared to initpar
    # setting absolute_sigma=True gives crazy low lsperror values
    lsperror = np.sqrt(np.diag(lscov))

    if verbose is True:
        logger.info('Least Squares RV and stellar parameters:')
        for k,n in enumerate(['Teff','logg','[Fe/H]','RV','Vsini']):
            logger.info('%s = %f' % (n,lspars[k]))
    lsmodel = spec_interp_vsini(spec.wave,teff=lspars[0],logg=lspars[1],feh=lspars[2],rv=lspars[3],vsini=lspars[4])
    lschisq = np.sqrt(np.sum(((spec.flux.flatten()-lsmodel)/spec.err.flatten())**2)/len(lsmodel))
    if verbose is True: logger.info('chisq = %5.2f' % lschisq)

    # Put it into the output structure
    npar = len(lspars)
    dtype = np.dtype([('pars',float,npar),('parerr',float,npar),('parcov',float,(npar,npar)),('chisq',float)])
    out = np.zeros(1,dtype=dtype)
    out['pars'] = lspars
    out['parerr'] = lsperror
    out['parcov'] = lscov
    out['chisq'] = lschisq

    return out, lsmodel

def fit_elem(spec,params,elem,verbose=0,alinefile=None,mlinefile=None,logger=None):
    """ Fit an individual element."""

    t0 = time.time()

    if logger is None:
        logger = dln.basiclogger()

    # Create fitparams
    #fitparams = [e+'_H' for e in elem]
    fitparams = elem.copy()

    if verbose>0:
        logger.info('Fitting: '+', '.join(fitparams))

    # Initialize the fitter
    spfitter = SpecFitter(spec,params,fitparams=fitparams,verbose=(verbose>=2),
                          alinefile=alinefile,mlinefile=mlinefile)
    spfitter.logger = logger
    spfitter.norm = True  # normalize the synthetic spectrum
    #spfitter.verbose = True
    bounds = mkbounds(elem)
    pinit = initpars(params,elem,bounds)

    # Initalize output
    npar = len(fitparams)
    dtyp = []
    for f in fitparams:
        dtyp += [(f,float)]
    dtyp += [('pars',float,npar),('chisq',float),('nsynfev',int)]
    dtype = np.dtype(dtyp)
    out = np.zeros(1,dtype=dtype)

    # Loop over elemental abundances
    flag = 0
    abund = -2.0
    dabund = 1.0
    count = 0
    abundarr = []
    chisq = []
    modelarr = []
    # Loop from -2 to +1 or until we get through the minimum
    while (flag==0):
        model = spfitter.model(spec.wave.flatten(),abund)
        chisq1 = spfitter.chisq(model)
        abundarr.append(abund)
        modelarr.append(model)
        chisq.append(chisq1)
        if verbose>0:
            logger.info('%f %f' % (abund,chisq1))
        # Are we done?
        if (abund>=1) and (chisq1 != np.min(np.array(chisq))):
            flag = 1
        if (abund >= 10):
            flag = 1
        # Increment the abundance
        abund += dabund
        count += 1

    # Best value is at the end, just return that value
    bestind = np.argmin(chisq)
    if (bestind==0) or (bestind==len(chisq)-1):
        bestabund = abundarr[bestind]
        for k,f in enumerate(fitparams):
            out[f] = bestabund
        out['pars'] = bestabund
        out['chisq'] = np.min(chisq)
        out['nsynfev'] = spfitter.nsynfev
        model = modelarr[bestind]
        if verbose>0:
            logger.info('%f %f' % (bestabund,np.min(chisq)))
            logger.info('nfev = %i' % spfitter.nsynfev)
            logger.info('dt = %.2f sec.' % (time.time()-t0))
            logger.info(' ')
        return out, model

    # Now refine twice
    for i in range(2):
        # Get best value
        bestind = np.argmin(np.array(chisq))
        # get values half-way to left and right
        # Left
        lftind = bestind-1
        lftabund = np.mean([abundarr[lftind],abundarr[bestind]])
        lftmodel = spfitter.model(spec.wave.flatten(),lftabund)
        lftchisq = spfitter.chisq(lftmodel)
        abundarr.append(lftabund)
        modelarr.append(lftmodel)
        chisq.append(lftchisq)
        if verbose>0:
            logger.info('%f %f' % (lftabund,lftchisq))
        # Right
        rgtind = bestind+1
        rgtabund = np.mean([abundarr[bestind],abundarr[rgtind]])
        rgtmodel = spfitter.model(spec.wave.flatten(),rgtabund)
        rgtchisq = spfitter.chisq(rgtmodel)
        abundarr.append(rgtabund)
        modelarr.append(rgtmodel)
        chisq.append(rgtchisq)
        if verbose>0:
            logger.info('%f %f' % (rgtabund,rgtchisq))
        # Sort arrays
        si = np.argsort(abundarr)
        abundarr = [abundarr[k] for k in si]
        chisq = [chisq[k] for k in si]
        modelarr = [modelarr[k] for k in si]

    # Now interpolate to find the best value
    abundarr2 = np.linspace(np.min(abundarr),np.max(abundarr),1000)
    chisq2 = interp1d(abundarr,chisq,kind='quadratic')(abundarr2)
    bestind = np.argmin(chisq2)
    bestabund = abundarr2[bestind]

    # Get the model at the best value
    model = spfitter.model(spec.wave.flatten(),bestabund)
    bestchisq = spfitter.chisq(model)
    # Populate output structure
    for k,f in enumerate(fitparams):
        out[f] = bestabund
    out['pars'] = bestabund
    out['chisq'] = bestchisq
    out['nsynfev'] = spfitter.nsynfev

    if verbose>0:
        logger.info('%f %f' % (bestabund,bestchisq))
        logger.info('nfev = %i' % spfitter.nsynfev)
        logger.info('dt = %.2f sec.' % (time.time()-t0))
        logger.info(' ')

    return out, model

def fit_lsq(spec,params,fitparams=None,fparamlims=None,verbose=0,alinefile=None,mlinefile=None,logger=None):
    """
    Fit a spectrum with a synspec synthetic spectrum and determine stellar parameters and
    abundances using least-squares.

    Parameters
    ----------
    spec : Spec1D object
         The observed spectrum to match.
    params : dict
         Dictionary of initial values to use or parameters/elements to hold fixed.
    fitparams : list, optional
         List of parameter names to fit (e.g., TEFF, LOGG, FE_H, RV).  By default all values
         in PARAMS are fit.
    fparamlims : dict, optional
         Dictionary of lower and upper limits for each of the fitparams.
    verbose : int, optional
         Verbosity level (0, 1, or 2).  The default is 0 and verbose=2 is for debugging.
    alinefile : str, optional
         The atomic linelist to use.  Default is None which means the default synple linelist is used.
    mlinefile : str, optional
         The molecular linelist to use.  Default is None which means the default synple linelist is used.
    logger : logging object, optional
         Logging object.

    Returns
    -------
    out : numpy structured array
        Catalog of best-fit values.
    model : numpy array
        The best-fit synthetic stellar spectrum.

    Example
    -------

    .. code-block:: python

         spec = doppler.read(file)
         params = {'teff':5500,'logg':3.0,'fe_h':-1.0,'rv':0.0,'ca_h':-1.0}
         fitparams = ['teff','logg','fe_h','rv','ca_h']
         out,model = specfit.fit_lsq(spec,params,fitparams=fitparams)

    """

    t0 = time.time()

    if logger is None:
        logger = dln.basiclogger()

    # Normalize the spectrum
    if spec.normalized==False:
        spec.normalize()

    # Capitalize the inputs
    # Make key names all CAPS
    params = dict((key.upper(), value) for (key, value) in params.items())

    # Fitting parameters
    if fitparams is None:
        fitparams = list(params.keys())
    fitparams = [v.upper() for v in fitparams]  # all CAPS
    npar = len(fitparams)

    # Initialize the fitter
    spfitter = SpecFitter(spec,params,fitparams=fitparams,verbose=(verbose>=2),
                          alinefile=alinefile,mlinefile=mlinefile)
    spfitter.logger = logger
    spfitter.norm = True  # normalize the synthetic spectrum
    bounds = mkbounds(fitparams,fparamlims)
    pinit = initpars(params,fitparams,bounds)

    if verbose>0:
        logger.info('Fitting: '+', '.join(fitparams))

    # Fit the spectrum using curve_fit
    dx_lim = mkdxlim(fitparams)
    pars, cov = curve_fit(spfitter.model,spfitter.wave,spfitter.flux,dx_lim=dx_lim,
                          sigma=spfitter.err,p0=pinit,bounds=bounds,jac=spfitter.jac)
    error = np.sqrt(np.diag(cov))

    if verbose>0:
        logger.info('Best values:')
        for k in range(npar):
            logger.info('%s = %.3f +/- %.3f' % (fitparams[k],pars[k],error[k]))
    model = spfitter.model(spfitter.wave,*pars)
    chisq = np.sqrt(np.sum(((spfitter.flux-model)/spfitter.err)**2)/len(model))
    if verbose>0:
        logger.info('chisq = %.2f' % chisq)
        logger.info('nfev = %i' % spfitter.nsynfev)
        logger.info('dt = %.2f sec.' % (time.time()-t0))

    # Put it into the output structure
    dtyp = []
    for f in fitparams:
        dtyp += [(f,float),(f+'_ERR',float)]
    dtyp += [('pars',float,npar),('parerr',float,npar),('parcov',float,(npar,npar)),('chisq',float),('nsynfev',int)]
    dtype = np.dtype(dtyp)
    out = np.zeros(1,dtype=dtype)
    for k,f in enumerate(fitparams):
        out[f] = pars[k]
        out[f+'_ERR'] = error[k]
    out['pars'] = pars
    out['parerr'] = error
    out['parcov'] = cov
    out['chisq'] = chisq
    out['nsynfev'] = spfitter.nsynfev

    # Reshape final model spectrum
    model = model.reshape(spec.flux.shape)

    return out, model

def fit(spec,params=None,elem=None,figfile=None,fitvsini=False,fitvmicro=False,
        fparamlims=None,verbose=1,alinefile=None,mlinefile=None,logger=None):
    """
    Fit a spectrum with a synspec synthetic spectrum and determine stellar parameters and
    abundances using a multi-step iterative method.

    Step 1: Fit Teff/logg/[Fe/H]/RV using Doppler
    Step 2: Fit Teff/logg/[Fe/H]/RV + vsini with Doppler model
    Step 3: Fit stellar parameters (Teff/logg/[Fe/H]/[alpha/H]), RV and broadening (Vrot/Vmicro)
    Step 4: Fit each element one at a time holding everything else fixed.
    Step 5: Fit everything simultaneously

    Parameters
    ----------
    spec : Spec1D object
         The observed spectrum to match.
    params : dict, optional
         Dictionary of initial values to use or parameters/elements to hold fixed.
    elem : list, optional
         List of elements to fit.  The default is:
         elem = ['C','N','O','NA','MG','AL','SI','K','CA','TI','V','CR','MN','CO','NI','CU','CE','ND']
         Input an empty list [] to fit no elements.
    figfile : string, optional
         The filename for a diagnostic plot showing the observed spectrum and model spectrum.
    fitvsini : bool, optional
         Fit rotational velocity (vsini).  By default, Vsini will be fit initially with a Doppler
           model, but only included in the final fit if it improved chisq.
    fitvmicro : bool, optional
         Fit Vmicro.  Default is False.  By default, Vmicro is set (if not included in PARAMS)
         logg>=3.8:  vmicro = 2.0
         logg<3.8:   vmicro = 10^(0.226−0.0228*logg+0.0297*(logg)^2−0.0113*(logg)^3 )
    fparamlims : dict, optional
         Dictionary of lower and upper limits for each of the fitted parameter.
         For example, if params is {'teff': 9000, 'logg': 4.00, 'rv': -16.124}, fparamlims
         could be {'teff': [8000,10000], 'logg': [3.50,4.50], 'rv': [-20.124,-12.124]}.
    verbose : int, optional
         Verbosity level (0, 1, or 2).  The default is 0 and verbose=2 is for debugging.
    alinefile : str, optional
         The atomic linelist to use.  Default is None which means the default synple linelist is used.
    mlinefile : str, optional
         The molecular linelist to use.  Default is None which means the default synple linelist is used.
    logger : logging object, optional
         Logging object.

    Returns
    -------
    out : numpy structured array
        Catalog of best-fit values.
    model : numpy array
        The best-fit synthetic stellar spectrum.

    Example
    -------

    .. code-block:: python

         spec = doppler.read(file)
         out,model = specfit.fit(spec)

    """

    t0 = time.time()

    if logger is None:
        logger = dln.basiclogger()
        logger.handlers[0].setFormatter(logging.Formatter("%(asctime)s [%(levelname)-5.5s]  %(message)s"))
        logger.handlers[0].setStream(sys.stdout)

    # Default set of elements
    if elem is None:
        elem = ['C','N','O','NA','MG','AL','SI','K','CA','TI','V','CR','MN','CO','NI','CU','SR','CE','ND']

    # Normalize the spectrum
    if spec.normalized==False:
        spec.normalize()

    # Print out inputs
    if verbose>0:
        logger.info('Inputs:')
        if params is not None:
            logger.info('PARAMS:')
            for k,n in enumerate(params.keys()):
                logger.info('%s = %f' % (n,params[n]))
        else:
            logger.info('PARAMS: None')
        if fitvmicro:
            logger.info('Fitting VMICRO')
        if fitvsini:
            logger.info('Fitting VSINI')
        if len(elem)>0:
            logger.info('Elements to fit: '+', '.join(elem))
        else:
            logger.info('No elements to fit')
        logger.info(' ')

    # Input linelists
    if verbose and alinefile is not None:
        logger.info('Using input atomic linelist: ',alinefile)
    if verbose and mlinefile is not None:
        logger.info('Using input molecular linelist: ',mlinefile)

    # 1) Doppler (Teff, logg, feh, RV)
    #---------------------------------
    t1 = time.time()
    if verbose>0:
        logger.info('Step 1: Running Doppler')
    # Use Doppler to get initial guess of stellar parameters and RV
    dopout, dopfmodel, dopspecm = doppler.fit(spec)
    if verbose>0:
        logger.info('Teff = %.2f +/- %.2f' % (dopout['teff'][0],dopout['tefferr'][0]))
        logger.info('logg = %.3f +/- %.3f' % (dopout['logg'][0],dopout['loggerr'][0]))
        logger.info('[Fe/H] = %.3f +/- %.3f' % (dopout['feh'][0],dopout['feherr'][0]))
        logger.info('Vrel = %.4f +/- %.4f' % (dopout['vrel'][0],dopout['vrelerr'][0]))
        logger.info('chisq = %.3f' % dopout['chisq'][0])
        logger.info('dt = %.2f sec.' % (time.time()-t1))
        # typically 5 sec

    # 2) Fit vsini as well with Doppler model
    #-----------------------------------------
    t2 = time.time()
    if verbose>0:
        logger.info(' ')
        logger.info('Step 2: Fitting vsini with Doppler model')
    # For APOGEE resolution you need vsini~4 km/s or greater to see an effect
    initpar2 = [dopout['teff'][0], dopout['logg'][0], dopout['feh'][0], dopout['vrel'][0], 5.0]
    out2, model2 = dopvrot_lsq(spec,initpar=initpar2,verbose=verbose,logger=logger)
    if verbose>0:
        logger.info('Teff = %.2f +/- %.2f' % (out2['pars'][0][0],out2['parerr'][0][0]))
        logger.info('logg = %.3f +/- %.3f' % (out2['pars'][0][1],out2['parerr'][0][1]))
        logger.info('[Fe/H] = %.3f +/- %.3f' % (out2['pars'][0][2],out2['parerr'][0][2]))
        logger.info('Vrel = %.4f +/- %.4f' % (out2['pars'][0][3],out2['parerr'][0][3]))
        logger.info('Vsini = %.3f +/- %.3f' % (out2['pars'][0][4],out2['parerr'][0][4]))
        logger.info('chisq = %.3f' % out2['chisq'][0])
        logger.info('dt = %.2f sec.' % (time.time()-t2))
        # typically 5 sec
    if out2['chisq'][0] > dopout['chisq'][0]:
        if verbose>0:
            logger.info('Doppler Vrot=0 chisq is better')
        out2['pars'][0] = [dopout['teff'][0],dopout['logg'][0],dopout['feh'][0],dopout['vrel'][0],0.0]

    # Initialize params
    if params is None:
        params = {}
    else:
        params = dict((key.upper(), value) for (key, value) in params.items())  # all CAPS
    # Using input values when possible, otherwise Doppler values
    for k,name in enumerate(['TEFF','LOGG','FE_H','RV','VROT']):
        if params.get(name) is None:
            params[name] = out2['pars'][0][k]

    # Get Vmicro using Teff/logg relation
    # APOGEE DR14 vmicro relation (Holtzman et al. 2018)
    # for stars with [M/H]>-1 and logg<3.8
    # vmicro = 10^(0.226−0.0228*logg+0.0297*(logg)^2−0.0113*(logg)^3 )
    # coef = [0.226,0.0228,0.0297,−0.0113]
    # only giants, was fit in dwarfs
    if params.get('VMICRO') is None:
        vmicro = 2.0  # default
        if params['LOGG']<3.8:
            vmcoef = [0.226,0.0228,0.0297,-0.0113]
            vmicro = 10**dln.poly(params['LOGG'],vmcoef[::-1])
        params['VMICRO'] = vmicro

    # for giants
    # vmacro = 10^(0.741−0.0998*logg−0.225[M/H])
    # maximum of 15 km/s

    # 3) Fit stellar parameters (Teff, logg, feh, alpha, RV, Vsini)
    #--------------------------------------------------------------
    t3 = time.time()
    if verbose>0:
        logger.info(' ')
        logger.info('Step 3: Fitting stellar parameters, RV and broadening')
    params3 = params.copy()
    fitparams3 = ['TEFF','LOGG','FE_H','ALPHA_H','RV']
    if params3['VROT']>0 or fitvsini is True:
        fitparams3.append('VROT')
    # Fit Vmicro as well if it's a dwarf
    if params3['LOGG']>3.8 or params3['TEFF']>8000 or fitvmicro is True:
        fitparams3.append('VMICRO')
    out3, model3 = fit_lsq(spec,params3,fitparams3,fparamlims,verbose=verbose,
                                    alinefile=alinefile,mlinefile=mlinefile,logger=logger)
    # typically 9 min.

    # Should we fit C_H and N_H as well??

    # Tweak the continuum
    if verbose is not None:
        logger.info('Tweaking continuum using best-fit synthetic model')
    tmodel = Spec1D(model3,wave=spec.wave.copy(),lsfpars=np.array(0.0))
    spec = doppler.rv.tweakcontinuum(spec,tmodel)

    # 4) Fit each element separately
    #-------------------------------
    t4 = time.time()
    if verbose>0:
        logger.info(' ')
        logger.info('Step 4: Fitting each element separately')
    params4 = params3.copy()
    for k in range(len(fitparams3)):
        params4[fitparams3[k]] = out3['pars'][0][k]
    nelem = len(elem)
    if nelem>0:
        if verbose>0:
            logger.info('Elements: '+', '.join(elem))
        elemcat = np.zeros(nelem,dtype=np.dtype([('name',np.str,10),('par',np.float64),('parerr',np.float64)]))
        elemcat['name'] = elem
        for k in range(nelem):
            t4b = time.time()
            parselem = params4.copy()
            if elem[k] in ['O','MG','SI','S','CA','TI']:
                parselem[elem[k]+'_H'] = params4['ALPHA_H']
            else:
                parselem[elem[k]+'_H'] = params4['FE_H']
            fitparselem = [elem[k]+'_H']
            #out4, model4 = fit_lsq(spec,parselem,fitparselem,verbose=verbose,logger=logger)
            out4, model4 = fit_elem(spec,parselem,fitparselem,verbose=verbose,
                                    alinefile=alinefile,mlinefile=mlinefile,logger=logger)
            elemcat['par'][k] = out4['pars'][0]
            #elemcat['parerr'][k] = out4['parerr'][0]
        if verbose>0:
            logger.info('dt = %f sec.' % (time.time()-t4))
            logger.info(' ')
    else:
        if verbose>0:
            logger.info('No elements to fit')
    # about 50 min.

    # 5) Fit all parameters simultaneously
    #---------------------------------------
    # if NO elements to fit, then nothing to do
    if nelem>0:
        t5 = time.time()
        if verbose>0:
            logger.info('Step 5: Fit all parameters simultaneously')
        params5 = params4.copy()
        for k in range(nelem):
            params5[elem[k]+'_H'] = elemcat['par'][k]
        if params5.get('ALPHA_H') is not None:
            del params5['ALPHA_H']
        fitparams5 = ['TEFF','LOGG','FE_H','RV']
        if 'VROT' in fitparams3 or fitvsini is True:
            fitparams5.append('VROT')
        if 'VMICRO' in fitparams3 or fitvmicro is True:
            fitparams5.append('VMICRO')
        fitparams5 = fitparams5+list(np.char.array(elem)+'_H')
        out5, model5 = fit_lsq(spec,params5,fitparams5,fparamlims,verbose=verbose,
                               alinefile=alinefile,mlinefile=mlinefile,logger=logger)
    else:
        out5 = out3
        model5 = model3
        fitparams5 = fitparams3

    # Make final structure and save the figure
    out = out5
    dtyp = []
    npar = len(fitparams5)
    for f in fitparams5:
        dtyp += [(f,float),(f+'_ERR',float)]
    dtyp += [('pars',float,npar),('parerr',float,npar),('parcov',float,(npar,npar)),('chisq',float),('vhelio',float)]
    dtype = np.dtype(dtyp)
    out = np.zeros(1,dtype=dtype)
    for k,f in enumerate(fitparams5):
        out[f] = out5['pars'][0][k]
        out[f+'_ERR'] = out5['parerr'][0][k]
    out['pars'] = out5['pars'][0]
    out['parerr'] = out5['parerr'][0]
    out['parcov'] = out5['parcov'][0]
    out['chisq'] = out5['chisq'][0]
    out['vhelio'] = out5['RV']+spec.barycorr()
    if verbose>0:
        logger.info('Vhelio = %.3f' % out['vhelio'])
    # Final model
    model = Spec1D(model5,wave=spec.wave.copy(),lsfpars=np.array(0.0))
    model.lsf = spec.lsf.copy()
    # Make figure
    if figfile is not None:
        specfigure(figfile,spec,model,out,verbose=(verbose>=2))

    if verbose>0:
        logger.info('dt = %.2f sec.' % (time.time()-t0))

    return out, model

max_stars_repo_name: dnidever/fraunhofer

max_issues_repo_name: dnidever/fraunhofer

max_forks_repo_name: dnidever/fraunhofer

reason: import numpy,from scipy,from astropy


text:
import os
import json
import numpy as np
import pandas as pd
import datetime as dt
from pathlib import Path
from collections import defaultdict
from dateutil.relativedelta import relativedelta

from fluxgapfill.models import EnsembleModel
from fluxgapfill.test import parse_model_dirs
from fluxgapfill.predictors import parse_predictors
from fluxgapfill.metrics import get_pred_interval
from fluxgapfill.models import get_model_class

def gapfill(
        data_dir,
        sites=None,
        models=None,
        predictors=None,
        predictors_paths=None,
        distribution="laplace",
        budget_date_ranges_path=None,
        **kwargs
):
    """
    Gapfill data with trained models.

    Args:
        data_dir (<str>): path to the project folder containing site data foloders to save
                          all raw, processed data and results.
        sites (list<str>): Comma-separated list of site IDs to train on.
                           Must match the name(s) of the data directories.
        models (list<str>): Comma-separated list of model names to train.
                            Options: ['rf', 'ann', 'lasso', 'xgb']
        predictors (list<str>): Comma-separated list of predictors. Ignored if
                                predictors_path is provided.
                                Certain keyword predictors are used to denote
                                specific sets of predictors:
                                ['temporal', 'all']
        predictors_paths (list<str>): Comma-separated list of paths to files
                                      containing predictor names. See
                                      predictors/metereological.txt for an
                                      example.
        distribution (str): Which distribution to use for prediction.
                            Options: ['laplace', 'normal']
        budget_date_ranges_path (str): Dictionary mapping site names to date
                                       ranges. See
                                           gapfill/budget_date_ranges.json
                                       for an example.

    Writes gapfilled data to
        data/{SiteID}/gapfilled/{model}_{predictors}_{distribution}.csv
    where the CSV has all the same columns as raw.csv, excluding any existing
    gapfilled columns, plus columns for the
        predicted mean (FCH4_F)
        95% uncertainty (FCH4_uncertainty)
        spread individual predictions (FCH4{1-N})
    and writes budget data to
        data/{SiteID}/gapfilled/{model}_{predictors}_{distribution}_budget.csv
    """
    # parse inputs
    if isinstance(sites, str):
        sites = sites.split(",")
    if isinstance(models, str):
        models = models.split(",")
        for model in models:
            try:
                get_model_class(model)
            except Exception as e:
                raise ValueError(f"Model {model} not supported.")

    if (predictors is not None) and (predictors_paths is not None):
        predictor_subsets = parse_predictors(predictors_paths, predictors)   # list?
    else:
        predictor_subsets = None

    if budget_date_ranges_path is not None:
        with Path(budget_date_ranges_path).open() as f:
            budget_date_ranges = json.load(f)
    else:
        budget_date_ranges = {}

    #
    data_dir = Path(data_dir)
    if sites is None:
        sites = os.listdir(data_dir)
        if "results" in sites:
            sites.remove("results")

    for site in sites:
        site_dir = data_dir/site
        site_gap_dir = site_dir / "gapfilled"
        site_gap_dir.mkdir(exist_ok=True)

        if models is None:
            models = os.listdir(site_dir/'models')

        for model in models:
            if predictor_subsets is None:
                predictor_subsets = os.listdir(site_dir/'models'/model)

            for predictor_subset in predictor_subsets:
                print(f'Gapfilling: {site}, {model}, {predictor_subset}, {distribution}')
                gap_df = gapfill_site_model_predictor(site,
                                                      model,
                                                      predictor_subset,
                                                      distribution,
                                                      site_dir)
                out_fn = site_gap_dir / f"{model}_{predictor_subset}_{distribution}.csv"
                print(f" - Writing gapfilled data to {out_fn}\n")
                gap_df.to_csv(out_fn, index=False)

                ## compute annual budget
                print(f'Computing Annual Budget: {site}, {model}, {predictor_subset}, {distribution}')
                # For budget estimates, convert from nmol m-2 s-1 to g C m-2 halfhour-1
                ch4_conversion = lambda fch4: fch4*60*30*12.0107*10**-9
                for col in [col for col in gap_df.columns if 'FCH4' in col]:
                    gap_df[col] = gap_df[col].apply(ch4_conversion)

                site_budget_date_ranges = get_site_budget_ranges(site, budget_date_ranges, gap_df)
                budget_df = compute_annual_budget(site, gap_df,
                                                  site_budget_date_ranges)

                # Write budget data to file
                out_fn = (
                    site_gap_dir / f"{model}_{predictor_subset}_{distribution}_budget.csv"
                )
                print(f" - Writing budget data to {out_fn}\n")
                budget_df.to_csv(out_fn, index=False)

def gapfill_site_model_predictor(site, model, predictor, distribution, site_dir):

    model_dir = site_dir/"models"/model/predictor
    Model = EnsembleModel(model_dir)

    gap_df = pd.read_csv(site_dir / 'gap.csv')
    scale_fn = model_dir / "scale.json"
    if not scale_fn.exists():
        raise ValueError(
            "Must run <python main.py test> with --distribution " +
            f"{distribution} to compute an uncertainty scale before " +
            "gapfilling."
        )
    with scale_fn.open() as f:
        uncertainty_scale = json.load(f)

    y_hat_individual = Model.predict_individual(gap_df)
    y_hat_mean = np.mean(y_hat_individual, axis=0)
    y_hat_spread = (
        (y_hat_individual - y_hat_mean) * uncertainty_scale + y_hat_mean
    )
    y_hat_dist = Model.predict_dist(
        gap_df,
        distribution=distribution,
        uncertainty_scale=uncertainty_scale
    )
    y_hat_uncertainty_95 = get_pred_interval(y_hat_dist)

    num_models = len(y_hat_individual)
    gapfill_columns = {
        'FCH4_F': y_hat_mean,
        'FCH4_F_UNCERTAINTY': y_hat_uncertainty_95[:, 1] - y_hat_mean,
        **{
            f"FCH4_F{i+1}": y_hat_spread[i]
            for i in range(num_models)
        }
    }
    for column, gapfill_data in gapfill_columns.items():
        gap_df[column] = gapfill_data

    # Merge with original data
    raw_fn = site_dir / 'raw.csv'
    raw_df = pd.read_csv(raw_fn)

    # Drop existing columns from the raw df
    existing_fch4_columns = [
        column for column in raw_df.columns
        if 'FCH4_' in column
    ]
    raw_df.drop(columns=existing_fch4_columns, inplace=True)
    gap_df = raw_df.merge(
        gap_df[list(gapfill_columns.keys()) + ['TIMESTAMP_END']],
        on='TIMESTAMP_END',
        how='left'
    )

    # Add in observed values to FCH4_F, FCH4_F{1-N}
    fch4_columns = [f'FCH4_F{i+1}' for i in range(num_models)]
    observed_rows = gap_df['FCH4_F'].isna()
    for fch4_column in ['FCH4_F'] + fch4_columns:
        gap_df.loc[observed_rows, fch4_column] = (
            gap_df.loc[observed_rows, 'FCH4']
        )

    # Set uncertainties around observed values to 0
    gap_df.loc[observed_rows, 'FCH4_F_UNCERTAINTY'] = 0
    return gap_df

def get_site_budget_ranges(site, budget_date_ranges, gap_df):
    # Get date ranges to compute annual budget
    if site in budget_date_ranges:
        site_budget_date_ranges = [
            [
                dt.datetime(
                    int(range_start.split('-')[1]),
                    int(range_start.split('-')[0]),
                    1
                ),
                dt.datetime(
                    int(range_end.split('-')[1]),
                    int(range_end.split('-')[0]),
                    1
                ) + relativedelta(months=1)
            ]
            for [range_start, range_end] in budget_date_ranges[site]
        ]
    else:
        annual_sum_year_start = gap_df['Year'].min()
        annual_sum_year_end = gap_df['Year'].max()
        site_budget_date_ranges = [
            [
                dt.datetime(year, 1, 1),
                dt.datetime(year+1, 1, 1)
            ]
            for year in range(
                annual_sum_year_start,
                annual_sum_year_end + 1
            )
        ]
    return site_budget_date_ranges

def compute_annual_budget(site, gap_df, site_budget_date_ranges):

    # Compute annual budgets between date ranges
    gap_df["TIMESTAMP_END"] = pd.to_datetime(
        gap_df["TIMESTAMP_END"],
        format='%Y%m%d%H%M'
    )
    budget_dict = defaultdict(list)
    for [range_start, range_end] in site_budget_date_ranges:
        annual_gap_df = gap_df[
            (gap_df['TIMESTAMP_END'] >= range_start) &
            (gap_df['TIMESTAMP_END'] < range_end)
        ]

        fch4_columns = [col for col in gap_df.columns if 'FCH4_F' in col]
        fch4_columns.remove('FCH4_F_UNCERTAINTY')
        annual_budgets = annual_gap_df[fch4_columns].sum()

        budget_mean = annual_budgets.mean()
        budget_uncertainty = annual_budgets.std() * 1.96

        date_range_str = f"{range_start.date()} and {range_end.date()}"
        budget_str = f"{budget_mean:.2f} +/- {budget_uncertainty:.2f}"
        print(f" - Budget between {date_range_str} is {budget_str}")

        budget_dict['range_start'].append(range_start)
        budget_dict['range_end'].append(range_end)
        budget_dict['budget_mean'].append(budget_mean)
        budget_dict['budget_uncertainty'].append(budget_uncertainty)

    budget_df = pd.DataFrame(budget_dict)

    # Combine budgets and uncertainties per year with an average
    mean_budget_dict = {
        'range_start': site_budget_date_ranges[0][0],
        'range_end': site_budget_date_ranges[-1][1],
        'budget_mean': budget_df['budget_mean'].mean(),
        'budget_uncertainty': budget_df['budget_uncertainty'].mean()
    }
    budget_df = budget_df.append(mean_budget_dict, ignore_index=True)
    return budget_df

ext: py

max_stars_repo_name: stanfordmlgroup/methane-gapfill-ml

max_issues_repo_name: stanfordmlgroup/methane-gapfill-ml

max_forks_repo_name: stanfordmlgroup/methane-gapfill-ml

reason: import numpy


text:
'''
@author: kris
'''
# import modules; set up logging
from gensim.models import Word2Vec
from gensim.models import KeyedVectors
from gensim.test.utils import datapath
import numpy as np
import logging, os, sys, gzip
import datetime

logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', filename='word2vec.out', level=logging.INFO)
# Path to a file that contains lines with the locations of files
# containing the sentences we want for our Word2Vec model
# Also works with entities that are just stacked line by line
pathsLocator = "./sentencesPaths.txt"
outputPath = "./entity_embeddings.txt"

# Model to load
to_load = '/vol2/cb/crunchbase-201806/embeddings/dim200-iter10-win5/CB_sg1_size200_mincount1_window5_neg15_iter10.wv.vectors.npy'
#'/home/faerberm/mag-training/MAG_sg1_size128_minCount5_window5_neg15_iter10_alpha_cbowMean.wv.vectors.npy'
#'/vol2/cb/crunchbase-201806/embeddings/dim200-iter10-win5/CB_sg1_size200_mincount1_window5_neg15_iter10'
#'MAG_sg1_size128_minCount5_window5_neg15_iter5'
loadKeyedVector = True

#'dbpedia_sg1_size200_mincount1_window5_neg15_iter10'
#'RDF2Vec_sg1_size200_mincount1_window5_neg15_iter20'
#'MAG_sg1_size200_mincount1_window5_neg15_iter15'

#What is the newline character on the machine
newline = '\n'
ignorePrefix = '#'
#What separates one walk from another (aka. one sentence from another)?
walkSeparator = "\t"
#What separates the single 'units' of a given walk?
hopSeparator = '->'

# Mapping dict
entity_mapping_dict = {}
# Mapping file
mapping_file = "/home/noulletk/prog/bmw/dbpedia_full/resources/data/walks/walk_entity_mapping.txt"
mapping_sep = "\t"
hasMapping = False

iterationCounter = {'val': 0}

#Load mappings if there are any
if hasMapping:
	for mapping_line in open(mapping_file, mode='rt'):
		mapping_tokens = mapping_line.rstrip(newline).split(mapping_sep)
		if len(mapping_tokens) == 2:
			entity_mapping_dict[mapping_tokens[0]] = mapping_tokens[1]
	print("Loaded %s mappings!" % (len(entity_mapping_dict)))

class MySentences:
	def __init__(self, iterationCounter):
		self.iterationCounter = iterationCounter

	def __iter__(self):
		print("Running Iteration #%s" % (iterationCounter['val']))
		iterationCounter['val'] += 1
		# Iterate to find which files are to be read
		for fname in open(pathsLocator, mode='rt'):  # os.listdir(self.dirname):
			sentencesPath = fname.rstrip(newline)
			# Ignore commented-out lines
			if sentencesPath.startswith(ignorePrefix):
				continue
			now = datetime.datetime.now()
			print("[%s] Grabbing sentences from: %s" % (now.strftime("%Y-%m-%d %H:%M"), sentencesPath))
			try:
				# Go through all paths
				for line in open(sentencesPath, mode='rt'):
					# If you're NOT grouping the walks and separating them by tabs
					sentence = line.rstrip(newline).split(hopSeparator)
					for tokenPos in range(len(sentence)):
						token = sentence[tokenPos]
						# Give the proper URL for the entity IF it exists, otherwise return the entity itself
						sentence[tokenPos] = entity_mapping_dict.get(token, token)
					#print(sentence)
					yield sentence

			except Exception:
				print("Failed reading file:")
				print(sentencesPath)

#load model
if loadKeyedVector:
	print("Loading [KeyedVectors] from: ",to_load)
	#model_wv = KeyedVectors.load(to_load, mmap='r')
	#model_wv = KeyedVectors.load_word2vec_format(to_load, binary=True)
	#model_wv = KeyedVectors.load_word2vec_format(to_load)
	model_wv = KeyedVectors.load(to_load)
	#model_wv = KeyedVectors.load_word2vec_format(datapath('word2vec_pre_kv_c'), binary=False)  # C text format
	#model_wv = KeyedVectors.load_word2vec_format(to_load, binary=True, unicode_errors='ignore')
else:
	print("Loading [MODEL] from: ",to_load)
	model_wv = Word2Vec.load(to_load).wv

print("Vocab keys size:",len(model_wv.vocab.keys()))
print("Outputting entity embeddings to: ",outputPath)
sentences = MySentences(iterationCounter)

#Open the output file for the entity embeddings
outFile = open(outputPath, "w")

#Make a dictionary for in-memory aggregation while going over sentences
default_val = None
entity_embeddings_dict = {}
vocab_keys = model_wv.vocab.keys()

displayCounter = 0
maxDisplay = 10
for voc in vocab_keys:
	print(voc)
	if displayCounter >= maxDisplay:
		break
	displayCounter+=1

print("Compute entity embeddings (through combination of word embeddings)...")
counter = 0

'''
for sentence in sentences:
	entity = sentence[0]
	entity_embedding = None
	#Sum over all words' embeddings and then output the resulting embedding
	for word in sentence:
		word_embedding = model.wv[word]
		if default_val is None:
			#Initialise default_val if it isn't yet
			default_val = np.zeros(word_embedding.shape)
		if entity_embedding is None:
			entity_embedding = np.zeros(word_embedding.shape)
		entity_embedding += word_embedding
	entity_embeddings_dict[entity] = entity_embeddings_dict.get(entity, default_val) + entity_embedding
	if (counter % 1000000 == 0):
		print("Combined word embeddings: ",counter)
		print("Last one completed: ",entity)
	counter+=1
'''

#Go through all sentences to see which entities we want
for sentence in sentences:
	# idea is that the entity is in the document, so we check what it is like and
	# since every entity has 'the same' treatment, that we can determine their probabilities based on that
	entity = sentence[0]
	if hasMapping:
		entity = entity_mapping_dict.get(entity, entity)
	entity_embedding = None
	dict_val = entity_embeddings_dict.get(entity, None)
	if (dict_val is None):
		if entity in vocab_keys:
			entity_embedding = model_wv[entity]
			entity_embeddings_dict[entity] = entity_embedding
			#Encountered first time, so output it
			outFile.write("%s" % entity)
			for number in entity_embedding:
				outFile.write("\t%s" % number)
			outFile.write("\n")

	if (counter % 1000000 == 0):
		print("Lines passed through: ",counter)
		print("Current line's entity: ",entity)
		print("Embeddings output: ",len(entity_embeddings_dict))
	counter+=1

#print("Output computed entity embeddings!")
#for (entity, entity_embedding) in entity_embeddings_dict.items():
#	#Output computed embedding
#	outFile.write("%s" % entity)
#	for number in entity_embedding:
#		outFile.write("\t%s" % number)
#	outFile.write("\n")

#Close the output file post finishing output operations
outFile.close()
print("Finished outputting entity embeddings")

max_stars_repo_name: michaelfaerber/Agnos

max_issues_repo_name: michaelfaerber/Agnos

max_forks_repo_name: michaelfaerber/Agnos