<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import division

""" 
Cross-correlate individual spectra against the AMBRE grid and provide initial
guesses. This is effectively what GUESS will do, but Ireland & Lin have yet to
respond.
"""

__author__ = "<NAME> <<EMAIL>>"

import cPickle as pickle
import matplotlib
import os
from collections import OrderedDict
from glob import glob
from warnings import simplefilter
matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np
from astropy.io import fits
from astropy.table import Table
from astropy.constants import c as speed_of_light

import specutils
from convert import motions

# Suppress "polyfit may be poorly conditioned" messages
simplefilter("ignore", np.RankWarning)
c = speed_of_light.to("km/s").value

# CONTROLS
UPDATE_FILES = False
SAVE_BEST_FIT_MODELS = False
continuum_degree = {
    1: 4,
    2: 4,
    3: 4,
    4: 3
}
ccf_mask = None
chi_sq_mask = [
    [7590, 7617],
    [7620, 7700]
]

vrad_limits = (-1000, 1000)
standardised_file_folder = "/home/acasey/miner3/galah/standardised_spectra/gap_v6.2/"
standardised_file_folder = ""
output_figures_folder = ""
output_models_folder = ""


# CODE
def get_unique_id(filename):
    base = filename.split("/")[-1]
    date_fco_ccd, fib_galah_id = base.split("_")
    date_fco_ccd = date_fco_ccd[:-1] # strip the ccd number from the right
    fib_galah_id = fib_galah_id[:-5] # strip the .fits text from the right
    return "?_".join([date_fco_ccd, fib_galah_id]) + ".fits"

# Identify unique stars
standardised_files = glob(os.path.join(standardised_file_folder, "*.fits"))
id_filename_masks =  map(get_unique_id, standardised_files)

unique_filename_masks = set(id_filename_masks)

# Load in the synthetic spectra from the pickled fluxes.
# Available from: https://zenodo.org/record/14977/files/galah-ambre-grid.pkl
with open("galah-ambre-grid.pkl", "rb") as fp:
    model_grid, model_wavelengths, model_intensities, pixels = pickle.load(fp)
    model_intensities = model_intensities.reshape(model_grid.size, sum(pixels))

# Create a set of rows to store tabular data in.
rows = []
for i, filename_mask in enumerate(unique_filename_masks):

    # get the spectra
    filenames = glob(os.path.join(standardised_file_folder, filename_mask))

    print("Found {0} spectra matching mask {1}".format(len(filenames),
        filename_mask))

    images = map(fits.open, filenames)

    # Check that MOON_DEG exists, otherwise put it in.
    if "MOON_DEG" not in images[0][0].header:
        moon_deg = motions.moon_distance(images[0][0].header)
        print("Inserting MOON_DEG ({0:.1f}) into image headers".format(moon_deg))
        for image in images:
            image[0].header["MOON_DEG"] = moon_deg
            image[0].header.comments["MOON_DEG"] \
                = "Distance to the moon during observing (deg)"

    # Get a list of all the CCDs we actually have.
    ccds = [image[0].header["CCD"] for image in images]

    # Do CCF against all of the arms.
    ccf_results = {}
    for filename, image in zip(filenames, images):

        obs_dispersion = image[0].header["CRVAL1"] + image[0].header["CDELT1"] \
            * (np.arange(image[0].data.size) - image[0].header["CRPIX1"])
        ccd = image[0].header["CCD"]
        si, ei = map(sum, (pixels[:ccd - 1], pixels[:ccd]))

        v, v_err, R = specutils.ccf.cross_correlate(
            obs_dispersion, image[0].data, model_wavelengths[si:ei],
            model_intensities[:, si:ei], rebin="observed", rebin_method="fast",
            ccf_method="fast", continuum_degree=continuum_degree[ccd])<commit_msg>Apply CCF mask<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import division

""" 
Cross-correlate individual spectra against the AMBRE grid and provide initial
guesses. This is effectively what GUESS will do, but Ireland & Lin have yet to
respond.
"""

__author__ = "<NAME> <<EMAIL>>"

import cPickle as pickle
import matplotlib
import os
from collections import OrderedDict
from glob import glob
from warnings import simplefilter
matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np
from astropy.io import fits
from astropy.table import Table
from astropy.constants import c as speed_of_light

import specutils
from convert import motions

# Suppress "polyfit may be poorly conditioned" messages
simplefilter("ignore", np.RankWarning)
c = speed_of_light.to("km/s").value

# CONTROLS
UPDATE_FILES = False
SAVE_BEST_FIT_MODELS = False
continuum_degree = {
    1: 4,
    2: 4,
    3: 4,
    4: 3
}
ccf_mask = [
    [7500, 7700]
]
chi_sq_mask = [
    [7590, 7617],
    [7620, 7700]
]

vrad_limits = (-1000, 1000)
standardised_file_folder = "/home/acasey/miner3/galah/standardised_spectra/gap_v6.2/"
standardised_file_folder = ""
output_figures_folder = ""
output_models_folder = ""


# CODE
def get_unique_id(filename):
    base = filename.split("/")[-1]
    date_fco_ccd, fib_galah_id = base.split("_")
    date_fco_ccd = date_fco_ccd[:-1] # strip the ccd number from the right
    fib_galah_id = fib_galah_id[:-5] # strip the .fits text from the right
    return "?_".join([date_fco_ccd, fib_galah_id]) + ".fits"

# Identify unique stars
standardised_files = glob(os.path.join(standardised_file_folder, "*.fits"))
id_filename_masks =  map(get_unique_id, standardised_files)

unique_filename_masks = set(id_filename_masks)

# Load in the synthetic spectra from the pickled fluxes.
# Available from: https://zenodo.org/record/14977/files/galah-ambre-grid.pkl
with open("galah-ambre-grid.pkl", "rb") as fp:
    model_grid, model_wavelengths, model_intensities, pixels = pickle.load(fp)
    model_intensities = model_intensities.reshape(model_grid.size, sum(pixels))

# Create a set of rows to store tabular data in.
rows = []
for i, filename_mask in enumerate(unique_filename_masks):

    # get the spectra
    filenames = glob(os.path.join(standardised_file_folder, filename_mask))

    print("Found {0} spectra matching mask {1}".format(len(filenames),
        filename_mask))

    images = map(fits.open, filenames)

    # Check that MOON_DEG exists, otherwise put it in.
    if "MOON_DEG" not in images[0][0].header:
        moon_deg = motions.moon_distance(images[0][0].header)
        print("Inserting MOON_DEG ({0:.1f}) into image headers".format(moon_deg))
        for image in images:
            image[0].header["MOON_DEG"] = moon_deg
            image[0].header.comments["MOON_DEG"] \
                = "Distance to the moon during observing (deg)"

    # Get a list of all the CCDs we actually have.
    ccds = [image[0].header["CCD"] for image in images]

    # Do CCF against all of the arms.
    ccf_results = {}
    for filename, image in zip(filenames, images):

        obs_dispersion = image[0].header["CRVAL1"] + image[0].header["CDELT1"] \
            * (np.arange(image[0].data.size) - image[0].header["CRPIX1"])
        ccd = image[0].header["CCD"]
        si, ei = map(sum, (pixels[:ccd - 1], pixels[:ccd]))

        # Apply the CCF mask.
        ccf_flux_data = image[0].data.copy()
        for mask_region in ccf_mask:
            sj, ej = np.searchsorted(obs_dispersion, mask_region)
            ccf_flux_data = np.nan

        v, v_err, R = specutils.ccf.cross_correlate(
            obs_dispersion, ccf_flux_data, model_wavelengths[si:ei],
            model_intensities[:, si:ei], rebin="observed", rebin_method="fast",
            ccf_method="fast", continuum_degree=continuum_degree[ccd])
{'fraction_non_alphanumeric': 0.117286418765827, 'fraction_numerical': 0.018392642942822873, 'mean_word_length': 3.782664117272148, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6472772', 'n_tokens_mistral': 2682, 'n_tokens_neox': 2557, 'n_words': 646}