Periodogram functionality, periodogram-built period priors, and a new parameterization

.pre-commit-config.yaml

@@ -52,7 +52,9 @@ repos:
     hooks:
       - id: nbstripout
         files: .ipynb
-        exclude: ^docs/tutorials/case-studies/population-binary-fraction\.ipynb$
+        # Case studies ship with stored outputs: docs/conf.py sets
+        # nb_execution_mode = "off", so a stripped notebook renders with no figures.
+        exclude: ^docs/tutorials/case-studies/.*\.ipynb$
   - repo: local
     hooks:
       - id: ty

conftest.py

@@ -19,3 +19,5 @@
 install_import_hook("harv.kepler", "beartype.beartype")
 install_import_hook("harv.likelihood", "beartype.beartype")
 install_import_hook("harv.models", "beartype.beartype")
+install_import_hook("harv.periodogram", "beartype.beartype")
+install_import_hook("harv.stats.grid_density", "beartype.beartype")

docs/sharp-bits.md

@@ -96,3 +96,25 @@ This rewrites the samples into the equivalent convention
 - `arg_peri` wrapped into $\[0, 2\\pi)$
 
 without changing the physical orbit.
+
+## Periodograms batch over sources only with a fixed frequency grid
+
+`harv.periodogram.periodogram` works under `jax.jit` and `jax.vmap`, so you can compute
+periodograms for a whole population of sources in one traced call:
+
+```python
+batched = jax.tree.map(lambda *xs: jnp.stack(xs), *sources)
+grid = hp.frequency_grid(t_span=Q(1000, "day"), period_min=Q(5, "day"), n_grid=1024)
+results = jax.jit(jax.vmap(lambda d: hp.periodogram(d, grid, prior=prior)))(batched)
+```
+
+Two things have to hold. The grid must be shape-fixed, which means passing an explicit
+`frequency_grid` as above, or giving `period_min`, `period_max`, and `n_grid` together.
+Letting the grid size come from each source's own time baseline (i.e., omitting
+`n_grid`) cannot be traced, since the number of grid points is then an array shape that
+JAX needs to know at trace time. The stacked sources also need a common number of
+observations, for the same reason; sources with different epoch counts will retrace.
+
+Both conditions point the same way as the advice in `frequency_grid`: use one shared
+grid across a population so the results are directly comparable and the sampler compiles
+once.

docs/tutorials/case-studies/index.md

@@ -6,6 +6,7 @@ End-to-end analyses of real systems with harv.
 :maxdepth: 1
 
 population-binary-fraction
+kepler-periodogram
 <!-- HD285968
 gaia-bh3 -->
 ::::

docs/tutorials/data/Gaia-DR4-preview.ipynb

@@ -0,0 +1,245 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pathlib\n",
+    "\n",
+    "import astropy.table as at\n",
+    "import astropy.units as u\n",
+    "import numpy as np\n",
+    "from astropy.time import Time\n",
+    "from astropy.units import Quantity as Q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "epoch_astrometry_file = pathlib.Path(\n",
+    "    \"~/data/Gaia/DR4/gaia-dr4-prerelease-epoch-astrometry_2026-06-26/GAIA_DR4_PRERELEASE_EPOCH_ASTROMETRY_RAW.xml\"\n",
+    ")\n",
+    "table = at.Table.read(epoch_astrometry_file, format=\"votable\")\n",
+    "first_mask = ~table[\"source_id\"].mask\n",
+    "\n",
+    "table = table[first_mask]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "TCB_REFERENCE_EPOCH = Time(\"2010-01-01T00:00:00\", format=\"isot\", scale=\"tcb\")\n",
+    "DR4_REFERENCE_EPOCH = Time(\"2017.5\", format=\"jyear\", scale=\"tcb\")\n",
+    "\n",
+    "table[\"relative_time_day\"] = (\n",
+    "    TCB_REFERENCE_EPOCH.jyear\n",
+    "    + table[\"obs_time_tcb\"] * (u.nanosecond.to(u.year))\n",
+    "    + table[\"obs_time_bary_corr\"] * (u.nanosecond.to(u.year))\n",
+    "    - DR4_REFERENCE_EPOCH.jyear\n",
+    ") * u.year.to(u.day)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "source_ids = np.unique(table[\"source_id\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "source_id_to_name = {\n",
+    "    1457486023639239296: \"Gaia 4\",\n",
+    "    4318465066420528000: \"Gaia BH3\",\n",
+    "    3937211745905473024: \"HD 114762\",\n",
+    "    435469040545191680: \"plx G14\",\n",
+    "    3926186255616949504: \"plx G19\",\n",
+    "    4181040337841125632: \"plx G9\",\n",
+    "    2237987199365376: \"qso 1\",\n",
+    "    10973744521070720: \"qso 2\",\n",
+    "    60730287810150016: \"qso 3\",\n",
+    "}\n",
+    "\n",
+    "# parallax in mas, period in days\n",
+    "truths = {\n",
+    "    \"Gaia 4\": {\"parallax\": 13.643, \"period\": 571.3},\n",
+    "    \"Gaia BH3\": {\"parallax\": 1.675, \"period\": 4194.7},  # from astrometric solution only\n",
+    "    \"HD 114762\": {\"parallax\": 24.855, \"period\": 83.9},\n",
+    "    \"plx G9\": {\"parallax\": 1.028},\n",
+    "    \"plx G14\": {\"parallax\": 1.021},\n",
+    "    \"plx G19\": {\"parallax\": 0.754},\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5",
+   "metadata": {},
+   "source": [
+    "Collapse CCD measurements into a single measurement per transit:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def get_data(rows):\n",
+    "    preproc_time_day = np.array(\n",
+    "        [\n",
+    "            np.average(\n",
+    "                row[\"relative_time_day\"][1:][row[\"used_by_agis_al\"][1:].filled(False)]\n",
+    "            )\n",
+    "            for row in rows\n",
+    "        ]\n",
+    "    )\n",
+    "    preproc_pos_al = np.array(\n",
+    "        [\n",
+    "            np.average(\n",
+    "                row[\"centroid_pos_al\"][1:][row[\"used_by_agis_al\"][1:].filled(False)],\n",
+    "                weights=1\n",
+    "                / row[\"centroid_pos_error_al\"][1:][\n",
+    "                    row[\"used_by_agis_al\"][1:].filled(False)\n",
+    "                ]\n",
+    "                ** 2,\n",
+    "            )\n",
+    "            for row in rows\n",
+    "        ]\n",
+    "    )\n",
+    "    preproc_pos_al_err = np.array(\n",
+    "        [\n",
+    "            np.sqrt(\n",
+    "                1\n",
+    "                / np.sum(\n",
+    "                    1\n",
+    "                    / row[\"centroid_pos_error_al\"][1:][\n",
+    "                        row[\"used_by_agis_al\"][1:].filled(False)\n",
+    "                    ]\n",
+    "                    ** 2\n",
+    "                )\n",
+    "            )\n",
+    "            for row in rows\n",
+    "        ]\n",
+    "    )\n",
+    "\n",
+    "    # preproc_pos_al_err = np.sqrt(\n",
+    "    #     preproc_pos_al_err**2 + 0.1**2\n",
+    "    # )\n",
+    "\n",
+    "    excess_noise = np.array(rows[\"agis_source_excess_noise\"].filled(np.nan))\n",
+    "\n",
+    "    preproc_scan_angle = np.array(\n",
+    "        [\n",
+    "            np.average(\n",
+    "                row[\"scan_pos_angle\"][1:][row[\"used_by_agis_al\"][1:].filled(False)]\n",
+    "            )\n",
+    "            for row in rows\n",
+    "        ]\n",
+    "    )\n",
+    "\n",
+    "    preproc_parallax_factor = np.array(rows[\"parallax_factor_al\"])\n",
+    "\n",
+    "    _mask = (\n",
+    "        np.isfinite(preproc_time_day)\n",
+    "        & np.isfinite(preproc_pos_al)\n",
+    "        & np.isfinite(preproc_pos_al_err)\n",
+    "        & np.isfinite(preproc_scan_angle)\n",
+    "        & np.isfinite(preproc_parallax_factor)\n",
+    "    )\n",
+    "    print(_mask.sum(), len(_mask))\n",
+    "\n",
+    "    return {\n",
+    "        \"relative_time\": Q(preproc_time_day[_mask].astype(\"f8\"), \"day\"),\n",
+    "        \"pos_al\": Q(preproc_pos_al[_mask].astype(\"f8\"), \"mas\"),\n",
+    "        \"pos_al_err\": Q(preproc_pos_al_err[_mask].astype(\"f8\"), \"mas\"),\n",
+    "        \"scan_angle\": Q(preproc_scan_angle[_mask].astype(\"f8\"), \"deg\"),\n",
+    "        \"parallax_factor\": preproc_parallax_factor[_mask].astype(\"f8\"),\n",
+    "        \"excess_noise\": Q(excess_noise[_mask].astype(\"f8\"), \"mas\"),\n",
+    "    }"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "this_path = pathlib.Path(\".\").resolve()\n",
+    "dr4_preview_path = this_path / \"gaia-dr4-prerelease\"\n",
+    "dr4_preview_path.mkdir(exist_ok=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for source_id in source_ids:\n",
+    "    rows = table[table[\"source_id\"] == source_id]\n",
+    "    name = source_id_to_name.get(source_id, str(source_id))\n",
+    "    print(source_id, name)\n",
+    "\n",
+    "    data = at.QTable(get_data(rows))\n",
+    "\n",
+    "    # put parallax, period into table metadata\n",
+    "    data.meta[\"parallax_mas\"] = truths.get(name, {}).get(\"parallax\", np.nan)\n",
+    "    data.meta[\"period_day\"] = truths.get(name, {}).get(\"period\", np.nan)\n",
+    "\n",
+    "    data.write(dr4_preview_path / f\"{name.replace(' ', '-')}.ecsv\", overwrite=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "harv (3.12.10)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}