<commit_before># Third-party
import astropy.units as u
import numpy as np
from thejoker.sampler.mcmc import TheJokerMCMCModel

__all__ = ['unimodal_P', 'max_likelihood_sample', 'MAP_sample']


def unimodal_P(samples, data):
    """Check whether the samples returned are within one period mode.

    Parameters
    ----------
    samples : `~thejoker.JokerSamples`

    Returns
    -------
    is_unimodal : bool
    """

    P_samples = samples['P'].to(u.day).value
    P_min = np.min(P_samples)
    T = np.ptp(data.t.mjd)
    delta = 4*P_min**2 / (2*np.pi*T)

    return np.ptp(P_samples) < delta


def max_likelihood_sample(data, samples):
    """Return the maximum-likelihood sample.

    Parameters
    ----------
    data : `~thejoker.RVData`
    samples : `~thejoker.JokerSamples`

    """
    chisqs = np.zeros(len(samples))

    for i in range(len(samples)):
        orbit = samples.get_orbit(i)
        residual = data.rv - orbit.radial_velocity(data.t)
        err = np.sqrt(data.stddev**2 + samples['jitter'][i]**2)
        chisqs[i] = np.sum((residual**2 / err**2).decompose())

    return samples[np.argmin(chisqs)]


def MAP_sample(data, samples, joker_params):
    """Return the maximum a posteriori sample.

    Parameters
    ----------
    data : `~thejoker.RVData`
    samples : `~thejoker.JokerSamples`
    joker_params : `~thejoker.JokerParams`

    """
    model = TheJokerMCMCModel(joker_params, data)

    ln_ps = np.zeros(len(samples))

    mcmc_p = model.pack_samples_mcmc(samples)        
    for i in range(len(samples)):
        ln_ps[i] = model.ln_posterior(mcmc_p[i])

    return samples[np.argmax(ln_ps)]
<commit_msg>option to return index of MAP sample<commit_after># Third-party
import astropy.units as u
import numpy as np
from thejoker.sampler.mcmc import TheJokerMCMCModel

__all__ = ['unimodal_P', 'max_likelihood_sample', 'MAP_sample']


def unimodal_P(samples, data):
    """Check whether the samples returned are within one period mode.

    Parameters
    ----------
    samples : `~thejoker.JokerSamples`

    Returns
    -------
    is_unimodal : bool
    """

    P_samples = samples['P'].to(u.day).value
    P_min = np.min(P_samples)
    T = np.ptp(data.t.mjd)
    delta = 4*P_min**2 / (2*np.pi*T)

    return np.ptp(P_samples) < delta


def max_likelihood_sample(data, samples):
    """Return the maximum-likelihood sample.

    Parameters
    ----------
    data : `~thejoker.RVData`
    samples : `~thejoker.JokerSamples`

    """
    chisqs = np.zeros(len(samples))

    for i in range(len(samples)):
        orbit = samples.get_orbit(i)
        residual = data.rv - orbit.radial_velocity(data.t)
        err = np.sqrt(data.stddev**2 + samples['jitter'][i]**2)
        chisqs[i] = np.sum((residual**2 / err**2).decompose())

    return samples[np.argmin(chisqs)]


def MAP_sample(data, samples, joker_params, return_index=False):
    """Return the maximum a posteriori sample.

    Parameters
    ----------
    data : `~thejoker.RVData`
    samples : `~thejoker.JokerSamples`
    joker_params : `~thejoker.JokerParams`

    """
    model = TheJokerMCMCModel(joker_params, data)

    ln_ps = np.zeros(len(samples))

    mcmc_p = model.pack_samples_mcmc(samples)
    for i in range(len(samples)):
        ln_ps[i] = model.ln_posterior(mcmc_p[i])

    if return_index:
        idx = np.argmax(ln_ps)
        return samples[idx], idx
    else:
        return samples[np.argmax(ln_ps)]

{'fraction_non_alphanumeric': 0.14714452214452214, 'fraction_numerical': 0.004079254079254079, 'mean_word_length': 2.857303370786517, 'pattern_counts': {'":': 0, '<': 5, '<?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': '1048315', 'n_tokens_mistral': 1275, 'n_tokens_neox': 1213, 'n_words': 239}