<commit_before># Standard library
from os.path import abspath, expanduser, join

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

# Project
from ..config import TWOFACE_CACHE_PATH
from ..util import Timer
from ..log import log as logger
from .connect import db_connect, Base
from .model import (AllStar, AllVisit, AllVisitToAllStar, Status, RedClump,
                    StarResult)
from .query_helpers import paged_query

__all__ = ['initialize_db', 'load_red_clump', 'load_nessrg']

def tblrow_to_dbrow(tblrow, colnames, varchar_cols=[]):
    row_data = dict()
    for k in colnames:
        if k in varchar_cols:
            row_data[k.lower()] = tblrow[k].strip()

        # special case the array columns
        elif k.lower().startswith('fparam') and 'cov' not in k.lower():
            i = int(k[-1])
            row_data[k.lower()] = tblrow[k[:-1]][i]

        elif k.lower().startswith('fparam') and 'cov' in k.lower():
            i = int(k[-2])
            j = int(k[-1])
            row_data[k.lower()] = tblrow[k[:-2]][i,j]

        else:
            row_data[k.lower()] = tblrow[k]

    return row_data

def initialize_db(allVisit_file, allStar_file, database_file,
                  drop_all=False, batch_size=4096):
    """Initialize the database given FITS filenames for the APOGEE data.

    Parameters
    ----------
    allVisit_file : str
        Full path to APOGEE allVisit file.
    allStar_file : str
        Full path to APOGEE allStar file.
    database_file : str
        Filename (not path) of database file in cache path.
    drop_all : bool (optional)
        Drop all existing tables and re-create the database.
    batch_size : int (optional)
        How many rows to create before committing.
    """

    database_path = join(TWOFACE_CACHE_PATH, database_file)

    norm = lambda x: abspath(expanduser(x))
    allvisit_tbl = fits.getdata(norm(allVisit_file))
    allstar_tbl = fits.getdata(norm(allStar_file))

    # Remove bad velocities and flagged bad visits:
    skip_mask = np.sum(2 ** np.array([9, 12, 13, # PERSIST_HIGH, PERSIST_JUMP_POS, PERSIST_JUMP_NEG
                                      3, 4])) # VERY_BRIGHT_NEIGHBOR, LOW_SNR
    allvisit_tbl = allvisit_tbl[np.isfinite(allvisit_tbl['VHELIO']) &
                                np.isfinite(allvisit_tbl['VRELERR']) &
                                (allvisit_tbl['VRELERR'] < 100.) & # MAGIC NUMBER
                                ((allvisit_tbl['STARFLAG'] & skip_mask) == 0)]

    v_apogee_ids, counts = np.unique(allvisit_tbl['APOGEE_ID'],
                                     return_counts=True)
    star_mask = np.isin(allstar_tbl['APOGEE_ID'], v_apogee_ids[counts >= 3])

    # Remove STAR_BAD stars:
    skip_mask = 2 ** np.array([23])
    star_mask &= ((allstar_tbl['ASPCAPFLAG'] & skip_mask) == 0)

    # Remove stars with bad logg measurements, or dwarfs
    star_mask &= (allstar_tbl['logg'] > 0) & (allstar_tbl['logg'] < 3.5)
    allstar_tbl = allstar_tbl[star_mask]

    # Only load visits for stars that we're loading
    allvisit_tbl = allvisit_tbl[np.isin(allvisit_tbl['APOGEE_ID'],
                                        allstar_tbl['APOGEE_ID'])]

    allvisit_tbl = Table(allvisit_tbl)
    allstar_tbl = Table(allstar_tbl)
<commit_msg>change logg cut<commit_after># Standard library
from os.path import abspath, expanduser, join

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

# Project
from ..config import TWOFACE_CACHE_PATH
from ..util import Timer
from ..log import log as logger
from .connect import db_connect, Base
from .model import (AllStar, AllVisit, AllVisitToAllStar, Status, RedClump,
                    StarResult)
from .query_helpers import paged_query

__all__ = ['initialize_db', 'load_red_clump', 'load_nessrg']

def tblrow_to_dbrow(tblrow, colnames, varchar_cols=[]):
    row_data = dict()
    for k in colnames:
        if k in varchar_cols:
            row_data[k.lower()] = tblrow[k].strip()

        # special case the array columns
        elif k.lower().startswith('fparam') and 'cov' not in k.lower():
            i = int(k[-1])
            row_data[k.lower()] = tblrow[k[:-1]][i]

        elif k.lower().startswith('fparam') and 'cov' in k.lower():
            i = int(k[-2])
            j = int(k[-1])
            row_data[k.lower()] = tblrow[k[:-2]][i,j]

        else:
            row_data[k.lower()] = tblrow[k]

    return row_data

def initialize_db(allVisit_file, allStar_file, database_file,
                  drop_all=False, batch_size=4096):
    """Initialize the database given FITS filenames for the APOGEE data.

    Parameters
    ----------
    allVisit_file : str
        Full path to APOGEE allVisit file.
    allStar_file : str
        Full path to APOGEE allStar file.
    database_file : str
        Filename (not path) of database file in cache path.
    drop_all : bool (optional)
        Drop all existing tables and re-create the database.
    batch_size : int (optional)
        How many rows to create before committing.
    """

    database_path = join(TWOFACE_CACHE_PATH, database_file)

    norm = lambda x: abspath(expanduser(x))
    allvisit_tbl = fits.getdata(norm(allVisit_file))
    allstar_tbl = fits.getdata(norm(allStar_file))

    # Remove bad velocities and flagged bad visits:
    skip_mask = np.sum(2 ** np.array([9, 12, 13, # PERSIST_HIGH, PERSIST_JUMP_POS, PERSIST_JUMP_NEG
                                      3, 4])) # VERY_BRIGHT_NEIGHBOR, LOW_SNR
    allvisit_tbl = allvisit_tbl[np.isfinite(allvisit_tbl['VHELIO']) &
                                np.isfinite(allvisit_tbl['VRELERR']) &
                                (allvisit_tbl['VRELERR'] < 100.) & # MAGIC NUMBER
                                ((allvisit_tbl['STARFLAG'] & skip_mask) == 0)]

    v_apogee_ids, counts = np.unique(allvisit_tbl['APOGEE_ID'],
                                     return_counts=True)
    star_mask = np.isin(allstar_tbl['APOGEE_ID'], v_apogee_ids[counts >= 3])

    # Remove STAR_BAD stars:
    skip_mask = 2 ** np.array([23])
    star_mask &= ((allstar_tbl['ASPCAPFLAG'] & skip_mask) == 0)

    # Remove stars with bad logg measurements, or dwarfs
    star_mask &= (allstar_tbl['logg'] > 0) & (allstar_tbl['logg'] < 4.)
    allstar_tbl = allstar_tbl[star_mask]

    # Only load visits for stars that we're loading
    allvisit_tbl = allvisit_tbl[np.isin(allvisit_tbl['APOGEE_ID'],
                                        allstar_tbl['APOGEE_ID'])]

    allvisit_tbl = Table(allvisit_tbl)
    allstar_tbl = Table(allstar_tbl)

{'fraction_non_alphanumeric': 0.11116131947582467, 'fraction_numerical': 0.00858563036601898, 'mean_word_length': 2.5603217158176945, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3665656', 'n_tokens_mistral': 2272, 'n_tokens_neox': 2133, 'n_words': 528}