summary(results) takes a long time to execute #23
-
|
Hello, I attach the python script I'm using and the data. I'm doing some basic tests for exoplanet detection with radial velocities. Thanks in advance |
Beta Was this translation helpful? Give feedback.
Replies: 2 comments 4 replies
-
|
Small update: I've noticed that the time it gets stuck is proportional to the number of live points I use. With only a few live points used in the nested sampling, |
Beta Was this translation helpful? Give feedback.
-
|
Hi @nicochunger I'm not comfortable unzipping things locally (as a general policy), but I'd love to help you. Could you please just paste your code in github's very nice markdown code block? Like this: def some_code(...): |
Beta Was this translation helpful? Give feedback.
-
|
Sure sorry, I also have a data file that's why I zipped it. The code: from jaxns.nested_sampling import NestedSampler, save_results, load_results
from jaxns.prior_transforms import PriorChain, UniformPrior
from jaxns.plotting import plot_cornerplot, plot_diagnostics
from jaxns.utils import summary, resample
from jax import random, jit
from jax import numpy as jnp
import pandas as pd
from datetime import datetime
import matplotlib.pyplot as plt
# Load data
data_51peg = pd.read_csv('51Peg_HIRES_Pub-2017.txt', sep='\t', skiprows=(1,))
time = jnp.array(data_51peg['rjd'])
vrad = jnp.array(data_51peg['vrad'])
svrad = jnp.array(data_51peg['svrad'])
# Center radial velocities
vrad -= jnp.median(vrad)
# Plot raw timeseries
# plt.title('51 Pegasi')
# plt.errorbar(time, vrad, yerr=svrad, fmt='.', color='b')
# plt.xlabel('Date [BJD]')
# plt.ylabel('Radial Velocity [m/s]')
# Error bars are so small they can't be seen in the plot
def keplerian(C_ms, K_ms, P_day, ecc, omega_rad, ml0_rad):
""" Keplerian model """
epoch = 55000
ma0_rad = ml0_rad - omega_rad
# Compute mean anomaly at current time (or array)
ma = 2*jnp.pi/P_day * (time - epoch) + ma0_rad
# Compute true anomaly
nu = true_anomaly(ma, ecc)
return C_ms + K_ms * (jnp.cos(nu + omega_rad) + ecc * jnp.cos(omega_rad))
def true_anomaly(ma, ecc):
E = jnp.array(ma, copy=True) #ma.copy()
E0 = jnp.zeros_like(E)
# deltaE = jnp.array([1])
# This is usually a dynamic algorithm with a while loop unitl a stopping
# criteria is met. Not sure yet if I can do that wiht jax.
# For now I just fix the number of loops
for i in range(5):
E0 = jnp.array(E, copy=True)
diff = E - ecc*jnp.sin(E) - ma
deriv = 1 - ecc*jnp.cos(E)
deltaE = diff/deriv
E = E0 - deltaE
v = 2*jnp.arctan(jnp.sqrt((1+ecc)/(1-ecc))*jnp.tan(E/2))
return v
def lnlike(C_ms, K_ms, P_day, ecc, omega_rad, ml0_rad, jitter, **kwargs):
"""
gaussian likelihood
"""
kep = keplerian(C_ms, K_ms, P_day, ecc, omega_rad, ml0_rad)
N = time.size
# On top of the intrisic error in the data we add another white noise term
# to account for unknown noise sources.
var = svrad**2 + jitter**2
residuals = vrad - kep
cte = -0.5 * N * jnp.log(2 * jnp.pi)
return cte - jnp.sum(jnp.log(jnp.sqrt(var))) - jnp.sum(residuals**2 / (2 * var))
# Define priors
prior_C = UniformPrior('C_ms', -50., 50.)
prior_K = UniformPrior('K_ms', 0.1, 100.)
prior_P = UniformPrior('P_day', 0.7, 100)
prior_ecc = UniformPrior('ecc', 0., 0.99)
prior_omega = UniformPrior('omega_rad', 0., 2*jnp.pi)
prior_ml0 = UniformPrior('ml0_rad', 0., 2*jnp.pi)
prior_jitter = UniformPrior('jitter', 0., 10.)
prior_chain = PriorChain().push(prior_C).push(prior_K).push(prior_P).push(prior_ecc).push(prior_omega).push(prior_ml0).push(prior_jitter)
ndims = prior_chain.U_ndims
print(f'Number of free parameters: {ndims}')
# Create ns object
ns = NestedSampler(lnlike, prior_chain)
print('Starting run...')
t0 = datetime.now()
results = jit(ns)(random.PRNGKey(35111651), termination_frac=0.001)
print(f"Time to run: {datetime.now() - t0}")
t0 = datetime.now()
# General results
# # can always save results to play with later
# save_results(results, 'save.npz')
# # loads results that you may have saved
# results = load_results('save.npz')
summary(results)
print(f"Time to print results: {datetime.now() - t0}")
def plot_hist(data, title):
plt.figure()
plt.hist(data, bins='auto', label='jaxns Posterior')
plt.axvline(4.23078, c='r', ls='--', label='True value')
plt.title(title)
plt.legend(loc='upper right')
plt.show()
samples = resample(random.PRNGKey(34252), results.samples, results.log_p, S=int(results.ESS))
plot_hist(samples['P_day'], 'Period of 51Peg b')And here is the data which should be saved as "51Peg_HIRES_Pub-2017.txt" Thanks! |
Beta Was this translation helpful? Give feedback.
-
|
Thanks @nicochunger, neat model! So this problem here is due to the way that JAX works. When you run a JAX op it immediately passes control back to python and lets python continue until you actually need the result of the computation. So what's happening is that you're calling Then To make sure a computation is done, you could block on any of the arrays in the results = jit(ns)(random.PRNGKey(35111651), termination_frac=0.001)
results.num_likelihood_evaluations.block_until_ready()I personally accomplish this with plotting (which blocks until results are ready): results = jit(ns)(random.PRNGKey(35111651), termination_frac=0.001)
plot_diagnostics(results)
plot_cornerplot(results)
summary(results)For me when I do this complie+run takes 4 minutes, and summary takes 1 second. Note, you can rewrite your def true_anomaly(ma, ecc):
E = jnp.array(ma, copy=True) # ma.copy()
E0 = jnp.zeros_like(E)
# deltaE = jnp.array([1])
# This is usually a dynamic algorithm with a while loop unitl a stopping
# criteria is met. Not sure yet if I can do that wiht jax.
# For now I just fix the number of loops
def body(state):
(i, E) = state
E0 = E
diff = E - ecc * jnp.sin(E) - ma
deriv = 1 - ecc * jnp.cos(E)
deltaE = diff / deriv
E = E0 - deltaE
return (i + 1, E)
def cond(state):
(i, E) = state
return i < 5
(_, E) = while_loop(
cond,
body,
init_val=(jnp.asarray(0), E))
v = 2 * jnp.arctan(jnp.sqrt((1 + ecc) / (1 - ecc)) * jnp.tan(E / 2))
return v |
Beta Was this translation helpful? Give feedback.
-
|
Great thank you for the explanation. Yeah, I have similar run times, around 4 minutes. Thanks for the tip on how to do while loops. I'm also comparing it with PolyChord, which I've been using until now. So faster run times could be very valuable. I will continue with my tests and report back :) |
Beta Was this translation helpful? Give feedback.
-
|
Yep exactly. Please report back on differences in runtime with PolyChord, that would be wonderful to share! |
Beta Was this translation helpful? Give feedback.
Thanks @nicochunger, neat model! So this problem here is due to the way that JAX works. When you run a JAX op it immediately passes control back to python and lets python continue until you actually need the result of the computation. So what's happening is that you're calling
results = jit(ns)(random.PRNGKey(35111651), termination_frac=0.001)and then it passes back control to python once it's done compiling, but only when it's started running. For me compilation takes 40 seconds.Then
summaryblocks until theresultsis actually done, which gives you the apparent wait time.To make sure a computation is done, you could block on any of the arrays in the
resultslike this: