<commit_before>"""Simulate some toy model light curves."""
import numpy as np
import matplotlib.pyplot as plt
import zachopy.units as u
from zachopy.Talker import Talker


rotation_table = None


def get_rotationtable():
    """Lazily load the McQuillan rotation periods table"""
    import astropy.io.ascii as ascii
    import pkgutil
    global rotation_table
    if rotation_table is None:
        print "[lightcurve.py] reading McQuillan rotation table"
        rotation_table = ascii.read(pkgutil.get_data(__name__, 'data/rotation_McQuillan.txt'))
    return rotation_table


transit_table = None


def get_transit_table():
    """Lazily load the Kepler TCE table"""
    # from Seader et al.
    # "The search includes a total of $198,646$ targets,
    # of which $112,001$ were observed in every quarter
    # and $86,645$ were observed in a subset of the 17 quarters."
    import astropy.io.ascii as ascii
    import pkgutil
    global transit_table
    if transit_table is None:
        print "[lightcurve.py] reading Kepler TCE table"
        transit_table = ascii.read(pkgutil.get_data(__name__, 'data/keplerTCE_DR24.txt'))
    return transit_table


def draw_rotation(seed=None):
    """Return a random sine curve, drawn from McQuillan rotation periods"""
    np.random.seed(seed)
    row = np.random.choice(get_rotationtable())
    # noinspection PyTypeChecker
    return Sinusoid(P=row['PRot'],
                    E=np.random.uniform(0, row['PRot']),
                    A=row['Rper'] / 1.0e6)


def draw_transit(seed=None):
    """Draw a random transit from the Kepler TCE list"""
    np.random.seed(seed)
    row = np.random.choice(get_transit_table())
    T14 = row['tce_duration'] / 24.0
    return Trapezoid(P=row['tce_period'],
                     E=row['tce_time0bk'] + 2545833.0,
                     T14=T14,
                     T23=max(T14 - 2 * row['tce_ingress'] / 24.0, 0),
                     D=row['tce_depth'] / 1.0e6)


def parseCode(code):
    """extract name and traits from a lightcurve code"""

    traits = {}
    name, traitstrings = code.split('|')
    for t in traitstrings.split(','):
        k, v = t.split('=')
        traits[k] = np.float(v)
    return name, traits


def generate(code):
    """generate a lightcurve object from a lightcurve code string"""
    name, traits = parseCode(code)
    return globals()[name](**traits)


# TODO: Support Custom
# TODO: This should not have a **kw splat
def random(options=('trapezoid', 'sin'),
           fractionwithextremelc=0.01, fractionwithrotation=None,
           fractionwithtrapezoid=None, fractionwithcustom=0.0, **kw):
    """
    random() returns random Lightcurve.

    random() makes use of these keyword arguments:

        options=['trapezoid', 'sin'] (a list of the kinds of a variability to choose from)
        fractionwithextremelc=0.01 (what fraction should be drawn from a separate "extreme" population of light curves?)
        fractionwithtrapezoid=None (what fraction should get trapezoid? if None, defaults to 18%, from Kepler)
        fractionwithrotation=None (what fraction should get rotation? if None, defaults to 26%, from Kepler)

        when injecting, a star gets (in this order):
            a chance to be extreme
            a chance to be a trapezoid
            a chance to be a rotator
    """

    # first of all, give preference to try to be extreme
    if np.random.uniform(0, 1) < fractionwithextremelc:
        return cartoonrandom(options=options, extreme=True)
    else:
        # by default, set the fraction that get rotation to that from Kepler
        if fractionwithrotation is None:
            fractionwithrotation = 34030.0 / 133030.0
        # by default, set the fraction that get transits to that from Kepler
        if fractionwithtrapezoid is None:
            fractionwithtrapezoid = 20152 / 112001.0

        # if 'custom' in options:
        #     if np.random.uniform(0,1) < fractionwithcustom:
        #         return drawCustom()

        # give trapezoids preference over sin curves
        if 'trapezoid' in options:
            if np.random.uniform(0, 1) < fractionwithtrapezoid:
                return draw_transit()

        # then, try to include a sine curve
        if 'sin' in options:
            if np.random.uniform(0, 1) < fractionwithrotation:
                return draw_rotation()


    # if nothing else, make the light curve a constant
    return constant()


def cartoonrandom(options=('trapezoid', 'sin'), extreme=False, seed=None):
    """generate a random lightcurve, from a cartoonish population"""
    np.random.seed(seed)
    if extreme:
        opens = ['sin']
    name = np.random.choice(options)
    if name == 'sin':
        if extreme:
            p = [0.1, 30.0]
            a = [0.1, 1]
        else:
            p = [0.1, 30.0]
            a = [0.0001, 0.02]

        P = 10 ** np.random.uniform(*np.log10(p))
        E = np.random.uniform(0, P)
        A = 10 ** np.random.uniform(*np.log10(a))

        # noinspection PyTypeChecker
        return Sinusoid(P=P, E=E, A=A)

    if name == 'trapezoid':
        if extreme:
            p = [0.1, 30.0]
            d = [0.1, 1]
        else:
            p = [0.1, 30.0]
            d = [0.0001, 0.01]
        P = 10 ** np.random.uniform(*np.log10(p))
        E = np.random.uniform(0, P)

        mass = np.random.uniform(0.1, 1.5)
        radius = mass
        stellar_density = 3 * mass * u.Msun / (4 * np.pi * (radius * u.Rsun) ** 3)
        rsovera = (3 * np.pi / u.G / (P * u.day) ** 2 / stellar_density) ** (1.0 / 3.0)
        T14 = rsovera * P / np.pi
        # noinspection PyTypeChecker
        T23 = np.random.uniform(0, T14)
        D = 10 ** np.random.uniform(*np.log10(d))

        # noinspection PyTypeChecker
        return Trapezoid(P=P, E=E, D=D, T23=T23, T14=T14)


def testplot(n=100, step=0.01, **kw):
    """make a plot testing random light curves,
        keywords will be passed to "random()\""""
    plt.ion()
    plt.cla()
    f = plt.figure(figsize=(5, 10))
    ax = plt.subplot()
    o = 0
    for i in range(n):
        lc = random(**kw)
        print lc
        lc.demo(offset=o, ax=ax)
        o += step
    ax.set_ylim(step * n, 0)
    plt.draw()


class LightCurve(Talker):
    """The LightCurve class defines the basics of a light curve object, which can<commit_msg>Deterministic light curve randomness with seeds<commit_after>"""Simulate some toy model light curves."""
import numpy as np
import matplotlib.pyplot as plt
import zachopy.units as u
from zachopy.Talker import Talker


rotation_table = None


def get_rotationtable():
    """Lazily load the McQuillan rotation periods table"""
    import astropy.io.ascii as ascii
    import pkgutil
    global rotation_table
    if rotation_table is None:
        print "[lightcurve.py] reading McQuillan rotation table"
        rotation_table = ascii.read(pkgutil.get_data(__name__, 'data/rotation_McQuillan.txt'))
    return rotation_table


transit_table = None


def get_transit_table():
    """Lazily load the Kepler TCE table"""
    # from Seader et al.
    # "The search includes a total of $198,646$ targets,
    # of which $112,001$ were observed in every quarter
    # and $86,645$ were observed in a subset of the 17 quarters."
    import astropy.io.ascii as ascii
    import pkgutil
    global transit_table
    if transit_table is None:
        print "[lightcurve.py] reading Kepler TCE table"
        transit_table = ascii.read(pkgutil.get_data(__name__, 'data/keplerTCE_DR24.txt'))
    return transit_table


def draw_rotation(seed=None):
    """Return a random sine curve, drawn from McQuillan rotation periods"""
    prng = np.random.RandomState(seed)
    row = prng.choice(get_rotationtable())
    # noinspection PyTypeChecker
    return Sinusoid(P=row['PRot'],
                    E=prng.uniform(0, row['PRot']),
                    A=row['Rper'] / 1.0e6)


def draw_transit(seed=None):
    """Draw a random transit from the Kepler TCE list"""
    prng = np.random.RandomState(seed)
    row = prng.choice(get_transit_table())
    T14 = row['tce_duration'] / 24.0
    return Trapezoid(P=row['tce_period'],
                     E=row['tce_time0bk'] + 2545833.0,
                     T14=T14,
                     T23=max(T14 - 2 * row['tce_ingress'] / 24.0, 0),
                     D=row['tce_depth'] / 1.0e6)


def parseCode(code):
    """extract name and traits from a lightcurve code"""

    traits = {}
    name, traitstrings = code.split('|')
    for t in traitstrings.split(','):
        k, v = t.split('=')
        traits[k] = np.float(v)
    return name, traits


def generate(code):
    """generate a lightcurve object from a lightcurve code string"""
    name, traits = parseCode(code)
    return globals()[name](**traits)


# TODO: Support Custom
# TODO: This should not have a **kw splat
def random(options=('trapezoid', 'sin'),
           fractionwithextremelc=0.01, fractionwithrotation=None,
           fractionwithtrapezoid=None, fractionwithcustom=0.0, **kw):
    """
    random() returns random Lightcurve.

    random() makes use of these keyword arguments:

        options=['trapezoid', 'sin'] (a list of the kinds of a variability to choose from)
        fractionwithextremelc=0.01 (what fraction should be drawn from a separate "extreme" population of light curves?)
        fractionwithtrapezoid=None (what fraction should get trapezoid? if None, defaults to 18%, from Kepler)
        fractionwithrotation=None (what fraction should get rotation? if None, defaults to 26%, from Kepler)

        when injecting, a star gets (in this order):
            a chance to be extreme
            a chance to be a trapezoid
            a chance to be a rotator
    """

    # first of all, give preference to try to be extreme
    if np.random.uniform(0, 1) < fractionwithextremelc:
        return cartoonrandom(options=options, extreme=True)
    else:
        # by default, set the fraction that get rotation to that from Kepler
        if fractionwithrotation is None:
            fractionwithrotation = 34030.0 / 133030.0
        # by default, set the fraction that get transits to that from Kepler
        if fractionwithtrapezoid is None:
            fractionwithtrapezoid = 20152 / 112001.0

        # if 'custom' in options:
        #     if np.random.uniform(0,1) < fractionwithcustom:
        #         return drawCustom()

        # give trapezoids preference over sin curves
        if 'trapezoid' in options:
            if np.random.uniform(0, 1) < fractionwithtrapezoid:
                return draw_transit()

        # then, try to include a sine curve
        if 'sin' in options:
            if np.random.uniform(0, 1) < fractionwithrotation:
                return draw_rotation()


    # if nothing else, make the light curve a constant
    return constant()


def cartoonrandom(options=('trapezoid', 'sin'), extreme=False, seed=None):
    """generate a random lightcurve, from a cartoonish population"""

    prng = np.random.RandomState(seed)
    if extreme:
        options = ['sin']
    name = prng.choice(options)
    if name == 'sin':
        if extreme:
            p = [0.1, 30.0]
            a = [0.1, 1]
        else:
            p = [0.1, 30.0]
            a = [0.0001, 0.02]

        P = 10 ** prng.uniform(*np.log10(p))
        E = prng.uniform(0, P)
        A = 10 ** prng.uniform(*np.log10(a))

        # noinspection PyTypeChecker
        return Sinusoid(P=P, E=E, A=A)

    if name == 'trapezoid':
        if extreme:
            p = [0.1, 30.0]
            d = [0.1, 1]
        else:
            p = [0.1, 30.0]
            d = [0.0001, 0.01]
        P = 10 ** prng.uniform(*np.log10(p))
        E = prng.uniform(0, P)

        mass = prng.uniform(0.1, 1.5)
        radius = mass
        stellar_density = 3 * mass * u.Msun / (4 * np.pi * (radius * u.Rsun) ** 3)
        rsovera = (3 * np.pi / u.G / (P * u.day) ** 2 / stellar_density) ** (1.0 / 3.0)
        T14 = rsovera * P / np.pi
        # noinspection PyTypeChecker
        T23 = prng.uniform(0, T14)
        D = 10 ** prng.uniform(*np.log10(d))

        # noinspection PyTypeChecker
        return Trapezoid(P=P, E=E, D=D, T23=T23, T14=T14)


def testplot(n=100, step=0.01, **kw):
    """make a plot testing random light curves,
        keywords will be passed to "random()\""""
    plt.ion()
    plt.cla()
    f = plt.figure(figsize=(5, 10))
    ax = plt.subplot()
    o = 0
    for i in range(n):
        lc = random(**kw)
        print lc
        lc.demo(offset=o, ax=ax)
        o += step
    ax.set_ylim(step * n, 0)
    plt.draw()


class LightCurve(Talker):
    """The LightCurve class defines the basics of a light curve object, which can
{'fraction_non_alphanumeric': 0.11322984792372548, 'fraction_numerical': 0.031360806871010956, 'mean_word_length': 2.3782273090231567, 'pattern_counts': {'":': 0, '<': 11, '<?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': '7607174', 'n_tokens_mistral': 4353, 'n_tokens_neox': 3992, 'n_words': 1225}