Discovery

Discovery is a next-generation pulsar-timing-array data-analysis package, built for speed on a JAX backend that supports GPU execution and autodifferentiation. If Enterprise is Spock, logical and elegant, Discovery is all Scotty, fast, efficient, and not above a hack if it gets you to warp speed.

Requirements

Discovery needs a modern Python with numpy, scipy, jax, pyarrow. It will be happier running on an Nvidia GPU with CUDA-enabled JAX.

Discovery's subpackages (such as discovery.flow and the packages under discovery.samplers) require additional dependencies.

Examples

The folder examples contains a growing set of usage examples.

Data model

The Discovery data model consists of Kernel objects (think of a noise matrix N, which can be inverted and applied to a vector, N^{-1} y, or even sandwiched with it, y^T N^{-1} y) and of GP objects, consisting of a basis F (sized ntoas x ngp) and a prior/kernel Phi. The kernel WoodburyKernel(N, F, P) combines a noise kernel and a GP.

Pulsar data

Discovery uses lightweight Pulsar objects saved as Arrow Feather files. To convert an Enterprise Pulsar objects use discovery.Pulsar.save_feather(psr, filename, noisedict), where noisedict is an optional dictionary of default parameter values. Some example datasets (based on NG 15yr) are included in the data folder.

Noise and signal components (signals.py)

• residuals(psr): just a numpy array of residuals for pulsar psr.

• makenoise_measurement(psr, [noisedict]): returns a Kernel object that implements EFAC + EQUAD measurement noise for pulsar psr. Parameters are multiplexed to pulsar and backend (e.g., B1855+09_430_ASP_efac), Enterprise-style. If those parameters appear is noisedict, their values will be frozen to those specifications.

• makenoise_measurement_simple(psr, [noisedict]): same, but no backends.

• makegp_ecorr(psr, [noisedict]): returns a GP object that implements ECORR measurement noise for pulsar psr. The quantization style. Parameters are multiplexed to pulsar and backend, and frozen if included in noisedict.

• makegp_ecorr_simple(psr, [noisedict]): same, but no backends.

• makegp_fourier(psr, prior, components, [T, fourierbasis, common, name]): returns a GP object that implements a finite GP over a vector basis. Here prior must be a JAX-ready function with signature priorfunc(f, df, arg1, [...]), where f is a vector of frequencies, and df a vector of integration weights associated with the frequencies. The resulting GP parameters are named {psrname}_{name}_{argX}, unless they are included in the list common. The Fourier basis is generated by calling fourierbasis(psr, components, T) (see below), which must return (f, df, F). In this function T defaults to the pulsar span.

• makegp_fourier_global(psrs, prior, orf, components, T, [fourierbasis, name]): returns a GlobalVariableGP object that implements a GP for all pulsars together, with a matched Fourier basis. Here prior is the same as for makegp_fourier (with GP parameters named {name}_{argX}), and orf must be a function orf(pos1, pos2) of two numpy vectors, such as hd_orf, monopole_orf, dipole_orf, and the diagonal uncorrelated_orf. The Fourier basis is generated by calling fourierbasis(psr, components, T) (see below), which must return (f, df, F). In this function T does not have a default (but see getspan). If prior and orf are given as equal-length lists of functions, the object implements the composite process sum_i GP(prior[i], orf[i]) (with GP parameters named {name}_{orf[i].__name__}_{argX}).

• makegp_improper(psr, fmat, [constant, name]): returns a GP object with improper prior (formally, a constant vector equal to constant) and basis matrix fmat. Here constant defaults to 1e40.
• makegp_timing(psr, [constant], [variance]): convenience function to call makegp_improper with the pulsar design matrix psr.Mmat and prior diagonal value constant. Design-matrix columns are normalized. To obtain "physical" priors, use variance to a value in s^2 (e.g., 1e-12``) instead of setting constant`.

GP basis and priors (signals.py)

• fourierbasis(psr, components, [T]): returns (f, df, F) for a basis of interleaved sines and cosines evaluated over psr TOAs with frequencies k/T, with k = 1, ..., components. Again T defaults to the pulsar span.

• dmfourierbasis(psr, components, [T, fref]): same, but rescales the basis by (fref / psr.freqs)**2, useful to define DMGP. Here T defaults to the pulsar span and fref to 1400.

• getspan(psr or psrs): returns the TOA span of the pulsar or iterable of pulsars. Useful for makegp_fourier and makegp_fourier_global.

• powerlaw(f, df, log10_A, gamma): implements the standard red-noise/GW spectrum 10^(2 log10_A) f^(-gamma) (yr^(gamma - 3) / pi^2 / 12) df. This is a JAX-ready function, so one would pass powerlaw to makegp_fourier.

• freespectrum(f, df, log10_rho: typing.Sequence): implements 10^(2 log10_rho). Note that makegp_fourier uses the log10_rho annotation to treat it as a vector; the resulting parameter name would be, e.g., B1855+09_fourierGP_log10_rho(10) if components = 10.

• makepowerlaw_crn(components): makes the prior for a combined red-noise/common-process GP, limiting the common-process to components frequencies. Returns a function with the signature powerlaw_crn(f, df, log10_A, gamma, crn_log10_A, crn_gamma); calling makegp_fourier(..., common=['crn_log10_A, crn_gamma'], name='rednoise') would then result in parameter names ['B1855+09_rednoise_log10_A', 'B1855+09_rednoise_gamma', ..., 'crn_log10_A', 'crn_gamma'].

Delays (signals.py)

• makedelay(psr, delayfunc, [common, name]): returns a JAX-ready function that implements a deterministic delay for psr. Here delayfunc must be a JAX-ready function with signature delayfunc(arg1, ...). The resulting delay parameters are named {psrname}_{name}_{argX}, unless they are included in the list common. If the first few arguments are defined attributes of psr (e.g., toas or freqs), they are automatically passed to the function and excluded from params; however they must come before all the other variable parameters.
• make_solardm(psr) [in solar.py]: calculates the DM delay in a 1/r^2 solar wind density model. Returns a function with signature solardm(n_earth).
• make_chromaticdecay(psr) [in solar.py]: calculates a chromatic exponential-dip delay. Returns a function with signature decay(t0, log10_Amp, log10_tau, idx).

Likelihood (likelihood.py)

• PulsarLikelihood(signals, concat=True): returns a PulsarLikelihood object, with a logL property that implements the single-pulsar likelihood as a JAX-ready function. The likelihood takes as a single argument a dictionary of parameters named as discussed above. Here signals is an iterable that must contain exactly one residual vector, exactly one noise Kernel, any number of GP objects, and any number of deterministic delays (any callable). If concat=False, the likelihood is built by nesting WoodburyKernels, first consuming ConstantGP objects (those without parameters) and then VariableGP, but otherwise respecting the order in signals. If concat=True, the ConstantGPs and VariableGPs are separately concatenated, and then nested.

• GlobalLikelihood(psls, globalgp=None): returns GlobalLikelihood object, with a logL property that implements the multi-pulsar likelihood as a JAX-ready function. The likelihood takes as a single argument a dictionary of parameters. Here psls is an iterable that may contain any number of PulsarLikelihood objects, and globalgp is a GlobalVariableGP object, such as returned by makegp_fourier_global, encoding a joint GP for all pulsars.

• The two likelihood objects have a sample property—a JAX-ready function that generates a random realization of the data if given a JAX key and a dictionary of parameters. (According to JAX protocol, sample actually returns a tuple consisting of the "split" key and the data realization.)

• The PulsarLikelihood and GlobalLikelihood objects have a sample_conditional(key, params) -> (key, dict) property–a JAX-ready function that returns a dictionary of GP coefficients drawn from the conditional (normal) probability distribution of the parameter-dependent GPs (for PulsarLikelihood) or of the globalgp(s), given the params. This function calls GlobalLikelihood.conditional(params), which returns a concatenated vector mu of the mean GP coefficients and the result cf of running scipy.linalg.cho_factor on the matrix Phi^-1 + T^t K T, where Phi and T are prior and basis of the relevant GP.

Priors (prior.py)

• makelogprior_uniform(params, [priordict]): returns a function logprior(params) that evaluates the total log prior according to priordict (given, e.g., as {'FourierGP_log10_A': [-18, -11]'}). Some standard enterprise_extensions priors are included by default (e.g., crn_log10_A, crn_gamma, gw_log10_A, gw_gamma, ...). Parameters that are in the list params but are not in priordict and have no default are not included in the computation.
• sample_uniform(params, [priordict]): creates a dictionary of random values for the parameters in the list params, using the uniform priors in priordict or in the system default. Fails if a parameter in params has no specification.

Optimal statistic (os.py)

• OS(gbl): creates an optimal statistic object, where gbl is a GlobalLikelihood in which each PulsarLikelihood has a makegp_fourier component named gw with at least the common parameter gw_log10_A. The prior and Fourier basis from that GP will be used to build the OS. If the GlobalLikelihood has a globalgp, it is unused.
• OS.os(params, [orfa]): returns the dictionary {'os': ..., 'os_sigma': ..., 'snr': ..., 'log10_A': ...} as computed for params. orfa defaults to hd_orfa but can also be dipole_orfa or monopole_orfa—these are all functions of the cosine between pulsars. OS.os can be jitted and vmapped over params. (If one wishes to use the orfa argument, the jitting/vmapping syntax is more complicated: jax.jit(os.os, static_argnums=1) and jax.vmap(os.os, in_axes=(0, None)).)
• OS.scramble(params, pos, [orfa]): returns the same dictionary as OS.os, using the pulsar positions in pos (either a list of three-vectors or an npsr x 3 array). OS.scramble can be jitted and vmapped over params (jax.vmap(os.scramble, (0, None))) or positions (jax.vmap(os.scramble, (None, 0))).
• OS.shift(params, phases, [orfa]): returns the same dictionary as OS.os, shifting the GW basis vectors by the phases in phases (an npsr x ngw array of floats in [0,2*pi]). OS.shift can be jitted and vmapped over params or phases.
• OS.gx2cdf(params, xs, cutoff=1e-6, limit=100, epsabs=1e-6): returns a vector of the GX2 OS CDF evaluated at the SNRs xs. The CDF is computed by considering only GPs (i.e., not white noise), using Imhof's method. The parameters limit and epsabs are passed to scipy.integrate.quad. If cutoff is a float, GX2 eigenvalues are limited to those larger than cutoff; if cutoff is an integer, only the cutoff largest eigenvalues are used. Currently this function cannot be jitted or vmapped.

Experimental batched interface (likelihood.py and signals.py)

• ArrayLikelihood(psls, [commongp=], [globalgp=]): returns an ArrayLikelihood object analog to GlobalLikelihood, but set up especially for "batched" GPs (given as commongp) that implement the same process for all pulsars (i.e., red noise with the same number of components and the same power-law prior, but different amplitude and gamma parameters). Here psls is an iterable that may contain any number of PulsarLikelihood objects, commongp is a vectorized GP (as returned, e.g., by makecommongp_fourier) or a list of vectorized GPs, and globalgp is a GlobalVariableGP object as for GlobalLikelihood. In addition to logL, ArrayLikelihood implements a hierarchical clogL that takes the commongp and globalgp coefficients as additional parameters (e.g., 'B1855+09_red_noise_coefficients(60)', 'B1937+21_gw_coefficients(28)', etc.).
• makecommongp_fourier(psrs, prior, components, [T, fourierbasis, common, name)]): returns a vector GP object that implements a finite GP over a vector basis for a set of pulsars. This function is analog to makegp_fourier but takes a list of pulsars, and is meant to be used with ArrayLikelihood.

Example usage

The following implements a standard NANOGrav likelihood + prior for pulsar psr, with free parameters ['{psrname}_rednoise_gamma', '{psrname}_rednoise_log10_A', 'crn_gamma', 'crn_log10_A']:

import discovery as ds

psl = ds.PulsarLikelihood([psr.residuals,
                           ds.makenoise_measurement(psr, noisedict),
                           ds.makegp_ecorr(psr, noisedict),
                           ds.makegp_timing(psr),
                           ds.makegp_fourier(psr, ds.powerlaw, 30, name='rednoise'),
                           ds.makegp_fourier(psr, ds.powerlaw, 14, common=['crn_log10_A', 'crn_gamma'], name='crn')])
logl = psl.logL
logp = dp.makelogprior_uniform(logl.params)

Then you would use jax.jit(logl) (and perhaps jax.jit(jax.grad(logl)))) and jax.jit(logp) in inference.

The following implements a standard NANOGrav HD likelihood:

Tspan = ds.getspan(psrs)
gbl = ds.GlobalLikelihood((ds.PulsarLikelihood([psr.residuals,
                                                ds.makenoise_measurement(psr, noisedict),
                                                ds.makegp_ecorr(psr, noisedict),
                                                ds.makegp_timing(psr),
                                                ds.makegp_fourier(psr, ds.powerlaw, 30, T=Tspan, name='rednoise')
                                                ]) for psr in psrs),
                          ds.makegp_fourier_global(psrs, ds.powerlaw, ds.hd_orf, 14, T=Tspan, name='gw'))
logl = gbl.logL
logp = dp.makelogprior_uniform(logl.params)

Here are ArrayLikelihood versions, which will be especially fast on a GPU:

curn = ds.ArrayLikelihood([ds.PulsarLikelihood([psr.residuals,
                                                ds.makenoise_measurement(psr, psr.noisedict),
                                                ds.makegp_ecorr(psr, psr.noisedict),
                                                ds.makegp_timing(psr, svd=True, constant=1e-6)]) for psr in psrs],
                          ds.makecommongp_fourier(psrs, ds.makepowerlaw_crn(14), 30, T=Tspan, common=['crn_log10_A', 'crn_gamma'], name='red_noise'))

hd = ds.ArrayLikelihood([ds.PulsarLikelihood([psr.residuals,
                                              ds.makenoise_measurement(psr, psr.noisedict),
                                              ds.makegp_ecorr(psr, psr.noisedict),
                                              ds.makegp_timing(psr, svd=True, constant=1e-6)]) for psr in psrs],
                         ds.makecommongp_fourier(psrs, ds.powerlaw, 30, T=Tspan, name='red_noise'),
                         ds.makegp_fourier_global(psrs, ds.powerlaw, ds.hd_orf, 14, T=Tspan, name='gw'))

To create a random realization of a model (e.g., gbl) as a function of its parameters:

sampler = gbl.sampler

p0 = ds.sample_uniform(sampler.params)
key = jax.random.PRNGKey(42)

key, y = sampler(key, p0)

To simulate residuals for a pulsar based on an existing dataset:

psr = ds.Pulsar.read_feather('data/NG15yr-B1855+09.feather')

psl = ds.PulsarLikelihood([psr.residuals,
                           ds.makenoise_measurement(psr, psr.noisedict),
                           ds.makegp_ecorr(psr, psr.noisedict),
                           ds.makegp_timing(psr, variance=1e-12),
                           ds.makegp_fourier(psr, ds.powerlaw, 30, name='rednoise'),
                           ds.makegp_fourier(psr, ds.powerlaw, 14, common=['crn_log10_A', 'crn_gamma'], name='crn')
                           ])

sampler = psl.sample

# fix parameters to reasonable value
noisedict_red = {p: v for p,v in psr.noisedict.items() if 'rednoise' in p}
params = {**ds.sample_uniform(sampler.params), **noisedict_red, 'crn_gamma': 4.3, 'crn_log10_A': -14.5}

# should jit the sampler if simulating many datasets
key = ds.rngkey(42)
key, res = sampler(key, params)

# update Pulsar object if desired (make PulsarLikelihood anew to analyze)
psr.residuals = res

To simulate residuals for an array:

import glob
psrs = [ds.Pulsar.read_feather(f) for f in glob.glob('data/*.feather')[:5]]

Tspan = ds.getspan(psrs)

gbl = ds.GlobalLikelihood((ds.PulsarLikelihood([psr.residuals,
                                                ds.makenoise_measurement(psr, psr.noisedict),
                                                ds.makegp_ecorr(psr, psr.noisedict),
                                                ds.makegp_timing(psr, variance=1e-14),
                                                ds.makegp_fourier(psr, ds.powerlaw, 30, T=Tspan, name='rednoise')
                                                ]) for psr in psrs),
                          ds.makegp_fourier_global(psrs, ds.powerlaw, ds.hd_orf, 14, T=Tspan, name='gw'))

sampler = gbl.sample

noisedict_red = {p: v for psr in psrs for p,v in psr.noisedict.items() if 'rednoise' in p}
params = {**ds.sample_uniform(sampler.params), **noisedict_red, 'crn_gamma': 4.3, 'crn_log10_A': -14.5}

key = ds.rngkey(43)
key, res = sampler(key, params)

To load and save PTMCMC chain files into Pandas tables:

chain = ds.read_chain('PTMCMC_example_output')
ds.save_chain(chain, 'example_chain.feather')
chain = ds.read_chain('example_chain.feather')

Note that the resulting Pandas tables will contain useful attributes chain.attr['noisedict'] (a Python dict), chain.attr['priors'] (a list of strings), and chain.attr['runtime_info'] (the runtime_info.txt file rendered as a list of strings).