Red giant example

Bayesian inference of stellar parameters using an emulator trained on MESA/GYRE models with global asteroseismic data of red giants from the NASA Kepler mission

#### inference of a population of observed red giant's stellar parameters 
#### using an emulator neural network trained on a grid of MESA/GYRE models 

import seistron

Observations

## load the observed data 
observations = seistron.data.red_giants.yu_2018() # Ji Yu 2018 data 

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[2], line 2
      1 ## load the observed data 
----> 2 observations = seistron.data.red_giants.yu_2018() # Ji Yu 2018 data 

AttributeError: module 'seistron.data.red_giants' has no attribute 'yu_2018'

seistron.visualize.hr_diagram(observations, cbar='[Fe/H]') # plot an HRD 

seistron.visualize.kiel_diagram(observations) 

seistron.visualize.jcd_diagram(observations, cbar='Teff') 

Theoretical models

## load the theoretical models 
model_grid = seistron.models.red_giants.li_rg()   # Yaguang Li's red giants 

seistron.visualize.hr_diagram(models=model_grid, cbar='M') 

seistron.visualize.hr_diagram(observations, models=model_grid) 

Emulator

## train or load the emulator 
pretrained_model_filename = 'transformer_rg.tron' # save or load here 
# emulator = seistron.emulate.transformer.train( 
#    model_grid, 
#    inputs=model_grid.inputs, 
#    outputs=model_grid.outputs) # need a GPU or computing cluster 
# emulator.save_pretrained(emulator, save_filename=pretrained_model_filename)
emulator = seistron.emulator.load_pretrained(pretrained_model_filename) 

seistron.visualize.show_crossval_score(emulator) # visualize emulator accuracy 

seistron.visualize.compare_emulators([
    emulator, 
    seistron.emulator.linear(models=model_grid), 
    seistron.emulator.random_forest(models=model_grid, cv=True)
]) 

# for each output column, plot 1-to-1 line between the predicted and actual, and the residuals as a function of the actual 
seistron.visualize.show_residuals(emulator, 
    only=['Delta_nu', 'nu_max', 'Teff']) 

emulator.summarize() # print std of residuals, which is representative of systematic interpolation errors that should be propagated 

Inference

## apply the emulator to do inference on real data 
## first, do it for just one star to obtain posteriors for its mass, age, etc 
priors = {'M':   seistron.sample.priors.kroupa_IMF, 
          'age': seistron.sample.priors.flat([0, 13.8])}

example_star = observations.iloc[0]
MAP_one = seistron.sample.maximum_a_posteriori(example_star, priors=priors, 
    method='nelder-mead') # need a sensible starting point 

samples = seistron.sample.do_one( # better to run this on a cluster like below 
    data=example_star, 
    emulator=emulator, 
    method='HMC', 
    num_samples=1e7, 
    priors=priors, 
    starting_values=MAP_one) 

seistron.visualize.corner(samples) # show a nice corner plot 

seistron.sample.summarize(samples) # print means, standard deviations, and χ2

Mock data test

## do a test using mock data based on one of the models 
# user-specifiable noise levels for observables
noise = {'Delta_nu': 0.1,  # in muHz
         'nu_max':   2.0,  # in muHz
         'Teff':     50.0, # in K
         'logg':     0.05}

# apply noise to the first model in the data set and try to recover its params 
mock_test = emulator.generate_mock_data(model_grid.iloc[0], noise)
MAP_test = seistron.sample.maximum_a_posteriori(mock_test, 
    priors=priors, method='nelder-mead') 

samples_test = seistron.sample.do_one(mock_first, emulator=emulator,
    num_samples=1e4, priors=priors, starting_values=MAP_test)

seistron.visualize.summarize(samples_test)

# see how well the inferences scale with increased noise on the observables 
noise_factors = {'Delta_nu': [0.1, 0.3, 1, 3, 10]} 
noise_test = seistron.sample.noise_test(mock_first, 
    emulator=emulator, num_samples=1e4, priors=priors, 
    starting_values=MAP_test, noise=noise, noise_factors=noise_factors) 

seistron.visualize.noise_test(noise_test) 

# perhaps we want some wrapper that automates the MAP/sampling/visualizing/summarizing/etc so that we can do it all in one go 

Hierarchical Bayes

## hierarchical Bayes for one correlated (shared) parameter while characterizing all the observations of all the stars simultaneously 
MAP_all = seistron.sample.maximum_a_posteriori(observations, priors=priors, 
    method='nelder-mead')

samples = seistron.sample.hierarchical( # with access to a cluster 
    data=observations, 
    emulator=emulator, 
    method='HMC', 
    num_samples=1e10, 
    priors=priors, 
    correlate={'Y': ['Z', 'age'], # galactic chemical enrichment 
           'alpha_MLT': ['Teff', 'logg', 'Z'], # convection 
           'alpha_ov': ['M', 'age', 'Z']}, 
    correlation_type={'Y': 'linear', 
                      'alpha_MLT': 'linear', 
                      'alpha_ov': 'gaussian_process'}, 
    starting_values=MAP_all, 
    nproc=128, # number of processes for parallelism 
    cluster='grace') 

seistron.visualize.corner(samples, 
    only=['Teff', 'logg', 'Z', 'alpha_MLT']) 

## visualize draws from the posterior for one star in the data set 
seistron.visualize.posterior_samples(example_star, samples) 

seistron.visualize.hr_diagram(example_star, 
    models=model_grid, samples=samples) 

Hierarchical Bayes mock data test

## mock data test: make mock data out of each of the models in the data set and test our ability to recover correlated parameters 
# still need to brainstorm a good way to implement this 
cov_params = lambda inputs: inputs['Y'] = 1.4 * inputs['Z'] + 0.2473 

mock_correlated = generate_correlated_mock_data(n_stars=50, 
    cov_params=cov_params, noise_dict=noise)

#samples = seistron.sample.hierarchical...