<commit_before>__author__ = "adrn <<EMAIL>>"

# Standard library
import os
from os.path import abspath, expanduser, join, dirname

# Third-party
from astropy.table import Table

# Project
import twoface
from twoface.util import Timer
from twoface.log import log as logger
from twoface.db import Session, db_connect, AllStar, AllVisit, Status
from twoface.db.core import Base
from twoface.config import TWOFACE_CACHE_PATH

def main(allVisit_file=None, allStar_file=None, test=False, **kwargs):

    # if running in test mode, get test files
    if test:
        base_path = abspath(dirname(twoface.__file__))
        allVisit_file = os.path.join(base_path, 'db', 'tests', 'test-allVisit.fits')
        allStar_file = os.path.join(base_path, 'db', 'tests', 'test-allStar.fits')

        database_path = join(TWOFACE_CACHE_PATH, 'test.sqlite')

    else:
        assert allVisit_file is not None
        assert allStar_file is not None

        database_path = join(TWOFACE_CACHE_PATH, 'apogee.sqlite')

    norm = lambda x: abspath(expanduser(x))
    allvisit_tbl = Table.read(norm(allVisit_file), format='fits', hdu=1)
    allstar_tbl = Table.read(norm(allStar_file), format='fits', hdu=1)

    engine = db_connect(database_path)
    # engine.echo = True
    logger.debug("Connected to database at '{}'".format(database_path))

    # this is the magic that creates the tables based on the definitions in twoface/db/model.py
    Base.metadata.drop_all()
    Base.metadata.create_all()

    # load allVisit and allStar files
    # allvisit_tbl

    session = Session()

    logger.debug("Loading allStar, allVisit tables...")

    allstar_colnames = [str(x).split('.')[1].upper() for x in AllStar.__table__.columns]
    i = allstar_colnames.index('ID')
    allstar_colnames.pop(i)

    allvisit_colnames = [str(x).split('.')[1].upper() for x in AllVisit.__table__.columns]
    i = allvisit_colnames.index('ID')
    allvisit_colnames.pop(i)

    batch_size = 4000
    stars = []
    all_visits = dict()
    with Timer() as t:
        for i,row in enumerate(allstar_tbl):
            row_data = dict([(k.lower(), row[k]) for k in allstar_colnames])
            star = AllStar(**row_data)
            stars.append(star)

            if star.apogee_id not in all_visits:
                visits = []
                for j,visit_row in enumerate(allvisit_tbl[allvisit_tbl['APOGEE_ID'] == row['APOGEE_ID']]):
                    _data = dict([(k.lower(), visit_row[k]) for k in allvisit_colnames])
                    visits.append(AllVisit(**_data))
                all_visits[star.apogee_id] = visits

            else:
                visits = all_visits[star.apogee_id]

            star.visits = visits

            if i % batch_size == 0 and i > 0:
                session.add_all(stars)
                session.add_all([item for sublist in all_visits.values() for item in sublist])
                session.commit()
                logger.debug("Loaded batch {} ({:.2f} seconds)".format(i*batch_size, t.elapsed()))
                t.reset()

                all_visits = dict()
                stars = []

    if len(stars) > 0:
        session.add_all(stars)
        session.add_all([item for sublist in all_visits.values() for item in sublist])
        session.commit()

    logger.debug("tables loaded in {:.2f} seconds".format(t.elapsed()))

    # Load the status table
    logger.debug("Populating Status table...")
    statuses = list()
    statuses.append(Status(id=0, message='untouched'))
    statuses.append(Status(id=1, message='pending'))
    statuses.append(Status(id=2, message='needs more prior samples'))
    statuses.append(Status(id=3, message='needs mcmc'))
    statuses.append(Status(id=4, message='error'))
    statuses.append(Status(id=5, message='completed'))

    session.add_all(statuses)
    session.commit()
    logger.debug("...done")

    session.close()

if __name__ == "__main__":
    from argparse import ArgumentParser
    import logging

    # Define parser object
    parser = ArgumentParser(description="Initialize the TwoFace project database.")

    vq_group = parser.add_mutually_exclusive_group()
    vq_group.add_argument('-v', '--verbose', action='count', default=0, dest='verbosity')
    vq_group.add_argument('-q', '--quiet', action='count', default=0, dest='quietness')
    # parser.add_argument('-o', '--overwrite', action='store_true', dest='overwrite',
    #                     default=False, help='Destroy everything.')

    parser.add_argument("--allstar", dest="allStar_file", default=None,
                        type=str, help="Path to APOGEE allStar FITS file.")
    parser.add_argument("--allvisit", dest="allVisit_file", default=None,
                        type=str, help="Path to APOGEE allVisit FITS file.")
<commit_msg>only load finite vhelio values<commit_after>__author__ = "adrn <<EMAIL>>"

# Standard library
import os
from os.path import abspath, expanduser, join, dirname

# Third-party
from astropy.table import Table
import numpy as np

# Project
import twoface
from twoface.util import Timer
from twoface.log import log as logger
from twoface.db import Session, db_connect, AllStar, AllVisit, Status
from twoface.db.core import Base
from twoface.config import TWOFACE_CACHE_PATH

def main(allVisit_file=None, allStar_file=None, test=False, **kwargs):

    # if running in test mode, get test files
    if test:
        base_path = abspath(dirname(twoface.__file__))
        allVisit_file = os.path.join(base_path, 'db', 'tests', 'test-allVisit.fits')
        allStar_file = os.path.join(base_path, 'db', 'tests', 'test-allStar.fits')

        database_path = join(TWOFACE_CACHE_PATH, 'test.sqlite')

    else:
        assert allVisit_file is not None
        assert allStar_file is not None

        database_path = join(TWOFACE_CACHE_PATH, 'apogee.sqlite')

    norm = lambda x: abspath(expanduser(x))
    allvisit_tbl = Table.read(norm(allVisit_file), format='fits', hdu=1)
    allstar_tbl = Table.read(norm(allStar_file), format='fits', hdu=1)

    engine = db_connect(database_path)
    # engine.echo = True
    logger.debug("Connected to database at '{}'".format(database_path))

    # this is the magic that creates the tables based on the definitions in twoface/db/model.py
    Base.metadata.drop_all()
    Base.metadata.create_all()

    # load allVisit and allStar files
    allvisit_tbl = allvisit_tbl[np.isfinite(allvisit_tbl['VHELIO'])]

    session = Session()

    logger.debug("Loading allStar, allVisit tables...")

    allstar_colnames = [str(x).split('.')[1].upper() for x in AllStar.__table__.columns]
    i = allstar_colnames.index('ID')
    allstar_colnames.pop(i)

    allvisit_colnames = [str(x).split('.')[1].upper() for x in AllVisit.__table__.columns]
    i = allvisit_colnames.index('ID')
    allvisit_colnames.pop(i)

    batch_size = 4000
    stars = []
    all_visits = dict()
    with Timer() as t:
        for i,row in enumerate(allstar_tbl):
            row_data = dict([(k.lower(), row[k]) for k in allstar_colnames])
            star = AllStar(**row_data)
            stars.append(star)

            if star.apogee_id not in all_visits:
                visits = []
                for j,visit_row in enumerate(allvisit_tbl[allvisit_tbl['APOGEE_ID'] == row['APOGEE_ID']]):
                    _data = dict([(k.lower(), visit_row[k]) for k in allvisit_colnames])
                    visits.append(AllVisit(**_data))
                all_visits[star.apogee_id] = visits

            else:
                visits = all_visits[star.apogee_id]

            star.visits = visits

            if i % batch_size == 0 and i > 0:
                session.add_all(stars)
                session.add_all([item for sublist in all_visits.values() for item in sublist])
                session.commit()
                logger.debug("Loaded batch {} ({:.2f} seconds)".format(i*batch_size, t.elapsed()))
                t.reset()

                all_visits = dict()
                stars = []

    if len(stars) > 0:
        session.add_all(stars)
        session.add_all([item for sublist in all_visits.values() for item in sublist])
        session.commit()

    logger.debug("tables loaded in {:.2f} seconds".format(t.elapsed()))

    # Load the status table
    logger.debug("Populating Status table...")
    statuses = list()
    statuses.append(Status(id=0, message='untouched'))
    statuses.append(Status(id=1, message='pending')) # OLD - don't use
    statuses.append(Status(id=2, message='needs more prior samples'))
    statuses.append(Status(id=3, message='needs mcmc'))
    statuses.append(Status(id=4, message='error'))
    statuses.append(Status(id=5, message='completed'))

    session.add_all(statuses)
    session.commit()
    logger.debug("...done")

    session.close()

if __name__ == "__main__":
    from argparse import ArgumentParser
    import logging

    # Define parser object
    parser = ArgumentParser(description="Initialize the TwoFace project database.")

    vq_group = parser.add_mutually_exclusive_group()
    vq_group.add_argument('-v', '--verbose', action='count', default=0, dest='verbosity')
    vq_group.add_argument('-q', '--quiet', action='count', default=0, dest='quietness')
    # parser.add_argument('-o', '--overwrite', action='store_true', dest='overwrite',
    #                     default=False, help='Destroy everything.')

    parser.add_argument("--allstar", dest="allStar_file", default=None,
                        type=str, help="Path to APOGEE allStar FITS file.")
    parser.add_argument("--allvisit", dest="allVisit_file", default=None,
                        type=str, help="Path to APOGEE allVisit FITS file.")

{'fraction_non_alphanumeric': 0.12732919254658384, 'fraction_numerical': 0.004347826086956522, 'mean_word_length': 2.987205943045811, 'pattern_counts': {'":': 2, '<': 7, '<?xml version=': 0, '>': 11, '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': '64509', 'n_tokens_mistral': 3147, 'n_tokens_neox': 2974, 'n_words': 692}