Changes for .kilo/plan.md: 96 added lines, 75 removed lines.
Original line number
Diff line number
Diff line
# Task Plan: Implementing Gap-Aware Horizontal Layout in `plot_models.py`
## Purpose
The goal is to modify the visual representation of model fits in `PyAETNA/reporting/plot_models.py`. Currently, the plot shows a single continuous time series (with potential gaps) in one large left column. The new requirement is to split this left column into multiple **sub-columns**, where each sub-column represents one contiguous cluster of data points identified by `detect_data_clusters` from `plot_model_helpers.py`.
Crucially, we must ensure that the **horizontal scale** (time units per pixel) is identical across all sub-columns so that a specific duration looks visually the same regardless of which subplot it appears in. The rightmost KDE column must remain unchanged in its width ratio (fixed at 1 unit relative to the total left side).
## Strategy & Methodology
1.**Cluster Detection**: Use `detect_data_clusters` on the sorted time array to identify $N$ contiguous blocks of data.
2.
3.**Proportional Width Calculation**:
* Calculate the time span $\Delta T_i$ for each cluster $i$.
* Total left-side width is fixed at **6 units**.
* The weight for each sub-column $i$ will be: $W_i = \left( \frac{\Delta T_i}{\sum_{j=1}^{N} \Delta T_j} \right) \times 6$.
* The width of the KDE column is fixed at **1 unit**.
4.**Grid Construction**: Use `matplotlib.gridspec.GridSpec` with `width_ratios = [W_1, W_2, ..., W_N, 1]`.
5.**Adaptive Plotting Loop**: Iterate through the detected clusters. For each cluster, slice the underlying arrays (`t`, `y`, `dy`, etc.) and populate the corresponding column in the `GridSpec` (e.g., Overlay, Detrended, Residuals rows).
6.**Consistency Checks**: Ensure all sub-columns share a common Y-axis within their respective rows to allow comparison across time gaps.
## Implementation Steps
### Step 1: Import and Detection
Import the cluster detection utility and run it on the masked/sorted time array.
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.
width_ratios=list(cluster_widths)+[kde_weight]
```
---
### Step 3: Grid Construction
Configure the `GridSpec` with the new column count and width ratios. The number of rows remains unchanged (3 or 4 depending on non-time trends).
```python
frommatplotlib.gridspecimportGridSpec
## 2️⃣ Identify Functional Areas (themes)
The current implementation can be grouped into the following logical sections:
num_rows=4if_show_panel0else3# Assume standard logic for panel counting
| **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. |
### Step 4: The Core Plotting Loop
Iterate through `clusters`. For each cluster $i$, create the vertical stack of axes (Overlay $\to$ Detrended $\to$ Residuals).
```python
# Example for one row of columns (e.g. Overlay)
axes_overlay=[]
foriinrange(n_clusters):
ax=fig.add_subplot(gs[row_index,i],sharey=...)# sharey with previous cluster in same row
axes_overlay.append(ax)
# Inside the loop: slice data for current cluster index
start,end=clusters[i]
t_c=t[start:end+1]
y_c=y[start:end+1]
# ... plot on ax ...
```
These themes map cleanly onto separate smaller Python modules.
### Step 5: Axis labels
---
Make sure that the plots with the data clusters have one shared y-axis label to the left of the left-most subplot and one shared x-axis label conveniently placed at the bottom-center of the series of clusters.
## 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
├── overlay_plots.py # functions that draw the overlay panel components
└── detrended_residuals.py # detrending, residual calculation and KDE inset
```
### Step 6: Final Verification
1.**Verify Imports**: Ensure `detect_data_clusters` is correctly imported and the name matches the alias used.
2.**Compilation Check**: Run `python -m py_compile PyAETNA/reporting/plot_models.py` to ensure no syntax errors or missing dependencies were introduced during refactoring.
3.**Visual Inspection**: Confirm that clusters with longer time spans appear physically wider on the screen and that the $x$ scale is perceptually identical across all columns.
* 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.*
