<jupyter_start><jupyter_text>Mode fittingHere we will make a simple hierarchical model that encodes some knowledge of quasi-equally spaced modes of oscillation into the prior. Using data from papers:<jupyter_code>import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
import seaborn as sns
import pystan<jupyter_output><empty_output><jupyter_text>Spectral analysis of a signal generated by a sum of 14 cosine waves whose frequencies follow: $f_{n, {\rm true}} = (n + 0.5) \Delta \nu + \mathcal{N}(0, 0.02)$.Then spectral analysis identifies the component frequencies and from this the frequency spacing can be found.<jupyter_code>def gaussian(f, f0, h, w):
    return h * np.exp(-0.5 * (f - f0)**2 / w**2)

fs = 10e3
N = 1e5
dnu = 2.0
numax = 14.0
time = np.arange(N) / fs
f0s = (np.arange(0, 14, 1) + 0.5) * dnu 
#f0s += np.random.randn(len(f0s)) * 0.02
x = 0

for n in f0s:
    #print(gaussian(n, numax, 25.0, 5.0))
    x += gaussian(n, numax, 25.0, 5.0)*np.cos(2*np.pi*n*time)
    
x += 200*np.random.randn(len(x))
x = x/3.5

f, Pxx_den = signal.periodogram(x, fs, scaling='spectrum')
                
fig, axs = plt.subplots(1, 2, figsize=(14, 4))
axs[0].plot(time,x)
axs[0].set_ylabel('Luminosity (Arbitrary units)')
axs[0].set_xlabel('Time (s)')
axs[0].set_xlim([0, 3])
axs[1].plot(f, Pxx_den)
#axs[1].set_ylim([1e-7, 1e2])
axs[1].set_xlim([0, 28])
axs[1].set_xlabel('Frequency [Hz]')
axs[1].set_ylabel('PSD')<jupyter_output><empty_output><jupyter_text>Can find $\Delta\nu$ by calculating the distance between 2 adjacent peaks:<jupyter_code>peaks = signal.find_peaks(Pxx_den, 3)
for i in range(len(peaks[0])-1):
    print(f[peaks[0][i+1]]-f[peaks[0][i]])
bin_width = f[1] - f[0]
w = int(dnu / bin_width)
s = 0
h = int(np.floor(len(Pxx_den[s:]) / w))
print(len(Pxx_den[s:]))
print(h,w)
ladder_p = np.reshape(Pxx_den[s:h*w+s], [h, w])
ladder_f = np.reshape(f[s:h*w+s], [h, w])
from astropy.convolution import Gaussian1DKernel, convolve
fig, ax = plt.subplots()
ax.plot(ladder_f[0,:] / dnu, np.mean(ladder_p, axis=0))
# Create kernel
g = Gaussian1DKernel(stddev=2)
# Convolve data
z = convolve(np.mean(ladder_p, axis=0), g)
ax.plot(ladder_f[0,:] / dnu, z, 'k-', lw=2)
# Create kernel
g = Gaussian1DKernel(stddev=10)
# Convolve data
z = convolve(np.mean(ladder_p, axis=0), g)
ax.plot(ladder_f[0,:] / dnu, z, 'r-', lw=1)
n_max = numax / dnu
nn = 14
d_f = ladder_f[0,:] / dnu
fig, ax = plt.subplots(nn)
for i in range(int(nn)):
        ax[i].plot(d_f, ladder_p[i,:], label=f'Index: {i}')
        ax[i].set_yticks([])        
plt.subplots_adjust(hspace=0.0, wspace=0.0)

#print(len(ladder_f[0,:]))
#print(len(ladder_p[:,0]))
#print(ladder_f) 
#print(ladder_p)
code = '''
functions {
    real gaussian(real f, real numax, real width, real height){
        return height * exp(-0.5 * (f - numax)^2 / width^2);
    }
}
data {
    int N; // Data points per order
    int M; // Number of orders
    real freq[M, N];
    real snr[M, N];
    real dnu_est;
    real numax_est;
}
parameters {
    real<lower = 0> dnu;
    real<lower = 0> numax;
    real<lower = 0> envwidth;
    real<lower = 0> envheight;
    real modewidth;
    real delta[M];
    real<lower = 0> scale;
}
transformed parameters {
    real mode_freqs[M];
    real mode_heights[M];
    for (i in 1:M){
        mode_freqs[i] = (i-0.5) * dnu + delta[i] * scale;
        mode_heights[i] = gaussian(mode_freqs[i], numax, envwidth, envheight);
    }
}
model {
    real beta[N];
    for (i in 1:M){
        for (j in 1:N){
            beta[j] =  (1.0 + lor(freq[i, j], mode_freqs[i], modewidth, mode_heights[i]));
        }
        snr[i, :] ~ normal(beta, 1.0);
    }
    dnu ~ normal(dnu_est, dnu_est * 0.01);
    dnu ~ uniform(dnu_est * 0.9, dnu_est * 1.1);
    numax ~ normal(numax_est, numax_est * 0.01);
    numax ~ uniform(numax_est * 0.9, numax_est * 1.1);
    envwidth ~ normal(5.0, 1.0);
    envheight ~ normal(25.0, 3.0);
    modewidth ~ lognormal(log(0.1), 0.05);
    modewidth ~ uniform(0, 0.5);
    delta ~ normal(0, 1); 
    scale ~ normal(0, 0.01);
}
'''
import pystan
sm = pystan.StanModel(model_code=code)
stan_data = {'N': len(ladder_f[0,:]), 'M': len(ladder_p[:,0]), 
             'freq': ladder_f, 'snr': ladder_p,
            'dnu_est': dnu, 'numax_est': numax}
nchains = 4
start = {'dnu': dnu, 'numax': numax}
fitsm = sm.sampling(data=stan_data, iter=20000, chains=nchains, init=[start for n in range(nchains)])
fitsm.plot()
plt.show()
print(fitsm)<jupyter_output><empty_output><jupyter_text>Guy's CodeLet's set up the data first.  We will define a bunch of Lorentzian modes that are nearly equally spaced in frequency.  The mode heights will be controlled by a Gaussian function. So the mode frequencies will be defined as:$f_{n, {\rm true}} = (n + 0.5) \Delta \nu + \mathcal{N}(0, 0.02)$.and the envelope (height of lorentzian) will just be:$H_n = 25.0 \exp(-0.5 (f_{n, {\rm true}} - \nu_{\rm max})^2 / 5^2)$The data are then just the model of the sum of the Lorentzians, evaluated a frequencies $f_{i}$, plus 1.  To the data we add Gaussian noise (to keep things simple) with mean zero and standard deviation of unity.<jupyter_code>def lor(f, f0, w, h):
    return h / (1.0 + 4.0 * ((f - f0)/w)**2)

def gaussian(f, f0, h, w):
    return h * np.exp(-0.5 * (f - f0)**2 / w**2)

np.random.seed(53)
f = np.linspace(0, 28, 1000)
dnu = 2.0
numax = 14.0
f0s = (np.arange(0, 14, 1) + 0.5) * dnu 
f0s += np.random.randn(len(f0s)) * 0.02  #f_n,true
true = np.ones(len(f))
true += np.sum([lor(f, n, 0.1, gaussian(n, numax, 25.0, 5.0)) for n in f0s], axis=0)
#data = true * np.random.chisquare(2, len(true)) / 2.0
data = true + np.random.randn(len(true))
fig, ax = plt.subplots()
ax.plot(f, true, label='True', zorder=0)
ax.plot(f, data, label='Data', alpha=0.4)
ax.legend()
ax.set_ylabel('SNR')
ax.set_xlabel('Frequnency')<jupyter_output><empty_output><jupyter_text>We now transform the data into a ladder, or echelle, with one mode in each segment.<jupyter_code>bin_width = f[1] - f[0]
print(bin_width)
w = int(dnu / bin_width)
print(len(data[s:]), w)
s = 0
h = int(np.floor(len(data[s:]) / w))
ladder_p = np.reshape(data[s:h*w+s], [h, w])
ladder_f = np.reshape(f[s:h*w+s], [h, w])<jupyter_output>0.1
1000 20<jupyter_text>We can collapse the echelle and combine with some smoothing techniques to show you the data:<jupyter_code>from astropy.convolution import Gaussian1DKernel, convolve
fig, ax = plt.subplots()
ax.plot(ladder_f[0,:] / dnu, np.mean(ladder_p, axis=0))
# Create kernel
g = Gaussian1DKernel(stddev=5)
# Convolve data
z = convolve(np.mean(ladder_p, axis=0), g)
ax.plot(ladder_f[0,:] / dnu, z, 'k-', lw=2)
# Create kernel
g = Gaussian1DKernel(stddev=15)
# Convolve data
z = convolve(np.mean(ladder_p, axis=0), g)
ax.plot(ladder_f[0,:] / dnu, z, 'r-', lw=1)
n_max = numax / dnu
nn = 14
d_f = ladder_f[0,:] / dnu<jupyter_output><empty_output><jupyter_text>Now we can plot the uncollapsed echelle:<jupyter_code>fig, ax = plt.subplots(nn)
for i in range(int(nn)):
        ax[i].plot(d_f, ladder_p[i,:], label=f'Index: {i}')
        ax[i].set_yticks([])        
plt.subplots_adjust(hspace=0.0, wspace=0.0)

#print(len(ladder_f[0,:]))
#print(len(ladder_p[:,0]))
#print(ladder_f) 
#print(ladder_p)<jupyter_output>71
14
[[ 0.          0.02802803  0.05605606  0.08408408  0.11211211  0.14014014
   0.16816817  0.1961962   0.22422422  0.25225225  0.28028028  0.30830831
   0.33633634  0.36436436  0.39239239  0.42042042  0.44844845  0.47647648
   0.5045045   0.53253253  0.56056056  0.58858859  0.61661662  0.64464464
   0.67267267  0.7007007   0.72872873  0.75675676  0.78478478  0.81281281
   0.84084084  0.86886887  0.8968969   0.92492492  0.95295295  0.98098098
   1.00900901  1.03703704  1.06506507  1.09309309  1.12112112  1.14914915
   1.17717718  1.20520521  1.23323323  1.26126126  1.28928929  1.31731732
   1.34534535  1.37337337  1.4014014   1.42942943  1.45745746  1.48548549
   1.51351351  1.54154154  1.56956957  1.5975976   1.62562563  1.65365365
   1.68168168  1.70970971  1.73773774  1.76576577  1.79379379  1.82182182
   1.84984985  1.87787788  1.90590591  1.93393393  1.96196196]
 [ 1.98998999  2.01801802  2.04604605  2.07407407  2.1021021   2.13013013
   2.15815816  2.18618619  2.21421421  2.24[...]<jupyter_text>And now we build the example Pystan model:<jupyter_code>code = '''
functions {
    real lor(real freq, real f0, real w, real h){
        return h / (1 + 4 * ((freq - f0)/w)^2);
    }
    real gaussian(real f, real numax, real width, real height){
        return height * exp(-0.5 * (f - numax)^2 / width^2);
    }
}
data {
    int N; // Data points per order
    int M; // Number of orders
    real freq[M, N];
    real snr[M, N];
    real dnu_est;
    real numax_est;
}
parameters {
    real<lower = 0> dnu;
    real<lower = 0> numax;
    real<lower = 0> envwidth;
    real<lower = 0> envheight;
    real modewidth;
    real delta[M];
    real<lower = 0> scale;
}
transformed parameters {
    real mode_freqs[M];
    real mode_heights[M];
    for (i in 1:M){
        mode_freqs[i] = (i-0.5) * dnu + delta[i] * scale;
        mode_heights[i] = gaussian(mode_freqs[i], numax, envwidth, envheight);
    }
}
model {
    real beta[N];
    for (i in 1:M){
        for (j in 1:N){
            beta[j] =  (1.0 + lor(freq[i, j], mode_freqs[i], modewidth, mode_heights[i]));
        }
        snr[i, :] ~ normal(beta, 1.0);
    }
    dnu ~ normal(dnu_est, dnu_est * 0.01);
    dnu ~ uniform(dnu_est * 0.9, dnu_est * 1.1);
    numax ~ normal(numax_est, numax_est * 0.01);
    numax ~ uniform(numax_est * 0.9, numax_est * 1.1);
    envwidth ~ normal(5.0, 1.0);
    envheight ~ normal(25.0, 3.0);
    modewidth ~ lognormal(log(0.1), 0.05);
    modewidth ~ uniform(0, 0.5);
    delta ~ normal(0, 1); 
    scale ~ normal(0, 0.01);
}
'''
import pystan
sm = pystan.StanModel(model_code=code)<jupyter_output>INFO:pystan:COMPILING THE C++ CODE FOR MODEL anon_model_ae1e6452f5130395d3f50a4387192288 NOW.<jupyter_text>The code takes a while to converge.  We run for 20000 iterations and check the results.<jupyter_code>stan_data = {'N': len(ladder_f[0,:]), 'M': len(ladder_p[:,0]), 
             'freq': ladder_f, 'snr': ladder_p,
            'dnu_est': dnu, 'numax_est': numax}
nchains = 4
start = {'dnu': dnu, 'numax': numax}
fitsm = sm.sampling(data=stan_data, iter=20000, chains=nchains, init=[start for n in range(nchains)])
fitsm.plot()
plt.show()
print(fitsm)<jupyter_output>Inference for Stan model: anon_model_ae1e6452f5130395d3f50a4387192288.
4 chains, each with iter=20000; warmup=10000; thin=1; 
post-warmup draws per chain=10000, total post-warmup draws=40000.

                   mean se_mean     sd   2.5%    25%    50%    75%  97.5%  n_eff   Rhat
dnu                 2.0  9.9e-6 9.1e-4    2.0    2.0    2.0    2.0    2.0   8355    nan
numax             13.97  3.5e-4   0.07  13.84  13.92  13.97  14.02  14.11  40000    1.0
envwidth           5.07  4.1e-4   0.08   4.91   5.01   5.07   5.13   5.23  40000    1.0
envheight         24.97  2.5e-3   0.43  24.13  24.68  24.97  25.26  25.81  29869    1.0
modewidth           0.1  1.0e-5 2.0e-3    0.1    0.1    0.1    0.1    0.1  37423    1.0
delta[0]          -0.36  4.3e-3   0.85   -2.0  -0.92  -0.37   0.19   1.37  40000    1.0
delta[1]           0.43  2.8e-3   0.57  -0.73   0.07   0.45   0.81   1.51  40000    1.0
delta[2]           -2.0  3.1e-3   0.44  -2.93  -2.29  -1.99   -1.7  -1.18  21163    1.0
delta[3]       [...]<jupyter_text>The convergence is good! (The nan Rhat is because the dnu value is so well constrained).  We can check the inferred frequencies with respect to the true frequencies:<jupyter_code>fig, ax = plt.subplots()
ax.scatter(f0s, fitsm['mode_freqs'].mean(axis=0) - f0s)
ax.errorbar(f0s, fitsm['mode_freqs'].mean(axis=0) - f0s, yerr=fitsm['mode_freqs'].std(axis=0))
print(fitsm['mode_freqs'].shape)
print(f0s.shape)<jupyter_output>(40000, 14)
(14,)<jupyter_text>Here is a corner plot of the results:<jupyter_code>import corner
post = np.vstack([fitsm['dnu'], fitsm['numax'], fitsm['envheight'], fitsm['envwidth'], fitsm['modewidth']]).T
corner.corner(post)
plt.show()<jupyter_output><empty_output><jupyter_text>We can now compare the true model with the estimated model:<jupyter_code>best = np.ones(len(f))
best += np.sum([lor(f, n, fitsm['modewidth'].mean(), gaussian(n, fitsm['numax'].mean(),
                            fitsm['envheight'].mean(), fitsm['envwidth'].mean())
                   ) for n in fitsm['mode_freqs'].mean(axis=0)], axis=0)
plt.plot(f, best, 'k-.', label='best', zorder=99)
plt.plot(f, true, label='true', zorder=98)
plt.plot(f, data, label='Data')
plt.legend()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16341579448144625, 'fraction_numerical': 0.11203615604186488, 'mean_word_length': 3.3299004462753174, 'pattern_counts': {'":': 0, '<': 54, '<?xml version=': 0, '>': 54, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '667035', 'license': 'MIT', 'n_tokens_mistral': 6160, 'n_tokens_neox': 5415, 'n_words': 1246, 'path': 'JoshFiles/Stan Practice/Mode_Fitting.ipynb', 'repo_name': 'daw538/y4project', 'size': '12612'}