Skip to content

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

Merged
adrn merged 38 commits into
mainfrom
periodogram
Sep 14, 2026
Merged

adrn merged 38 commits into
mainfrom
periodogram

Conversation

@adrn

@adrn adrn commented Jul 31, 2026 •

Copy link
Copy Markdown
Owner

Adds harv.periodogram, which builds adds support around computing custom keplerian periodograms for RV and astrometry data. The periodogram statistic is a log Bayes factor: marginal likelihood of a Fourier model at trial period P, minus that of the null model (null = constant RV for RV data, or sky position, parallax, + proper motion for astrometry, i.e. 5 parameter solution). So for astrometry, parallax/proper-motion/scan-law power cancels out of the periodogram instead of showing up as peaks at a year.

To make this work, there are two new parameterizations: FourierRV and FourierGaiaAstrometry: a truncated Fourier series with linear amplitudes, so period is the only nonlinear parameter and there's no Kepler solve. n_terms=0 is the null model.

Two prior builders ingest the periodogram to build period priors: tempered_period_prior (exp(beta*delta) plus floor) and peak_period_prior (equal mass per peak, amplitude-agnostic).

For hierarchical use, attach_ln_pint records the per-sample interim prior density as an extra Samples column so it can be divided out later.

Separately: the rejection sampler now warns when the evidence ESS is < 3, and Samples.acceptance_diagnostics() reports some statistics. This came up while testing the above functionality.

One thing left open: what amplitude scale to recommend for the Fourier priors. A data-RMS scale under-estimates the amplitude for partial arcs and pulls the peak short. Needs a comparison against converged posteriors before picking a default, so for now the user has to choose.

TODO:

  • Resolve guidance for default amplitude scale priors
  • Update customize-prior tutorial and (maybe) add new tutorial docs/tutorials/rv/6-periodogram-prior.ipynb
  • Case study: Kepler periodograms. Compare kepmodel to this one - profile vs. marginalize likelihood

@adrn adrn mentioned this pull request Aug 31, 2026
4 tasks
Comment thread src/harv/samplers/samples.py

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Prior normalization and domain handling contain correctness bugs, and several documentation targets and examples are broken.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds Fourier-model periodograms, periodogram-informed priors, persistence and hierarchical bookkeeping, plus rejection-sampling diagnostics.

Changes:

  • Introduces RV/Gaia Fourier parameterizations and periodogram APIs.
  • Adds tempered/peak priors, HDF5 persistence, and interim-prior metadata.
  • Expands diagnostics, documentation, and automated tests.
File summaries
File Description
tests/unit/test_samplers/test_acceptance_diagnostics.py Tests rejection diagnostics and warnings.
tests/unit/periodogram/test_priors_builders.py Tests prior builders.
tests/unit/periodogram/test_prior_interface.py Tests prior validation and extensions.
tests/unit/periodogram/test_periodogram_rv.py Tests RV periodograms.
tests/unit/periodogram/test_periodogram_gaia.py Tests Gaia and joint periodograms.
tests/unit/periodogram/test_io.py Tests prior persistence.
tests/unit/periodogram/test_grid.py Tests frequency grids and harmonic caps.
tests/unit/periodogram/test_distribution.py Tests grid-density distribution.
tests/unit/periodogram/test_attach_ln_pint.py Tests interim-density attachment.
tests/unit/periodogram/__init__.py Defines the test package.
tests/unit/models/test_fourier.py Tests Fourier parameterizations.
tests/integration/test_periodogram_prior.py Tests the end-to-end workflow.
src/harv/samplers/samples.py Adds acceptance diagnostics and Fourier compatibility.
src/harv/samplers/rejection.py Emits under-resolution warnings.
src/harv/periodogram/priors.py Builds periodogram-informed priors.
src/harv/periodogram/io.py Persists period priors.
src/harv/periodogram/grid.py Constructs frequency grids.
src/harv/periodogram/distribution.py Implements LogGridDensity.
src/harv/periodogram/core.py Implements periodogram computation.
src/harv/periodogram/__init__.py Exports periodogram APIs.
src/harv/models/rv.py Supports Fourier RV models.
src/harv/models/parameterizations/fourier.py Defines Fourier parameterizations.
src/harv/models/parameterizations/__init__.py Exports Fourier parameterizations.
src/harv/models/astrometry.py Supports Fourier Gaia models.
src/harv/models/__init__.py Exports new model types.
src/harv/custom_types.py Adds frequency quantity types.
src/harv/__init__.py Exposes the periodogram module.
pyproject.toml Configures warning filtering.
docs/tutorials/rv/index.md Adds a tutorial index entry.
docs/tutorials/rv/1-customize-prior.ipynb Adds periodogram-prior guidance.
docs/spec.md Documents the new APIs and diagnostics.
conftest.py Enables runtime type checking in tests.
Review details

Suppressed comments (1)

src/harv/periodogram/priors.py:93

  • When the requested prior domain lies entirely above the periodogram grid, this transition knot precedes u_lo, producing a non-monotonic knot array. Only insert it when it is inside the requested domain.
    if u_hi > u[-1] + 2 * eps:
        knots.append(np.array([u[-1] + eps]))
        vals.append(np.array([0.0]))
  • Files reviewed: 31/32 changed files
  • Comments generated: 12
  • Review effort level: Balanced

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment thread src/harv/periodogram/priors.py Outdated
Comment thread src/harv/periodogram/priors.py Outdated
Comment thread src/harv/periodogram/priors.py Outdated
Comment thread src/harv/periodogram/io.py Outdated
Comment thread src/harv/periodogram/core.py
Comment thread src/harv/periodogram/grid.py
Comment thread docs/tutorials/rv/index.md Outdated
Comment thread docs/spec.md Outdated
Comment thread docs/spec.md Outdated
Comment thread docs/tutorials/rv/1-customize-prior.ipynb Outdated
adrn and others added 24 commits September 13, 2026 22:06
All 18 ty errors were in code new on this branch:

- LogGridDensity: numpyro declares arg_constraints / reparametrized_params /
  pytree_data_fields as instance-level attributes, so ClassVar overrides are
  LSP violations; annotate them plainly and silence ruff's RUF012 instead.
  Widen log_prob / log_prob_ln / cdf / icdf to ArrayLike to match
  Distribution. sample() keeps key: jax.Array with a narrow ty: ignore --
  numpyro's jax.dtypes.prng_key is a dtype class, not the runtime type of a
  PRNG key, and beartype rejects real keys against it.
- low / high read the knots directly; Distribution._support is typed
  Constraint | None upstream, so its bounds are unresolved.
- periodogram(): accumulate the per-dataset terms in a list and sum, instead
  of None-initialized running totals; cast .period to NTime.
- _write_group(): take the already-narrowed LogGridDensity and unit rather
  than re-deriving them from a QuantityDistribution field typed Distribution.
- frequency_grid(): branch on `data is not None` so _data_t_span narrows.
- RVModel / GaiaAstrometryModel: ty: ignore for .eccentricity / .sky_orbit,
  which the Fourier parameterizations joined those unions without defining.

Annotations and narrowing only -- no behavior change. Test results are
identical before and after (same 32 pre-existing failures, 1084 passed).

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Four cleanups, spec updated for each (docs/spec.md is authoritative):

1. LogGridDensity moves to harv.stats.grid_density, next to the other
   numpyro extensions, and is exported as harv.stats.LogGridDensity. It gets
   its own module rather than joining numpyro_ext.py so it stays outside the
   vendored-code lint exemption: ruff's per-file ignore is narrowed from
   src/harv/stats/*.py to the two vendored modules by name, and conftest.py
   adds the beartype hook for harv.stats.grid_density. Spec grows a
   "Statistical utilities (harv.stats)" section, which the periodogram
   chapter now cross-references.

2. harv.periodogram.io is deleted. Nothing outside its own tests ever called
   save_period_prior / load_period_prior: the per-sample interim-prior
   density that population reweighting actually reads already rides along in
   Samples via attach_interim_period_prior -> to_hdf5, so the HDF5 prior-spec
   group was a second, unused representation of the same information.

3. priors.py naming and explanation:
   - "pint" (prior, interim -- and a name collision with the pint units
     package) is gone. attach_ln_pint -> attach_interim_period_prior, column
     ln_pint_period -> ln_interim_period_prior, LN_PINT_PERIOD_KEY ->
     LN_INTERIM_PERIOD_PRIOR_KEY. The docstring now says what an interim
     prior is and why it has to be divided back out.
   - _EDGE_EPS_FACTOR is no longer a module-level magic number: it is a local
     in _assemble_knots with a comment explaining that knots are joined by
     straight lines, so extending the domain past the grid without a
     transition knot would ramp Delta across the whole extension -- and that
     the offset only has to be far below the knot spacing.
   - Locals renamed out of one-letter shorthand (u -> ln_period, pw ->
     peak_width, rho -> density) and the module TODO resolved.

4. MIN_EVIDENCE_ESS is no longer a hard-coded bar. It stays as the single
   documented default (with the reasoning: ESS = 3 is where the MC error on
   logZ_int reaches ~0.6 nats), but the threshold is now user-controlled via
   the static RejectionSampler.min_evidence_ess field and
   Samples.acceptance_diagnostics(min_evidence_ess=...). 0.0 silences the
   check, inf always flags, and the warning text and the diagnostics dict
   both report the bar that was applied.

Full suite: 1114 passed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
5e18cd2 raised frequency_grid's samples_per_peak default to 8, but
periodogram() kept forwarding its own default of 5, so the new value only
took effect for direct frequency_grid() calls. Align periodogram(), the
frequency_grid docstring, and the two spec signature blocks on 8.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1. The prior domain may no longer reach past the periodogram grid.

   _assemble_knots treated the un-evaluated region outside the grid as
   Delta = 0. Since exp(beta * 0) outranks a periodogram that is negative
   everywhere -- the ordinary no-signal case, where the Occam factor beats the
   fit -- a source with no detectable signal got an interim period prior
   concentrated on periods nobody looked at: measured 94% of the mass in
   1000-20000 d for a flat-RV source whose grid stopped at 1000 d.

   The same machinery appended its transition knots unconditionally, so a
   domain lying entirely outside the grid produced a non-monotonic knot array
   and LogGridDensity died with "ln_grid must be strictly increasing".

   Both go away together: the domain now raises unless it lies within the grid,
   and the eps/transition-knot block is deleted. A strict subset is still
   supported and renormalized. Bounds are compared with a 1e-6 ln-period
   tolerance and then clipped, because the natural call passes back the same
   period_min/period_max that built the grid and those round-trip through 1/f
   and log() -- landing within ~1e-7 in float32, exact in x64. The same
   tolerance keeps a grid knot sitting a hair from an endpoint from duplicating
   it, which would trip the same monotonicity check.

2. The base model is only period-independent when its own priors are.

   lnL_base was always evaluated once, at the longest trial period, and
   subtracted from every grid point. With a LinearPriorCallable on a base
   column (e.g. PeriodDependentKPrior on v_sys -- an input the API accepts) the
   baseline genuinely varies: 1.3 nats between P=10 d and P=1000 d, silently
   tilting every Delta. Detect callable base-column priors and evaluate the
   base model across the grid in that case only; the single evaluation stays
   the fast path. ln_likelihood_base is correspondingly scalar or
   per-frequency, and the cross-dataset reduction now broadcasts rather than
   stacks, so a container mixing the two shapes adds up instead of raising (or
   silently collapsing over the grid axis).

   Reuses the existing _is_callable_prior predicate, moved from models/joint.py
   to models/_helpers.py beside _needs_explicit_sampling so the periodogram is
   not importing a private name out of joint.py.

3. periodogram(n_terms=0) was silently promoted to 1 -- the warning guard
   1 < 0 can never fire -- and then failed with "prior.linear_priors is missing
   entries for ['cos_amp_1', 'sin_amp_1']", naming parameters the caller never
   asked for. Negative values behaved the same. Validate n_terms >= 1 up front.

Full suite: 1128 passed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
- arg_constraints declared log_density as real_vector, whose _Real check
  rejects every non-finite value -- but -inf is a documented, deliberately
  produced input (a zero-density knot, which _to_prior creates from
  np.log(0)). Under numpyro.enable_validation() or validate_args=True,
  constructing the very thing peak_period_prior returns raised
  "LogGridDensity distribution got invalid log_density parameter". Use
  independent(less_than(inf), 1), which admits -inf, still rejects +inf and
  NaN, and keeps event_dim == 1.

- Mark LogGridDensity @Final, per CLAUDE.md's abstract-final rule. Nothing
  subclasses it, and every other concrete class added on this branch already
  carries it.

- peak_period_prior documented each peak carrying exactly (1 - floor)/n_peaks,
  but the top-hats were normalized by their analytic width while
  LogGridDensity integrates them as sampled on the knots -- so a peak clipped
  by the domain edge, or one narrower than the knot spacing, quietly carried
  less than its share (~12% at the default samples_per_peak=8). Normalize each
  top-hat by its own trapezoid mass on the knots instead. Measured per-peak
  mass now matches the documented share to ~1e-6 relative, so the tests assert
  it at rel=1e-4 rather than the rel=0.25 that hid the discrepancy.

Full suite: 1133 passed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
- pyproject: declare numpy>=2.0. It was never listed, arriving only
  transitively via jax/numpyro/h5py, while 11 modules under src/harv import it
  directly -- and priors.py calls np.trapezoid, which exists only in NumPy 2.x
  (1.x spells it trapz). An environment resolving numpy 1.26, still supported
  by jax, would have died in tempered_period_prior.

- The under-resolution warning blamed harv rather than the caller, so a user's
  filterwarnings(module=...) could never match it. stacklevel alone cannot fix
  this: the warning is raised in _finalize_posterior, reached from either run()
  or run_with_samples(), with equinox method wrappers interleaved -- measured
  depths differ per entry point, and stacklevel=2 landed inside equinox. Use
  Python 3.12's skip_file_prefixes to skip harv's and equinox's frames
  entirely. Tested from both entry points.

- periodogram(): samples_per_peak was silently ignored when an explicit
  frequency_grid was passed. Make it None-defaulted, include it in the
  mutual-exclusion check, and forward it only when set so frequency_grid keeps
  sole ownership of the default.

- Delete the dead branch in _resolve_per_dataset -- `if is_container and name
  == "prior": return value` was followed by an unconditional `return value`,
  implying a distinction the function does not make. The is_container
  parameter falls out with it.

- periodogram() recomputed the multi-dataset time baseline inline that
  frequency_grid had just computed; reuse grid._data_t_span instead.

Full suite: 1135 passed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Found while verifying the tutorial-6 call pattern: under JAX's default float32,
the periodogram in that tutorial's own configuration (16 epochs, high SNR,
marginal log-likelihoods reaching ~5e4 nats) returns a delta_ln_likelihood with
non-finite entries. tempered_period_prior then died inside LogGridDensity with
"log_density must have positive total mass", which says nothing about the
actual problem. The tutorial itself is fine -- it enables x64 in its first cell,
and the same run is clean in float64 (delta 1.5e4..4.8e4, peak at 34.4 d against
a true 35 d) -- but nothing pointed a float32 user at the cause.

_assemble_knots now checks Delta up front and raises a ValueError naming the
count of bad grid points and the likely fix (jax_enable_x64). This is a
diagnosis, not a numerical fix: the underlying float32 fragility of the
marginal likelihood on high-SNR data is untouched.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
- The "Interpreting acceptance" section had been inserted into the middle of
  the RejectionSampler.run keyword list, orphaning "Forced on by top_k." from
  the return_logprobs bullet it belongs to and leaving return_evidence_stats
  documented twice with different wording. Restore the continuation line, merge
  the two bullets into the one that carries the keys table, and move the
  section below the completed list.

- The Prior builders usage snippet called hp.periodogram without the required
  prior= argument -- it raised TypeError as written, and contradicted the
  "Priors are explicit" section three subsections above. Build the Fourier
  trial prior in the snippet, as the API-sketch example already does.

- Export PeriodDependentKPrior from harv.models.priors alongside its two
  siblings, so the :class: cross-reference in periodogram/core.py resolves.

- 1-customize-prior.ipynb shipped a markdown cell reading "TODO: bad example -
  should do for a source with more observations", which rendered verbatim on
  the published page. The concern is real -- that source is sparse -- so make
  it a reader-facing note that says so and points at tutorial 6, rather than
  deleting it.

Full suite: 1138 passed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The TODO pointed at docs/spec.md, "Amplitude and nuisance priors", which does
not exist. The open question it refers to lives under "Priors are explicit".

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Every Keplerian parameterization in harv already defaults to a
parameter-dependent amplitude prior; only the Fourier ones got a flat
Normal(0, sigma_amp). That asymmetry was the standing
TODO(default-amplitude-prior), and a mini-grid settles it for RV: a tilted
prior makes the true mode dominant where flat does not, provided sigma_K0
matches the companion being searched for.

The proviso drives the API. The Occam factor only reaches d*ln sigma(P) once
sigma^2 lambda(P) >> 1, so sigma_0 is the tilt's *gain*, not a nuisance width
-- an over-wide scale, harmless under a flat prior, tilts the periodogram
toward long periods here.

- FourierRV.default_prior gains sigma_K0/P0, FourierGaiaAstrometry gains
  sigma_a0/P0, each mutually exclusive with sigma_amp. Both scales are still
  required explicitly: no data-driven default.
- periodogram(..., prior_params={"parallax": ...}) supplies values a
  LinearPriorCallable needs but the scan does not own. Bound into the
  callables, deliberately *not* merged into nl_values: log_prob's auto mode
  reclassifies any linear name out of nl_values as an explicit column, which
  would silently fix the parallax in both trial and base models. A regression
  test asserts the marginalized column set is unchanged.
- Callable priors are probed once, eagerly, so a missing or misspelled key
  raises a TypeError naming prior_params rather than a KeyError from inside a
  jit/vmap trace. Unresolved keys are ignored downstream, so a typo would
  otherwise fail identically to an omission.
- Gaia's path needs an external parallax whose uncertainty a point value
  discards; recorded as TODO(parallax-marginalization) with both routes.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Classical periodograms (Lomb-Scargle, kepmodel) maximize over the linear
amplitudes rather than integrating them out, reporting
z0(f) = 1/2 (chi2_base - chi2_trial(f)). There was no way to get that number
out of harv, so comparing the two statistics meant rebuilding the design
matrix by hand in numpy.

This is a second statistic, not a limit of the first: as the amplitude priors
widen, Delta_marginal -> z0 - (1/2)(d_trial - d_base) ln(Lambda) -> -inf,
because the trial model has more columns than the base. The periodogram
docstring previously claimed the opposite; corrected here.

- AbstractComponentModel._log_prob_profile reuses the same design matrix and
  the same extension-modified covariance as the marginal path, whitens through
  the existing to_linear_op, and solves by SVD so it stays finite where the
  trial columns go collinear with the base ones at long periods.
- periodogram(prior=False) skips every prior-consulting step. prior_params
  alongside it is a TypeError rather than silently ignored, and False inside a
  per-dataset mapping is rejected: a log Bayes factor and a (1/2)dchi2 are not
  commensurable, so summing them across datasets is meaningless.
- PeriodogramResult.statistic records which one was computed; plot() labels
  the axis accordingly.
- The n_terms cap matters more here, not less: the marginal likelihood stays
  finite when columns outnumber observations because the prior regularizes,
  while the least-squares solve would drive chi2 to zero at every frequency.

Tested against a numpy lstsq oracle, and via the bridge identity
Delta_marginal = z0 - Occam - shrinkage, which is the only check that both
modes see the same design matrix and noise model.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
adrn and others added 13 commits September 13, 2026 22:08
The marginal statistic stays well-posed when the linear columns outnumber the
observations: M = I + B^T B is at least the identity however rank-deficient the
design matrix is. Only the unregularized least-squares solve has a breakdown
point there, where chi^2 hits zero at every trial period and the statistic goes
identically flat. So profile mode still reduces n_terms, while marginal mode now
keeps the requested value and only warns.

Both modes still warn at the same two-observations-per-column bar. That bar is a
convention sitting well above the rank limit, not a rank condition: it is placed
where recovery of the true period empirically falls off, which tracks the
observations-per-column ratio rather than the absolute column count. On
simulated RV data recovery drops ~40% as the ratio crosses 2, at 10 epochs
(H=2 -> H=3) and again at 14 (H=3 -> H=4).

The warning no longer claims the trial model "overfits", which is false at 7
columns against 10 observations, and the profile branch no longer claims chi^2
has already hit zero when it has not. Both now report the measured ratio and
the bar being applied, so the message survives a reader checking the arithmetic.

_resolve_linear_priors loses its eff_terms argument: it is only reached on the
marginal path, where the requested and effective counts now always agree.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
@adrn
adrn enabled auto-merge September 14, 2026 02:12
@adrn
adrn merged commit 9e459e3 into main Sep 14, 2026
13 checks passed
@adrn
adrn deleted the periodogram branch September 14, 2026 02:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants