plot_models.py refactored

# Fix JAX Segmentation Fault in PyAETNA `relabel_array` — Action Plan

## Problem Statement

**User Context:**  

G (a researcher at INAF Catania) recently updated JAX, JAXlib, and NumPyro (originally installed March 2025) on the ExoREC-Ganymede H100 server. A previously working PyAETNA analysis script (`runFit.py` in TOI5572 directory) now crashes with a segmentation fault at the JAX/XLA layer.

**Crash Signature:**  

```

Fatal Python error: Segmentation fault

  File ".../jax/_src/numpy/setops.py", line 856 in unique

  File ".../PyAETNA/sampler.py", line 1749 in relabel_array

```

**Root Cause:**  

JAX 0.4.31+ flipped the default XLA CPU backend from the legacy StreamExecutor to the new **Thunk runtime** (IFRT/PjRt). The Thunk runtime has a known bug when handling `jnp.unique` operations with **dynamic output shapes** — it crashes at the compiled kernel dispatch level. The `relabel_array` function in `sampler.py` calls `jnp.unique` without specifying an output size, forcing XLA to infer the size dynamically.

**Workaround (Not Sustainable):**  

Setting `XLA_FLAGS="--xla_cpu_use_thunk_runtime=false"` reverts to the old backend and the script runs. However, this flag is deprecated and will be removed in a future JAX release.

---

## Solution Overview

Replace all calls to `jnp.unique()` in the `relabel_array` function with **static-size** `jnp.unique(..., size=...)` calls. This forces XLA to compile with a fixed output shape, eliminating the dynamic shape inference that triggers the Thunk runtime bug.

**Benefits:**

- ✅ Removes dependency on deprecated `XLA_FLAGS` workaround

- ✅ Future-proof (works with all XLA backends, current and future)

- ✅ Potentially faster (static shapes = better XLA optimizations)

- ✅ JAX best practice for JIT/pmap/vmap environments

---

## Step-by-Step Action Plan

### Step 1: Locate the `relabel_array` function

**File:** `/home/gscandariato/software/pyLib/PyAETNA/PyAETNA/sampler.py`  

**Line:** ~1749

**Action:**  

- Open `sampler.py` in an editor

- Search for `def relabel_array` to find the function definition

- Read the full function body to identify all `jnp.unique(...)` calls

- Note the input array type and expected output range

**Expected output:** Full function code snippet with all `jnp.unique` invocations highlighted

---

### Step 2: Analyze `jnp.unique` usage patterns in `relabel_array`

**Context:**  

The `relabel_array` function likely:

1. Takes an input array of class/label assignments (e.g., MCMC chain cluster labels from mixture models)

2. Calls `jnp.unique()` to find distinct values

3. Uses this to relabel or validate cluster identity across chains

**Action:**  

- For **each** `jnp.unique(...)` call in the function, determine:

  - **Input array:** What is being passed? (e.g., `relabeling_data`, `raw_labels`)

  - **Max possible unique values:** What is the upper bound on output size?

    - For cluster labels: typically `n_components` or `n_clusters` (get from function parameters or parent context)

    - For array elements: typically `len(input_array)`

  - **Is `sorted=True` used?** (preserve existing behavior)

  - **Are there any assumptions** about the output being "unsorted"? (Note for preservation)

**Deliverable:** Annotated list of all `jnp.unique` calls with their context and max size constraints

---

### Step 3: Modify each `jnp.unique` call

Very important: any modification should be compliant with jitting and must not affect the perfomances of the code.

**Pattern:**  

```python

# BEFORE (dynamic output size → segfault on Thunk runtime)

result = jnp.unique(input_array)

# AFTER (static output size → safe on all XLA backends)

result = jnp.unique(input_array, size=max_unique_values, fill_value=-1)

```

**Action:**  

For each `jnp.unique` call identified in Step 2:

1. **Determine `size=` parameter:**

   - Conservative approach: Use `len(input_array)` (always safe, worst-case one unique value per element)

   - Optimized approach: Use domain knowledge (e.g., `n_components`, `n_chains * max_labels_per_chain`)

2. **Choose `fill_value`:**

   - Use a sentinel value that will **never** appear in real data

   - Common choices: `-1` (if labels are non-negative), or `np.iinfo(arr.dtype).min`

3. **Preserve other arguments:**

   - Keep `sorted=`, `return_index=`, `return_counts=`, etc. as they were

   - Example: `jnp.unique(arr, size=100, sorted=True, return_counts=True, fill_value=-1)`

4. **Update the code:**

   - Replace the old `jnp.unique(...)` with the new static-size variant

   - Add a comment explaining the `size=` parameter for future maintainers

**Pseudo-example:**

```python

# Line ~1760 in relabel_array

# BEFORE:

unique_labels = jnp.unique(raw_chain_labels)

# AFTER:

# Use size=n_components to enforce static shape (JAX Thunk runtime requires this)

unique_labels = jnp.unique(raw_chain_labels, size=self.n_components, fill_value=-1)

```

---

### Step 4: Validate the fix

**File:** `/DATA/ExoREC/TOI5572/runFit.py`

**Action:**  

1. **Remove any workaround flags** from the script (if present):

   - Delete `os.environ["XLA_FLAGS"] = ...` lines

   - Delete any shell-level `export XLA_FLAGS=...` calls

2. **Run the fixed script:**

   ```bash

   cd /DATA/ExoREC/TOI5572

   /home/gscandariato/cuda12/bin/python -X faulthandler runFit.py

   ```

3. **Check for:**

   - ✅ No segmentation fault

   - ✅ No deprecation warnings about `xla_cpu_use_thunk_runtime`

   - ✅ Output is numerically identical to previous successful runs

   - ✅ Runtime is comparable (or slightly faster)

4. **If segfault persists:**

   - There may be additional `jnp.unique` calls in other functions called by `runFit.py`

   - Search the entire codebase: `grep -rn "jnp.unique" /home/gscandariato/software/pyLib/PyAETNA/`

   - Apply the same fix to any other occurrences

---

### Step 5: Document and commit

**Action:**  

1. **Add docstring note** to `relabel_array`:

   ```python

   def relabel_array(...):

       """

       [existing docstring]

       Notes:

       - All jnp.unique() calls use explicit size= parameter to enforce static shapes.

       This is required for JAX 0.4.31+ Thunk runtime compatibility (XLA CPU backend).

       See: https://github.com/jax-ml/jax/issues/... [optional: add GitHub issue link]

       """

   ```

2. **Commit to version control** (if applicable):

   ```bash

   git add sampler.py

   git commit -m "Fix: Replace dynamic-shape jnp.unique with static-size variant (JAX Thunk runtime segfault)"

   ```

3. **Update project notes** (e.g., CHANGELOG, deployment docs):

   - Record that PyAETNA now requires JAX 0.4.31+ (no longer needs `XLA_FLAGS` workaround)

---

## Success Criteria

- [ ] Script runs without segmentation fault

- [ ] No JAX deprecation warnings appear

- [ ] All MCMC chains complete successfully

- [ ] Output file produced (sampler chains, posteriors, etc.)

- [ ] Numerical results match expected posteriors (spot-check 2–3 key parameters)

- [ ] No performance regression observed

---

## Rollback Plan (if needed)

If the fix introduces unexpected behavior:

1. **Quick rollback:** Revert to the deprecated flag in `runFit.py`:

   ```python

   import os

   os.environ["XLA_FLAGS"] = "--xla_cpu_use_thunk_runtime=false"

   import jax

   import numpyro

   # ... rest of script

   ```

2. **Long-term rollback:** Pin JAX/JAXlib to pre-Thunk version:

   ```bash

   pip install "jax==0.4.30" "jaxlib==0.4.30[cuda12]" --force-reinstall

   ```

---

## Estimated Effort

- **Step 1–2:** 10–15 min (locate and analyze)

- **Step 3:** 5–20 min (modify, depending on number of `jnp.unique` calls)

- **Step 4:** 5–10 min (validate + debug if needed)

- **Step 5:** 5 min (document + commit)

**Total:** ~30–50 minutes

---

## References

- JAX GitHub Issues: XLA Thunk Runtime CPU backend migration (late 2024–2025)

- JAX NumPy docs: [`jnp.unique(size=...)`](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.unique.html)

- JAX best practices: Static shapes in JIT-compiled code

 No newline at end of file

**Version:** 3.6.0

PyAETNA provides a flexible, configuration-driven pipeline for fitting

transit, phase‑curve light curves and radial velocities using a Bayesian

framework with high-performance JAX/NumPyro sampling.  It handles

instrument-specific systematics via polynomials, Fourier series, or

Gaussian processes, and produces publication-quality diagnostic reports

with minimal user effort.

  [![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)

---

## Features

## What's PyAETNA?

PyAETNA is a flexible, configuration-driven pipeline for fitting

transit, phase‑curve light curves and radial velocities using a Bayesian

framework with high-performance JAX/NumPyro sampling.  It handles

instrument-specific systematics via polynomials, Fourier series, or

Gaussian processes, and produces high-quality diagnostic reports

with minimal user effort.

### Features

- **JAX/NumPyro Backend**: JIT-compiled, GPU-accelerated MCMC sampling

- **Flexible MCMC Sampling**:

  - Automatic multimodal detection and restart strategies

- **Comprehensive Light Curve Modeling**:

  - Primary transits 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

  - Doppler beaming

  - Correlated noise modeling

  - Confidence bands from posterior samples

## Why Use PyAETNA?

- **Single YAML, complete analysis.** Define data paths, priors,

  systematics, MCMC strategy, and post-processing physics in one

  human-readable configuration file — no boilerplate Python scripts required.

- **GPU-accelerated end-to-end.** The entire forward model, likelihood, and

  gradient computation runs inside JIT-compiled JAX, delivering order-of-

  magnitude speedups on GPU/TPU hardware without code changes.

- **Adaptive three-phase sampling.** The pipeline manages warmup and

  production automatically, with ensemble warmup (ESS or AIES) for

  exploration and NUTS for efficient local convergence. Warmup and

  production behavior is still highly configurable in the YAML file (sampler

  choice, number of chains, adaptation schedule, etc.).

- **Checkpointed, restartable runs.** MCMC state is periodically pickled so

  that long runs can survive crashes, and subsequent launches resume exactly

  where they left off — including the ensemble warmup phase.

- **Mixed-precision ephemeris firewall.** Transit mid-times are folded in

  float64 while the rest of the model evaluates in float32 for speed,

  eliminating the catastrophic cancellation that plagues naïve

  single-precision transit fits.

- **Rich systematics toolkit.** Polynomial trends, Fourier series, and

  scalable Gaussian Processes (O(N) quasiseparable or O(N³) dense kernels

  via tinygp) can be combined freely, scoped per-visit or per-instrument,

  and stacked on the same dataset.

- **Joint multi-instrument, multi-type fitting.** Photometric light curves

  and radial velocities from different instruments are modelled

  simultaneously with fully shared orbital parameters and independent

  per-instrument systematics.

- **High-quality and informative plots.** Diagnostic PDFs include phase-folded fits

  with residuals, best-fit models per visit, corner plots, trace plots,

  Lomb–Scargle periodograms, and an optional physics report with Cowan &

  Agol energy-balance, brightness temperatures, and passband-integrated

  albedos to clearly check the outcome of model exploration.

- **Physics post-processing in one step.** A dedicated `physics:` stage

  derives limb-darkening coefficients, orbital geometry, equilibrium

  temperatures, Bond albedos, and heat redistribution efficiencies —

  propagating full MCMC uncertainties via Monte Carlo resampling.

- **Model-atmosphere support.** BT-Settl and PHOENIX synthetic spectra are

  downloaded on the fly (with local caching) for passband-integrated

  brightness-temperature inversion, replacing the default blackbody

  approximation with no additional libraries to install.

## Quick Start

### Installation

  - Detrending options (polynomials, Fourier, Gaussian Processes)

  - Best practices and troubleshooting

  A complete **online configurator** is available at http://pyaetna-e92eeb.pages.ict.inaf.it/.

- **[Dependencies Guide](docs/dependencies.md)** - Detailed information about software requirements

  - Installation instructions (CPU and GPU)

  - Platform-specific notes (Linux, macOS, Windows)

```

PyAETNA/

├── runPyAETNA.py            # Main CLI entrypoint (run the pipeline)

├── pyproject.toml           # Packaging metadata (version, deps)

├── requirements.txt         # Optional dependency pins / fallback

├── runPyAETNA.py              # CLI entrypoint

├── pyproject.toml             # Packaging metadata and dependencies

├── requirements.txt           # Dependency pins (fallback)

├── README.md                  # This file

├── docs/                    # Documentation (configuration, dependencies)

├── LICENSE                    # MIT licence

│

├── 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

│   └── scripts/               #   Auxiliary CLI scripts

│

├── docs/                      # Documentation

│   ├── configuration.md       #   Full YAML config reference

│   ├── configuration_keys.md# Config key reference

│   └── dependencies.md      # Installation and dependency notes

├── examples/                # Example configs and run setups

│   ├── 501.yaml

│   ├── example_lc_rv_joint.yaml

│   └── reference_config.yaml

├── localDocs/               # Internal investigations, notes, benchmarks

├── tests/                   # Test datasets, configs, and generated outputs

│   ├── test_phase_curve.yaml

│   ├── test_lc_rv_joint.yaml

│   └── test_rv_only.yaml

└── PyAETNA/                 # Core library package

    ├── __init__.py           # Package exports & version string

    ├── run.py                # CLI entrypoint (dispatches sampling → reports → physics)

    ├── sampler.py            # MCMC sampler (NumPyro/PL model)

    ├── reports.py            # Plotting & report generation (PDFs, LC/RV fits)

    ├── planet_phys.py        # Derived physical quantities (eclipse, phase-curve, albedo)

    ├── gp.py                 # GP kernel utilities (tinygp-based)

    ├── model_utils.py        # Shared model assembly + detrending helpers

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

    └── other modules and helper files

│   ├── configuration_keys.md  #   Quick-reference key table

│   ├── dependencies.md        #   Installation guide (CPU & GPU)

│   ├── 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

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

│

├── examples/                  # Example configurations and data

│   ├── reference_config.yaml  #   Annotated reference with EVERY supported key

│   ├── 501.yaml               #   CHEOPS dayside emission example

│   ├── transits.yaml          #   TESS transit-only example

│   ├── example_lc_rv_joint.yaml # Joint LC + RV example

│   └── *.dat                  #   Accompanying data files

│

├── tests/                     # Test suite

│   ├── generate_test_datasets.py # Synthetic data generator

│   ├── test_lc_only.yaml      #   LC-only test config

│   ├── test_rv_only.yaml      #   RV-only test config

│   ├── test_lc_rv_joint.yaml  #   Joint LC + RV test config

│   ├── test_lc_gp_rv.yaml    #   LC (GP) + RV test config

│   ├── test_gp_per_visit.yaml #   GP per-visit block-diagonal test

│   ├── test_phase_curve.yaml  #   Phase curve + RV test config

│   ├── test_exptime_integration.py # Exposure-time supersampling tests

│   └── *.dat                  #   Synthetic test data

│

└── localDocs/                 # Internal development notes (not shipped)

```

## Citation

If you use PyAETNA in your research, please cite:

- The jaxoplanet paper (Foreman-Mackey et al.)

- The NumPyro paper (Phan et al. 2019)

- This repository (citation details coming soon)

If you use PyAETNA in your research, please add the following citation in your publication:

#TODO add citation to PyAETNA

Please also cite the relevant works

listed in [docs/CITATIONS.md](docs/CITATIONS.md). For your convenience, at the end of every run the

log prints a tailored list of references based on the modules and

features activated in your analysis, together with a pointer to the

BibTeX entries.

## License

[License information to be added]

This project is licensed under the MIT License — see the [LICENSE](LICENSE) file.

## TODO

- Test in pleiadi

- stellar spectrum input by the user

- parse config file checking for the correctness of the priors

- parse config file checking for the correctness of the priors and add logging if priors are not well defined

- check citations and include citation to PyAETNA

- 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

## Contact

PyAETNA.egg-info/dependency_links.txt

PyAETNA.egg-info/requires.txt

PyAETNA.egg-info/top_level.txt

PyAETNA/data/passbands/CHEOPS.dat

PyAETNA/data/passbands/TESS.dat

PyAETNA/kernels/factory.py

PyAETNA/oldCode/PyAETNA_JAXdistributions.py

PyAETNA/oldCode/PyAETNA_MCMCFunctions.py

PyAETNA/oldCode/PyAETNA_NormalizedGLS.py

PyAETNA/oldCode/gsMCMCFunctions_v2.py

PyAETNA/oldCode/runPyAETNA.py

PyAETNA/oldCode/transitModelHANNU_v2.py

PyAETNA/reporting/__init__.py

PyAETNA/reporting/corner_plot.py

PyAETNA/reporting/diagnostics.py

PyAETNA/reporting/gp.py

PyAETNA/reporting/lightcurve_plots.py

PyAETNA/reporting/logger.py

PyAETNA/reporting/map_search.py

PyAETNA/reporting/model_components.py

PyAETNA/reporting/noise_plots.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_models.py

PyAETNA/reporting/ttv_plots.py

PyAETNA/scripts/normalize_flux_by_visit.py

PyAETNA/scripts/plot_random_gp.py

tests/test_exptime_integration.py

tests/test_gp_per_visit.py

tests/test_rm_effect.py

tests/test_spectral_mixture.py

tests/test_ttv.py

 No newline at end of file

# -*- coding: utf-8 -*-

"""

data_validation.py

~~~~~~~~~~~~~~~~~~

Input validation and data preparation for plot_models.

Responsibilities:

- Validate plot inputs (visit/instrument selection)

- Extract and normalize data arrays

- Determine data type from configuration

- Apply visit/instrument masking

"""

import numpy as np

def validate_plot_inputs(time, visit, inst, visit_to_plot, instrument_to_plot):

    """

    Validate and construct data mask.

    Args:

        time: Time array

        visit: Visit array

        inst: Instrument array

        visit_to_plot: Optional visit ID to filter

        instrument_to_plot: Optional instrument ID to filter

    Returns:

        mask: Boolean array selecting valid data points

    Raises:

        ValueError: If no data points match selection

    """

    if visit_to_plot is not None:

        mask = visit == visit_to_plot

    elif instrument_to_plot is not None:

        mask = inst == instrument_to_plot

    else:

        mask = np.ones(len(time), dtype=bool)

    if np.sum(mask) == 0:

        raise ValueError(

            f"No data points found for visit={visit_to_plot}, "

            f"instrument={instrument_to_plot}"

        )

    return mask

def determine_data_type(inst_array, data_types_config):

    """

    Determine dominant data type (RV or LC) for a selection.

    Args:

        inst_array: Array of instrument IDs

        data_types_config: Dict mapping instrument ID → 'RV' or 'LC'

    Returns:

        str: Data type ('RV' or 'LC'), defaults to 'LC' if not found

    """

    unique_insts = np.unique(inst_array)

    data_types = [data_types_config.get(iid, 'LC') for iid in unique_insts]

    # Return first type (warn if mixed)

    return data_types[0] if data_types else 'LC'

def extract_masked_data(time, flux, err, visit, inst, mask):

    """

    Extract and convert masked data to NumPy arrays.

    Args:

        time, flux, err: Data arrays

        visit, inst: Visit and instrument arrays

        mask: Boolean selection mask

    Returns:

        tuple: (t, y, dy, vis, insts) as NumPy arrays

    """

    return (

        np.asarray(time[mask]),

        np.asarray(flux[mask]),

        np.asarray(err[mask]),

        np.array(visit[mask]),

        np.array(inst[mask])

    )

def sort_by_time(t, *arrays):

    """

    Sort all arrays by time.

    Args:

        t: Time array

        *arrays: Other arrays to sort in same order

    Returns:

        tuple: (t_sorted, array1_sorted, array2_sorted, ...)

    """

    sort_idx = np.argsort(t)

    return tuple(arr[sort_idx] for arr in (t, *arrays))