Fix jnp.unique usage in relabel_array and update config file for test cases

# Refactoring Plan for ``plot_models.py``

## 1️⃣ Goal

Split the monolithic **`PyAETNA/reporting/plot_models.py`** into a set of focused, smaller thematic sub‑modules while:

- Removing *ghost variables* (imports or locals that are never used).

- Eliminating duplicated code blocks.

- Simplifying complex logic where possible.

The end result should be efficient, JIT compliant, easier to read, test and maintain, with each file having a single responsibility. Also `PyAETNA/reporting/plot_models.py` must be smaller after moving parts of the monolithic code to delegated submodules.

---

## 2️⃣ Identify Functional Areas (themes)

The current implementation can be grouped into the following logical sections:

| Theme | Typical responsibilities |

|------|---------------------------|

| **Model component computation** | Functions that build model pieces (`compute_model_components`, `build_mean_function`, `detrend_data`, etc.). |

| **Confidence‑band calculations** | Sampling, percentile extraction, non‑time trend handling, GP‑process variance logic. |

| **Plot utilities & styling** | Matplotlib figure/axes creation, colour maps, legend helpers, axis formatting. |

| **Overlay panel rendering** | Code that draws the individual model components (normalisation, planet, trend, GP) on the overlay axes. |

| **Detrended / residual panels** | Logic for detrending data, plotting the detrended series and residuals, KDE inset. |

These themes map cleanly onto separate smaller Python modules.

---

## 3️⃣ Proposed Sub‑module Layout

```

PyAETNA/reporting/

│   plot_models.py          # thin façade that orchestrates the workflow

│

├── model_components.py    # core model building helpers

├── confidence_bands.py    # sampling, percentile & GP variance utilities

├── plot_utils.py          # generic Matplotlib helpers (fig size, gridspec, colours)

├── overlay_plots.py       # functions that draw the overlay panel components

└── detrended_residuals.py # detrending, residual calculation and KDE inset

```

### File responsibilities

* **`model_components.py`**

	* `compute_model_components`, `build_mean_function`, `detrend_data`, `build_partial_for_plot`, `build_composite_for_plot`.

	* Any low‑level helpers that operate on `plData`, `plModel` and the parameter dictionary.

* **`confidence_bands.py`**

	* Sampling logic (`samples_dict`, `_sample_cache`).

	* Percentile extraction (`np.percentile`) for each component.

	* GP‑specific helpers: `_precompute_gp_grid_context`, `_gp_predict_on_grid`, and the law‑of‑total‑variance widening.

* **`plot_utils.py`**

	* Figure creation, `GridSpec` handling, width‑ratio calculation, cluster detection (`detect_data_clusters`).

	* Common colour palettes (`_get_distinct_colors`).

	* Helper to add reference uncertainty bars (`_add_unc_ref_bars`).

* **`overlay_plots.py`**

	* Functions that receive the prepared model arrays and draw the four component lines (normalisation, planet, trend, GP) with their confidence bands.

	* Separate implementations for RV vs LC where necessary.

* **`detrended_residuals.py`**

	* `detrend_data` usage to produce the detrended series.

	* Residual calculation and plotting, including the KDE inset logic.

* **`plot_models.py` (facade)**

	* Import the above modules.

	* Validate arguments, decide which panels to draw, orchestrate the workflow.

	* Keep only high‑level orchestration – no heavy numerical code.

---

## 4️⃣ Refactoring Steps

1. **Create the new files** with minimal stub content and import statements. ✅ Completed

2. **Move functions** and code blocks from `plot_models.py` to `model_components.py` according to the logic stated in sections 2 and 3 (copy‑paste, then delete from the original). Keep the public API identical and update the imports in `plot_models.py`.

3. **Move functions** and code blocks from `plot_models.py` to `confidence_bands.py` according to the logic stated in sections 2 and 3 (copy‑paste, then delete from the original). Keep the public API identical and update the imports in `plot_models.py`.

4. **Run the test** executing a quick end‑to‑end run of `runPyAETNA.py` to ensure nothing broke.

5. **Clean** `plot_models.py` by deleting the code moved to the submodules and adding the calls to the delegated functions.

6. **Search for ghost variables**:

	 * Unused imports (`jax`, `warnings`, etc.).

	 * Locally defined variables that are never read after assignment (e.g., temporary caches).

7. **Detect duplicate code blocks** – many sections repeat the same error‑bar plotting logic for different panels; extract a helper like `plot_errorbars(ax, t, y, dy, style)` into `plot_utils.py`.

8. **Simplify complex conditionals**:

	 * Replace nested `if … else …` chains that only toggle a boolean with a single expression (`show_gp = show_gp and np.any(gp != 0)`).

	 * Use list comprehensions for building colour maps instead of explicit loops where possible.

9. **Add docstrings** to the newly created functions, describing inputs/outputs – this will aid future maintenance.

10. **Remove** test scripts, stubs and code for backward compatibility produced during the refactoring.

11. **Run** final test by executing `runPyAETNA.py` to catch regressions. If the run aborts, debug the code and rerun the script iteratively.

---

## 5️⃣ Checklist after Refactor

- [ ] All original functionality works (visual output unchanged).

- [ ] No `ImportError` when running the package.

- [ ] Linter (`flake8`/`ruff`) reports zero unused imports or variables.

- [ ] Duplicate code sections have been replaced by shared helpers.

- [ ] Unit tests (if present) pass; otherwise, manually verify a few plots.

---

## 6️⃣ Next Actions for the User

1. **Run the refactor** using the steps above or let me generate the skeleton files automatically.

2. After moving code, run `runPyAETNA.py` to catch regressions. If the run aborts, debug the code and rerun the script iteratively.

3. If you encounter any ambiguous sections (e.g., a variable whose purpose is unclear), ask for clarification – I can help decide whether it should stay or be removed.

---

*This plan is now stored in `/.kilo/plan.md` and can be referenced later while performing the actual code moves.*

# 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

def relabel_array(arr):

    # 'inverse' returns an array of the same shape as 'arr'

    # where each element is the index of that value in the unique array

    unique_vals, inverse = jnp.unique(arr, return_inverse=True)

    # Use a static size for jnp.unique to avoid dynamic-output-shape bugs

    # in JAX's Thunk runtime on newer XLA CPU backends.

    unique_vals, inverse = jnp.unique(

        arr,

        return_inverse=True,

        size=arr.size,

        fill_value=-1,

    )

    # JAX indices are 0-based, so add 1 to start from 1

    new_labels = inverse + 1

- 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

- implement different limb darkening laws

- refactor reporting and add timeZeroPoint

- refactor sampling and consider planThunk

- add TIMEREF to configuration and show it in the plots (update generator and documentation)

- refactor sampling 

## Contact

        # configFile = 'tests/test_lc_rv_joint.yaml'

        # configFile = 'tests/test_lc_gp_rv.yaml'

        # configFile = 'tests/test_phase_curve.yaml'

        # configFile = 'tests/test_ttv.yaml'

        configFile = 'tests/test_spectral_mixture.yaml'

        configFile = 'tests/test_ttv.yaml'

        # configFile = 'tests/test_spectral_mixture.yaml'

    # Run with the config file (data.working_dir and data.file_path are read from YAML)

    run(configFile)