# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Mode fitting
#
# Here we will make a simple hierarchical model that encodes some knowledge of quasi-equally spaced modes of oscillation into the prior.

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import pystan

# +
dnu_sol = 136  # book value of large freq separation for the sun

df = pd.read_table('broomhall2009.txt', delim_whitespace=True, names=('n', 'l', 'nu', 'nu_unc'))  #import data

lmodes = []
dnu_avg = []
f_mod = []

for i in range(4):
    lmodes.append(df[df.l == i])
    dnu_avg.append(np.median(np.diff(lmodes[i].nu)))
    f_mod.append(lmodes[i].nu % dnu_avg[i])
    label = "l="+str(i)
    plt.scatter(f_mod[i], lmodes[i].nu, label = label)
    
plt.xlabel(r'Frequency modulo ($\mu Hz$)')
plt.ylabel(r'Frequency ($\mu Hz$)')
plt.legend()
# -

# The lower departure from a straight line is due to the BCZ, so the lower frequencies are not useful in analysing the response from the HeII ionization zone. Frequencies below 1700 $\mu$Hz are omitted.

# +
for i in range(4):
    lmodes[i] = lmodes[i].loc[lmodes[i].nu > 2200]
    lmodes[i] = lmodes[i].set_index(np.arange(0,len(lmodes[i]),1))

dnu_avg = []
f_mod = []
for i in range(4):
    dnu_avg.append(np.median(np.diff(lmodes[i].nu)))
    f_mod.append(lmodes[i].nu % dnu_avg[i])
    label = "l="+str(i)
    plt.scatter(f_mod[i], lmodes[i].nu, label = label)
    
plt.xlabel(r'Frequency modulo ($\mu Hz$)')
plt.ylabel(r'Frequency ($\mu Hz$)')
plt.legend()


# -

# To find a ballpark figure before defining priors, I will apply the model from Vrard:
#
# $$\nu_{UP}(n,0)=\bigg(n+\epsilon+\frac{\alpha}{2}(n-n_{max})^2+\frac{\mathcal{A}\mathcal{G}}{2\pi}sin\bigg(\frac{2\pi(n-n_{max})}{\mathcal{G}}+\phi\bigg)\bigg)\langle\Delta\nu\rangle$$

def model(n, dnu, nmax, epsilon, alpha, A, G, phi):
    freqs = (n + epsilon + (alpha/2)*(n-nmax)**2 + (A*G/2*np.pi)*np.sin(2*np.pi*(n-nmax)/G + phi))*dnu #Vrard universal pattern
    return freqs


n = lmodes[0].n
dnu = dnu_avg[0]
numax = 2900
epsilon = 1.48 #1.437
nmax = 21.9
alpha = 0.003 #from vrard paper
A = 0.0009
print(dnu_avg[0], A)
G = 2 #3
phi = 0
f = model(n, dnu, nmax, epsilon, alpha, A, G, phi)

plt.scatter(f_mod[0], lmodes[0].nu, label = 'Data')
plt.plot(f % dnu_avg[0], f, label = 'Model', color = 'r')
plt.ylabel(r'Frequency($\mu Hz$)')
plt.xlabel(r'Frequency modulo ($\mu Hz$)')
plt.legend()
#plt.xlim(0, 135.2)

code = '''
functions {
    real model(real n, real dnu, real nmax, real epsilon, real alpha, real A, real G, real phi){
        real freqs = (n + epsilon + (alpha/2)*(n-nmax)^2 + (A*G/2*pi())*sin(2*pi()*(n-nmax)/G + phi))*dnu;
        return freqs;
    }
}
data {
    int N; // Data points
    real n[N];
    real fobs[N];
}
parameters {
    real<lower=19> nmax;
    real dnu;
    real epsilon;
    real<lower=0> alpha;
    real A;
    real<lower=1.5, upper=8> G;
    real<lower=-pi(), upper=pi()> phi;
}
model {
    vector[N] mod;
    nmax ~ normal(22, 2);
    dnu ~ normal(135, 0.001*135);
    epsilon ~ normal(1.480, 0.05);
    alpha ~ normal(0.0025, 0.0010);
    A ~ normal(0.001, 0.001);
    G ~ normal(2, 0.2);
    phi ~ normal(0, pi()/4);
    for (i in 1:N) mod[i] = model(n[i], dnu, nmax, epsilon, alpha, A, G, phi);
    mod ~ normal(fobs, 10);
}
'''
import pystan
sm = pystan.StanModel(model_code=code)

stan_data = {'N': len(lmodes[0]), 'n': lmodes[0].n,
             'fobs': lmodes[0].nu
            }
nchains = 4
start = {'A': 0, 'G': 1.6, 'phi': -0.9*np.pi}
fitsm = sm.sampling(data=stan_data, iter=20000, chains=nchains, init=[start for n in range(nchains)])

fitsm.plot()
plt.show()

print(fitsm)

# +
stanfit = model(lmodes[0].n, fitsm['dnu'].mean(), fitsm['nmax'].mean(), fitsm['epsilon'].mean(), fitsm['alpha'].mean(), fitsm['A'].mean(), fitsm['G'].mean(), fitsm['phi'].mean())

plt.subplots()
plt.scatter(f_mod[0], lmodes[0].nu, label = 'Data')
plt.plot(stanfit % dnu_avg[0], stanfit, label = 'Stan model', color = 'r')
plt.legend()
# -

import corner
data = np.vstack([fitsm['nmax'], fitsm['dnu'], fitsm['epsilon'], fitsm['alpha'], fitsm['A'], fitsm['G'], fitsm['phi'],]).T
corner.corner(data, labels=[r'$n_{max}$', r'$\Delta\nu$', r'$\epsilon$', r'$\alpha$', r'$A$', r'$G$', r'$\phi$'])
plt.show()

# This analysis will use the methods detailed in Vrard 2015 and, as such, will use its definitions.
#
# The definition for the local frequency spacing is:
#
# $\Delta\nu(n) = \frac{\nu_{n+1,0}-\nu_{n-1,0}}{2}$.
#
# At the edges of the measured radial modes, we cannot use this equation and replace it by the frequency difference between two consecutive radial modes.

lmodes[0]['dnu_n'] = (lmodes[0]['nu'].diff(2).shift(-1)/2)
lmodes[0]['dnu_n'][0] = lmodes[0]['nu'].diff(1)[1]
lmodes[0]['dnu_n'][16] = (lmodes[0]['nu'].diff(1)[16])

# The stan model determined that the maximum frequency is $\nu_{max} = 2777.27 \mu Hz$. This is then used in the determination of the universal pattern to which the frequency separation will be compared:
#
# $\Delta\nu_{UP}(n)=(1+\alpha(n-n_{max}))\langle\Delta\nu\rangle$
#
# We compute the difference between the observed local large separation and the theoretical local large separation predicted by the universal pattern:
#
# $\delta_{g,obs}=\Delta\nu(n)-\Delta\nu_{UP}(n)$

# +
nmax = 22
numax = 3168.618
alpha = 0.015*dnu_avg[0]**(-0.32)
dnu_UP = (1+alpha*(lmodes[0].n-nmax))*dnu_avg[0]

plt.scatter(lmodes[0].nu, lmodes[0].dnu_n, label = r'$\Delta\nu(n)$')
plt.xlabel(r'Frequency($\mu Hz$)')
plt.ylabel(r'$(\nu_{n+1,0}-\nu_{n-1,0})/2$ ($\mu Hz$)')
plt.scatter(lmodes[0].nu, dnu_UP, label = r'$\Delta\nu_{UP}(n)$')
plt.legend()
# -

# We now subtract the universal pattern from the data in order to remove the curvature term.

# +
deltag = lmodes[0].dnu_n - dnu_UP

plt.scatter(lmodes[0].nu, deltag, label='Data')
plt.xlabel(r'Frequency($\mu Hz$)')
plt.ylabel(r'$\delta_{g,obs}$ ($\mu Hz$)')
plt.legend()
# -

# The amplitude can be found by fitting an oscillatory component to the resultant frequency variations obtained after removal of the curvature term from the measurements:
#
# $\delta_{g,obs}=\mathcal{A}\langle\Delta\nu\rangle cos\big(\frac{2\pi(\nu-\nu_{max})}{\mathcal{G}\langle\Delta\nu\rangle}+\phi\big)$
#
# where $\mathcal{G}$ is the period of the oscillation expressed in units of $\langle\Delta\nu\rangle$, $\mathcal{A}$ is the amplitude of the oscillation in units of $\langle\Delta\nu\rangle$ and $\phi$ is the phase of the oscillation centered on $\nu_{max}$.

code = '''
functions {
    real dgobs(real A, real numax, real nu, real G, real phi){
        return A * cos((2*pi()*(nu-numax))/G + phi);
    }
}
data {
    int N; // Data points
    real nu[N];
    real numax;
    real dnu_avg;
    real dg[N];
}
parameters {
    real<lower = 0> G;
    real<lower = 0> A;
    real<lower = 0> phi;
}
model {
    vector[N] mod;
    A ~ normal(4, 1.5);
    G ~ normal(7, 2);
    phi ~ normal(1.5*pi(), 1);
    for (i in 1:N) mod[i] = dgobs(A, numax, nu[i], G, phi);
    mod ~ normal(dg, 1);
}
'''
import pystan
sm = pystan.StanModel(model_code=code)

# The code takes a while to converge.  We run for 20000 iterations and check the results.

stan_data = {'N': len(lmodes[0]), 
             'nu': lmodes[0].nu, 'dg': deltag,
             'numax': numax, 'dnu_avg': dnu_avg[0]}
nchains = 4
start = {'A': 3, 'G': 4, 'phi': 1}
fitsm = sm.sampling(data=stan_data, iter=20000, chains=nchains, init=[start for n in range(nchains)])

fitsm.plot()
plt.show()

print(fitsm)

# +
dgstan = fitsm['A'].mean() * np.cos((2*np.pi*(lmodes[0].nu-numax))/fitsm['G'].mean() + fitsm['phi'].mean())

plt.subplots()
plt.plot(lmodes[0].nu, dgstan, label = 'Stan model')
plt.scatter(lmodes[0].nu, deltag, label = 'Data')
# -

# Here is a corner plot of the results:

import corner
data = np.vstack([fitsm['A'], fitsm['G'], fitsm['phi']]).T
corner.corner(data, labels=[r'$\mathcal{A}$', r'$\mathcal{G}$', r'$\phi$'])
plt.show()

# From Nature volume 215, pages 43–44 (01 July 1967) the helium mass fraction of the sun is between 0.20 and 0.27.



StarcoderdataJupyterScriptsDedupFiltered