<jupyter_start><jupyter_text>Make fake data:<jupyter_code># np.random.seed(123)
np.random.seed(456)

N = 5  # number of data points
# t = np.sort(np.random.uniform(0, 0.1, size=N))
t = np.linspace(0, 1, N)
# sig_y = np.random.uniform(0.1, 0.5, size=N)  # noise! MAGIC NUMBER
sig_y = np.full(len(t), 0.1)
C = np.diag(sig_y**2)
# y = 10 + np.random.normal(0., sig_y, size=N)
y = 0 + np.zeros(N)
data = RVData(t=t, rv=y*u.km/u.s, stddev=sig_y*u.km/u.s, 
              t0=Time(0., format='mjd', scale='tcb'))

# P_samples = np.exp(np.random.uniform(np.log(P_lim[0]), 
#                                      np.log(P_lim[1]),
#                                      size=nsamples))

# M0_samples = np.random.uniform(0, 2*np.pi, size=nsamples)

P_samples = prior_samples['P'].value #* 3.14
M0_samples = prior_samples['M0'].to_value(u.radian)

# hyperprior on linear parameters:
mu_th = np.array([0., 0])
_ = data.plot()
M = np.stack((np.ones((N, nsamples)),
              np.cos(2*np.pi*t[:, None] / P_samples[None] - M0_samples[None]))).T
def get_aA_bB(M, y, C, mu, V):
    Cinv = np.linalg.inv(C)
    Vinv = np.linalg.inv(V)
    
    B = C + M @ V @ M.T
    b = M.dot(mu)
    
    Ainv = Vinv + M.T @ Cinv @ M
    A = np.linalg.inv(Ainv)
    a = A.dot(Vinv @ mu + M.T @ Cinv @ y)

    return a, A, b, B
lls = np.zeros(nsamples)
linear_samples = np.zeros((nsamples, M.shape[-1]))
for i in range(nsamples):
    # varK = (P_samples[i] / 1.)**(-2/3) * 45**2
    # VV = np.diag([1e3**2, varK])
    VV = V
    a, A, b, B = get_aA_bB(M[i], y, C, mu_th, VV)
    lls[i] = multivariate_normal.logpdf(y, mean=b, cov=B)
    linear_samples[i] = np.random.multivariate_normal(a, A)
samples = prior_samples.copy()
samples['v0'] = linear_samples[:, 0] * data.rv.unit
samples['K'] = linear_samples[:, 1] * data.rv.unit

samples['M0'][samples['K'] < 0] += np.pi*u.rad
samples['K'][samples['K'] < 0] = np.abs(samples['K'][samples['K'] < 0])
samples.t0 = data._t0_bmjd
uu = np.random.uniform(size=nsamples)
good_samples_mask = uu < np.exp(lls - lls.max())
good_samples_mask.sum()
_ = plot_rv_curves(samples[good_samples_mask][:256],
                   data=data,
                   t_grid=np.linspace(-1, 10, 1024))
plt.ylim(-1, 1)
fig, ax = plt.subplots(1, 1, figsize=(7, 5))
_, bins, *_ = ax.hist(prior_samples['P'].value, 
                      bins=np.geomspace(P_lim[0], P_lim[1], 64),
                      label='prior')
ax.hist(samples['P'].value[good_samples_mask], bins=bins, label='posterior')
ax.set_xscale('log')

baseline = np.ptp(data._t_bmjd)
ax.set_title(r'data baseline: {:.1f},    $\Lambda =$ diag({:.0e}, {:.0e})$^2$'
             .format(baseline, *np.sqrt(Vdiag)), 
             fontsize=16)
ax.set_xlabel('$P$')
ax.legend(loc='upper left', fontsize=14)

fig.tight_layout()
fig.savefig('baseline{:.1f}_V{}.png'
            .format(baseline, '_'.join(['{:.0e}'.format(x) for x in np.sqrt(Vdiag)])),
            dpi=250)<jupyter_output><empty_output><jupyter_text>---<jupyter_code>from thejoker.sampler.likelihood import marginal_ln_likelihood, design_matrix, cy_rv_from_elements
import h5py
from schwimmbad import SerialPool
with h5py.File('/Users/apricewhelan/Downloads/test_joker.hdf5', 'r') as f:
    nonlinear_ps = f['samples'][:]
joker_lls = np.array([marginal_ln_likelihood(nonlinear_ps[i], data, joker.params) 
                      for i in range(len(nonlinear_ps))])
joker_good_samples_mask = uu < np.exp(joker_lls - joker_lls.max())
joker_samples = prior_samples[joker_good_samples_mask]
len(joker_samples)
# lnlike = compute_likelihoods(nsamples, '/Users/apricewhelan/Downloads/test_joker.hdf5', 0, data,
#                              pars, SerialPool())
# joker_samples, lnprior, lnlike = joker.rejection_sample(
#     data, n_prior_samples=nsamples,
#     return_logprobs=True,
#     prior_cache_file='/Users/apricewhelan/Downloads/test_joker.hdf5')
fig, ax = plt.subplots(1, 1, figsize=(8, 5))
_, bins, *_ = ax.hist(prior_samples['P'].value, 
                      bins=np.geomspace(P_lim[0], P_lim[1], 64),
                      label='prior')
ax.hist(joker_samples['P'].value, bins=bins, label='posterior')
ax.set_xscale('log')

baseline = np.ptp(data._t_bmjd)
ax.set_title(r'data baseline: {:.1f},    $\Lambda =$ diag({:.0e}, {:.0e})$^2$'
             .format(baseline, *np.sqrt(Vdiag)), 
             fontsize=16)
ax.set_xlabel('$P$')
ax.legend(loc='upper left', fontsize=14)

fig.tight_layout()
fig.savefig('baseline{:.1f}_V{}.png'
            .format(baseline, '_'.join(['{:.0e}'.format(x) for x in np.sqrt(Vdiag)])),
            dpi=250)<jupyter_output><empty_output><jupyter_text>---<jupyter_code>phase_samples = np.arctan2(samples[:, 2], samples[:, 1])
amp_samples = np.sqrt(samples[:, 1]**2 + samples[:, 2]**2)
from corner import corner
flat = np.stack((samples[:, 0], amp_samples, phase_samples, P_samples)).T[good_samples_mask]
_ = corner(flat, plot_contours=False, plot_density=False, labels=['const', 'amp', 'phase', 'period'])
model_t = np.linspace(-1, 100., 1024)

fig, axes = plt.subplots(2, 1, figsize=(10, 8), sharex=True)

for ax in axes:
    ax.errorbar(t, y, sig_y, marker='o', ls='none')

    for j in np.where(good_samples_mask)[0][:128]:
        ax.plot(model_t, 
                (samples[j][0] + samples[j][1] * np.cos(2*np.pi*model_t / P_samples[j]) + 
                 samples[j][2] * np.sin(2*np.pi*model_t / P_samples[j])),
                marker='', lw=0.5, alpha=0.2)

axes[0].set_ylim(0, 20)
axes[1].set_xlim(-1, 2)
axes[1].set_ylim(np.mean(y) - 2, np.mean(y) + 2)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.19357168652943302, 'fraction_numerical': 0.02925243770314193, 'mean_word_length': 3.420590582601756, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '403462', 'license': 'MIT', 'n_tokens_mistral': 2406, 'n_tokens_neox': 2259, 'n_words': 373, 'path': 'ConstCosSin-model.ipynb', 'repo_name': 'adrn/TheSimplestJoker', 'size': '5538'}