def __getitem__(self, item):
        if item not in self.valid_keys:
            raise KeyError('{} not a valid item for CCFContainer!'.format(item))

        if item == 'ml' and self.ml is not None:
            return DataStructures.xypoint(x=self.x, y=self.ml)
        elif item == 'dc' and self.dc is not None:
            return DataStructures.xypoint(x=self.x, y=self.dc)
        elif item == 'simple' and self.simple is not None:
            return DataStructures.xypoint(x=self.x, y=self.simple)
        elif item == 'weighted' and self.weighted is not None:
            return DataStructures.xypoint(x=self.x, y=self.weighted)
        elif item == 'simple-weighted' and self.simple_weighted is not None:
            return DataStructures.xypoint(x=self.x, y=self.simple_weighted)

        return None  # We should never get here...
    def Interpolate_Old(self, dictionary, SpT):
        # First, we must convert the relations above into a monotonically increasing system
        # Just add ten when we get to each new spectral type
        relation = DataStructures.xypoint(len(dictionary))

        # Strip the spectral type of the luminosity class information
        SpT = re.search("[A-Z]([0-9]\.?[0-9]*)", SpT).group()

        xpoints = []
        ypoints = []
        for key, index in zip(dictionary, range(len(dictionary))):
            # Convert key to a number
            xpoints.append(self.SpT_To_Number(key))
            ypoints.append(dictionary[key])

        sorting_indices = [i[0] for i in sorted(enumerate(xpoints), key=lambda x: x[1])]
        for index in range(len(dictionary)):
            i = sorting_indices[index]
            relation.x[index] = xpoints[i]
            relation.y[index] = ypoints[i]

        RELATION = UnivariateSpline(relation.x, relation.y, s=0)

        spnum = self.SpT_To_Number(SpT)
        if spnum > 0:
            return RELATION(spnum)
        else:
            return np.nan
def astropy_smooth(data, vel, linearize=False, kernel=convolution.Gaussian1DKernel, **kern_args):
    """
    Smooth using a gaussian filter, using astropy.

    Parameters:
    ===========
    - data:          kglib.utils.DataStructures.xypoint instance
                     The data to smooth.

    - vel:           float
                     The velocity scale to smooth out.
                     Can either by an astropy quantity or a float in km/s

    - linearize:     boolean
                     If True, we will put the data in a constant
                     log-wavelength spacing grid before smoothing.
                     The output has the same spacing as the input
                     regardless of this variable.

    - kernel:        astropy.convolution kernel
                     The astropy kernel to use. The default is the
                     Gaussian1DKernel.

    - kern_args:     Additional kernel arguments beyond width

    Returns:
    ========
    A smoothed version of the data, on the same wavelength grid as the data
    """

    if linearize:
        original_data = data.copy()
        datafcn = spline(data.x, data.y, k=3)
        linear = DataStructures.xypoint(data.x.size)
        linear.x = np.logspace(np.log10(data.x[0]), np.log10(data.x[-1]), linear.size())
        linear.y = datafcn(linear.x)
        data = linear

    # Figure out feature size in pixels
    if not isinstance(vel, u.quantity.Quantity):
        vel *= u.km / u.second

    featuresize = (vel / constants.c).decompose().value
    dlam = np.log(data.x[1] / data.x[0])
    Npix = featuresize / dlam

    # Make kernel and smooth
    kern = kernel(Npix, **kern_args)
    smoothed = convolution.convolve(data.y, kern, boundary='extend')

    if linearize:
        fcn = spline(data.x, smoothed)
        return fcn(original_data.x)
    return smoothed
def IterativeLowPass(data, vel, numiter=100, lowreject=3, highreject=3, width=5, linearize=False):
    """
    An iterative version of LowPassFilter.
    It will ignore outliers in the low pass filter. New parameters that are
    not described in the docstring fro LowPassFilter are:

    Parameters:
    ===========
    - numiter:       integer
                     The maximum number of iterations to take

    - lowreject:     integer
                     How many sigma below the current filtered curve do we
                     count as bad and ignore in the next iteration?

    - highreject:    integer
                     How many sigma above the current filtered curve do we
                     count as bad and ignore in the next iteration?
    """

    datacopy = data.copy()
    if linearize:
        datafcn = spline(datacopy.x, datacopy.y, k=3)
        errorfcn = spline(datacopy.x, datacopy.err, k=1)
        contfcn = spline(datacopy.x, datacopy.cont, k=1)
        linear = DataStructures.xypoint(datacopy.x.size)
        linear.x = np.linspace(datacopy.x[0], datacopy.x[-1], linear.size())
        linear.y = datafcn(linear.x)
        linear.err = errorfcn(linear.x)
        linear.cont = contfcn(linear.x)
        datacopy = linear.copy()

    done = False
    iter = 0
    datacopy.cont = FittingUtilities.Continuum(datacopy.x, datacopy.y, fitorder=9, lowreject=2.5, highreject=5)
    while not done and iter < numiter:
        done = True
        iter += 1
        smoothed = LowPassFilter(datacopy, vel, width=width)
        residuals = datacopy.y / smoothed
        mean = np.mean(residuals)
        std = np.std(residuals)
        badpoints = np.where(np.logical_or((residuals - mean) < -lowreject * std, residuals - mean > highreject * std))[
            0]
        if badpoints.size > 0:
            done = False
            datacopy.y[badpoints] = smoothed[badpoints]
    if linearize:
        return linear.x, smoothed
    else:
        return smoothed
    def __call__(self, T, metal, vsini=0.0, return_xypoint=True, **kwargs):
        """
        Given parameters, return an interpolated spectrum

        If return_xypoint is False, then it will only return
          a numpy.ndarray with the spectrum

        Before interpolating, we will do some error checking to make
        sure the requested values fall within the grid
        """

        # Scale the requested values
        T = (T - self.T_scale[0]) / self.T_scale[1]
        metal = (metal - self.metal_scale[0]) / self.metal_scale[1]

        # Get the minimum and maximum values in the grid
        T_min = min(self.grid[:, 0])
        T_max = max(self.grid[:, 0])
        metal_min = min(self.grid[:, 1])
        metal_max = max(self.grid[:, 1])
        input_list = (T, metal)

        # Check to make sure the requested values fall within the grid
        if (T_min <= T <= T_max and
                        metal_min <= metal <= metal_max):

            y = self.interpolator(input_list)
        else:
            if self.debug:
                warnings.warn("The requested parameters fall outside the model grid. Results may be unreliable!")
            print(T, T_min, T_max)
            print(metal, metal_min, metal_max)
            y = self.NN_interpolator(input_list)

        # Test to make sure the result is valid. If the requested point is
        # outside the Delaunay triangulation, it will return NaN's
        if np.any(np.isnan(y)):
            if self.debug:
                warnings.warn("Found NaNs in the interpolated spectrum! Falling back to Nearest Neighbor")
            y = self.NN_interpolator(input_list)

        model = DataStructures.xypoint(x=self.xaxis, y=y)
        vsini *= units.km.to(units.cm)
        model = Broaden.RotBroad(model, vsini, linear=self.rebin)


        # Return the appropriate object
        if return_xypoint:
            return model
        else:
            return model.y
def process_model(model, data, vsini_model=None, resolution=None, vsini_primary=None,
                  maxvel=1000.0, debug=False, logspace=True):
    """
    Process a stellar model to prepare it for cross correlation

    Parameters:
    - model:          string, or kglib.utils.DataStructures.xypoint instance
                      If a string, should give the path to an ascii file with the model
                      Otherwise, should hold the model data

    - data:           list of kglib.utils.DataStructures.xypoint instances
                      The already-processed data.

    - vsini_model:    float
                      The rotational velocity to apply to the model spectrum

    - vsini_primary:  float
                      The rotational velocity of the primary star

    - resolution:     float
                      The detector resolution in $\lambda / \Delta \lambda$

    - maxvel:         float
                      The maximum velocity to include in the eventual CCF.
                      This is used to trim the data appropriately for each echelle order.

    - debug:          boolean
                      Print some extra stuff?

    - logspace:       boolean
                      Rebin the model to constant log-spacing?

    Returns:
    ========
    A list of kglib.utils.DataStructures.xypoint instances with the processed model.
    """
    # Read in the model if necessary
    if isinstance(model, str):
        if debug:
            print("Reading in the input model from %s" % model)
        x, y = np.loadtxt(model, usecols=(0, 1), unpack=True)
        x = x * u.angstrom.to(u.nm)
        y = 10 ** y
        left = np.searchsorted(x, data[0].x[0] - 10)
        right = np.searchsorted(x, data[-1].x[-1] + 10)
        model = DataStructures.xypoint(x=x[left:right], y=y[left:right])
    elif not isinstance(model, DataStructures.xypoint):
        raise TypeError(
            "Input model is of an unknown type! Must be a DataStructures.xypoint or a string with the filename.")


    # Linearize the x-axis of the model (in log-spacing)
    if logspace:
        if debug:
            print("Linearizing model")
        xgrid = np.logspace(np.log10(model.x[0]), np.log10(model.x[-1]), model.size())
        model = FittingUtilities.RebinData(model, xgrid)

    # Broaden
    if vsini_model is not None and vsini_model > 1.0 * u.km.to(u.cm):
        if debug:
            print("Rotationally broadening model to vsini = %g km/s" % (vsini_model * u.cm.to(u.km)))
        model = Broaden.RotBroad(model, vsini_model, linear=True)


    # Reduce resolution
    if resolution is not None and 5000 < resolution < 500000:
        if debug:
            print("Convolving to the detector resolution of %g" % resolution)
        model = FittingUtilities.ReduceResolutionFFT(model, resolution)

    # Divide by the same smoothing kernel as we used for the data
    if vsini_primary is not None:
        smoothed = HelperFunctions.astropy_smooth(model, vel=SMOOTH_FACTOR * vsini_primary, linearize=False)
        model.y += model.cont.mean() - smoothed
        model.cont = np.ones(model.size()) * model.cont.mean()


    # Rebin subsets of the model to the same spacing as the data
    model_orders = []
    model_fcn = spline(model.x, model.y)
    if debug:
        model.output("Test_model.dat")
    for i, order in enumerate(data):
        if debug:
            sys.stdout.write("\rGenerating model subset for order %i in the input data" % (i + 1))
            sys.stdout.flush()
        # Find how much to extend the model so that we can get maxvel range.
        dlambda = order.x[order.size() / 2] * maxvel * 1.5 / 3e5
        left = np.searchsorted(model.x, order.x[0] - dlambda)
        right = np.searchsorted(model.x, order.x[-1] + dlambda)
        right = min(right, model.size() - 2)

        # Figure out the log-spacing of the data
        logspacing = np.log(order.x[1] / order.x[0])

        # Finally, space the model segment with the same log-spacing
        start = np.log(model.x[left])
        end = np.log(model.x[right])
        xgrid = np.exp(np.arange(start, end + logspacing, logspacing))

        segment = DataStructures.xypoint(x=xgrid, y=model_fcn(xgrid))
        segment.cont = FittingUtilities.Continuum(segment.x, segment.y, lowreject=1.5, highreject=5, fitorder=2)
        model_orders.append(segment)

    print("\n")
    return model_orders
def HighPassFilter(data, vel, width=5, linearize=False):
    """
    Function to apply a high-pass filter to data.

    Parameters:
    ===========
    - data:        kglib.utils.DataStructures.xypoint instance
                   The data to filter.

    - vel:         float
                   The width of the features to smooth out, in velocity space (in cm/s).

    - width:       integer
                   How long it takes the filter to cut off, in units of wavenumber.

    - linearize:   boolean
                   Do we need to resample the data to a constant wavelength spacing
                   before filtering? Set to True if the data is unevenly sampled.

    Returns:
    ========
    If linearize = True:  Returns a numpy.ndarray with the new x-axis, and the filtered data
    Otherwise:  Returns just a numpy.ndarray with the filtered data.
    """

    if linearize:
        original_data = data.copy()
        datafcn = spline(data.x, data.y, k=3)
        errorfcn = spline(data.x, data.err, k=3)
        contfcn = spline(data.x, data.cont, k=3)
        linear = DataStructures.xypoint(data.x.size)
        linear.x = np.linspace(data.x[0], data.x[-1], linear.size())
        linear.y = datafcn(linear.x)
        linear.err = errorfcn(linear.x)
        linear.cont = contfcn(linear.x)
        data = linear

    # Figure out cutoff frequency from the velocity.
    featuresize = 2 * data.x.mean() * vel / constants.c.cgs.value  # vel MUST be given in units of cm
    dlam = data.x[1] - data.x[0]  # data.x MUST have constant x-spacing
    Npix = featuresize / dlam

    nsamples = data.size()
    sample_rate = 1.0 / dlam
    nyq_rate = sample_rate / 2.0  # The Nyquist rate of the signal.
    width /= nyq_rate
    cutoff_hz = min(1.0 / featuresize, nyq_rate - width * nyq_rate / 2.0)  # Cutoff frequency of the filter

    # The desired attenuation in the stop band, in dB.
    ripple_db = 60.0

    # Compute the order and Kaiser parameter for the FIR filter.
    N, beta = kaiserord(ripple_db, width)
    if N % 2 == 0:
        N += 1

    # Use firwin with a Kaiser window to create a lowpass FIR filter.
    taps = firwin(N, cutoff_hz / nyq_rate, window=('kaiser', beta), pass_zero=False)

    # Extend data to prevent edge effects
    y = np.r_[data.y[::-1], data.y, data.y[::-1]]

    # Use lfilter to filter data with the FIR filter.
    smoothed_y = lfilter(taps, 1.0, y)

    # The phase delay of the filtered signal.
    delay = 0.5 * (N - 1) / sample_rate
    delay_idx = np.searchsorted(data.x, data.x[0] + delay)
    smoothed_y = smoothed_y[data.size() + delay_idx:-data.size() + delay_idx]
    if linearize:
        fcn = spline(data.x, smoothed_y)
        return fcn(original_data.x)
    else:
        return smoothed_y
def ReadFits(datafile, errors=False, extensions=False, x=None, y=None, cont=None, return_aps=False, debug=False):
    """
    Read a fits file. If extensions=False, it assumes IRAF's multispec format.
    Otherwise, it assumes the file consists of several fits extensions with
    binary tables, with the table names given by the x,y,cont, and errors keywords.

    See ReadExtensionFits for a convenience function that assumes my standard names

    Parameters:
    ===========
    - datafile:       string
                      The name of the file to read

    - errors:         boolean, integer, or string
                      If False, indicates there are no errors in the datafile
                      If an integer AND extensions = False, indicates the index
                      that the errors are found in.
                      If a string AND extenstions = True, indicates the name
                      of the error field in the binary table.

    - extensions:     boolean
                      Is the data stored in several fits extensions? If not,
                      we assume it is in multispec format

    - x:              string
                      The name of the field with the x-coordinate (wavelength, etc).
                      Ignored if extensions = False

    - y:              string
                      The name of the field with the y-coordinate (flux, etc).
                      Ignored if extensions = False

    - cont:           string
                      The name of the field with the continuum estimate.
                      Ignored if extensions = False

    - return_aps:     boolean
                      Return the aperture wavelength fields as well as the
                      extracted orders. The wavelength fields define the
                      wavelengths in multispec format. Ignored if extensions = True.

    - debug:          boolean
                      Print some extra information to screen

    Returns:
    ========
    A list of kglib.utils.DataStructures.xypoint instances with the data for each
    echelle order.
    """
    if debug:
        print( "Reading in file %s: " % datafile)

    if extensions:
        # This means the data is in fits extensions, with one order per extension
        # At least x and y should be given (and should be strings to identify the field in the table record array)
        if type(x) != str:
            x = raw_input("Give name of the field which contains the x array: ")
        if type(y) != str:
            y = raw_input("Give name of the field which contains the y array: ")
        orders = []
        hdulist = pyfits.open(datafile)
        if cont == None:
            if not errors:
                for i in range(1, len(hdulist)):
                    data = hdulist[i].data
                    xypt = DataStructures.xypoint(x=data.field(x), y=data.field(y))
                    orders.append(xypt)
            else:
                if type(errors) != str:
                    errors = raw_input("Give name of the field which contains the errors/sigma array: ")
                for i in range(1, len(hdulist)):
                    data = hdulist[i].data
                    xypt = DataStructures.xypoint(x=data.field(x), y=data.field(y), err=data.field(errors))
                    orders.append(xypt)
        elif type(cont) == str:
            if not errors:
                for i in range(1, len(hdulist)):
                    data = hdulist[i].data
                    xypt = DataStructures.xypoint(x=data.field(x), y=data.field(y), cont=data.field(cont))
                    orders.append(xypt)
            else:
                if type(errors) != str:
                    errors = raw_input("Give name of the field which contains the errors/sigma array: ")
                for i in range(1, len(hdulist)):
                    data = hdulist[i].data
                    xypt = DataStructures.xypoint(x=data.field(x), y=data.field(y), cont=data.field(cont),
                                                  err=data.field(errors))
                    orders.append(xypt)

    else:
        # Data is in multispec format rather than in fits extensions
        # Call Rick White's script
        try:
            retdict = multispec.readmultispec(datafile, quiet=not debug)
        except ValueError:
            warnings.warn("Wavelength not found in file %s. Using a pixel grid instead!" % datafile)
            hdulist = pyfits.open(datafile)
            data = hdulist[0].data
            hdulist.close()
            numpixels = data.shape[-1]
            numorders = data.shape[-2]
            wave = np.array([np.arange(numpixels) for i in range(numorders)])
            retdict = {'flux': data,
                       'wavelen': wave,
                       'wavefields': np.zeros(data.shape)}

        # Check if wavelength units are in angstroms (common, but I like nm)
        hdulist = pyfits.open(datafile)
        header = hdulist[0].header
        hdulist.close()
        wave_factor = 1.0  #factor to multiply wavelengths by to get them in nanometers
        for key in sorted(header.keys()):
            if "WAT1" in key:
                if "label=Wavelength" in header[key] and "units" in header[key]:
                    waveunits = header[key].split("units=")[-1]
                    if waveunits == "angstroms" or waveunits == "Angstroms":
                        # wave_factor = u.nm/u.angstrom
                        wave_factor = u.angstrom.to(u.nm)
                        if debug:
                            print( "Wavelength units are Angstroms. Scaling wavelength by ", wave_factor)

        if errors == False:
            numorders = retdict['flux'].shape[0]
        else:
            numorders = retdict['flux'].shape[1]
        orders = []
        for i in range(numorders):
            wave = retdict['wavelen'][i] * wave_factor
            if errors == False:
                flux = retdict['flux'][i]
                err = np.ones(flux.size) * 1e9
                err[flux > 0] = np.sqrt(flux[flux > 0])
            else:
                if type(errors) != int:
                    errors = int(raw_input("Enter the band number (in C-numbering) of the error/sigma band: "))
                flux = retdict['flux'][0][i]
                err = retdict['flux'][errors][i]
            cont = FittingUtilities.Continuum(wave, flux, lowreject=2, highreject=4)
            orders.append(DataStructures.xypoint(x=wave, y=flux, err=err, cont=cont))
        if return_aps:
            # Return the aperture wavefields too
            orders = [orders, retdict['wavefields']]
    return orders
def plot_expected(orders, prim_spt, Tsec, instrument, vsini=None, rv=0.0, twoaxes=False):
    """
    Plot the data orders, with a model spectrum added at appropriate flux ratio

    Parameters
    ==========
    - orders:        A list of kglib.utils.Datastructures.xypoint instances
                     The observed spectra, split into echelle orders

    - prim_spt:      string
                     The primary star spectral type

    - Tsec:          float
                     The secondary temperature

    - instrument:    string
                     The name of the instrument the observation came from

    - vsini:         float
                     The vsini of the companion, in km/s

    - rv:            float
                     The rv shift of the companion
    """

    sns.set_context("paper", font_scale=2.0)
    sns.set_style("white")
    sns.set_style("ticks")

    # First, get the model
    dir_prefix = "/media/ExtraSpace"
    if not os.path.exists(dir_prefix):
        dir_prefix = "/Volumes/DATADRIVE"
    inst2hdf5 = {
        "TS23": "{}/PhoenixGrid/TS23_Grid.hdf5".format(dir_prefix),
        "HRS": "{}/PhoenixGrid/HRS_Grid.hdf5".format(dir_prefix),
        "CHIRON": "{}/PhoenixGrid/CHIRON_Grid.hdf5".format(dir_prefix),
        "IGRINS": "{}/PhoenixGrid/IGRINS_Grid.hdf5".format(dir_prefix),
    }
    hdf5_int = StellarModel.HDF5Interface(inst2hdf5[instrument])
    wl = hdf5_int.wl
    pars = {"temp": Tsec, "logg": 4.5, "Z": 0.0, "alpha": 0.0}
    fl = hdf5_int.load_flux(pars)

    # Broaden, if requested
    if vsini is not None:
        m = DataStructures.xypoint(x=wl, y=fl)
        m = Broaden.RotBroad(m, vsini * units.km.to(units.cm))
        wl, fl = m.x, m.y

    # get model continuum
    c = FittingUtilities.Continuum(wl, fl, fitorder=5, lowreject=2, highreject=10)

    # Interpolate the model
    x = wl * units.angstrom
    plt.plot(wl, fl)
    plt.plot(wl, c)
    plt.show()
    modelfcn = interp(x.to(units.nm), fl / c)

    # Get the wavelength-specific flux ratio between the primary and secondary star
    MS = SpectralTypeRelations.MainSequence()
    Tprim = MS.Interpolate("temperature", prim_spt)
    Rprim = MS.Interpolate("radius", prim_spt)
    sec_spt = MS.GetSpectralType("temperature", Tsec, prec=1e-3)
    Rsec = MS.Interpolate("radius", sec_spt)
    flux_ratio = blackbody_lambda(x, Tprim) / blackbody_lambda(x, Tsec) * (Rprim / Rsec) ** 2
    fluxratio_fcn = interp(x.to(units.nm), 1.0 / flux_ratio)

    # Loop over the orders:
    if twoaxes:
        fig, axes = plt.subplots(2, 1, sharex=True)
        top, bottom = axes
        for order in orders:
            order.cont = FittingUtilities.Continuum(order.x, order.y, fitorder=3, lowreject=2, highreject=5)
            top.plot(order.x, order.y, "k-", alpha=0.4)
            top.plot(order.x, order.cont, "r--")

            total = order.copy()

            xarr = total.x * (1 + rv / constants.c.to(units.km / units.s).value)
            model = (modelfcn(xarr) - 1.0) * fluxratio_fcn(xarr)
            total.y += total.cont * model
            top.plot(total.x, total.y, "g-", alpha=0.4)

            bottom.plot(total.x, total.y - order.y, "k-", alpha=0.4)

        return fig, [top, bottom], orders

    fig, ax = plt.subplots(1, 1)
    for order in orders:
        order.cont = FittingUtilities.Continuum(order.x, order.y, fitorder=3, lowreject=2, highreject=5)
        ax.plot(order.x, order.y, "k-", alpha=0.4)

        total = order.copy()

        xarr = total.x * (1 + rv / constants.c.to(units.km / units.s).value)
        model = (modelfcn(xarr) - 1.0) * fluxratio_fcn(xarr)
        total.y += total.cont * model
        ax.plot(total.x, total.y, "g-", alpha=0.4)

    # Label
    ax.set_xlabel("Wavelength (nm)")
    ax.set_ylabel("Flux (arbitrary units)")
    xlim = ax.get_xlim()
    ylim = ax.get_ylim()
    ax.plot([9e9], [9e9], "k-", label="Actual data")
    ax.plot([9e9], [9e9], "g-", label="Expected data")
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)
    leg = ax.legend(loc="best", fancybox=True)

    return fig, ax, orders
def MakeModelDicts(model_list, vsini_values=[10, 20, 30, 40], type='phoenix',
                   vac2air=True, logspace=False, hdf5_file=HDF5_FILE, get_T_sens=False):
    """
    This will take a list of models, and output two dictionaries that are
    used by GenericSearch.py and Sensitivity.py

    Parameters:
    ===========
    - model_list:        iterable
                         A list of model filenames

    - vsini_values:     iterable
                        A list of vsini values to broaden
                        the spectrum by (we do that later!)

    - type:             string
                        The type of models. Currently,
                        phoenix, kurucz, and hdf5 are implemented

    - vac2air:          boolean
                        If true, assumes the model is in
                        vacuum wavelengths and converts to air

    - logspace:         boolean
                        If true, it will rebin the
                        data to a constant log-spacing

    - hdf5_file:        string
                        The absolute path to the HDF5 file
                        with the models. Only used if type=hdf5

    - get_T_sens:       boolean
                        Flag for getting the temperature sensitivity.
                        If true, it finds the derivative of each pixel dF/dT

    Returns:
    ========
     A dictionary containing the model with keys of temperature, gravity,
     metallicity, and vsini,and another one with a processed flag with the same keys
    """
    vsini_values = np.atleast_1d(vsini_values)
    if type.lower() == 'phoenix':
        modeldict = defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(DataStructures.xypoint))))
        processed = defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(bool))))
        for fname in model_list:
            temp, gravity, metallicity = ClassifyModel(fname)
            print("Reading in file %s" % fname)
            data = pandas.read_csv(fname,
                                   header=None,
                                   names=["wave", "flux"],
                                   usecols=(0, 1),
                                   sep=' ',
                                   skipinitialspace=True)
            x, y = data['wave'].values, data['flux'].values
            if vac2air:
                n = 1.0 + 2.735182e-4 + 131.4182 / x ** 2 + 2.76249e8 / x ** 4
                x /= n
            model = DataStructures.xypoint(x=x * units.angstrom.to(units.nm), y=10 ** y)
            if logspace:
                xgrid = np.logspace(np.log(model.x[0]), np.log(model.x[-1]), model.size(), base=np.e)
                model = FittingUtilities.RebinData(model, xgrid)
            for vsini in vsini_values:
                modeldict[temp][gravity][metallicity][vsini] = model
                processed[temp][gravity][metallicity][vsini] = False

    elif type.lower() == 'kurucz':
        modeldict = defaultdict(
            lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(DataStructures.xypoint)))))
        processed = defaultdict(
            lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(bool)))))
        for fname in model_list:
            temp, gravity, metallicity, a = ClassifyModel(fname)
            print("Reading in file %s" % fname)
            data = pandas.read_csv(fname,
                                   header=None,
                                   names=["wave", "flux"],
                                   usecols=(0, 1),
                                   sep=' ',
                                   skipinitialspace=True)
            x, y = data['wave'].values, data['flux'].values
            if vac2air:
                n = 1.0 + 2.735182e-4 + 131.4182 / x ** 2 + 2.76249e8 / x ** 4
                x /= n
            model = DataStructures.xypoint(x=x * units.angstrom.to(units.nm), y=10 ** y)
            if logspace:
                xgrid = np.logspace(np.log(model.x[0]), np.log(model.x[-1]), model.size(), base=np.e)
                model = FittingUtilities.RebinData(model, xgrid)
            for vsini in vsini_values:
                modeldict[temp][gravity][metallicity][a][vsini] = model
                processed[temp][gravity][metallicity][a][vsini] = False

    elif type.lower() == 'hdf5':
        hdf5_int = HDF5Interface(hdf5_file)
        x = hdf5_int.wl
        wave_hdr = hdf5_int.wl_header
        if vac2air:
            if not wave_hdr['air']:
                n = 1.0 + 2.735182e-4 + 131.4182 / x ** 2 + 2.76249e8 / x ** 4
                x /= n
        elif wave_hdr['air']:
            raise GridError(
                'HDF5 grid is in air wavelengths, but you requested vacuum wavelengths. You need a new grid!')
        modeldict = defaultdict(
            lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(DataStructures.xypoint)))))
        processed = defaultdict(
            lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(bool)))))
        for pars in model_list:
            temp, gravity, metallicity, a = pars['temp'], pars['logg'], pars['Z'], pars['alpha']
            y = hdf5_int.load_flux(pars)
            model = DataStructures.xypoint(x=x * units.angstrom.to(units.nm), y=y)
            for vsini in vsini_values:
                modeldict[temp][gravity][metallicity][a][vsini] = model
                processed[temp][gravity][metallicity][a][vsini] = False

    else:
        raise NotImplementedError("Sorry, the model type ({:s}) is not available!".format(type))

    if get_T_sens:
        # Get the temperature sensitivity. Warning! This assumes the wavelength grid is the same in all models.
        sensitivity = defaultdict(
            lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(DataStructures.xypoint)))))
        Tvals = sorted(modeldict.keys())
        for i, T in enumerate(Tvals):
            gvals = sorted(modeldict[T].keys())
            for gravity in gvals:
                metal_vals = sorted(modeldict[T][gravity].keys())
                for metal in metal_vals:
                    alpha_vals = sorted(modeldict[T][gravity][metal].keys())
                    for alpha in alpha_vals:
                        # get the temperature just under this one
                        lower, l_idx = get_model(modeldict, Tvals, i, gravity, metal, vsini_values[0], alpha, mode='lower')
                        upper, u_idx = get_model(modeldict, Tvals, i, gravity, metal, vsini_values[0], alpha, mode='upper')
                        T_low = Tvals[l_idx]
                        T_high = Tvals[u_idx]
                        slope = (upper.y - lower.y) / (T_high - T_low)
                        for vsini in vsini_values:
                            sensitivity[T][gravity][metal][alpha][vsini] = slope**2
        return modeldict, processed, sensitivity

    return modeldict, processed
def Correlate(data, model_orders, debug=False, addmode="ML",
              orderweights=None, get_weights=False, prim_teff=10000.0):
    """
    This function does the actual correlation. The interface is slightly less useful than GetCCF,
    but can still be called by the user.

    Parameters:
    ===========
    - data:                list of kglib.utils.DataStructures.xypoint instances
                           The data we are cross-correlating against. Each element in the list
                           is treated like an echelle order.

    - model_orders:        list of kglib.utils.DataStructures.xypoint instances
                           Models relevant for each data orders. Must have the same
                           length as data

    - debug:               boolean
                           Prints debugging info to the screen, and saves various files.

    - addmode:             string
                           The CCF addition mode. The default is Maximum Likelihood
                           (from Zucker 2003, MNRAS, 342, 1291). The other valid options
                           are "simple", which will just do a straight addition,
                           "dc", which weights by the CCF value itself, and "weighted",
                           which weights each order. Maximum Likelihood is better for
                           finding weak signals, but simple is better for determining
                           parameters from the CCF (such as vsini)

    - orderweights:        list of floats
                           Weights to apply to each order. Only used if addmode="weighted".
                           Must have the same length as the data list

    - get_weights:         boolean
                           If true, attempts to determine weights from the information content
                           and flux ratio of the companion spectra.
                           The weights are only used if addmode="weighted"

    - prim_teff:           float
                           The effective temperature of the primary star. Used to determine the
                           flux ratio, which in turn is used to make the weights. Ignored if
                           addmode is not "weighted" or get_weights is False.

    Returns:
    ========
    A CCFContainer object
    """

    # Error checking
    if "weighted" in addmode.lower() and orderweights is None and not get_weights:
        raise ValueError("Must give orderweights if addmode == weighted")

    corrlist = []
    normalization = 0.0
    info_content = []
    flux_ratio = []
    snr = []
    for ordernum, order in enumerate(data):
        model = model_orders[ordernum]
        if get_weights:
            slopes = [(model.y[i + 1] / model.cont[i + 1] - model.y[i - 1] / model.cont[i - 1]) /
                      (model.x[i + 1] - model.x[i - 1]) for i in range(1, model.size() - 1)]
            prim_flux = Planck(model.x * units.nm.to(units.cm), prim_teff)
            lines = FittingUtilities.FindLines(model)
            sec_flux = np.median(model.y.max() - model.y[lines])
            flux_ratio.append(np.median(sec_flux) / np.median(prim_flux))
            info_content.append(np.sum(np.array(slopes) ** 2))
            snr.append(1.0 / np.std(order.y))

        reduceddata = order.y / order.cont
        reducedmodel = model.y / model.cont

        # Get the CCF for this order
        l = np.searchsorted(model.x, order.x[0])
        if l > 0:
            if order.x[0] >= model.x[l]:
                dl = (order.x[0] - model.x[l]) / (model.x[l + 1] - model.x[l])
                l += dl
            else:
                logging.debug('Less!')
                dl = (model.x[l] - order.x[0]) / (model.x[l] - model.x[l - 1])
                l -= dl
            logging.debug('dl = {}'.format(dl))
        ycorr = Normalized_Xcorr.norm_xcorr(reduceddata, reducedmodel, trim=False)
        N = ycorr.size
        distancePerLag = np.log(model.x[1] / model.x[0])
        v1 = -(order.size() + l - 0.5) * distancePerLag
        vf = v1 + N * distancePerLag
        offsets = np.linspace(v1, vf, N)
        velocity = -offsets * constants.c.cgs.value * units.cm.to(units.km)
        corr = DataStructures.xypoint(velocity.size)
        corr.x = velocity[::-1]
        corr.y = ycorr[::-1]

        # Only save part of the correlation
        left = np.searchsorted(corr.x, minvel)
        right = np.searchsorted(corr.x, maxvel)
        corr = corr[left:right]

        # Make sure that no elements of corr.y are > 1!
        if max(corr.y) > 1.0:
            corr.y /= max(corr.y)


        # Save correlation
        if np.any(np.isnan(corr.y)):
            warnings.warn("NaNs found in correlation from order %i\n" % (ordernum + 1))
            continue
        normalization += float(order.size())
        corrlist.append(corr.copy())

    if get_weights:
        if debug:
            print("Weight components: ")
            print("lam_0  info  flux ratio,  S/N")
            for i, f, o, s in zip(info_content, flux_ratio, data, snr):
                print(np.median(o.x), i, f, s)
        info_content = (np.array(info_content) - min(info_content)) / (max(info_content) - min(info_content))
        flux_ratio = (np.array(flux_ratio) - min(flux_ratio)) / (max(flux_ratio) - min(flux_ratio))
        snr = (np.array(snr) - min(snr)) / (max(snr) - min(snr))
        orderweights = (1.0 * info_content ** 2 + 1.0 * flux_ratio ** 2 + 1.0 * snr ** 2)
        orderweights /= orderweights.sum()
        logging.debug('Weights:')
        logging.debug(orderweights)

    # Add up the individual CCFs
    total = corrlist[0].copy()
    total_ccfs = CCFContainer(total.x)

    if addmode.lower() == "ml" or addmode.lower() == 'all':
        # use the Maximum Likelihood method from Zucker 2003, MNRAS, 342, 1291
        total.y = np.ones(total.size())
        for i, corr in enumerate(corrlist):
            correlation = spline(corr.x, corr.y, k=1)
            N = data[i].size()
            total.y *= np.power(1.0 - correlation(total.x) ** 2, float(N) / normalization)
        total_ccfs['ml'] = np.sqrt(1.0 - total.y)
    if addmode.lower() == "simple" or addmode.lower() == 'all':
        # do a simple addition
        total.y = np.zeros(total.size())
        for i, corr in enumerate(corrlist):
            correlation = spline(corr.x, corr.y, k=1)
            total.y += correlation(total.x)
        total_ccfs['simple'] = total.y / float(len(corrlist))
    if addmode.lower() == "dc" or addmode.lower() == 'all':
        total.y = np.zeros(total.size())
        for i, corr in enumerate(corrlist):
            N = data[i].size()
            correlation = spline(corr.x, corr.y, k=1)
            total.y += float(N) * correlation(total.x) ** 2 / normalization
        total_ccfs['dc'] = np.sqrt(total.y)
    if addmode.lower() == "weighted" or (addmode.lower() == 'all' and orderweights is not None):
        total.y = np.zeros(total.size())
        for i, corr in enumerate(corrlist):
            w = orderweights[i] / np.sum(orderweights)
            correlation = spline(corr.x, corr.y, k=1)
            total.y += w * correlation(total.x) ** 2 
        total_ccfs['weighted'] = np.sqrt(total.y)
    if addmode.lower() == 'simple-weighted' or (addmode.lower() == 'all' and orderweights is not None):
        total.y = np.zeros(total.size())
        for i, corr in enumerate(corrlist):
            w = orderweights[i] / np.sum(orderweights)
            correlation = spline(corr.x, corr.y, k=1)
            total.y += correlation(total.x) * w
        total_ccfs['simple-weighted'] = total.y / float(len(corrlist))

    if addmode.lower() == 'all':
        return (total_ccfs, corrlist) if debug else total_ccfs
    return (total_ccfs[addmode], corrlist) if debug else total_ccfs[addmode]
def Process(model, data, vsini, resolution, debug=False):
    """
    Process the model to prepare for cross-correlation.

    Parameters:
    ===========
    - data:          list of kglib.utils.DataStructures.xypoint instances
                     The data we are cross-correlating against. Each element in the list
                     is treated like an echelle order.

    - vsini:         float
                     The rotational velocity for the model template, in km/s

    - resolution:    float
                     The detector resolution. The model is convolved with a gaussian of appropriate width to get
                     to this resolution.


    - debug:         boolean
                     Print a bit more information and output the broadened model to './Test_model.dat'

    Returns:
    =========
    - model_orders:  list of kglib.utils.DataStructures.xypoint instances
                     The broadened models, resampled to the same spacing as the data and ready for cross-correlation.
    """
    # Read in the model if necessary
    if isinstance(model, str):
        logging.debug("Reading in the input model from {0:s}".format(model))
        x, y = np.loadtxt(model, usecols=(0, 1), unpack=True)
        x = x * units.angstrom.to(units.nm)
        y = 10 ** y
        left = np.searchsorted(x, data[0].x[0] - 10)
        right = np.searchsorted(x, data[-1].x[-1] + 10)
        model = DataStructures.xypoint(x=x[left:right], y=y[left:right])
    elif not isinstance(model, DataStructures.xypoint):
        raise TypeError(
            "Input model is of an unknown type! Must be a DataStructures.xypoint or a string with the filename.")


    # Linearize the x-axis of the model
    logging.debug('Linearizing model')
    xgrid = np.linspace(model.x[0], model.x[-1], model.size())
    model = FittingUtilities.RebinData(model, xgrid)


    # Broaden
    logging.debug("Rotationally broadening model to vsini = {0:g} km/s".format(vsini * units.cm.to(units.km)))
    if vsini > 1.0 * units.km.to(units.cm):
        model = RotBroad.Broaden(model, vsini, linear=True)


    # Reduce resolution
    logging.debug(u"Convolving to the detector resolution of {}".format(resolution))
    if resolution is not None and  5000 < resolution < 500000:
        model = FittingUtilities.ReduceResolution(model, resolution)


    # Rebin subsets of the model to the same spacing as the data
    model_orders = []
    if debug:
        model.output("Test_model.dat")

    for i, order in enumerate(data):
        if debug:
            sys.stdout.write("\rGenerating model subset for order %i in the input data" % (i + 1))
            sys.stdout.flush()
        # Find how much to extend the model so that we can get maxvel range.
        dlambda = order.x[order.size() / 2] * maxvel * 1.5 / 3e5
        left = np.searchsorted(model.x, order.x[0] - dlambda)
        right = np.searchsorted(model.x, order.x[-1] + dlambda)
        right = min(right, model.size() - 2)

        # Figure out the log-spacing of the data
        logspacing = np.log(order.x[1] / order.x[0])

        # Finally, space the model segment with the same log-spacing
        start = np.log(model.x[left])
        end = np.log(model.x[right])
        xgrid = np.exp(np.arange(start, end + logspacing, logspacing))

        segment = FittingUtilities.RebinData(model[left:right + 1].copy(), xgrid)
        segment.cont = FittingUtilities.Continuum(segment.x, segment.y, lowreject=1.5, highreject=5, fitorder=2)
        model_orders.append(segment)

    return model_orders