Metadata-Version: 2.4
Name: qpe-fit
Version: 0.2.1
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy==2.3.5
Requires-Dist: scipy
Requires-Dist: ultranest==4.4.0
Requires-Dist: emcee==3.1.6
Requires-Dist: corner==2.2.3
Requires-Dist: matplotlib==3.10.8
Requires-Dist: kerrgeopy
Requires-Dist: h5py==3.15.1
Provides-Extra: gpu
Requires-Dist: cupy-cuda12x; extra == "gpu"
Dynamic: license-file

# QPE-FIT (Fast Inference with Timing) v0.2.0

This repository contains the code for `QPE-FIT`, a Bayesian parameter estimation tool to fit Quasi-Periodic Eruption (QPE) timings. Under the model assumption that QPEs are generated by an extreme mass-ratio inspiral (EMRI) orbiter colliding with the accretion disk around a massive black hole, QPE timings encode information about the EMRI-disk system, and in principle, can be used to learn their underlying parameters.

`QPE-FIT` provides two forward models, selected at the command line with `qpe-fit --model {geodesic,analytic}`:

- **`analytic`** (default): a fast closed-form nodal + apsidal eruption-timing model, fit directly in time space with [emcee](https://emcee.readthedocs.io/) (`--sampler mcmc`). It is ~10⁴× cheaper than the geodesic integration and self-seeds (period scan → MAP → covariance), so no initial guess is required.
- **`geodesic`**: the full post-Newtonian Kerr EMRI-disk trajectory, fit by [nested sampling](https://en.wikipedia.org/wiki/Nested_sampling_algorithm) with [UltraNest](https://johannesbuchner.github.io/UltraNest) (`--sampler ns`).

The [original paper](https://arxiv.org/abs/2508.20162) used Markov Chain Monte Carlo on the geodesic model.

Please email joheen@mit.edu with any requests, bug reports, or comments!

## Installation

```bash
conda create --name myenv python=3.12
conda activate myenv
conda install pip

# CPU only
pip install qpe-fit

# With GPU support (requires CUDA)
pip install qpe-fit[gpu]
```
The current version of QPE-FIT was only tested on CuPy/CUDA v12, which supports Python 3.8-3.12 (hence the pickiness regarding python version above). If you need a different CuPy version, ping me and I'll try to update the package accordingly.

## Dependencies

- [UltraNest](https://johannesbuchner.github.io/UltraNest) : Nested-sampling inference (geodesic model)
- [emcee](https://emcee.readthedocs.io/) : Ensemble-MCMC inference (analytic model)
- [kerrgeopy](https://kerrgeopy.readthedocs.io/) : Exact Kerr geodesic orbits (mock-data generation)
- [SciPy](https://scipy.org/) : MAP optimization for the analytic auto-seed
- [CuPy](https://cupy.dev/) : GPU-accelerated trajectory/residual computations (optional)
- [NumPy](https://numpy.org/) : Numerical computations
- [Corner](https://corner.readthedocs.io/) : Posterior distribution plots
- [h5py](https://www.h5py.org/) : HDF5 file I/O

## Command Overview

### `qpe-gen`
Generate mock QPE timings for a given set of EMRI/disk parameters. A non-evolving orbit uses the **exact Kerr geodesic** ([kerrgeopy](https://kerrgeopy.readthedocs.io/)); when secular evolution is switched on (`adot`/`edot` ≠ 0) it falls back to the adiabatic post-Newtonian trajectory (a `StableOrbit` has fixed constants of motion and cannot drift). `--pn` forces the PN trajectory in all cases.

**Usage:**
```bash
qpe-gen --params params.json --windows windows.txt --output-timings timings.txt
```

**Arguments:**
- `--params`: JSON file with orbital parameters
- `--windows`: Observation windows file (start/stop pairs in seconds)
- `--output-timings`: Output file for QPE timings
- `--dt`: Time step resolution (default: 10s)
- `--one-per-orbit`: Flag to generate only one QPE per orbit (not two)
- `--pn`: Force the PN trajectory even for a non-evolving orbit (default: exact geodesic when `adot = edot = 0`)

### `qpe-fit`
Fit QPE timings to estimate system parameters. The model and sampler are chosen with `--model {geodesic,analytic}` and `--sampler {ns,mcmc}` (default `--model analytic --sampler mcmc`; the `geodesic` model defaults to `--sampler ns`). All remaining options are forwarded to the selected backend, e.g. `qpe-fit --model geodesic --sampler ns --help`.

#### `--model geodesic --sampler ns`
Nested-sampling inference on the full post-Newtonian Kerr EMRI-disk trajectory.

**Usage:**
```bash
qpe-fit --model geodesic --gpu    # geodesic model + nested sampling
```

**Arguments:**
| Argument | Default | Description |
|----------|---------|-------------|
| `--output` | `output` | Sampling output directory |
| `--timings` | `timings.txt` | .txt file containing QPE timings (one per line, in seconds) |
| `--windows` | `windows.txt` | .txt file containing observation windows (one start-stop pair per line separated by a space, in seconds) |
| `--errors` | `errors.txt` | .txt file containing QPE timing errors (one per line, in seconds) |
| `--priors` | `priors.json` | .json file containing sampling priors |
| `--dt` | `10.0` | Time step size for likelihood evaluations (seconds) |
| `--gpu` | `False` | Flag to use GPU-accelerated likelihood evaluation |
| `--stepsampler` | `slice` | Step sampler: `none`, `slice`, `harm`, `rwalk` |
| `--direction` | `region` | Direction function: `region`, `random`, `mixture`, `cube` |
| `--region` | `simple` | Region class: `mlfriends`, `simple`, `ellipsoid` |
| `--popsize` | `256` | Number of walkers for step sampling |
| `--nsteps` | `256` | Number of steps for step sampling |
| `--nlive` | `600` | Minimum number of live points |
| `--dkl` | `0.5` | Target posterior uncertainty (KL divergence in nats) |
| `--frac-remain` | `0.01` | Integrate until this fraction of the integral remains |
| `--min-ess` | `400` | Minimum effective sample size |
| `--one-per-orbit` | `False` | Flag to generate only one QPE per orbit (not two) |

**Inferred Parameters:**
The script fits for up to 15 parameters. The disk-geometry parameters (`theta_d`, `P_d`, `phi_d`, `r_warp`) and the secular-evolution parameters (`adot`, `edot`) are optional: each is sampled only if it appears in `priors.json`, and omitting it recovers the simpler model exactly (e.g. `r_warp = 0` is a flat, unwarped disk; `adot = edot = 0` is a non-evolving orbit).
- `sma`: EMRI semimajor axis (gravitational radii $R_g$)
- `e`: EMRI eccentricity
- `incl`: EMRI inclination (degrees)
- `phi_r0`: Initial EMRI radial phase (radians)
- `phi_theta0`: Initial EMRI polar phase (radians)
- `phi_phi0`: Initial EMRI azimuthal phase (radians)
- `spin`: MBH spin parameter (0-0.998)
- `logMbh`: log(MBH mass / solar masses)
- `theta_obs`: Observer viewing angle (radians)
- `theta_d`: Disk inclination (degrees)
- `P_d`: Disk precession period (multiples of EMRI orbital period)
- `phi_d`: Initial disk azimuthal phase (radians)
- `r_warp`: Disk warp radius (gravitational radii $R_g$). Sets the radial scale of a Bardeen–Petterson-like warp (a radius-dependent twist + alignment layered on the rigid precession); `0` gives a flat, unwarped disk.
- `adot`: Rate of change of the semimajor axis $da/dt$ ($R_g$/s); `0` = no semimajor-axis evolution.
- `edot`: Rate of change of the eccentricity $de/dt$ (1/s); `0` = no eccentricity evolution.

#### `--model analytic --sampler mcmc` (default)
Fit the fast closed-form nodal + apsidal eruption-timing model directly in time space with `emcee`. The recurrence, precession periods, and phases are sampled directly, and the integer cycle number of each eruption is fixed by phase connection. The fit **self-seeds** (period scan → multi-start MAP → Laplace-covariance walker ball), so no initial guess is required.

**Usage:**
```bash
qpe-fit --priors priors_analytic.json    # analytic + mcmc is the default
```

**Key arguments:** `--timings`, `--errors`, `--windows`, `--priors`, `--output` (default `analytic_fit`), `--nwalkers` (64), `--nsteps` (20000), `--burnin-frac` (0.5), `--thin` (10), `--sys-err` (fixed systematic error [s]; ignored if `log_fsys` is sampled), `--scan-pts` / `--n-nod-scan` / `--map-starts` (auto-seed controls), `--init` (optional JSON initial guess).

**Sampled parameters** (set bounds/`type` in the priors JSON, same schema as the geodesic model):
- `t0`: reference eruption time (s)
- `P`: mean eruption recurrence (s)
- `e`: eccentricity
- `N_aps`: argument-of-pericentre precession period (orbits)
- `omega0`: initial argument of pericentre (radians)
- `beta`: local disk tilt (degrees)
- `incl`: orbital inclination (degrees)
- `N_nod`: nodal (Lense–Thirring) precession period (orbits)
- `psi0`: initial nodal phase (radians)
- `log_fsys` *(optional)*: log₁₀ of a fitted white-noise systematic $\sigma_\mathrm{sys} = f_\mathrm{sys}\,P$; include it in the priors to *fit* the systematic instead of fixing it with `--sys-err`.
- `chi` *(optional; a constraint, not a sampled parameter)*: a joint Kerr-bound prior on the weak-field spin $\chi(N_\mathrm{aps}, N_\mathrm{nod})$, e.g. `"chi": {"bounds": [0, 1]}`.

Outputs (in `analytic_fit/`): `samples.npy`, `logprob.npy`, `summary.json`, `physical_summary.json` (weak-field $M$, spin, semimajor axis), plus corner, O-C, and timing-residual plots.

## Example Workflow

```bash
# Step 1: Create parameter file (params.json)
cat > params.json << EOF
{
    "sma": 100, "e": 0.15, "incl": 45,
    "phi_r0": 0, "phi_theta0": 1, "phi_phi0": 2,
    "spin": 0.9, "logMbh": 5.8, "theta_obs": 0.5,
    "theta_d": 15, "P_d": 500, "phi_d": 0
}
EOF

# Step 2: Create observation windows (two 100ks windows)
echo "0 100000
500000 600000" > windows.txt

# Step 3: Generate mock QPE timings (exact Kerr geodesic, since the orbit is non-evolving)
qpe-gen --params params.json --windows windows.txt --output-timings timings.txt

# Step 4: Create mock timing uncertainties (100s errors)
awk '{print 100}' timings.txt > errors.txt

# Step 5: Create prior file (priors.json)
cat > priors.json << EOF
{
    "sma": {"bounds": [50, 200], "wrapped": false},
    "e": {"bounds": [0.0, 0.3], "wrapped": false},
    "incl": {"bounds": [0, 90], "wrapped": false},
    "phi_r0": {"bounds": [0, 6.283], "wrapped": true},
    "phi_theta0": {"bounds": [0, 6.283], "wrapped": true},
    "phi_phi0": {"bounds": [0, 6.283], "wrapped": true},
    "spin": {"bounds": [0.5, 0.998], "wrapped": false},
    "logMbh": {"bounds": [5.5, 6.0], "wrapped": false},
    "theta_obs": {"bounds": [0, 3.142], "wrapped": true},
    "theta_d": {"bounds": [0, 30], "wrapped": false},
    "P_d": {"bounds": [300, 700], "wrapped": false},
    "phi_d": {"bounds": [0, 6.283], "wrapped": true}
}
EOF

# Step 6: Run inference on the geodesic model with nested sampling (with GPU)
qpe-fit --model geodesic --gpu

# Results will be in the output/ directory
```
Of course, if you have _real_ QPE timing data you want to fit, you can skip Steps 1-4.
From here, you can follow the (excellent and informative) [UltraNest tutorials](https://johannesbuchner.github.io/UltraNest/using-ultranest.html) to interpret the results!

The fast **analytic** model is the **default**, so to fit the same `timings.txt`/`errors.txt`/`windows.txt` you just need an analytic priors file — it self-seeds, no initial guess required:
```bash
qpe-fit --priors priors_analytic.json
```

## Citation

If you use QPE-FIT in your research, please cite [J. Chakraborty & L.V. Drummond et al., ApJ, 992, 120 (2025)](https://ui.adsabs.harvard.edu/abs/2025ApJ...992..120C/abstract) and the sampler you used ([UltraNest](https://ui.adsabs.harvard.edu/abs/2021JOSS....6.3001B/abstract) or [emcee](https://ui.adsabs.harvard.edu/abs/2013PASP..125..306F/abstract))!
