Discovery

Sphinx Documentation is here. The old README user guide is below as well.

<img src="discovery.png" alt="Logo" width="300" align="left" style="margin-right: 20px;"/>

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_fourier_allpsr(psrs, prior, components, T, [fourierbasis, common, name])`: returns a `GlobalVariableGP` that implements uncorrelated GPs for each pulsar, which matched Fourier bases. Here `prior` is the same as for `makegp_fourier` (with GP parameters named `{psrname}_{name}_{argX}`). This object is useful, e.g., to obtain conditional RN coefficients in a CURN or HD model. -->

• 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.

<!-- To get the actual covariance matrix, one would use `cho_solve(cf, identity)`. -->

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](https: