<commit_before>import numpy as np
from scipy.interpolate import interp1d

def one_comp_model(wav, model1, gammas):
    """ Make 1 component simulations, broadcasting over gamma values.
    """
    # Enable single scalar inputs (turn to 1d np.array)
    if not hasattr(gammas, "__len__"):
        gammas = np.asarray(gammas)[np.newaxis]
        print(len(gammas))

    m1 = model1
    print(model1.shape)

    m1g = np.empty(model1.shape + (len(gammas),))   # am2rvm1g = am2rvm1 with gamma doppler-shift
    print(m1g.shape)
    for j, gamma in enumerate(gammas):
        wav_j = (1 + gamma / 299792.458) * wav
        m1g[:, j] = interp1d(wav_j, m1, axis=0, bounds_error=False)(wav)

    return interp1d(w, m1g, axis=0)    # pass it the wavelength values to return


def two_comp_model(wav, model1, model2, alphas, rvs, gammas):
    # Make 2 component simulations, broadcasting over alpha, rv, gamma values.

    # Enable single scalar inputs (turn to 1d np.array)
    if not hasattr(alphas, "__len__"):
        alphas = np.asarray(alphas)[np.newaxis]
    if not hasattr(rvs, "__len__"):
        rvs = np.asarray(rvs)[np.newaxis]
    if not hasattr(gammas, "__len__"):
        gammas = np.asarray(gammas)[np.newaxis]
        print(len(gammas))

    am2 = model2[:, np.newaxis] * alphas  # alpha * Model2 (am2)
    print(am2.shape)

    am2rv = np.empty(am2.shape + (len(rvs),))  # am2rv = am2 with rv doppler-shift
    print(am2rv.shape)
    for i, rv in enumerate(rvs):
        # nflux, wlprime = dopplerShift(wav, am2, rv)
        # am2rv[:, :, i] = nflux
        wav_i = (1 - rv / c) * wav
        am2rv[:, :, i] = interp1d(wav_i, am2, axis=0, bounds_error=False)(wav)

    # Normalize by (1 / 1 + alpha)
    am2rv = am2rv / (1 + alphas)[np.newaxis, :, np.newaxis]

    am2rvm1 = h[:, np.newaxis, np.newaxis] + am2rv  # am2rvm1 = am2rv + model_1
    print(am2rvm1.shape)

    am2rvm1g = np.empty(am2rvm1.shape + (len(gammas),))  # am2rvm1g = am2rvm1 with gamma doppler-shift
    for j, gamma in enumerate(gammas):
        wav_j = (1 - gamma / 299792.458) * wav
        am2rvm1g[:, :, :, j] = interp1d(wav_j, am2rvm1, axis=0, bounds_error=False)(wav)

    return interp1d(w, am2rvm1g, axis=0)  # pass it the wavelength values to return<commit_msg>Tidy broadcasted models to stop printing<commit_after>import numpy as np
from scipy.interpolate import interp1d

def one_comp_model(wav, model1, gammas):
    """ Make 1 component simulations, broadcasting over gamma values.
    """
    # Enable single scalar inputs (turn to 1d np.array)
    if not hasattr(gammas, "__len__"):
        gammas = np.asarray(gammas)[np.newaxis]
        # print(len(gammas))

    m1 = model1
    # print(model1.shape)

    m1g = np.empty(model1.shape + (len(gammas),))   # am2rvm1g = am2rvm1 with gamma doppler-shift
    print(m1g.shape)
    for j, gamma in enumerate(gammas):
        wav_j = (1 + gamma / 299792.458) * wav
        m1g[:, j] = interp1d(wav_j, m1, axis=0, bounds_error=False)(wav)

    assert m1g.shape == (len(model1), len(gammas)), "Dimensions of broadcast output not correct"
    return interp1d(wav, m1g, axis=0)    # pass it the wavelength values to return


def two_comp_model(wav, model1, model2, alphas, rvs, gammas):
    # Make 2 component simulations, broadcasting over alpha, rv, gamma values.

    # Enable single scalar inputs (turn to 1d np.array)
    if not hasattr(alphas, "__len__"):
        alphas = np.asarray(alphas)[np.newaxis]
    if not hasattr(rvs, "__len__"):
        rvs = np.asarray(rvs)[np.newaxis]
    if not hasattr(gammas, "__len__"):
        gammas = np.asarray(gammas)[np.newaxis]
        print(len(gammas))

    am2 = model2[:, np.newaxis] * alphas  # alpha * Model2 (am2)
    # print(am2.shape)

    am2rv = np.empty(am2.shape + (len(rvs),))  # am2rv = am2 with rv doppler-shift
    # print(am2rv.shape)
    for i, rv in enumerate(rvs):
        # nflux, wlprime = dopplerShift(wav, am2, rv)
        # am2rv[:, :, i] = nflux
        wav_i = (1 - rv / c) * wav
        am2rv[:, :, i] = interp1d(wav_i, am2, axis=0, bounds_error=False)(wav)

    # Normalize by (1 / 1 + alpha)
    am2rv = am2rv / (1 + alphas)[np.newaxis, :, np.newaxis]

    am2rvm1 = model1[:, np.newaxis, np.newaxis] + am2rv  # am2rvm1 = am2rv + model_1
    # print(am2rvm1.shape)

    am2rvm1g = np.empty(am2rvm1.shape + (len(gammas),))  # am2rvm1g = am2rvm1 with gamma doppler-shift
    for j, gamma in enumerate(gammas):
        wav_j = (1 - gamma / 299792.458) * wav
        am2rvm1g[:, :, :, j] = interp1d(wav_j, am2rvm1, axis=0, bounds_error=False)(wav)

    assert m1g.shape == (len(model1), len(alphas), len(rvs), len(gammas)), "Dimensions of broadcast output not correct"
    return interp1d(wav, am2rvm1g, axis=0)  # pass it the wavelength values to return
{'fraction_non_alphanumeric': 0.13505566057550936, 'fraction_numerical': 0.04053770216341105, 'mean_word_length': 3.094582975064488, 'pattern_counts': {'":': 0, '<': 3, '<?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': '6585274', 'n_tokens_mistral': 1948, 'n_tokens_neox': 1895, 'n_words': 452}