<jupyter_start><jupyter_text>We're going to try fitting a full asymptotic relation to some simulated dataWe'll do Gaussian noise cos it makes my life easier<jupyter_code>import numpy as np
import matplotlib.pyplot as plt
import theano.tensor as tt

import lightkurve as lk
from astropy.units import cds
from astropy import units as u
import seaborn as sns

import corner
import pystan
import pandas as pd
import pickle
import glob
from astropy.io import ascii
import os

import pymc3 as pm
import arviz<jupyter_output><empty_output><jupyter_text>Build the model<jupyter_code>class model():
    def __init__(self, f, n0_, n1_, n2_):
        self.f = f
        self.n0 = n0_
        self.n1 = n1_
        self.n2 = n2_
        self.npts = len(f)
        self.M = [len(n0_), len(n1_), len(n2_)]
        
    def epsilon(self, i, theano=True):
        eps = tt.zeros((3,3))
        eps0 = tt.set_subtensor(eps[0][0], 1.)
        eps1 = tt.set_subtensor(eps[1][0], tt.cos(i)**2)
        eps1 = tt.set_subtensor(eps1[1], 0.5 * tt.sin(i)**2)
        eps2 = tt.set_subtensor(eps[2][0], 0.25 * (3. * tt.cos(i)**2 - 1.)**2)
        eps2 = tt.set_subtensor(eps2[1], (3./8.)*tt.sin(2*i)**2)
        eps2 = tt.set_subtensor(eps2[2], (3./8.) * tt.sin(i)**4)

        eps = tt.set_subtensor(eps[0], eps0)
        eps = tt.set_subtensor(eps[1], eps1)
        eps = tt.set_subtensor(eps[2], eps2)
        
        if not theano:
            return eps.eval()
        return eps
            

    def lor(self, freq, h, w):
        return h / (1.0 + 4.0/w**2*(self.f - freq)**2)
    
    def mode(self, l, freqs, hs, ws, eps, split=0):
        for idx in range(self.M[l]):
            for m in range(-l, l+1, 1):
                self.modes += self.lor(freqs[idx] + (m*split), 
                                     hs[idx] * eps[l,abs(m)], 
                                     ws[idx])
  
    def model(self, p, theano=True):
        f0, f1, f2, g0, g1, g2, h0, h1, h2, split, i, b = p
               
        # Calculate the modes
        eps = self.epsilon(i, theano)        
        self.modes = np.zeros(self.npts)
        self.mode(0, f0, h0, g0, eps)
        self.mode(1, f1, h1, g1, eps, split)
        self.mode(2, f2, h2, g2, eps, split)
               
        #Create the model
        self.mod = self.modes + b
        return self.mod

    def asymptotic(self, n, numax, deltanu, alpha, epsilon):       
        nmax = (numax / deltanu) - epsilon
        over = (n + epsilon + ((alpha/2)*(nmax - n)**2))
        return over * deltanu
    
    def f0(self, p):
        numax, deltanu, alpha, epsilon, d01, d02 = p
        
        return self.asymptotic(self.n0, numax, deltanu, alpha, epsilon)
    
    def f1(self, p):
        numax, deltanu, alpha, epsilon, d01, d02 = p
        
        f0 = self.asymptotic(self.n1, numax, deltanu, alpha, epsilon)
        return f0 + d01
    
    def f2(self, p):
        numax, deltanu, alpha, epsilon, d01, d02 = p
        
        f0 = self.asymptotic(self.n2+1, numax, deltanu, alpha, epsilon)
        return f0 - d02
nmodes = 2   # Two overtones
nbase = 18   # Starting at n = 18
n0_ = np.arange(nmodes)+nbase
n1_ = np.copy(n0_)
n2_ = np.copy(n0_) - 1.
fs = .05     # Data has a frequency spacing of 0.05 microhertz
nyq = (0.5 * (1./58.6) * u.hertz).to(u.microhertz).value  # Setting a sensible nyquist value
ff = np.arange(fs, nyq, fs)  # Generating the full frequency range<jupyter_output><empty_output><jupyter_text>Set the asymptotic parameters<jupyter_code>deltanu_  =  60.
numax_= 1150.
alpha_ = 0.
epsilon_ = 0.
d01_ = deltanu_/2.
d02_ = 6.<jupyter_output><empty_output><jupyter_text>Generate the model for the full range<jupyter_code>mod = model(ff, n0_, n1_, n2_)<jupyter_output><empty_output><jupyter_text>Calculate the predicted mode frequencies<jupyter_code>init_f = [numax_, deltanu_, alpha_, epsilon_, d01_, d02_]
f0_ = mod.f0(init_f)
f1_ = mod.f1(init_f)
f2_ = mod.f2(init_f)<jupyter_output><empty_output><jupyter_text>Slice up the data to just be around the mode frequencies<jupyter_code>lo = f2_.min() - .25*deltanu_
hi = f1_.max() + .25*deltanu_

sel = (ff > lo) & (ff < hi)
f = ff[sel]<jupyter_output><empty_output><jupyter_text>And now lets reset the model for the new frequency range...<jupyter_code>mod = model(f, n0_, n1_, n2_)<jupyter_output><empty_output><jupyter_text>Set up initial guesses for the model parameters<jupyter_code>def gaussian(locs, l, numax, Hmax0):
    fwhm = 0.25 * numax
    std = fwhm/2.355

    Vl = [1.0, 1.22, 0.71, 0.14]

    return Hmax0 * Vl[l] * np.exp(-0.5 * (locs - numax)**2 / std**2)
init_m =[f0_,                         # l0 modes
           f1_,                         # l1 modes
           f2_,                         # l2 modes
           np.ones(len(f0_)) * 2.0,     # l0 widths
           np.ones(len(f1_)) * 2.0,     # l1 widths
           np.ones(len(f2_)) * 2.0,     # l2 widths
           np.sqrt(gaussian(f0_, 0, numax_, 1000.) * 2.0 * np.pi / 2.0) ,# l0 heights
           np.sqrt(gaussian(f1_, 1, numax_, 1000.) * 2.0 * np.pi / 2.0) ,# l1 heights
           np.sqrt(gaussian(f2_, 2, numax_, 1000.) * 2.0 * np.pi / 2.0) ,# l2 heights
           1.,       # splitting
           np.pi/4.,                    # inclination angle
           1.                           # background parameters
        ]

# Add on the chisquare 2 dof noise
p = mod.model(init_m, theano=False)*np.random.chisquare(2., size=len(f))/2<jupyter_output><empty_output><jupyter_text>Plot what our data looks like<jupyter_code>plt.plot(f, p)
plt.plot(f, mod.model(init_m, theano=False), lw=3)
plt.show()

plt.show()<jupyter_output><empty_output><jupyter_text>Fitting the model:<jupyter_code>pm_model = pm.Model()

with pm_model:   
    epsilon = pm.Normal('epsilon', 0, 1)
    dnu = pm.Normal('dnu', 60, 0.01)
    d02 = pm.Normal('d02', 0.1, 0.01)
    d01 = pm.Normal('d01', 0.5, 0.01)
    
    f0 = pm.Normal('f0', (n0_ + epsilon) * dnu, 1.0, shape=2)
    f1 = pm.Normal('f1', (n0_ + epsilon + d01) * dnu, 1.0, shape=2)
    f2 = pm.Normal('f2', (n0_ + epsilon - d02) * dnu, 1.0, shape=2)
    
    split = pm.Lognormal('split', np.log(0.4), 0.1)
    i = 70.0
    g = pm.Lognormal('g', np.log(1.0), 0.1, shape=2)
    h = pm.Lognormal('h', np.log(50.0), 0.1, shape=2)
    b = 1.0
    
    fit = mod.model([f0, f1, f2, g, g, g, h, h, h, split, i, b])
    
    like = pm.Gamma('like', alpha=1., beta=1./fit, observed=p)
    '''
    
    # I've left these below for you to use if you want because I trust these parameterizations
    xsplit = pm.HalfNormal('xsplit', sigma=2.0, testval=init_m[9] * np.sin(init_m[10]))
    cosi = pm.Uniform('cosi', 0., 1., testval=np.cos(init_m[10]))
    
    i = pm.Deterministic('i', tt.arccos(cosi))
    split = pm.Deterministic('split', xsplit/tt.sin(i))
    
    b = pm.Bound(pm.Normal, lower=0.)('b', mu=1., sigma=.1, testval=1.)
    
    fit = mod.model([f0, f1, f2, g0, g1, g2, h0, h1, h2, split, i, b])
    
    like = pm.Gamma('like', alpha=1., beta=1./fit, observed=p)
    '''
with pm_model:
    inference = pm.ADVI()
    approx = pm.fit(n=90000, method=inference)
start = dict(pm.summary(approx.sample(draws=1000)).mean(axis=1))
with pm_model:
    trace = pm.sample(1000, start=start)
pm.summary(trace)
pm.summary(approx.sample(draws=1000))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15146110577582658, 'fraction_numerical': 0.050898614350391, 'mean_word_length': 2.4225352112676055, 'pattern_counts': {'":': 0, '<': 42, '<?xml version=': 0, '>': 42, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '172796', 'license': 'MIT', 'n_tokens_mistral': 3104, 'n_tokens_neox': 2906, 'n_words': 681, 'path': 'code/tests/completed_tests/pb_issues_guy.ipynb', 'repo_name': 'ojhall94/malatium', 'size': '7289'}