<commit_before>import sqlite3
import numpy as np
import logging
import matplotlib.pyplot as plt
from scipy.integrate import simps
from scipy.interpolate import griddata
from config import *
import os
import datetime
from star import Star


logger = logging.getLogger(__name__)

class SolvePars:
    def __init__(self, key_parameter_known='plx', db='yy02.sql3',
                 nsigma=5, window_len_age=13):
        self.key_parameter_known = key_parameter_known
        self.get_isochrone_points_db = db
        self.get_isochrone_points_nsigma = nsigma
        self.smooth_window_len_age = window_len_age
        self.smooth_window_len_mass = 0
        self.smooth_window_len_logl = 0
        self.smooth_window_len_mv = 0
        self.smooth_window_len_r = 0
        self.smooth_window_len_logg = 0

class PlotPars:
    def __init__(self, figure_format='png', directory=""):
        self.age_xlim = [0, 14]
        self.mass_xlim = None
        self.logl_xlim = None
        self.mv_xlim = None
        self.r_xlim = None
        self.logg_xlim = None
        self.directory = directory
        self.figure_format = figure_format

def pdf(pdf_x, ips, prob, par, smooth_window_len):
    '''Calculates a probability distribution function (PDF) for parameter par
    given the x-values for the PDF, the isochrone points ips, and their 
    probability. Return PDF and smoothed PDF (using smooth_window_len) if
    possible (otherwise returns two non-smoothed PDFs), as well as a stats
    dictionary with mean, std, most probable value, etc.
    '''
    dx = 0.5*(pdf_x[1] - pdf_x[0])
    pdf_y = []
    for x in pdf_x:
        pdf_y.append(sum(prob[np.logical_and(ips[par] >= x-dx,
                                             ips[par] <  x+dx)]))
    pdf_y = np.array(pdf_y)
    pdf_y = pdf_y/simps(pdf_y, pdf_x)

    try:
        pdf_y_smooth = smooth(pdf_y, smooth_window_len)
        pdf_y_smooth = pdf_y_smooth/simps(pdf_y_smooth, pdf_x)
    except:
        pdf_y_smooth = pdf_y
        logger.warning('Unable to smooth '+par+' PDF.')

    stats = get_stats(pdf_x, pdf_y_smooth)

    if stats['most_probable']:
        print("{0:10s} = {1:6.3f} | {2:6.3f} - {3:6.3f} | "\
              "{4:6.3f} - {5:6.3f} | {6:6.3f} +/- {7:6.3f}"\
              .format(par,
                      stats['most_probable'],
                      stats['lower_limit_1sigma'],
                      stats['upper_limit_1sigma'],
                      stats['lower_limit_2sigma'],
                      stats['upper_limit_2sigma'],
                      stats['mean'], stats['std'])
              )
    else:
        print("{0:10s} =        |        -        |  "\
              "      -        | {1:6.3f} +/- {2:6.3f}"\
              .format(par, stats['mean'], stats['std'])
              )
        logger.warning("--- Unable to calculate PDF stats for "+par)

    return pdf_y, pdf_y_smooth, stats

def get_stats(pdf_x, pdf_y_smooth):
    stats = {}
    stats['most_probable'] = \
      np.mean(np.array(pdf_x)[pdf_y_smooth == max(pdf_y_smooth)])
    stats['mean'] = simps(pdf_y_smooth*pdf_x, pdf_x)
    stats['std'] = np.sqrt(simps(pdf_y_smooth*(pdf_x-stats['mean'])**2,\
                                 pdf_x))

    k = pdf_x <= stats['most_probable']
    pdf_y_left = 0.5*pdf_y_smooth[k]/simps(pdf_y_smooth[k], pdf_x[k])
    pdf_x_left = pdf_x[k]
    areas_left = []
    for x in pdf_x_left:
        areas_left.append(simps(pdf_y_left[pdf_x_left <= x],
                                pdf_x_left[pdf_x_left <= x]))
    areas_left = np.array(areas_left)
    if np.mean(areas_left) == 0:
        logger.warning("Left side of distribution is empty")
        stats['most_probable'] = None
        stats['lower_limit_1sigma'] = None
        stats['lower_limit_2sigma'] = None
        stats['upper_limit_1sigma'] = None
        stats['upper_limit_2sigma'] = None
        return stats

    k = pdf_x >= stats['most_probable']
    pdf_y_right = 0.5*pdf_y_smooth[k]/simps(pdf_y_smooth[k], pdf_x[k])
    pdf_x_right = pdf_x[k]
    areas_right = []
    for x in pdf_x_right:
        areas_right.append(simps(pdf_y_right[pdf_x_right <= x],
                                 pdf_x_right[pdf_x_right <= x]))
    areas_right = np.array(areas_right)

    stats['lower_limit_1sigma'] = \
      np.mean(griddata(areas_left, pdf_x_left, 0.158))
    stats['lower_limit_2sigma'] = \
      np.mean(griddata(areas_left, pdf_x_left, 0.022))
    stats['upper_limit_1sigma'] = \
      np.mean(griddata(areas_right, pdf_x_right, 0.341))
    stats['upper_limit_2sigma'] = \
      np.mean(griddata(areas_right, pdf_x_right, 0.477))

    return stats

def solve_one(Star, SolvePars, PlotPars):
    '''Calculates most likely parameters of Star using YY isochrone points
    '''
    ips = get_isochrone_points(Star,
                               SolvePars.get_isochrone_points_db,
                               SolvePars.get_isochrone_points_nsigma,
                               SolvePars.key_parameter_known)
    if ips == None:
        logger.warning('Could not get any isochrone points.')
        return None
    logger.info('Using {0} Y2 isochrone points'.format(len(ips['age'])))
    Star.yykeyparameterknown = SolvePars.key_parameter_known
    Star.yynpoints = len(ips['age'])
    ips['t'] = 10**ips['logt']
    ips['r'] = 10**(0.5*(np.log10(ips['mass'])-ips['logg']+4.437))

    if SolvePars.key_parameter_known == 'logg':
        prob = np.exp(-1*((ips['t']-Star.teff)/          \
                          (1.414214*Star.err_teff))**2)* \
               np.exp(-1*((ips['logg']-Star.logg)/       \
                          (1.414214*Star.err_logg))**2)* \
               np.exp(-1*((ips['feh']-Star.feh)/         \
                          (1.414214*Star.err_feh))**2)
    if SolvePars.key_parameter_known == 'plx':
        prob = np.exp(-1*((ips['t']-Star.teff)/          \
                          (1.414214*Star.err_teff))**2)* \<commit_msg>print # of isochrone points<commit_after>import sqlite3
import numpy as np
import logging
import matplotlib.pyplot as plt
from scipy.integrate import simps
from scipy.interpolate import griddata
from config import *
import os
import datetime
from star import Star


logger = logging.getLogger(__name__)

class SolvePars:
    def __init__(self, key_parameter_known='plx', db='yy02.sql3',
                 nsigma=5, window_len_age=13):
        self.key_parameter_known = key_parameter_known
        self.get_isochrone_points_db = db
        self.get_isochrone_points_nsigma = nsigma
        self.smooth_window_len_age = window_len_age
        self.smooth_window_len_mass = 0
        self.smooth_window_len_logl = 0
        self.smooth_window_len_mv = 0
        self.smooth_window_len_r = 0
        self.smooth_window_len_logg = 0

class PlotPars:
    def __init__(self, figure_format='png', directory=""):
        self.age_xlim = [0, 14]
        self.mass_xlim = None
        self.logl_xlim = None
        self.mv_xlim = None
        self.r_xlim = None
        self.logg_xlim = None
        self.directory = directory
        self.figure_format = figure_format

def pdf(pdf_x, ips, prob, par, smooth_window_len):
    '''Calculates a probability distribution function (PDF) for parameter par
    given the x-values for the PDF, the isochrone points ips, and their 
    probability. Return PDF and smoothed PDF (using smooth_window_len) if
    possible (otherwise returns two non-smoothed PDFs), as well as a stats
    dictionary with mean, std, most probable value, etc.
    '''
    dx = 0.5*(pdf_x[1] - pdf_x[0])
    pdf_y = []
    for x in pdf_x:
        pdf_y.append(sum(prob[np.logical_and(ips[par] >= x-dx,
                                             ips[par] <  x+dx)]))
    pdf_y = np.array(pdf_y)
    pdf_y = pdf_y/simps(pdf_y, pdf_x)

    try:
        pdf_y_smooth = smooth(pdf_y, smooth_window_len)
        pdf_y_smooth = pdf_y_smooth/simps(pdf_y_smooth, pdf_x)
    except:
        pdf_y_smooth = pdf_y
        logger.warning('Unable to smooth '+par+' PDF.')

    stats = get_stats(pdf_x, pdf_y_smooth)

    if stats['most_probable']:
        print("{0:10s} = {1:6.3f} | {2:6.3f} - {3:6.3f} | "\
              "{4:6.3f} - {5:6.3f} | {6:6.3f} +/- {7:6.3f}"\
              .format(par,
                      stats['most_probable'],
                      stats['lower_limit_1sigma'],
                      stats['upper_limit_1sigma'],
                      stats['lower_limit_2sigma'],
                      stats['upper_limit_2sigma'],
                      stats['mean'], stats['std'])
              )
    else:
        print("{0:10s} =        |        -        |  "\
              "      -        | {1:6.3f} +/- {2:6.3f}"\
              .format(par, stats['mean'], stats['std'])
              )
        logger.warning("--- Unable to calculate PDF stats for "+par)

    return pdf_y, pdf_y_smooth, stats

def get_stats(pdf_x, pdf_y_smooth):
    stats = {}
    stats['most_probable'] = \
      np.mean(np.array(pdf_x)[pdf_y_smooth == max(pdf_y_smooth)])
    stats['mean'] = simps(pdf_y_smooth*pdf_x, pdf_x)
    stats['std'] = np.sqrt(simps(pdf_y_smooth*(pdf_x-stats['mean'])**2,\
                                 pdf_x))

    k = pdf_x <= stats['most_probable']
    pdf_y_left = 0.5*pdf_y_smooth[k]/simps(pdf_y_smooth[k], pdf_x[k])
    pdf_x_left = pdf_x[k]
    areas_left = []
    for x in pdf_x_left:
        areas_left.append(simps(pdf_y_left[pdf_x_left <= x],
                                pdf_x_left[pdf_x_left <= x]))
    areas_left = np.array(areas_left)
    if np.mean(areas_left) == 0:
        logger.warning("Left side of distribution is empty")
        stats['most_probable'] = None
        stats['lower_limit_1sigma'] = None
        stats['lower_limit_2sigma'] = None
        stats['upper_limit_1sigma'] = None
        stats['upper_limit_2sigma'] = None
        return stats

    k = pdf_x >= stats['most_probable']
    pdf_y_right = 0.5*pdf_y_smooth[k]/simps(pdf_y_smooth[k], pdf_x[k])
    pdf_x_right = pdf_x[k]
    areas_right = []
    for x in pdf_x_right:
        areas_right.append(simps(pdf_y_right[pdf_x_right <= x],
                                 pdf_x_right[pdf_x_right <= x]))
    areas_right = np.array(areas_right)

    stats['lower_limit_1sigma'] = \
      np.mean(griddata(areas_left, pdf_x_left, 0.158))
    stats['lower_limit_2sigma'] = \
      np.mean(griddata(areas_left, pdf_x_left, 0.022))
    stats['upper_limit_1sigma'] = \
      np.mean(griddata(areas_right, pdf_x_right, 0.341))
    stats['upper_limit_2sigma'] = \
      np.mean(griddata(areas_right, pdf_x_right, 0.477))

    return stats

def solve_one(Star, SolvePars, PlotPars):
    '''Calculates most likely parameters of Star using YY isochrone points
    '''
    ips = get_isochrone_points(Star,
                               SolvePars.get_isochrone_points_db,
                               SolvePars.get_isochrone_points_nsigma,
                               SolvePars.key_parameter_known)
    if ips == None:
        logger.warning('Could not get any isochrone points.')
        return None
    print('Using {0} Y2 isochrone points'.format(len(ips['age'])))
    logger.info('Using {0} Y2 isochrone points'.format(len(ips['age'])))
    Star.yykeyparameterknown = SolvePars.key_parameter_known
    Star.yynpoints = len(ips['age'])
    ips['t'] = 10**ips['logt']
    ips['r'] = 10**(0.5*(np.log10(ips['mass'])-ips['logg']+4.437))

    if SolvePars.key_parameter_known == 'logg':
        prob = np.exp(-1*((ips['t']-Star.teff)/          \
                          (1.414214*Star.err_teff))**2)* \
               np.exp(-1*((ips['logg']-Star.logg)/       \
                          (1.414214*Star.err_logg))**2)* \
               np.exp(-1*((ips['feh']-Star.feh)/         \
                          (1.414214*Star.err_feh))**2)
    if SolvePars.key_parameter_known == 'plx':
        prob = np.exp(-1*((ips['t']-Star.teff)/          \
                          (1.414214*Star.err_teff))**2)* \
{'fraction_non_alphanumeric': 0.13006629185197616, 'fraction_numerical': 0.022992363849962238, 'mean_word_length': 2.232438296718199, 'pattern_counts': {'":': 0, '<': 15, '<?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': '1612352', 'n_tokens_mistral': 4323, 'n_tokens_neox': 4097, 'n_words': 724}