feat: Enhance PSD plotting and update GP kernel parameters

Metadata-Version: 2.4

Name: PyAETNA

Version: 3.6.0

Version: 3.7.0

Summary: Python Astronomical Exoplanet Transit and Nuisance Analysis - JAX-accelerated exoplanet light curve modeling

Author-email: Gaetano Scandariato <gaetano.scandariato@inaf.it>

License: MIT

  A JAX-based Bayesian framework for modeling and analyzing exoplanet light curves and radial velocities with advanced systematics treatment.

**Version:** 3.6.0

**Version:** 3.7.0

**Release note (v3.7.0):** named priors and cross-section parameter tying

implemented via a top-level `priors:` section and `@name` references.

  [![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/)

  [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)

  - NUTS (gradient-based, default)

  - AIES and ESS ensemble samplers (gradient-free)

  - Automatic multimodal detection and restart strategies

- **Named Priors & Global Tying**: tie parameters across different model components (e.g., sharing a stellar rotation period between multiple instrument and corresponding detrending schemes)

- **Comprehensive Light Curve Modeling**:

  - Primary transits and secondary eclipses (via `jaxoplanet`)

  - Primary transits, phase curves and secondary eclipses (via `jaxoplanet`)

  - **Transit Timing Variations (TTV)**: per-epoch free transit times with O−C posterior

  - Phase curves with hotspot offsets (Lambert sphere)

  - Ellipsoidal variations with tidal lag

├── PyAETNA/                   # Core library package

│   ├── __init__.py            #   Package exports and __version__

│   ├── run.py                 #   Pipeline orchestrator (sampling → reports → physics)

│   ├── sampler.py             #   MCMC sampler (NUTS / ESS / AIES), model builder

│   ├── reports.py             #   Observables report (PDF, LC/RV fits, corner plots)

│   ├── planet_phys.py         #   Derived physics (eclipse depths, albedo, Cowan & Agol)

│   ├── body_models.py         #   JAX planetary models (transit, eclipse, RV via jaxoplanet)

│   ├── gp.py                  #   GP kernel utilities (tinygp quasisep + dense)

│   ├── model_utils.py         #   Shared model assembly and detrending helpers

│   ├── _device_config.py      #   JAX device/platform auto-detection

│   ├── CHANGELOG.md           #   Version history

│   ├── sampling/              #   Modular sampling engine and model assembly

│   │   ├── __init__.py        #     Public sampling API

│   │   ├── data.py            #     Input data structures and validation

│   │   ├── precision.py       #     JAX precision/casting helpers

│   │   ├── utils.py           #     Shared sampling math utilities

│   │   ├── model/             #     PlModel core + detrending + GP builders

│   │   ├── config/            #     Config parsing and prior handling

│   │   │   ├── priors.py      #       Prior string parsing and dist mapping

│   │   │   └── parser/        #       Scoped parser modules (star/planet/gp/...)

│   │   └── engine/            #     Sampling loop and NumPyro integration

│   │       ├── runner.py      #       runSampler_NumPyro main runner

│   │       ├── numpyro_model.py #     Model closure and log-joint helpers

│   │       ├── logging.py     #       Sampling log/summary formatters

│   │       ├── data_loader.py #       Resume-safe data loading utilities

│   │       ├── checkpoint.py  #       Checkpoint save/restore logic

│   │       ├── sampler_utils.py #     Shared runner helpers

│   │       └── utils/         #       Engine utility modules

│   ├── reporting/             #   Modular report-generation package

│   │   ├── __init__.py        #     Reporting package exports

│   │   ├── plot_models.py     #     Main plotting/report orchestration

│   │   ├── confidence_bands.py #    Posterior confidence band computations

│   │   ├── grid_builder.py    #     Grids for phase/time-domain plotting

│   │   ├── lightcurve_plots.py #    Light-curve plotting routines

│   │   ├── periodogram_plots.py #   Periodogram and PSD diagnostics

│   │   ├── phase_plots.py     #     Phase-folded visualization helpers

│   │   ├── residual_plotting.py #   Residual panels and KDE inset helpers

│   │   ├── corner_plot.py     #     Posterior corner plot generation

│   │   ├── diagnostics.py     #     MCMC and model diagnostics

│   │   ├── pdf_report.py      #     Multi-page PDF report assembly

│   │   └── ...                #     Additional plotting/support modules

│   └── scripts/               #   Auxiliary CLI scripts

│

├── docs/                      # Documentation

│   ├── programmatic_usage.md  #   Calling PyAETNA from Python

│   ├── phase_curve_physics.md #   Phase-curve theory, derivations, references

│   ├── ARCH_DECISIONS.md      #   Architecture rationale (vectorisation, precision)

|   ├── docs/CITATIONS.md      # Full reference list with BibTeX entries

│   ├── CITATIONS.md           #   Full reference list with BibTeX entries

│   └── gp_per_visit_analysis/ #   GP per-visit block-diagonal documentation

│

├── examples/                  # Example configurations and data

- improve online configurator

- remove lazy split-index fallback in planet model path and require precomputed static indices for strict end-to-end JAX/JIT compliance

- check the checkpoint_interval logic

- split plots for multi-season datasets

- implement different limb darkening laws

- better arrangement of legends and text in plots

## Contact

PyAETNA/gp.py

PyAETNA/model_utils.py

PyAETNA/planet_phys.py

PyAETNA/reports.py

PyAETNA/run.py

PyAETNA/sampler.py

PyAETNA.egg-info/PKG-INFO

PyAETNA.egg-info/SOURCES.txt

PyAETNA.egg-info/dependency_links.txt

PyAETNA/oldCode/runPyAETNA.py

PyAETNA/oldCode/transitModelHANNU_v2.py

PyAETNA/reporting/__init__.py

PyAETNA/reporting/confidence_bands.py

PyAETNA/reporting/corner_plot.py

PyAETNA/reporting/data_validation.py

PyAETNA/reporting/diagnostics.py

PyAETNA/reporting/gp.py

PyAETNA/reporting/gp_plotting.py

PyAETNA/reporting/grid_builder.py

PyAETNA/reporting/grid_builder_utils.py

PyAETNA/reporting/lightcurve_plots.py

PyAETNA/reporting/logger.py

PyAETNA/reporting/map_search.py

PyAETNA/reporting/model_components.py

PyAETNA/reporting/model_math.py

PyAETNA/reporting/noise_plots.py

PyAETNA/reporting/panel_detrended.py

PyAETNA/reporting/panel_overlay.py

PyAETNA/reporting/panel_raw.py

PyAETNA/reporting/panel_residuals.py

PyAETNA/reporting/pdf_report.py

PyAETNA/reporting/periodogram_plots.py

PyAETNA/reporting/phase_plots.py

PyAETNA/reporting/plot_model_helpers.py

PyAETNA/reporting/plot_data_prep.py

PyAETNA/reporting/plot_models.py

PyAETNA/reporting/plot_styling.py

PyAETNA/reporting/residual_plotting.py

PyAETNA/reporting/trend_detection.py

PyAETNA/reporting/ttv_plots.py

PyAETNA/sampling/__init__.py

PyAETNA/sampling/data.py

PyAETNA/sampling/precision.py

PyAETNA/sampling/utils.py

PyAETNA/sampling/config/priors.py

PyAETNA/sampling/config/parser/__init__.py

PyAETNA/sampling/config/parser/detrending.py

PyAETNA/sampling/config/parser/gp.py

PyAETNA/sampling/config/parser/planets.py

PyAETNA/sampling/config/parser/prior_refs.py

PyAETNA/sampling/config/parser/scoping.py

PyAETNA/sampling/config/parser/star.py

PyAETNA/sampling/engine/checkpoint.py

PyAETNA/sampling/engine/data_loader.py

PyAETNA/sampling/engine/logging.py

PyAETNA/sampling/engine/numpyro_model.py

PyAETNA/sampling/engine/runner.py

PyAETNA/sampling/engine/sampler_utils.py

PyAETNA/sampling/engine/utils/logging_utils.py

PyAETNA/sampling/model/core.py

PyAETNA/sampling/model/detrending.py

PyAETNA/sampling/model/gp_models.py

PyAETNA/scripts/normalize_flux_by_visit.py

PyAETNA/scripts/plot_random_gp.py

tests/test_behavioral_ties.py

tests/test_exptime_integration.py

tests/test_gp_per_visit.py

tests/test_prior_refs.py

tests/test_rm_effect.py

tests/test_spectral_mixture.py

tests/test_ttv.py

 No newline at end of file

PyAETNA

PyAETNAouts

# Plan: Enhanced PSD Plotting in PyAETNA

## Objective

Improve the Power Spectral Density (PSD) diagnostic plots in the reports to provide a full variance budget analysis. This allows for a direct comparison between the observed data power, the nominal noise floor, the fitted white noise (jitter), and the GP model's correlated power.

## Proposed Plot Components

For the "GP + noise" panel (Row 1), the following components will be plotted:

1.  **Lomb-Scargle (LS) Periodogram**: The observed power of the detrended data (Black line).

2.  **Nominal Noise Baseline**: A flat line representing the PSD of the nominal data uncertainties $\langle \sigma_{\text{nom}}^2 \rangle$.

3.  **Total White Noise Baseline**: A flat line representing nominal uncertainties plus the MAP jitter $\langle \sigma_{\text{nom}}^2 \rangle + \sigma_{\text{jitter}}^2$. (Shown only if `log_jitter` is finite).

4.  **GP Kernel PSD**: The theoretical PSD of the GP kernel $S_{\text{Kernel}}(f)$. (Shown only if GP is active).

5.  **Total Model PSD**: The sum of the GP Kernel PSD and the Total White Noise Baseline: $S_{\text{Total}}(f) = S_{\text{Kernel}}(f) + S_{\text{WhiteNoise}}$. (Shown only if GP is active).

## Mathematical Logic

The LS PSD is normalized such that $\int S(f) df = \text{Var}(\text{data})$. 

To maintain consistency, the flat noise baselines are calculated as:

$$S_{\text{white}} = \frac{\sigma_{\text{white}}^2}{f_{\text{max}} - f_{\text{min}}}$$

where $\sigma_{\text{white}}^2$ is the variance of the white noise component.

## Implementation Steps

### 1. API Update

- Modify `plot_periodograms` in `PyAETNA/reporting/periodogram_plots.py` to accept `nominal_uncertainties` and `map_jitter`.

- Update the call site in `PyAETNA/reporting/pdf_report.py` to pass these values from the model/data objects.

### 2. Calculation Logic

- Compute $f_{\text{range}} = f_{\text{max}} - f_{\text{min}}$.

- Calculate $S_{\text{nom}} = \text{mean}(\sigma_{\text{nom}}^2) / f_{\text{range}}$.

- Calculate $S_{\text{jitter}} = (\text{mean}(\sigma_{\text{nom}}^2) + \sigma_{\text{jitter}}^2) / f_{\text{range}}$.

### 3. Plotting Updates

- Use `ax.axhline` for the flat baselines.

- Use `ax.plot` for the GP and Total Model curves.

- Implement conditional checks:

    - Jitter baseline $\rightarrow$ `if map_jitter is not None and not np.isinf(log_jitter)`.

    - GP lines $\rightarrow$ `if gp_kernel is not None and gp_x_variable == 'BJD'`.

### 4. Legend & Styling

- Update the legend to clearly label: "LS Periodogram", "Nominal Noise", "Total White Noise", "GP Kernel PSD", and "Total Model PSD".

- Use distinct colors and line styles (e.g., dotted for nominal, dashed for jitter).

## Verification

- Generate a report for a target with a known GP fit.

- Verify that the integral of the "Total Model PSD" visually matches the integral of the LS periodogram.

- Ensure that the noise floor of the LS periodogram aligns with the "Total White Noise" baseline.

        params: Dictionary of parameter values (e.g., from MCMC samples or MAP).

        kernel_specs: List of kernel specifications, each with:

            - 'type': Kernel type string (e.g., 'SHO', 'Matern32')

            - 'params': List of parameter names (e.g., ['log_gp_sigma', 'gp_freq', 'log_gp_lengthscale'])

            - 'params': List of parameter names (e.g., ['log_gp_sigma', 'gp_freq', 'log_gp_scale'])

            - 'fixed_params': (optional) Dict of fixed parameter values

            - 'subspace': (optional) Index for Subspace transform

            - 'section_name': (optional) GP section name for parameter lookup (e.g., 'GP_roll')