### Install Documentation Dependencies

Source: https://github.com/nanograv/discovery/blob/main/docs/README.md

Install the necessary dependencies to build the documentation. This command should be run from the project root.

```bash
pip install -e ".[docs]"
```

--------------------------------

### Complete Simulation Workflow Example

Source: https://github.com/nanograv/discovery/blob/main/docs/tutorials/simulations.md

Demonstrates a full simulation workflow, including data loading, model building, parameter sampling, data generation, and model re-evaluation. This example integrates multiple components of the Discovery library.

```python
import discovery as ds
import jax
import numpy as np
import glob
import os

# Find data
data_dir = os.path.join(ds.__path__[0], '..', '..', 'data')
data_pattern = os.path.join(data_dir, 'v1p1_de440_pint_bipm2019-*.feather')

# Load pulsars
psrs = [ds.Pulsar.read_feather(f) for f in glob.glob(data_pattern)[:5]]
Tspan = ds.getspan(psrs)

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

# Sample parameters from prior
params = ds.sample_uniform(gbl.sample.params)

# Override GW parameters
params['gw_gamma'] = 13/3
params['gw_log10_A'] = -14.5

# Generate data
key = ds.rngkey(42)
key, residuals = gbl.sample(key, params)

# Update pulsars
for psr, res in zip(psrs, residuals):
    psr.residuals = np.array(res)

# Rebuild and evaluate
new_gbl = ds.GlobalLikelihood(
    [ds.PulsarLikelihood([
        psr.residuals,
        ds.makenoise_measurement(psr, psr.noisedict),
        ds.makegp_ecorr(psr, psr.noisedict),
        ds.makegp_timing(psr, svd=True, variance=1e-20),
        ds.makegp_fourier(psr, ds.powerlaw, 30, T=Tspan,
                         name='rednoise')
    ]) for psr in psrs],
    globalgp=ds.makegp_fourier_global(
        psrs, ds.powerlaw, ds.hd_orf, 14, T=Tspan,
        name='gw'
    )
)

logl = jax.jit(new_gbl.logL)
print(f"Log-likelihood at true params: {logl(params)}")
```

--------------------------------

### Build Documentation Locally

Source: https://github.com/nanograv/discovery/blob/main/docs/installation.md

Install documentation dependencies and build the HTML documentation.

```bash
pip install -e ".[docs]"
cd docs
make html
```

--------------------------------

### Example Pulsar Data Directory Structure

Source: https://github.com/nanograv/discovery/blob/main/docs/guide/pulsar_data.md

Illustrates the directory structure for example pulsar datasets included with Discovery, typically found in the `data/` folder.

```bash
data/
├── v1p1_de440_pint_bipm2019-B1855+09.feather
├── v1p1_de440_pint_bipm2019-B1937+21.feather
├── v1p1_de440_pint_bipm2019-J0030+0451.feather
└── ...
```

--------------------------------

### signals.getstart

Source: https://github.com/nanograv/discovery/blob/main/docs/api/index.md

Gets the start time of signals.

```APIDOC
## getstart()

### Description
Gets the start time of signals.

### Method
N/A (Function Call)

### Endpoint
N/A

### Parameters
N/A

### Request Example
N/A

### Response
N/A
```

--------------------------------

### Setup Likelihood for OS Simulations

Source: https://github.com/nanograv/discovery/blob/main/examples/os_example.ipynb

Configures a global likelihood for simulations, including a delay model for residuals and fixed GPs. This setup differs from real-data analysis by incorporating simulated residuals.

```python
Tspan = ds.getspan(psrs)

sim = ds.GlobalLikelihood([ds.PulsarLikelihood([ds.makedelay(psr, ds.getresiduals, name='sim'), 
                                                ds.makenoise_measurement(psr, psr.noisedict),
                                                # ds.makegp_ecorr(psr, psr.noisedict), # cannot do fixed GP at this time 
                                                ds.makegp_timing(psr, variance=1e-12, variable=True),
                                                ds.makegp_fourier(psr, ds.powerlaw, 30, T=Tspan, name='red_noise'),
                                                ds.makegp_fourier(psr, ds.powerlaw, 14, T=Tspan,
                                                                  common=['gw_log10_A', 'gw_gamma'], name='gw')
                                                ]) for psr in psrs])
```

--------------------------------

### Load and Inspect Example Pulsar Data

Source: https://github.com/nanograv/discovery/blob/main/docs/guide/pulsar_data.md

Load an example pulsar from the included datasets and print basic information such as its name, number of TOAs, and data span.

```python
# Load example pulsar
psr = ds.Pulsar.read_feather('data/v1p1_de440_pint_bipm2019-B1855+09.feather')

print(f"Pulsar: {psr.name}")
print(f"TOAs: {len(psr.toas)}")
print(f"Span: {(psr.toas.max() - psr.toas.min()) / 365.25 / 86400:.1f} years")
```

--------------------------------

### Complete Multi-Pulsar HD Analysis Example

Source: https://github.com/nanograv/discovery/blob/main/docs/tutorials/basic_likelihood.md

This example demonstrates a full analysis pipeline, including loading data, building a global likelihood with HD correlation, defining priors, JIT-compiling functions, and performing basic evaluations.

```python
import discovery as ds
import jax
import glob

# Load data
psrs = [ds.Pulsar.read_feather(f)
        for f in glob.glob('data/v1p1_de440_pint_bipm2019-*.feather')[:5]]
Tspan = ds.getspan(psrs)

# Build global likelihood with HD correlation
gbl = ds.GlobalLikelihood(
    [
        ds.PulsarLikelihood([
            psr.residuals,
            ds.makenoise_measurement(psr, psr.noisedict),
            ds.makegp_ecorr(psr, psr.noisedict),
            ds.makegp_timing(psr, svd=True),
            ds.makegp_fourier(psr, ds.powerlaw, 30, T=Tspan,
                             name='rednoise')
        ]) for psr in psrs
    ],
    globalgp=ds.makegp_fourier_global(
        psrs, ds.powerlaw, ds.hd_orf, 14, T=Tspan,
        name='gw'
    )
)

# Define prior
logprior = ds.makelogprior_uniform(gbl.logL.params)

# JIT-compile
logl = jax.jit(gbl.logL)
logp = jax.jit(logprior)
grad_logl = jax.jit(jax.grad(logl))

# Sample from prior
params = ds.sample_uniform(gbl.logL.params)

# Evaluate
print(f"Log-likelihood: {logl(params)}")
print(f"Log-posterior: {logl(params) + logp(params)}")
```

--------------------------------

### Basic Setup for Pulsar Simulation

Source: https://github.com/nanograv/discovery/blob/main/docs/tutorials/simulations.md

Loads a real pulsar's data and sets up a likelihood with specified noise models. Ensure the data directory is correctly set.

```python
import discovery as ds
import jax
import numpy as np
import os

# Find data directory
data_dir = os.path.join(ds.__path__[0], '..', '..', 'data')

# Load a real pulsar to use its TOAs and observing setup
psr = ds.Pulsar.read_feather(
    os.path.join(data_dir, 'v1p1_de440_pint_bipm2019-B1855+09.feather')
)

# Build likelihood with desired noise model
psl = ds.PulsarLikelihood([
    psr.residuals,
    ds.makenoise_measurement(psr, psr.noisedict),
    ds.makegp_ecorr(psr, psr.noisedict),
    ds.makegp_timing(psr, svd=True, variance=1e-20),
    ds.makegp_fourier(psr, ds.powerlaw, 30, name='rednoise')
])
```

--------------------------------

### discovery.signals.getstart

Source: https://github.com/nanograv/discovery/blob/main/docs/api/signals.md

Retrieves the start time of pulsar data.

```APIDOC
## discovery.signals.getstart(psrs)

### Description
Retrieves the start time of the provided pulsar data.

### Parameters
#### Path Parameters
- **psrs** (list) - Required - A list of pulsar objects.

### Method
(Not specified, assumed to be a function call)

### Endpoint
(Not applicable for SDK functions)
```

--------------------------------

### Numpy Docstring Format Example

Source: https://github.com/nanograv/discovery/blob/main/docs/README.md

An example demonstrating the numpy docstring format, which is used for generating API documentation. This format includes sections for parameters, returns, notes, and references.

```python
def example_function(param1, param2):
    """
    Short description.

    Longer description with more details.

    Parameters
    ----------
    param1 : type
        Description of param1.
    param2 : type
        Description of param2.

    Returns
    -------
    return_type
        Description of return value.

    Notes
    -----
    Additional notes or mathematical equations:

    .. math:: 

        E = mc^2

    References
    ----------
    .. [1] Author et al. (Year), "Title", Journal.
    """

```

--------------------------------

### Setup Pulsar Array Model

Source: https://github.com/nanograv/discovery/blob/main/docs/tutorials/simulations.md

Loads pulsar data, defines the time span, and builds individual pulsar likelihoods including noise, Ecorr, timing, and red noise GPs. It then initializes a GlobalLikelihood with a global GW background.

```python
import glob

# Load pulsars
data_pattern = os.path.join(data_dir, 'v1p1_de440_pint_bipm2019-*.feather')
psrs = [ds.Pulsar.read_feather(f) for f in glob.glob(data_pattern)]

Tspan = ds.getspan(psrs)

# Build individual likelihoods
pulsarlikelihoods = []
for psr in psrs:
    psl = ds.PulsarLikelihood([
        psr.residuals,
        ds.makenoise_measurement(psr, psr.noisedict),
        ds.makegp_ecorr(psr, psr.noisedict),
        ds.makegp_timing(psr, svd=True, variance=1e-20),
        ds.makegp_fourier(psr, ds.powerlaw, 30, T=Tspan,
                         name='rednoise')
    ])
    pulsarlikelihoods.append(psl)

# Add GW background
gbl = ds.GlobalLikelihood(
    pulsarlikelihoods,
    globalgp=ds.makegp_fourier_global(
        psrs, ds.powerlaw, ds.hd_orf, 14, T=Tspan,
        name='gw'
    )
)
```

--------------------------------

### Global Likelihood Setup for OS Calculation

Source: https://github.com/nanograv/discovery/blob/main/examples/os_example.ipynb

Configures a global likelihood object with pulsar-specific noise models, timing GPs, and common red noise and gravitational wave signals for OS calculation. Requires `discovery` library and pulsar data.

```python
Tspan = ds.getspan(allpsrs)

gbl = ds.GlobalLikelihood([ds.PulsarLikelihood([psr.residuals,
                                                ds.makenoise_measurement(psr, psr.noisedict),
                                                ds.makegp_ecorr(psr, psr.noisedict),
                                                ds.makegp_timing(psr, svd=True),
                                                ds.makegp_fourier(psr, ds.powerlaw, 30, T=Tspan, name='red_noise'),
                                                ds.makegp_fourier(psr, ds.powerlaw, 14, T=Tspan,
                                                                  common=['gw_log10_A', 'gw_gamma'], name='gw')
                                                ]) for psr in allpsrs])
```

--------------------------------

### Setup CURN Likelihood and Transform

Source: https://github.com/nanograv/discovery/blob/main/examples/flow_example.ipynb

Configures a CURN likelihood and transforms it to ensure all parameters have infinite ranges, suitable for flow-based models.

```python
Tspan = ds.getspan(allpsrs)

m2 = ds.ArrayLikelihood((ds.PulsarLikelihood([psr.residuals,
                                              ds.makenoise_measurement(psr, psr.noisedict),
                                              ds.makegp_ecorr(psr, psr.noisedict),
                                              ds.makegp_timing(psr, svd=True)]) for psr in allpsrs),
                        commongp = ds.makecommongp_fourier(allpsrs, ds.makepowerlaw_crn(14), 30, T=Tspan,
                                                           common=['crn_log10_A', 'crn_gamma'], name='red_noise'))

logx = ds.makelogtransform_uniform(m2.logL)
```

--------------------------------

### Sample Initial Parameter Values

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Samples initial parameter values from their default uniform priors. This provides a starting point for the MCMC sampling process.

```python
p0 = ds.sample_uniform(logl.params)
```

--------------------------------

### Import slq from matfree.trace_estimation

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

This snippet shows the correct import statement for the slq function from the matfree.trace_estimation module. Ensure the module is installed and accessible.

```python
from matfree.trace_estimation import slq
```

--------------------------------

### Run NUTS Sampler and Get Results

Source: https://github.com/nanograv/discovery/blob/main/docs/tutorials/curn_example.ipynb

Runs the NUTS sampler and converts the resulting MCMC chain into a pandas DataFrame.

```python
npsampler.run(jax.random.key(0))
chain = npsampler.to_df()
# chain.to_csv('chain.feather', index=False)
```

--------------------------------

### Conditional Sampling Limitation Example

Source: https://github.com/nanograv/discovery/blob/main/docs/advanced/conditional_sampling.md

This code demonstrates a scenario where conditional sampling is not supported due to the presence of deterministic delays. Attempting to use it will raise a NotImplementedError. A workaround is to build a separate likelihood without delays for conditional sampling.

```python
# This will raise NotImplementedError
signals = [
    psr.residuals,
    ds.makenoise_measurement(psr, noisedict, ecorr=True),
    ds.makegp_timing(psr, svd=True),
    ds.makegp_fourier(psr, ds.powerlaw, 30, name='rn'),
    ds.makedelay(psr, delay_func, name='delay')  # Not supported
]

logl = ds.PulsarLikelihood(signals)
# logl.conditional(params)  # raises NotImplementedError
```

--------------------------------

### Sample initial parameters

Source: https://github.com/nanograv/discovery/blob/main/examples/os_example.ipynb

Generates a set of initial parameters for the OS calculation. It forces the 'gw_gamma' parameter to a specific value (13/3).

```python
p0 = ds.sample_uniform(os.params, priordict={'gw_(.*_)?gamma': [13/3,13/3]})
```

--------------------------------

### Build HTML Documentation

Source: https://github.com/nanograv/discovery/blob/main/docs/README.md

Build the HTML version of the documentation. Navigate to the 'docs' directory before running this command. The output will be in '_build/html/'.

```bash
cd docs
make html
```

--------------------------------

### signals.getresiduals

Source: https://github.com/nanograv/discovery/blob/main/docs/api/index.md

Gets residuals for signals.

```APIDOC
## getresiduals()

### Description
Gets residuals for signals.

### Method
N/A (Function Call)

### Endpoint
N/A

### Parameters
N/A

### Request Example
N/A

### Response
N/A
```

--------------------------------

### Initialize Variational Trainer

Source: https://github.com/nanograv/discovery/blob/main/examples/flow_example.ipynb

Sets up the VariationalFit trainer with the defined flow model and loss function. Includes parameters for learning rate, annealing schedule, and progress display.

```python
trainer = dsf.VariationalFit(dist=flow, loss_fn=loss, multibatch=1,
                             learning_rate=1e-2, annealing_schedule=lambda i: min(1.0, 0.5 + 0.5*i/500),
                             show_progress=True)
```

--------------------------------

### Import necessary libraries for JAX and discovery

Source: https://github.com/nanograv/discovery/blob/main/docs/tutorials/params_example.ipynb

Import JAX for numerical computation and `discovery` for pulsar timing analysis. Ensure 64-bit precision is enabled for JAX.

```python
import jax
jax.config.update('jax_enable_x64', True)

import jax.random
import jax.numpy as jnp

import discovery as ds
from discovery import prior
from discovery.params import Params, make_layout
```

--------------------------------

### signals.getspan

Source: https://github.com/nanograv/discovery/blob/main/docs/api/index.md

Gets the time span of signals.

```APIDOC
## getspan()

### Description
Gets the time span of signals.

### Method
N/A (Function Call)

### Endpoint
N/A

### Parameters
N/A

### Request Example
N/A

### Response
N/A
```

--------------------------------

### Initialize Gradient Log-Likelihood Functions

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Initializes and JIT-compiles gradient log-likelihood functions using different make_logdet methods. These functions are optimized for performance.

```python
glogl1 = jax.jit(jax.grad(model.cglogL(100, 5, 200, make_logdet='CG-Woodbury')))
glogl2 = jax.jit(jax.grad(model.cglogL(100, 5, 200, make_logdet='CG-MDL')))
glogl3 = jax.jit(jax.grad(model.cglogL(100, 5, 200, make_logdet='G-series')))
glogl4 = jax.jit(jax.grad(model.cglogL(100, 5, 200, make_logdet='D-series')))
```

--------------------------------

### Directly Sample OS Distribution

Source: https://github.com/nanograv/discovery/blob/main/examples/os_example.ipynb

Computes a single OS value by drawing a random vector from a normal distribution with degrees of freedom matching the OS count and calculating $1/2 x^T Q x$. The sampling is frozen to a specific parameter set.

```python
oss(ds.jnpnormal(ds.rngkey(42), oss.cnt))
```

--------------------------------

### Get active parameters for likelihood

Source: https://github.com/nanograv/discovery/blob/main/examples/likelihood_example.ipynb

Retrieves the list of currently active parameters for the PulsarLikelihood object.

```python
m.logL.params
```

--------------------------------

### Get pulsar positions

Source: https://github.com/nanograv/discovery/blob/main/examples/os_example.ipynb

Retrieves the positions of the pulsars as a JAX array. This is used for scrambling the OS calculation.

```python
jnp.array(os.pos)
```

--------------------------------

### Sample initial parameters

Source: https://github.com/nanograv/discovery/blob/main/examples/cw_example.ipynb

Samples initial uniform parameters for the time-domain likelihood model.

```python
p0 = ds.sample_uniform(tml.logL.params)
```

--------------------------------

### Get Shape of cfD[0] in Numpyro

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Retrieves the shape of the 'cfD[0]' object, useful for debugging and understanding data dimensions.

```python
cfD[0].shape
```

--------------------------------

### Evaluate Log-Likelihood

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Evaluates the log-likelihood function with the initial parameter values. This checks the model's fit to the data at the starting point.

```python
logl(p0)
```

--------------------------------

### Initialize PulsarLikelihood with nanograv backends

Source: https://github.com/nanograv/discovery/blob/main/examples/likelihood_example.ipynb

Initializes PulsarLikelihood with pulsar residuals and nanograv noise models, allowing for free parameters.

```python
m = ds.PulsarLikelihood([psr.residuals,
                         ds.makenoise_measurement(psr)])
```

--------------------------------

### Initialize Pulsar Array Likelihood

Source: https://github.com/nanograv/discovery/blob/main/examples/likelihood_example.ipynb

Create an ArrayLikelihood object for multiple pulsars, each with its own PulsarLikelihood instance. This setup is for analyzing a collection of pulsars independently.

```python
psrs = allpsrs[:3]
```

```python
m = ds.ArrayLikelihood([ds.PulsarLikelihood([psr.residuals,
                                             ds.makenoise_measurement(psr, psr.noisedict),
                                             ds.makegp_ecorr(psr, psr.noisedict),
                                             ds.makegp_timing(psr, svd=True),
                                             ds.makegp_fourier(psr, ds.powerlaw, components=30, name='rednoise')])
                        for psr in psrs])
```

```python
m.logL.params
```

```python
m.logL(ds.sample_uniform(m.logL.params))
```

--------------------------------

### Vectorized OS Calculation with JAX

Source: https://github.com/nanograv/discovery/blob/main/examples/os_example.ipynb

Creates a JIT-compiled, vectorized version of the `os.os` function using JAX for efficient computation across multiple parameter sets.

```python
vos = jax.jit(jax.vmap(os.os))
```

--------------------------------

### JIT Compile OS Function

Source: https://github.com/nanograv/discovery/blob/main/docs/tutorials/optimal_statistic.md

Use JAX's jit to compile the OS function for improved performance. Specify static arguments if needed for non-default configurations.

```python
import jax

# JIT-compile
os_jit = jax.jit(os_obj.os)

# For using non-default orfa, specify static argument
os_mono_jit = jax.jit(os_obj.os, static_argnums=1)
result = os_mono_jit(params, ds.monopole_orfa)
```

--------------------------------

### Import Conjugate Gradient Solver

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Imports the `solve_cg` function from `jaxopt.linear_solve` for solving linear systems using the conjugate gradient method.

```python
from jaxopt.linear_solve import solve_cg
```

--------------------------------

### Define Power-Law Prior Function

Source: https://github.com/nanograv/discovery/blob/main/docs/components/noise_signals.md

Example of a power-law prior function used with `makegp_fourier`. This function calculates the power spectral density (PSD) at given frequencies.

```python
def powerlaw(f, df, log10_A, gamma):
    """Power-law spectrum.

    Returns PSD at frequencies f.
    """
    return 10**(2*log10_A) * f**(-gamma) * (365.25*86400)**(-gamma+3) / (12*np.pi**2) * df
```

--------------------------------

### Use Custom Prior in Model

Source: https://github.com/nanograv/discovery/blob/main/docs/components/priors_spectra.md

Integrate a custom JAX prior function into a Discovery model using makegp_fourier. This example shows how to use the defined 'broken_powerlaw' function.

```python
# Use in model
gp = ds.makegp_fourier(psr, broken_powerlaw, components=30,
                       name='broken_rn')
```

--------------------------------

### Initialize PulsarLikelihood with noisedict parameters

Source: https://github.com/nanograv/discovery/blob/main/examples/likelihood_example.ipynb

Initializes PulsarLikelihood using parameters directly from the pulsar's noisedict, effectively fixing these parameters.

```python
m = ds.PulsarLikelihood([psr.residuals,
                         ds.makenoise_measurement(psr, psr.noisedict)])
```

--------------------------------

### Compute OS for Subchain Samples

Source: https://github.com/nanograv/discovery/blob/main/examples/os_example.ipynb

Calculates the OS for 1000 sampled parameter sets using the vectorized function. This is timed to show performance.

```python
%time oses = vos(p0s)
```

--------------------------------

### Complete OS Analysis with Significance Testing

Source: https://github.com/nanograv/discovery/blob/main/docs/tutorials/optimal_statistic.md

This snippet demonstrates a full Optimal Statistic analysis workflow. It loads pulsar data, constructs a global likelihood, computes the OS, and performs a scrambling test to determine the p-value.

```python
import discovery as ds
import jax
import numpy as np
import glob
import os

# Load data
data_dir = os.path.join(ds.__path__[0], '..', '..', 'data')
data_pattern = os.path.join(data_dir, 'v1p1_de440_pint_bipm2019-*.feather')
psrs = [ds.Pulsar.read_feather(f) for f in glob.glob(data_pattern)]
Tspan = ds.getspan(psrs)

# Build likelihood
gbl = ds.GlobalLikelihood([
    ds.PulsarLikelihood([
        psr.residuals,
        ds.makenoise_measurement(psr, psr.noisedict),
        ds.makegp_ecorr(psr, psr.noisedict),
        ds.makegp_timing(psr, svd=True),
        ds.makegp_fourier(psr, ds.powerlaw, 30, T=Tspan, name='rednoise'),
        ds.makegp_fourier(psr, ds.powerlaw, 14, T=Tspan,
                         common=['gw_log10_A', 'gw_gamma'], name='gw')
    ]) for psr in psrs
])

# Create OS
os_obj = ds.OS(gbl)

# Sample parameters
params = ds.sample_uniform(gbl.logL.params)

# Compute OS
result = os_obj.os(params)
print(f"OS: {result['os']:.3f}")
print(f"SNR: {result['snr']:.3f}")
print(f"log10_A: {result['log10_A']:.3f}")

# Scrambling test
npsr, ngw = len(psrs), 14
nscrambles = 1000

phi = np.random.uniform(0, 2*np.pi, (nscrambles, npsr))
theta = np.arccos(np.random.uniform(-1, 1, (nscrambles, npsr)))
positions = np.array([
    np.sin(theta) * np.cos(phi),
    np.sin(theta) * np.sin(phi),
    np.cos(theta)
]).transpose(1, 2, 0)

os_scramble = jax.vmap(os_obj.scramble, in_axes=(None, 0, None))
results_scrambled = os_scramble(params, positions, ds.hd_orfa)

p_value = (results_scrambled['snr'] > result['snr']).mean()
print(f"p-value (scrambling): {p_value:.4f}")
```

--------------------------------

### Get Observing Span

Source: https://github.com/nanograv/discovery/blob/main/docs/components/noise_signals.md

Retrieve the observing span for a single pulsar or multiple pulsars using the `getspan` function. Ensure the pulsar object(s) are properly defined before calling.

```python
# Single pulsar
Tspan = ds.getspan(psr)
```

```python
# Multiple pulsars
Tspan = ds.getspan(psrs)
```

--------------------------------

### Set up GlobalLikelihood object

Source: https://github.com/nanograv/discovery/blob/main/examples/os_example.ipynb

Initializes a GlobalLikelihood object for analyzing pulsar data. It configures noise models, Gaussian processes (GP) for timing and red noise, and a common 'gw' GP for gravitational wave signals across multiple pulsars.

```python
Tspan = ds.getspan(psrs)
t0 = ds.getstart(psrs)

gbl = ds.GlobalLikelihood([ds.PulsarLikelihood([psr.residuals,
                                                ds.makenoise_measurement(psr, psr.noisedict),
                                                ds.makegp_ecorr(psr, psr.noisedict),
                                                ds.makegp_timing(psr, svd=True, variable=True),
                                                ds.makegp_fftcov(psr, ds.powerlaw, 61, oversample=6, T=Tspan, name='red_noise'),
                                                ds.makegp_fftcov(psr, ds.brokenpowerlaw, 61, oversample=6, T=Tspan, t0=t0,
                                                                 common=['gw_log10_A', 'gw_gamma', 'gw_log10_fb'], name='gw')
                                                ]) for psr in psrs])
```

--------------------------------

### Add Vectorized Common Red Noise with ArrayLikelihood

Source: https://github.com/nanograv/discovery/blob/main/docs/tutorials/basic_likelihood.md

Constructs an ArrayLikelihood with a common Gaussian Process for red noise. This setup allows for vectorized operations across pulsars, optimizing performance.

```python
curn = ds.ArrayLikelihood(
    pulsarlikelihoods,
    commongp=ds.makecommongp_fourier(
        psrs, ds.makepowerlaw_crn(14), 30, T=Tspan,
        common=['crn_log10_A', 'crn_gamma'],
        name='red_noise'
    )
)
```

--------------------------------

### Compute and Plot OS Correlation

Source: https://github.com/nanograv/discovery/blob/main/examples/os_example.ipynb

Calls the `oscorr` function to get binned OS data and plots it with error bars. It also overlays the theoretical Hellings-Downs curve scaled by the GW amplitude.

```python
iotas, oses, osigs = oscorr(os, p0)
```

```python
pp.errorbar(iotas, oses, yerr=osigs, fmt='.')

a = np.linspace(1e-6, 180)
hd = ost.os(p0)['os'] * ds.hd_orfa(np.cos(a * (np.pi/180.0)))

pp.plot(a, hd)
pp.xlabel('iota'); pp.ylabel(r'$A_\mathrm{gw}^2$');
```

--------------------------------

### Add Red Noise (Powerlaw) to Pulsar Likelihood

Source: https://github.com/nanograv/discovery/blob/main/examples/likelihood_example.ipynb

Configures a PulsarLikelihood object with red noise modeled as a power-law. This snippet is useful for initial model setup when the spectral index (gamma) is unknown and needs to be fitted.

```python
m = ds.PulsarLikelihood([psr.residuals,
                         ds.makenoise_measurement(psr, psr.noisedict),
                         ds.makegp_ecorr(psr, psr.noisedict),
                         ds.makegp_timing(psr, svd=True),
                         ds.makegp_fourier(psr, ds.powerlaw, components=30, name='rednoise')])
```

```python
m.logL.params
```

```python
m.logL(ds.sample_uniform(m.logL.params))
```

--------------------------------

### JIT Compile Log-Likelihood and Gradient

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Applies JAX's JIT compilation to the log-likelihood function and its gradient for significant performance improvements.

```python
jlogl = jax.jit(jax.value_and_grad(logl))
```

--------------------------------

### Create NUTS sampler

Source: https://github.com/nanograv/discovery/blob/main/examples/cw_example.ipynb

Initializes a No-U-Turn Sampler (NUTS) for MCMC sampling using the prepared numpyro model.

```python
sampler = ds_numpyro.makesampler_nuts(flogl)
```

--------------------------------

### Build Likelihood and Get Conditional Distribution

Source: https://github.com/nanograv/discovery/blob/main/docs/advanced/conditional_sampling.md

Builds a PulsarLikelihood object and then uses the `conditional` method to obtain the mean and Cholesky factor of the conditional distribution of GP coefficients. Requires prior and noise dictionaries, and data from a pulsar.

```python
import discovery as ds

# Build likelihood
psr = ds.Pulsar.read_feather('data/v1p1_de440_pint_bipm2019-B1855+09.feather')

signals = [
    psr.residuals,
    ds.makenoise_measurement(psr, noisedict, ecorr=True),
    ds.makegp_timing(psr, svd=True),
    ds.makegp_fourier(psr, ds.powerlaw, 30, name='rn')
]

logl = ds.PulsarLikelihood(signals)

# Sample hyperparameters from priors
params = ds.sample_uniform(logl.params, priordict)

# Get conditional distribution
mu, cf = logl.conditional(params)
```

--------------------------------

### Run MCMC Sampler with Numpyro

Source: https://github.com/nanograv/discovery/blob/main/docs/tutorials/params_example.ipynb

Initialize and run the MCMC sampler using the defined `numpyro_model`. This step performs the sampling process to obtain posterior distributions.

```python
sampler = infer.MCMC(infer.NUTS(numpyro_model),
                     num_warmup=200, num_samples=200)
sampler.run(jax.random.PRNGKey(42))
```

--------------------------------

### Generate Sampled Parameters and Likelihood Values

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Generates 1000 sets of uniformly sampled parameters and calculates various log-likelihood values (JIT-compiled and conditional) for each set. Uses tqdm for progress visualization.

```python
p0s, cls = [], []
for i in tqdm.tqdm(range(1000)):
    p0 = ds.sample_uniform(model.logL.params)
    p0s.append(p0)
    cls.append([jlogl(p0), clogl1(p0), clogl2(p0), clogl3(p0), clogl4(p0)])
```

--------------------------------

### Create Optimal Statistic Object

Source: https://github.com/nanograv/discovery/blob/main/docs/tutorials/optimal_statistic.md

Initializes an OS object by creating a GlobalLikelihood with pulsar data and defining a 'gw' Gaussian Process component for each pulsar. Ensure the 'gw' component includes common parameters like 'gw_log10_A'.

```python
import discovery as ds
import glob
import os

# Find and load data
data_dir = os.path.join(ds.__path__[0], '..', '..', 'data')
data_pattern = os.path.join(data_dir, 'v1p1_de440_pint_bipm2019-*.feather')
psrs = [ds.Pulsar.read_feather(f) for f in glob.glob(data_pattern)]

Tspan = ds.getspan(psrs)

# Build likelihood with 'gw' component
gbl = ds.GlobalLikelihood(
    [ds.PulsarLikelihood([
        psr.residuals,
        ds.makenoise_measurement(psr, psr.noisedict),
        ds.makegp_ecorr(psr, psr.noisedict),
        ds.makegp_timing(psr, svd=True),
        ds.makegp_fourier(psr, ds.powerlaw, 30, T=Tspan,
                         name='rednoise'),
        ds.makegp_fourier(psr, ds.powerlaw, 14, T=Tspan,
                         common=['gw_log10_A', 'gw_gamma'],
                         name='gw')
    ]) for psr in psrs]
)

# Create OS object
os_obj = ds.OS(gbl)
```

--------------------------------

### Evaluate Likelihood with Parameters

Source: https://github.com/nanograv/discovery/blob/main/docs/tutorials/basic_likelihood.md

Create a parameter dictionary and evaluate the likelihood function. JIT-compilation is recommended for performance.

```python
import jax

params = {
    'B1855+09_rednoise_log10_A': -14.0,
    'B1855+09_rednoise_gamma': 4.33
}

# JIT-compile for performance
logl_jit = jax.jit(logl)
logL_value = logl_jit(params)
```

--------------------------------

### Evaluate likelihoods with initial parameters

Source: https://github.com/nanograv/discovery/blob/main/examples/cw_example.ipynb

Calculates the log-likelihood for both the time-domain (tml) and frequency-domain (fml) models using the initial parameters p0.

```python
tml.logL(p0), fml.logL(p0)
```

--------------------------------

### JIT-Compile Sampler for Faster Generations

Source: https://github.com/nanograv/discovery/blob/main/docs/tutorials/simulations.md

Use JAX's JIT compilation to significantly speed up repeated calls to a sampler function. The first call incurs a compilation overhead, while subsequent calls are much faster.

```python
sampler_jit = jax.jit(sampler)

# First call compiles (slower)
key, residuals = sampler_jit(key, params)

# Subsequent calls are fast
for i in range(100):
    key, residuals = sampler_jit(key, params)
```

--------------------------------

### Initialize Monte Carlo Loss Function

Source: https://github.com/nanograv/discovery/blob/main/examples/flow_example.ipynb

Sets up the Evidence Lower Bound (ELBO) loss function for Monte Carlo estimation, specifying the number of samples for parallelism.

```python
num_samples = 1024
loss = dsf.value_and_grad_ElboLoss(logx, num_samples=num_samples)
```

--------------------------------

### Vectorize OS for Parameter Samples

Source: https://github.com/nanograv/discovery/blob/main/docs/tutorials/optimal_statistic.md

Vectorize the OS function using JAX's vmap to compute OS for multiple parameter samples in parallel. Ensure parameters are batched correctly.

```python
import jax.numpy as jnp

# Create batch of parameters (each key has array of values)
nsamples = 100
params_batch = {}
for key in gbl.logL.params:
    param_samples = jnp.array([ds.sample_uniform([key])[key]
                                for _ in range(nsamples)])
    params_batch[key] = param_samples

# Vectorize over parameters (axis 0 of each dict value)
os_vmap = jax.vmap(os_obj.os, in_axes=(0, None))

# Compute for all samples
results = os_vmap(params_batch, ds.hd_orfa)

print(f"SNRs: {results['snr']}")
print(f"Mean SNR: {results['snr'].mean()}")
```

--------------------------------

### Initialize Optimal Statistic object

Source: https://github.com/nanograv/discovery/blob/main/examples/os_example.ipynb

Initializes the Optimal Statistic (OS) object using the configured GlobalLikelihood. This operation is timed to assess performance.

```python
%time os = ds.OS(gbl)
```

--------------------------------

### Initialize PulsarLikelihood with ECORR noise (GP)

Source: https://github.com/nanograv/discovery/blob/main/examples/likelihood_example.ipynb

Initializes PulsarLikelihood including ECORR noise modeled as a Gaussian Process (GP), with free parameters.

```python
m = ds.PulsarLikelihood([psr.residuals,
                         ds.makenoise_measurement(psr),
                         ds.makegp_ecorr(psr)])
```

--------------------------------

### Initialize PulsarLikelihood with GP timing model

Source: https://github.com/nanograv/discovery/blob/main/examples/likelihood_example.ipynb

Initializes PulsarLikelihood including a Gaussian Process (GP) timing model, with parameters fixed by the noisedict.

```python
m = ds.PulsarLikelihood([psr.residuals,
                         ds.makenoise_measurement(psr, psr.noisedict),
                         ds.makegp_ecorr(psr, psr.noisedict),
                         ds.makegp_timing(psr, svd=True)])
```

--------------------------------

### JIT Compile Conditional Log-Likelihood Functions

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Applies JAX JIT compilation to the conditional log-likelihood functions ('cglogL') with specific parameters and different determinant calculation methods.

```python
jlogl = jax.jit(model.logL)
clogl1 = jax.jit(model.cglogL(100, detmatvecs=5, detsamples=200, clip=1e-6, make_logdet='CG-Woodbury'))
clogl2 = jax.jit(model.cglogL(100, detmatvecs=5, detsamples=200, clip=1e-6, make_logdet='CG-MDL'))
clogl3 = jax.jit(model.cglogL(100, detmatvecs=5, detsamples=200, make_logdet='G-series'))
clogl4 = jax.jit(model.cglogL(100, detmatvecs=5, detsamples=200, make_logdet='D-series'))
```

--------------------------------

### Import Discovery library

Source: https://github.com/nanograv/discovery/blob/main/examples/os_example.ipynb

Imports the Discovery library, which provides tools for pulsar data analysis and gravitational wave searches.

```python
import discovery as ds
```

--------------------------------

### Create Measurement Noise Kernel

Source: https://github.com/nanograv/discovery/blob/main/docs/guide/data_model.md

Demonstrates how to create a measurement noise kernel. This kernel represents noise components like white noise or measurement uncertainties.

```python
import discovery as ds

# Create measurement noise kernel
noise = ds.makenoise_measurement(psr, noisedict)

# The kernel can be applied in likelihood computations
# log L ∝ -0.5 * y^T N^{-1} y
```

--------------------------------

### Sample from Uniform Priors

Source: https://github.com/nanograv/discovery/blob/main/docs/tutorials/params_example.ipynb

Draws a parameter dictionary from the default uniform priors for specified parameter names. This is useful for obtaining initial parameter values for sampling or optimization.

```python
p0 = ds.sample_uniform(logL.params)
p0
```

--------------------------------

### Calculate Log-Likelihood with JIT Compilation

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Compiles the 'model.logL' function using JAX's JIT and then calculates the log-likelihood for parameters 'p0'.

```python
logl = jax.jit(model.logL)
logl(p0)
```

--------------------------------

### Sample OS Rho and Sigma

Source: https://github.com/nanograv/discovery/blob/main/examples/os_example.ipynb

Samples the rho and sigma values for the OS from a given parameter set. This is a precursor to direct sampling of the OS distribution.

```python
oss = os.sample_rhosigma(p0)
```

--------------------------------

### Import Discovery library components

Source: https://github.com/nanograv/discovery/blob/main/examples/cw_example.ipynb

Imports core modules from the discovery library, including models and samplers.

```python
import discovery as ds
import discovery.models.nanograv as ds_nanograv
import discovery.samplers.numpyro as ds_numpyro
```

--------------------------------

### Initialize PulsarLikelihood with ECORR noise (GP) from noisedict

Source: https://github.com/nanograv/discovery/blob/main/examples/likelihood_example.ipynb

Initializes PulsarLikelihood with ECORR noise (GP) and parameters fixed by the noisedict.

```python
m = ds.PulsarLikelihood([psr.residuals,
                         ds.makenoise_measurement(psr, psr.noisedict),
                         ds.makegp_ecorr(psr, psr.noisedict)])
```

--------------------------------

### Vectorized Simulations for Parallel Realizations

Source: https://github.com/nanograv/discovery/blob/main/docs/tutorials/simulations.md

Generate multiple simulation realizations in parallel using JAX's vmap for efficient batch processing. This requires preparing batches of input parameters and random keys.

```python
import jax.numpy as jnp

# Create a batch of parameters
params_batch = jax.tree_map(
    lambda x: jnp.repeat(x[None, ...], 100, axis=0),
    params
)

# Create a batch of keys
keys = jax.random.split(key, 100)

# Vectorize sampler over batch
sampler_vmap = jax.vmap(sampler, in_axes=(0, 0))

# Generate 100 realizations at once
keys, residuals_batch = sampler_vmap(keys, params_batch)

print(f"Generated {len(residuals_batch)} realizations")
print(f"Each realization has {len(residuals_batch[0])} pulsars")
```

--------------------------------

### JIT Compile Gradient of Log-Likelihood in Numpyro

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Compiles the gradient of the standard log-likelihood function 'model.logL' using JAX's JIT.

```python
glogl = jax.jit(jax.grad(model.logL))
```

--------------------------------

### Access Matfree SLQ Functionality

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Shows how to access the 'slq' (Stochastic Lanczos Quadrature) functionality within the 'matfree' library.

```python
matfree.slq
```

--------------------------------

### Import JAX Linear Algebra Modules

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Imports JAX's NumPy and SciPy linear algebra modules for efficient numerical computations.

```python
jnpa, jspa = jax.numpy.linalg, jax.scipy.linalg
```

--------------------------------

### Execution Time Output

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Shows the CPU time taken for the execution of the preceding JAXopt `solve_cg` command, indicating the performance of the linear solve operation.

```text
Output:
CPU times: user 159 ms, sys: 11.7 ms, total: 171 ms
Wall time: 168 ms
```

--------------------------------

### Define Uniform Priors for Parameters

Source: https://github.com/nanograv/discovery/blob/main/docs/tutorials/basic_likelihood.md

Use Discovery's default uniform priors or specify custom ranges for parameters. The prior is then JIT-compiled for efficiency.

```python
import discovery as ds
import jax

# Use Discovery's default priors
logprior = ds.makelogprior_uniform(logl.params)

# Or specify custom priors
priordict = {
    'gw_log10_A': [-18, -11],
    'gw_gamma': [0, 7]
}
logprior = ds.makelogprior_uniform(logl.params, priordict)

# JIT-compile
logp = jax.jit(logprior)
```

--------------------------------

### Import JAXopt Linear Solver

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Imports the `solve_normal_cg` function from JAXopt's linear_solve module. This function is used for solving linear systems.

```python
import jaxopt
jaxopt.linear_solve.solve_normal_cg
```

--------------------------------

### discovery.prior.makelogtransform_classic

Source: https://github.com/nanograv/discovery/blob/main/docs/api/prior.md

Creates a log-transformed prior using a classic distribution.

```APIDOC
## discovery.prior.makelogtransform_classic(func, priordict={})

### Description
Generates a log-transformed prior function, applying it to a given function and classic prior definitions.

### Parameters
- **func** (callable) - The function to transform.
- **priordict** (dict, optional) - A dictionary containing prior definitions. Defaults to an empty dictionary.
```

--------------------------------

### Create a population of parameters

Source: https://github.com/nanograv/discovery/blob/main/examples/os_example.ipynb

Generates multiple sets of parameters, simulating parameters that might be obtained from an MCMC run. Each set is sampled uniformly with 'gw_gamma' fixed.

```python
p0s = ds.sample_uniform(os.params, priordict={'gw_(.*_)?gamma': [13/3,13/3]}, n=5)
```

--------------------------------

### Run NUTS sampler

Source: https://github.com/nanograv/discovery/blob/main/examples/cw_example.ipynb

Executes the NUTS sampler with a specified random seed to generate a Markov chain.

```python
sampler.run(jax.random.PRNGKey(42))
```

--------------------------------

### discovery.matrix.make_uind

Source: https://github.com/nanograv/discovery/blob/main/docs/api/matrix.md

Creates an indexed matrix.

```APIDOC
## discovery.matrix.make_uind(U)

### Description
Creates an indexed matrix from a given matrix U.

### Parameters
#### Path Parameters
- **U**: Input matrix.
```

--------------------------------

### discovery.matrix.ConstantKernel

Source: https://github.com/nanograv/discovery/blob/main/docs/api/matrix.md

Represents a constant kernel.

```APIDOC
## class discovery.matrix.ConstantKernel

### Description
A kernel with constant properties.

### Bases
- discovery.matrix.Kernel
```

--------------------------------

### Compute OS for multiple parameter sets

Source: https://github.com/nanograv/discovery/blob/main/examples/os_example.ipynb

Calculates the optimal statistic for all parameter sets in p0s using the vectorized function.

```python
os_vpar(p0s)
```

--------------------------------

### Configure JAX for 64-bit precision

Source: https://github.com/nanograv/discovery/blob/main/examples/cw_example.ipynb

Enables 64-bit floating-point precision in JAX for improved accuracy in scientific computations.

```python
import jax
jax.config.update('jax_enable_x64', True)

import jax.random
import jax.numpy as jnp
```

--------------------------------

### Create Batched Likelihood for Gravitational Waves (GPU)

Source: https://github.com/nanograv/discovery/blob/main/README.md

Sets up an ArrayLikelihood for gravitational wave analysis using a common GP, designed for GPU acceleration. This configuration includes individual pulsar timing and common red noise.

```python
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'))
```

--------------------------------

### Call model.cglogL with JIT compilation

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

This code compiles the model.cglogL function using JAX's jit for performance optimization before calling it with initial parameters p0. This is useful for speeding up repeated computations.

```python
clogl = jax.jit(model.cglogL(1000))
clogl(p0)
```

--------------------------------

### Evaluate JIT-Compiled Log-Likelihood

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Evaluates the JIT-compiled log-likelihood and gradient function with initial parameters. '.block_until_ready()' ensures computation completion.

```python
%time jlogl(p0)[0].block_until_ready()
```

--------------------------------

### signals.makegp_ecorr_simple

Source: https://github.com/nanograv/discovery/blob/main/docs/api/index.md

Creates a simple Ecorr Gaussian Process for signals.

```APIDOC
## makegp_ecorr_simple()

### Description
Creates a simple Ecorr Gaussian Process for signals.

### Method
N/A (Function Call)

### Endpoint
N/A

### Parameters
N/A

### Request Example
N/A

### Response
N/A
```

--------------------------------

### Evaluate Log-Likelihood with Initial Parameters

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Evaluates the log-likelihood function with the sampled initial parameters. '.block_until_ready()' is used to ensure computation completion.

```python
%time logl(p0).block_until_ready()
```

--------------------------------

### Import Matfree and Submodules

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Imports the 'matfree' library and specific submodules 'decomp', 'funm', and 'stochtrace' for use in numerical computations.

```python
import matfree
from matfree import decomp, funm, stochtrace
```

--------------------------------

### Execute matvec2 and display partial result

Source: https://github.com/nanograv/discovery/blob/main/examples/numpyro_example.ipynb

Applies the JIT-compiled 'matvec2' function to the generated matrices and displays the first 10 elements of the second row of the resulting array. This is for inspecting the output of the inverse-matrix-based matvec.

```python
matvec2(invorf, invphi, FtNmF)(FtNmy)[1,:10]
```