kepler-lens

main.py

import os
import wget
import tarfile
import click
import kplr
from string import Template
from pyke import kepconvert
from subprocess import call
import json
import sys

ARCHIVE_URL = 'https://archive.stsci.edu/pub/{}/lightcurves/tarfiles/'

K2_TARFILES = {
    'c1': 2,
    'c102': 2,
    'c111': 1,
    'c112': 2,
    'c12': 3,
    'c13': 2,
    'c3': 3,
    'c4': 1,
    'c5': 2,
    'c6': 3,
    'c7': 1,
    'c8': 2
}

KEPLER_TARFILES = {
    'Q0': 1,
    'Q1': 6,
    'Q2': 10,
    'Q3': 10,
    'Q4': 10,
    'Q5': 10,
    'Q6': 10,
    'Q7': 10,
    'Q8': 8,
    'Q9': 11,
    'Q10': 10,
    'Q11': 10,
    'Q12': 9,
    'Q13': 10,
    'Q14': 11,
    'Q15': 10,
    'Q16': 10,
    'Q17': 4
}

META_UPDATE_MC2_TEMPLATE = """
<'
$attributes
'>
JSON-> 'data' STORE 

MARK

'$rt' '~.*' { 'id' '$filenames' } NOW -1 ] FETCH

<%
    DUP NAME RENAME
    $$data 'attributes' GET SETATTRIBUTES
%> FOREACH
COUNTTOMARK ->LIST '$wt' META
"""


@click.group()
def cli():
    pass

@cli.command()
@click.option('--wtoken', default='w', help='warp WRITE token')
@click.option('--endpoint', default='http://localhost:8080', help='warp endpoint')
@click.option('--path', default='/tmp/kepler', help='kepler download folder')
@click.option('--limit', default='all', help='comma separated list of compagne to download.')
@click.argument('dataset', type=click.Choice(['k2', 'kepler']))
@click.option('--lock', default='', help='LOCK file to suppress at the end of import')
def init(wtoken, endpoint, path, limit, dataset, lock):
    if dataset == "kepler":
        download_campagne(path, limit, "kepler", KEPLER_TARFILES, wtoken, endpoint, lock)
    if dataset == "k2":
        download_campagne(path, limit, "k2", K2_TARFILES, wtoken, endpoint, lock)


def download_campagne(path, limit, dataset, dictFiles, wtoken, endpoint, lock):
    click.echo('Initializing the database, downloading {} dataset...'.format(dataset))

    baseurl = ARCHIVE_URL.format(dataset)

    lightcurvesfolder = path + "/lightcurves/"
    csvfolder = path + "/csv/"

    if not os.path.exists(path):

        click.echo('Creating folder {}'.format(lightcurvesfolder))
        os.makedirs(lightcurvesfolder)
        click.echo('Creating folder {}'.format(csvfolder))
        os.makedirs(csvfolder)

    if limit != "all":
        limits = limit.split(",")
        click.echo('Filtering compagne, only using {}'.format(limit))
        compagnes = {k:v for (k, v) in dictFiles.items() if k in limits}
    else:
        compagnes = dictFiles

    for compagne, nbrfiles in compagnes.items():
        click.echo('downloading {} dataset, {} files'.format(compagne, nbrfiles))
        for nbrfile in range(1, nbrfiles + 1):
            dl_campagne(compagne, nbrfile, lightcurvesfolder, baseurl, dataset)
            generate_csv(compagne, nbrfile, lightcurvesfolder, csvfolder)

        click.echo('compagne {} done!'.format(compagne))
        os.mknod(lightcurvesfolder + compagne + ".done")

    if limit != "":
        click.echo('removing LOCK file')
        os.remove(lock)
    
    click.echo('all compagnes are fetched, bye')