<jupyter_start><jupyter_text>Fitting the Barnes Gyro relation using a latentfit<jupyter_code>import numpy as np
import matplotlib.pyplot as plt
import pystan
import warnings
warnings.filterwarnings('ignore')<jupyter_output><empty_output><jupyter_text>We'll add the uncertainty with a latent implementation.$P = \tau^a \times b(B - V - c)^c$where $P$ is rotation in days$\tau$ is age (in Myr),B-V are colour in magnitude,the rest are free parameters<jupyter_code>angus = [0.55, 0.40, 0.45, 0.31]

def gyro(tau, bv, a, b, c, d):
    return (tau*1000)**a * b*(bv - c)**d<jupyter_output><empty_output><jupyter_text>Instead of treating uncertainty in Period as emergent from uncertainty in colour and age, we are going to treat these as if each was obtained in an independent observation, therefore having uncorrelated levels on uncertainty.<jupyter_code>npts = 55

bprp_true = np.abs(np.sort(np.random.normal(1., 0.2, npts)))
age_true = np.abs(np.random.normal(5., 3, npts))
period_true = gyro(age_true, bprp_true, *angus)

sigma_p = 6.*np.ones(npts)
period_obs = np.abs(period_true + np.random.randn(npts) * sigma_p)

sigma_bprp = .01 * np.ones(npts)
bprp_obs = bprp_true + np.random.randn(npts) * sigma_bprp

sigma_age = np.random.uniform(.05, .30) * age_true
age_obs = age_true + np.random.randn(npts) * sigma_age
import matplotlib.gridspec as gridspec
import mystyle as ms

with plt.style.context(ms.ms):
    fig = plt.figure(figsize=(14,6))
    gs = gridspec.GridSpec(2, 2)
    ax0 = plt.subplot(gs[0,:1])
    ax0.scatter(bprp_obs, period_obs,zorder=1, s=20, label='observed')
    ax0.errorbar(bprp_obs, period_obs, yerr=sigma_p, xerr=sigma_bprp, fmt='|',zorder=0)
    ax0.set_xlabel('B-V (mag)')
    ax0.set_ylabel('period (day)')
    ax0.legend()
    
    ax1 = plt.subplot(gs[1,1])
    ax1.scatter(age_obs, period_obs,zorder=1, s=20, label='observed')
    ax1.errorbar(age_obs, period_obs, yerr=sigma_p, xerr=sigma_age, fmt='|',zorder=0)
    ax1.set_xlabel('Age (Gyr)')
    ax1.set_ylabel('period (day)')
    ax1.legend()    
    
    ax2 = plt.subplot(gs[1,0])
    ax2.scatter(age_obs, bprp_obs,zorder=1, s=20, label='observed')
    ax2.errorbar(age_obs, bprp_obs, yerr=sigma_bprp, xerr=sigma_age, fmt='|',zorder=0)
    ax2.set_xlabel('Age (Gyr)')
    ax2.set_ylabel('B-V (mag)')
    ax2.legend()    
    
    agex = np.linspace(age_true.min(), age_true.max(), 100)
    ax1.plot(agex, gyro(agex, bprp_true.min(), *angus), label='BPRP {:.2f}'.format(bprp_true.min()), 
             lw=3, alpha=.5)
    ax1.plot(agex, gyro(agex, np.mean(bprp_true), *angus), label='BPRP {:.2f}'.format(np.median(bprp_true)), 
             lw=3, alpha=.5)
    ax1.plot(agex, gyro(agex, bprp_true.max(), *angus), label='BPRP {:.2f}'.format(bprp_true.max()), 
             lw=3, alpha=.5)    
    ax1.legend()
    
    bprpx = np.linspace(bprp_true.min(), bprp_true.max(), 100)
    ax0.plot(bprpx, gyro(age_true.min(), bprpx, *angus), label='Age {:.2f}'.format(age_true.min()), 
             lw=3, alpha=.5)
    ax0.plot(bprpx, gyro(np.mean(age_true), bprpx, *angus), label='Age {:.2f}'.format(np.median(age_true)), 
             lw=3, alpha=.5)
    ax0.plot(bprpx, gyro(age_true.max(), bprpx, *angus), label='Age {:.2f}'.format(age_true.max()), 
             lw=3, alpha=.5)    
    ax0.legend()    
    
    fig.tight_layout()
    plt.show()<jupyter_output><empty_output><jupyter_text>Now lets fit for Period using a latent parameter model<jupyter_code>latentfit = '''
data {
    int<lower=1> N;
    real<lower=0> P[N];
    real<lower=0> sigma_p[N];
    
    real bprp[N];
    real<lower=0> sigma_bprp[N];
    
    real<lower=0> age[N];
    real<lower=0> sigma_age[N];
}
parameters {
    real<lower=0> a;
    real<lower=0> b;
    real<lower=0> c;
    real<lower=0> d;
    
    real raw_bprp[N];
    real raw_age[N];
}
transformed parameters {
    real<lower=0> true_age[N];
    real true_bprp[N];
    
    for (n in 1:N){
        true_age[n] = sigma_age[n] * raw_age[n] + age[n];
        true_bprp[n] = sigma_bprp[n] * raw_bprp[n] + bprp[n];
    }
}
model {
    real mod[N];
    
    for (n in 1:N){
        mod[n] = (true_age[n]*1000)^a * b * (true_bprp[n] - c)^d;
    }
    
    a ~ normal(0., 1.);
    b ~ normal(0., 1.);
    c ~ normal(0., 1.);
    d ~ normal(0., 1.);
    
    raw_bprp ~ std_normal(); // implies true_bprp ~ normal(bprp, sigma_bprp)
    raw_age ~ std_normal(); // implies true_age ~ normal(age, sigma_age)
    
    mod ~ normal(P, sigma_p);
}
'''
sm = pystan.StanModel(model_code = latentfit, model_name='latentfit')
data = {'N' : len(period_obs),
       'P' : period_obs,
       'sigma_p' : sigma_p,
       'bprp' : bprp_obs,
       'sigma_bprp' : sigma_bprp,
       'age' : age_obs,
       'sigma_age' : sigma_age}

init = {'a' : angus[0],
        'b' : angus[1],
        'c' : angus[2],
        'd' : angus[3],
        'true_age' : age_true,
        'true_bprp' : bprp_true}
fit = sm.sampling(data=data, iter=20000, chains=4, seed=11,
                  init = [init for n in range(4)])
print(fit)
fit.plot()
plt.show()
import corner
labels=['a','b','c','d']
chain = np.array([fit[label] for label in labels])
corner.corner(chain.T, labels=labels, truths=angus,
              quantiles=[0.16, 0.5, 0.84],show_titles=True)
plt.show()
import matplotlib.gridspec as gridspec
import mystyle as ms

gyrores = [np.median(fit[label]) for label in ['a','b','c','d']]

with plt.style.context(ms.ms):
    fig = plt.figure(figsize=(14,6))
    gs = gridspec.GridSpec(2, 2)
    ax0 = plt.subplot(gs[0,:1])
    ax0.scatter(bprp_obs, period_obs,zorder=1, s=20, label='observed')
    ax0.errorbar(bprp_obs, period_obs, yerr=sigma_p, xerr=sigma_bprp, fmt='|',zorder=0)
    ax0.set_xlabel('B-V (mag)')
    ax0.set_ylabel('period (day)')
    ax0.legend()
    
    ax1 = plt.subplot(gs[1,1])
    ax1.scatter(age_obs, period_obs,zorder=1, s=20, label='observed')
    ax1.errorbar(age_obs, period_obs, yerr=sigma_p, xerr=sigma_age, fmt='|',zorder=0)
    ax1.set_xlabel('Age (Gyr)')
    ax1.set_ylabel('period (day)')
    ax1.legend()    
    
    ax2 = plt.subplot(gs[1,0])
    ax2.scatter(age_obs, bprp_obs,zorder=1, s=20, label='observed')
    ax2.errorbar(age_obs, bprp_obs, yerr=sigma_bprp, xerr=sigma_age, fmt='|',zorder=0)
    ax2.set_xlabel('Age (Gyr)')
    ax2.set_ylabel('B-V (mag)')
    ax2.legend()    
    
    agex = np.linspace(age_true.min(), age_true.max(), 100)
    ax1.plot(agex, gyro(agex, bprp_true.min(), *gyrores), label='BPRP {:.2f}'.format(bprp_true.min()), 
             lw=3, alpha=.5)
    ax1.plot(agex, gyro(agex, np.mean(bprp_true), *gyrores), label='BPRP {:.2f}'.format(np.median(bprp_true)), 
             lw=3, alpha=.5)
    ax1.plot(agex, gyro(agex, bprp_true.max(), *gyrores), label='BPRP {:.2f}'.format(bprp_true.max()), 
             lw=3, alpha=.5)    
    ax1.legend()
    
    bprpx = np.linspace(bprp_true.min(), bprp_true.max(), 100)
    ax0.plot(bprpx, gyro(age_true.min(), bprpx, *gyrores), label='Age {:.2f}'.format(age_true.min()), 
             lw=3, alpha=.5)
    ax0.plot(bprpx, gyro(np.mean(age_true), bprpx, *gyrores), label='Age {:.2f}'.format(np.median(age_true)), 
             lw=3, alpha=.5)
    ax0.plot(bprpx, gyro(age_true.max(), bprpx, *gyrores), label='Age {:.2f}'.format(age_true.max()), 
             lw=3, alpha=.5)    
    ax0.legend()    
    
    fig.tight_layout()
    plt.show()
gyropred = gyro(age_obs, bprp_obs, *gyrores)
with plt.style.context(ms.ms):
    fig = plt.figure()
    ax= fig.gca()
    ax.scatter(age_obs, gyropred - period_true)
    ax.axhline(0., lw=3, c='r', alpha=.5)
    ax.set_ylabel('P Predicted / P True')
    ax.set_xlabel('Age')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.18004452009951552, 'fraction_numerical': 0.029461830561738903, 'mean_word_length': 3.4458672875436553, 'pattern_counts': {'":': 0, '<': 28, '<?xml version=': 0, '>': 28, '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': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '495878', 'license': 'MIT', 'n_tokens_mistral': 3420, 'n_tokens_neox': 3261, 'n_words': 562, 'path': 'code/tests/gyro_tests/gyro_latent_test.ipynb', 'repo_name': 'ojhall94/malatium', 'size': '7637'}