<commit_before>__author__ = 'yishay'

import numpy as np
import scipy.linalg
import scipy.interpolate


class ContinuumFitPCA:
    BLUE_START = 1020
    RED_END = 1600
    LY_A_PEAK_BINNED = 1216
    LY_A_PEAK_INDEX = (LY_A_PEAK_BINNED - BLUE_START) / 0.5
    NUM_BINS = (RED_END - LY_A_PEAK_BINNED) * 2 + 1

    def __init__(self, red_pc_text_file, full_pc_text_file, projection_matrix_file,
                 fit_function='dot_product'):
        self.red_pc_table = np.genfromtxt(red_pc_text_file, skip_header=23)
        self.full_pc_table = np.genfromtxt(full_pc_text_file, skip_header=23)
        self.projection_matrix = np.genfromtxt(projection_matrix_file, delimiter=',')
        self.red_pc = self.red_pc_table[:, 3:13]
        self.full_pc = self.full_pc_table[:, 3:13]
        self.red_mean = self.red_pc_table[:, 1]
        self.full_mean = self.full_pc_table[:, 1]
        self.fit_function = {'dot_product': self.project_red_spectrum,
                             'weighted_ls': self.least_squares_red_spectrum}[fit_function]

    def red_to_full(self, red_pc_coefficients):
        return np.dot(self.projection_matrix.T, red_pc_coefficients)

    def project_red_spectrum(self, ar_red_flux, ar_red_ivar):
        del ar_red_ivar  # suppress unused variable
        return np.dot(ar_red_flux - self.red_mean, self.red_pc)

    def least_squares_red_spectrum(self, ar_red_flux, ar_red_ivar):
        ar_red_flux_diff = ar_red_flux - self.red_mean
        ar_sqrt_weights = np.sqrt(ar_red_ivar)
        x = self.red_pc * ar_sqrt_weights[:, None]
        y = ar_red_flux_diff * ar_sqrt_weights
        coefficients = scipy.linalg.lstsq(x, y)
        return coefficients[0]


    def full_spectrum(self, full_pc_coefficients):
        return np.dot(self.full_pc, full_pc_coefficients) + self.full_mean

    def fit_rebin(self, ar_wavelength_rest, ar_flux, ar_ivar, normalized):
        # mask ROUGHLY at the useful spectrum range.
        # include some extra data from the edges for the nearest neighbor interpolation.
        red_spectrum_mask = [(self.LY_A_PEAK_BINNED - 1 <= ar_wavelength_rest) &
                             (ar_wavelength_rest <= self.RED_END + 1)]
        ar_red_wavelength_rest = ar_wavelength_rest[red_spectrum_mask]
        ar_red_flux = ar_flux[red_spectrum_mask]<commit_msg>pca fit: selectable number of principal components.<commit_after>__author__ = 'yishay'

import numpy as np
import scipy.linalg
import scipy.interpolate


class ContinuumFitPCA:
    BLUE_START = 1020
    RED_END = 1600
    LY_A_PEAK_BINNED = 1216
    LY_A_PEAK_INDEX = (LY_A_PEAK_BINNED - BLUE_START) / 0.5
    NUM_BINS = (RED_END - LY_A_PEAK_BINNED) * 2 + 1

    def __init__(self, red_pc_text_file, full_pc_text_file, projection_matrix_file,
                 fit_function='dot_product', num_components=8):
        assert 0 < num_components <= 10
        self.red_pc_table = np.genfromtxt(red_pc_text_file, skip_header=23)
        self.full_pc_table = np.genfromtxt(full_pc_text_file, skip_header=23)
        self.projection_matrix = np.genfromtxt(projection_matrix_file, delimiter=',')[:num_components, :num_components]
        self.num_components = num_components
        self.red_pc = self.red_pc_table[:, 3:3 + num_components]
        self.full_pc = self.full_pc_table[:, 3:3 + num_components]
        self.red_mean = self.red_pc_table[:, 1]
        self.full_mean = self.full_pc_table[:, 1]
        self.fit_function = {'dot_product': self.project_red_spectrum,
                             'weighted_ls': self.least_squares_red_spectrum}[fit_function]

    def red_to_full(self, red_pc_coefficients):
        return np.dot(self.projection_matrix.T, red_pc_coefficients)

    def project_red_spectrum(self, ar_red_flux, ar_red_ivar):
        del ar_red_ivar  # suppress unused variable
        return np.dot(ar_red_flux - self.red_mean, self.red_pc)

    def least_squares_red_spectrum(self, ar_red_flux, ar_red_ivar):
        ar_red_flux_diff = ar_red_flux - self.red_mean
        ar_sqrt_weights = np.sqrt(ar_red_ivar)
        x = self.red_pc * ar_sqrt_weights[:, None]
        y = ar_red_flux_diff * ar_sqrt_weights
        coefficients = scipy.linalg.lstsq(x, y)
        return coefficients[0]


    def full_spectrum(self, full_pc_coefficients):
        return np.dot(self.full_pc, full_pc_coefficients) + self.full_mean

    def fit_rebin(self, ar_wavelength_rest, ar_flux, ar_ivar, normalized):
        # mask ROUGHLY at the useful spectrum range.
        # include some extra data from the edges for the nearest neighbor interpolation.
        red_spectrum_mask = [(self.LY_A_PEAK_BINNED - 1 <= ar_wavelength_rest) &
                             (ar_wavelength_rest <= self.RED_END + 1)]
        ar_red_wavelength_rest = ar_wavelength_rest[red_spectrum_mask]
        ar_red_flux = ar_flux[red_spectrum_mask]
{'fraction_non_alphanumeric': 0.0856610800744879, 'fraction_numerical': 0.01324229257190151, 'mean_word_length': 3.323792486583184, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 3, 'https://': 0, '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': '72521', 'n_tokens_mistral': 1807, 'n_tokens_neox': 1695, 'n_words': 299}