Merge pull request #1426 from abhisrkckl/bayesian-wip

Bayesian interface for Wideband data (No correlated noise)

CHANGELOG-unreleased.md

@@ -18,6 +18,7 @@ the released changes.
 - `Parameter.as_latex` method for latex representation of a parameter.
 - `pint.output.publish` module and `pintpublish` script for generating publication (LaTeX) output.
 - Added radial velocity methods for binary models
+- Support for wideband data in `pint.bayesian` (no correlated noise).
 - Added `DMWaveX` model (Fourier representation of DM noise)
 ### Fixed
 - Wave model `validate()` can correctly use PEPOCH to assign WAVEEPOCH parameter

docs/conf.py

@@ -119,6 +119,8 @@
     "examples/compare_tempo2_B1855.py",
     "examples/example_dmx_ranges.py",
     "examples/example_pulse_numbers.py",
+    # "examples/bayesian-example-NGC6440E.py",
+    # "examples/bayesian-wideband-example.py",
     "conf.py",
     "_ext",
 ]

docs/examples/bayesian-example-NGC6440E.py

@@ -1,3 +1,18 @@
+# ---
+# jupyter:
+#   jupytext:
+#     formats: ipynb,py:percent
+#     text_representation:
+#       extension: .py
+#       format_name: percent
+#       format_version: '1.3'
+#       jupytext_version: 1.14.7
+#   kernelspec:
+#     display_name: Python 3 (ipykernel)
+#     language: python
+#     name: python3
+# ---
+
 # %% [markdown]
 # # PINT Bayesian Interface Examples
 
@@ -6,6 +21,7 @@
 from pint.bayesian import BayesianTiming
 from pint.config import examplefile
 from pint.models.priors import Prior
+from pint.logging import setup as setup_log
 from scipy.stats import uniform
 
 # %%
@@ -16,6 +32,10 @@
 import io
 import matplotlib.pyplot as plt
 
+# %%
+# Turn off log messages. They can slow down the processing.
+setup_log(level="WARNING")
+
 # %%
 # Read the par and tim files
 parfile = examplefile("NGC6440E.par.good")
@@ -70,36 +90,48 @@
     + np.random.randn(nwalkers, bt.nparams) * maxlike_errors
 )
 
+# %%
+# ** IMPORTANT!!! **
+# This is used to exclude some of the following time-consuming steps from the readthedocs build.
+# Set this to False while actually using this example.
+rtd = True
+
 # %%
 # Use longer chain_length for real runs. It is kept small here so that
 # the sampling finishes quickly (and because I know the burn in is short
 # because of the cheating priors above).
-print("Running emcee...")
-chain_length = 1000
-sampler.run_mcmc(
-    start_points,
-    chain_length,
-    progress=True,
-)
+if not rtd:
+    print("Running emcee...")
 
-# %%
-# Merge all the chains together after discarding the first 100 samples as 'burn-in'.
-# The burn-in should be decided after looking at the chains in the real world.
-samples_emcee = sampler.get_chain(flat=True, discard=100)
+    chain_length = 1000
+
+    sampler.run_mcmc(
+        start_points,
+        chain_length,
+        progress=True,
+    )
 
 # %%
-# Plot the MCMC chains to make sure that the burn-in has been removed properly.
-# Otherwise, go back and discard more points.
-for idx, param_chain in enumerate(samples_emcee.T):
-    plt.subplot(bt.nparams, 1, idx + 1)
-    plt.plot(param_chain, label=bt.param_labels[idx])
-    plt.legend()
-plt.show()
+if not rtd:
+    # Merge all the chains together after discarding the first 100 samples as 'burn-in'.
+    # The burn-in should be decided after looking at the chains in the real world.
+    samples_emcee = sampler.get_chain(flat=True, discard=100)
+
+    # %%
+if not rtd:
+    # Plot the MCMC chains to make sure that the burn-in has been removed properly.
+    # Otherwise, go back and discard more points.
+    for idx, param_chain in enumerate(samples_emcee.T):
+        plt.subplot(bt.nparams, 1, idx + 1)
+        plt.plot(param_chain, label=bt.param_labels[idx])
+        plt.legend()
+    plt.show()
 
 # %%
 # Plot the posterior distribution.
-fig = corner.corner(samples_emcee, labels=bt.param_labels)
-plt.show()
+if not rtd:
+    fig = corner.corner(samples_emcee, labels=bt.param_labels)
+    plt.show()
 
 # %% [markdown]
 # ## Nested sampling with nestle
@@ -118,28 +150,30 @@
 # `dlogz` is the target accuracy in the computed Bayesian evidence.
 # Increasing `npoints` or decreasing `dlogz` gives more accurate results,
 # but at the cost of time.
-print("Running nestle...")
-result_nestle_1 = nestle.sample(
-    bt.lnlikelihood,
-    bt.prior_transform,
-    bt.nparams,
-    method="multi",
-    npoints=150,
-    dlogz=0.5,
-    callback=nestle.print_progress,
-)
+if not rtd:
+    print("Running nestle...")
+    result_nestle_1 = nestle.sample(
+        bt.lnlikelihood,
+        bt.prior_transform,
+        bt.nparams,
+        method="multi",
+        npoints=150,
+        dlogz=0.5,
+        callback=nestle.print_progress,
+    )
 
 # %%
 # Plot the posterior
 # The nested samples come with weights, which must be taken into account
 # while plotting.
-fig = corner.corner(
-    result_nestle_1.samples,
-    weights=result_nestle_1.weights,
-    labels=bt.param_labels,
-    range=[0.999] * bt.nparams,
-)
-plt.show()
+if not rtd:
+    fig = corner.corner(
+        result_nestle_1.samples,
+        weights=result_nestle_1.weights,
+        labels=bt.param_labels,
+        range=[0.999] * bt.nparams,
+    )
+    plt.show()
 
 # %% [markdown]
 # Let us create a new model with an EFAC applied to all toas (all
@@ -170,36 +204,41 @@
 print(bt2.likelihood_method)
 
 # %%
-result_nestle_2 = nestle.sample(
-    bt2.lnlikelihood,
-    bt2.prior_transform,
-    bt2.nparams,
-    method="multi",
-    npoints=150,
-    dlogz=0.5,
-    callback=nestle.print_progress,
-)
+if not rtd:
+    result_nestle_2 = nestle.sample(
+        bt2.lnlikelihood,
+        bt2.prior_transform,
+        bt2.nparams,
+        method="multi",
+        npoints=150,
+        dlogz=0.5,
+        callback=nestle.print_progress,
+    )
 
 # %%
 # Plot the posterior.
 # The EFAC looks consistent with 1.
-fig2 = corner.corner(
-    result_nestle_2.samples,
-    weights=result_nestle_2.weights,
-    labels=bt2.param_labels,
-    range=[0.999] * bt2.nparams,
-)
-plt.show()
+if not rtd:
+    fig2 = corner.corner(
+        result_nestle_2.samples,
+        weights=result_nestle_2.weights,
+        labels=bt2.param_labels,
+        range=[0.999] * bt2.nparams,
+    )
+    plt.show()
 
 # %% [markdown]
 # Now let us look at the evidences and compute the Bayes factor.
 
 # %%
-print(f"Evidence without EFAC : {result_nestle_1.logz} +/- {result_nestle_1.logzerr}")
-print(f"Evidence with EFAC : {result_nestle_2.logz} +/- {result_nestle_2.logzerr}")
+if not rtd:
+    print(
+        f"Evidence without EFAC : {result_nestle_1.logz} +/- {result_nestle_1.logzerr}"
+    )
+    print(f"Evidence with EFAC : {result_nestle_2.logz} +/- {result_nestle_2.logzerr}")
 
-bf = np.exp(result_nestle_1.logz - result_nestle_2.logz)
-print(f"Bayes factor : {bf} (in favor of no EFAC)")
+    bf = np.exp(result_nestle_1.logz - result_nestle_2.logz)
+    print(f"Bayes factor : {bf} (in favor of no EFAC)")
 
 # %% [markdown]
 # The Bayes factor tells us that the EFAC is unnecessary for this dataset.

docs/examples/bayesian-wideband-example.py

@@ -0,0 +1,110 @@
+# %% [markdown]
+# # PINT Bayesian Interface Example (Wideband)
+
+# %%
+import corner
+import emcee
+import matplotlib.pyplot as plt
+import numpy as np
+
+from pint.bayesian import BayesianTiming
+from pint.config import examplefile
+from pint.fitter import WidebandDownhillFitter
+from pint.logging import setup as setup_log
+from pint.models import get_model_and_toas
+
+# %%
+# Turn off log messages. They can slow down the processing.
+setup_log(level="WARNING")
+
+# %%
+# This is a simulated dataset.
+m, t = get_model_and_toas(examplefile("test-wb-0.par"), examplefile("test-wb-0.tim"))
+
+# %%
+# Fit the model to the data to get the parameter uncertainties.
+ftr = WidebandDownhillFitter(t, m)
+ftr.fit_toas()
+m = ftr.model
+
+# %%
+# Now set the priors.
+# I am cheating here by setting the priors around the maximum likelihood estimates.
+# This is a bad idea for real datasets and can bias the estimates. I am doing this
+# here just to make everything finish faster. In the real world, these priors should
+# be informed by, e.g. previous (independent) timing solutions, pulsar search results,
+# VLBI localization etc. Note that unbounded uniform priors don't work here.
+prior_info = {}
+for par in m.free_params:
+    param = getattr(m, par)
+    param_min = float(param.value - 10 * param.uncertainty_value)
+    param_max = float(param.value + 10 * param.uncertainty_value)
+    prior_info[par] = {"distr": "uniform", "pmin": param_min, "pmax": param_max}
+
+# %%
+# Set the EFAC and DMEFAC priors and unfreeze them.
+# Don't do this before the fitting step. The fitter doesn't know
+# how to deal with noise parameters.
+prior_info["EFAC1"] = {"distr": "normal", "mu": 1, "sigma": 0.1}
+prior_info["DMEFAC1"] = {"distr": "normal", "mu": 1, "sigma": 0.1}
+
+m.EFAC1.frozen = False
+m.EFAC1.uncertainty_value = 0.01
+m.DMEFAC1.frozen = False
+m.DMEFAC1.uncertainty_value = 0.01
+
+# %%
+# The likelihood function behaves better if `use_pulse_numbers==True`.
+bt = BayesianTiming(m, t, use_pulse_numbers=True, prior_info=prior_info)
+
+# %%
+print("Number of parameters = ", bt.nparams)
+print("Likelihood method = ", bt.likelihood_method)
+
+# %%
+nwalkers = 25
+sampler = emcee.EnsembleSampler(nwalkers, bt.nparams, bt.lnposterior)
+
+# %%
+# Start the sampler close to the maximul likelihood estimate.
+maxlike_params = np.array([param.value for param in bt.params], dtype=float)
+maxlike_errors = [param.uncertainty_value for param in bt.params]
+start_points = (
+    np.repeat([maxlike_params], nwalkers).reshape(bt.nparams, nwalkers).T
+    + np.random.randn(nwalkers, bt.nparams) * maxlike_errors
+)
+
+# %%
+# ** IMPORTANT!!! **
+# This is used to exclude the following time-consuming steps from the readthedocs build.
+# Set this to False while actually using this example.
+rtd = True
+
+# %%
+if not rtd:
+    print("Running emcee...")
+    chain_length = 1000
+    sampler.run_mcmc(
+        start_points,
+        chain_length,
+        progress=True,
+    )
+
+    samples_emcee = sampler.get_chain(flat=True, discard=100)
+
+# %%
+# Plot the chains to make sure they have converged and the burn-in has been removed properly.
+if not rtd:
+    for idx, param_chain in enumerate(samples_emcee.T):
+        plt.subplot(bt.nparams, 1, idx + 1)
+        plt.plot(param_chain)
+        plt.ylabel(bt.param_labels[idx])
+        plt.autoscale()
+    plt.show()
+
+# %%
+if not rtd:
+    fig = corner.corner(
+        samples_emcee, labels=bt.param_labels, quantiles=[0.5], truths=maxlike_params
+    )
+    plt.show()

docs/tutorials.rst

@@ -67,6 +67,8 @@ are not included in the default build because they take too long, but you can do
    examples/check_phase_connection.ipynb
    examples/PINT_observatories.ipynb
    examples/solar_wind.ipynb
+   examples/bayesian-example-NGC6440E.py
+   examples/bayesian-wideband-example.py
    examples-rendered/paper_validation_example.ipynb
 
 .. _`Time a Pulsar`: examples/time_a_pulsar.html