Commit c3e520ba authored by Gaetano Scandariato's avatar Gaetano Scandariato
Browse files

improved reporting and bug fixes

parent fb32bfb9
Loading
Loading
Loading
Loading
Loading
+27 −0
Changes for .kilo/legend_relocation_plan.md: 27 added lines, 0 removed lines.
Original line number Diff line number Diff line
# Plan: Relocate Legends for Overlay and Detrended Rows

## Motivation
Currently, legends for the model overlay and detrended rows are placed within the first subpanel (leftmost column). In multi-cluster/multi-instrument plots, this can clutter the data visualization in the first panel while leaving the right side of the figure underutilized. 

To maintain consistency with the residual row—which uses a dedicated KDE inset axis on the right—the legends for the upper rows should be moved to a similar dedicated space. This improves visual balance and ensures that data in the first subpanel is not obscured by the legend.

## Objectives
1.  **Modify Figure Layout**: Update `FigureLayoutManager` to expand the `GridSpec` by one column. This new rightmost column will be dedicated to legends for the `overlay` and `detrended` rows.
2.  **Implement Legend Axes**: Create `legend_ax` for the overlay and detrended rows in the new column, ensuring `axis('off')` is called to keep them borderless.
3.  **Update Overlay Rendering**: Modify `_render_rv_overlay` and `_render_lc_overlay` in `panel_overlay.py` to:
    *   Remove labels from the primary subpanels (set all to `_nolegend_`).
    *   Use the new dedicated `legend_ax` to render the legend, using `bbox_to_anchor` for precise centering.
4.  **Update Detrended Rendering**: Apply similar logic to the detrended row rendering.
5.  **Verification**: Ensure that the shift in `GridSpec` does not cause overlap between the new legends and the dynamically positioned KDE insets.

## Technical Steps
- [ ] **`PyAETNA/reporting/plot_styling.py`**: 
    - Update `FigureLayoutManager.create()`:
        - Change `GridSpec(self.n_panels, self.n_clusters, ...)` to `GridSpec(self.n_panels, self.n_clusters + 1, ...)`.
        - Add `width_ratios` to the `GridSpec` (e.g., `[1]*self.n_clusters + [0.1]`) to keep the legend column narrow.
        - Instantiate `self.axes['overlay_legend']` and `self.axes['detrended_legend']` using `self.fig.add_subplot(gs[row, self.n_clusters])`.
        - Call `.axis('off')` on these new axes.
- [ ] **`PyAETNA/reporting/panel_overlay.py`**:
    - Update `_render_rv_overlay` and `_render_lc_overlay` to stop using `lbl = j == 0`.
    - Use the `overlay_legend` axis to call `.legend()` with `loc='center left'` and `bbox_to_anchor=(0, 0.5)`.
- [ ] **Detrended Row Logic**: Identify the rendering function for the detrended row and implement the legend relocation to `detrended_legend`.
+21 −23
Changes for PyAETNA/reporting/panel_detrended.py: 21 added lines, 23 removed lines.
Original line number Diff line number Diff line
@@ -19,13 +19,14 @@ from .plot_data_prep import PlotContext
from .plot_styling import add_unc_ref_bars


def render_panel_detrended(detrended_axes_raw, ctx: PlotContext) -> None:
def render_panel_detrended(detrended_axes_raw, ctx: PlotContext, ax_legend=None) -> None:
    """Render Panel 2: detrended data + mean-function overlay.

    Parameters
    ----------
    detrended_axes_raw : list[Axes] or Axes
    ctx : PlotContext
    ax_legend : Optional axis for the legend
    """
    detrended_axes = (detrended_axes_raw if isinstance(detrended_axes_raw, list)
                      else [detrended_axes_raw])
@@ -39,8 +40,7 @@ def render_panel_detrended(detrended_axes_raw, ctx: PlotContext) -> None:
                    ctx.t[sel], ctx.y_detrended[sel], yerr=ctx.dy[sel],
                    fmt='.', markersize=3, elinewidth=0.5, alpha=0.6,
                    color=c,
                    label=(f'Visit {uv} (trend+GP removed)' if j == 0
                           else '_nolegend_'),
                    label='_nolegend_',
                    zorder=1)
    else:
        for j, ax in enumerate(detrended_axes):
@@ -48,56 +48,50 @@ def render_panel_detrended(detrended_axes_raw, ctx: PlotContext) -> None:
                ctx.t, ctx.y_detrended, yerr=ctx.dy,
                fmt='.', markersize=3, elinewidth=0.5, alpha=0.6,
                color='gray',
                label='Data (trend + GP removed)' if j == 0 else '_nolegend_',
                label='_nolegend_',
                zorder=1)

    # ── Mean-function model line and band ─────────────────────────────────
    if ctx.has_planets:
        for j, ax in enumerate(detrended_axes):
            lbl = j == 0
            if ctx.data_type_in_mask == 'RV':
                ax.plot(ctx.t_mf, ctx.pm_mf, '-', color='C0', linewidth=1.5,
                ax.plot(ctx.t_mf, ctx.pm_mf, '-', color='green', linewidth=1.5,
                        zorder=2,
                        label=('Mean function — planet only (1σ band)'
                               if lbl else '_nolegend_'))
                        label='RV Model')
                if ctx.mf_q16_grid is not None:
                    ax.fill_between(ctx.t_grid, ctx.mf_q16_grid,
                                    ctx.mf_q84_grid,
                                    color='C0', alpha=0.25, zorder=1.5,
                                    color='green', alpha=0.25, zorder=1.5,
                                    label='_nolegend_')
                elif ctx.pm_q16_grid is not None:
                    ax.fill_between(ctx.t_grid, ctx.pm_q16_grid,
                                    ctx.pm_q84_grid,
                                    color='C0', alpha=0.25, zorder=1.5,
                                    color='green', alpha=0.25, zorder=1.5,
                                    label='_nolegend_')
            else:
                ax.plot(ctx.t_mf, 1.0 + ctx.pm_mf, '-', color='C0',
                ax.plot(ctx.t_mf, 1.0 + ctx.pm_mf, '-', color='green',
                        linewidth=1.5, zorder=2,
                        label=('Mean function — 1 + planet (1σ band)'
                               if lbl else '_nolegend_'))
                        label='LC Model')
                if ctx.mf_q16_grid is not None:
                    ax.fill_between(ctx.t_grid, ctx.mf_q16_grid,
                                    ctx.mf_q84_grid,
                                    color='C0', alpha=0.25, zorder=1.5,
                                    color='green', alpha=0.25, zorder=1.5,
                                    label='_nolegend_')
                elif ctx.pm_q16_grid is not None:
                    ax.fill_between(ctx.t_grid, 1.0 + ctx.pm_q16_grid,
                                    1.0 + ctx.pm_q84_grid,
                                    color='C0', alpha=0.25, zorder=1.5,
                                    color='green', alpha=0.25, zorder=1.5,
                                    label='_nolegend_')
    else:
        for j, ax in enumerate(detrended_axes):
            lbl = j == 0
            if ctx.data_type_in_mask == 'RV':
                ax.axhline(0.0, color='C0', linestyle='-', linewidth=1.5,
                ax.axhline(0.0, color='green', linestyle='-', linewidth=1.5,
                           zorder=2,
                           label=('Mean function — no planet (zero RV)'
                                  if lbl else '_nolegend_'))
                           label='_nolegend_')
            else:
                ax.axhline(1.0, color='C0', linestyle='-', linewidth=1.5,
                ax.axhline(1.0, color='green', linestyle='-', linewidth=1.5,
                           zorder=2,
                           label=('Mean function — no planet (flat baseline)'
                                  if lbl else '_nolegend_'))
                           label='_nolegend_')

    # ── X-limits per cluster ──────────────────────────────────────────────
    for j, ax in enumerate(detrended_axes):
@@ -116,4 +110,8 @@ def render_panel_detrended(detrended_axes_raw, ctx: PlotContext) -> None:
    add_unc_ref_bars(detrended_axes[0], ctx.x_ref, y_center,
                     ctx.med_raw, ctx.med_tot, ctx.fmt_unc,
                     ctx.is_per_inst_plot)
    detrended_axes[0].legend(loc='upper right', fontsize=7, framealpha=0.9)
    
    # Add legend to the dedicated axis using handles from the first panel
    if ax_legend is not None and len(detrended_axes) > 0:
        handles, labels = detrended_axes[0].get_legend_handles_labels()
        ax_legend.legend(handles, labels, loc='center left', fontsize=7, frameon=False)
+20 −42
Changes for PyAETNA/reporting/panel_overlay.py: 20 added lines, 42 removed lines.
Original line number Diff line number Diff line
@@ -30,7 +30,7 @@ def _set_cluster_xlims(ax_list: list, ctx: PlotContext) -> None:
            ax.set_xlim(c_start - buf, c_end + buf)


def _render_rv_overlay(overlay_axes: list, ctx: PlotContext) -> None:
def _render_rv_overlay(overlay_axes: list, ctx: PlotContext, ax_legend=None) -> None:
    """Render per-component lines for an RV overlay panel."""
    has_trend = np.any(ctx.tr != 0)
    has_gp    = ctx.show_gp and np.any(ctx.gp != 0)
@@ -40,12 +40,10 @@ def _render_rv_overlay(overlay_axes: list, ctx: PlotContext) -> None:
    tr_mf_ov = ctx.tr_time_mf if ctx.has_non_time_trends else ctx.tr_mf

    for j, ax in enumerate(overlay_axes):
        lbl = j == 0  # only label lines on the first sub-panel

        # 1. Normalization γ
        ax.plot(ctx.t_mf, ctx.norm_mf, linestyle='--', color='C0',
                linewidth=1.2, zorder=2, alpha=0.8,
                label='γ (systemic RV)' if lbl else '_nolegend_')
                label='Normalization' if j == 0 else '_nolegend_')
        if ctx.norm_q16_grid is not None:
            ax.fill_between(ctx.t_grid, ctx.norm_q16_grid, ctx.norm_q84_grid,
                            color='C0', alpha=0.12, zorder=1.5,
@@ -55,7 +53,7 @@ def _render_rv_overlay(overlay_axes: list, ctx: PlotContext) -> None:
        if ctx.has_planets:
            ax.plot(ctx.t_mf, ctx.norm_mf + ctx.pm_mf, '-', color='C1',
                    linewidth=1.5, zorder=3,
                    label='Planet + γ' if lbl else '_nolegend_')
                    label='Planet + Norm' if j == 0 else '_nolegend_')
            if ctx.pn_q16_grid is not None:
                ax.fill_between(ctx.t_grid, ctx.pn_q16_grid, ctx.pn_q84_grid,
                                color='C1', alpha=0.15, zorder=2.5,
@@ -66,7 +64,7 @@ def _render_rv_overlay(overlay_axes: list, ctx: PlotContext) -> None:
            tlbl = ctx.trend_label
            ax.plot(ctx.t_mf, ctx.norm_mf + tr_mf_ov, '-', color='C4',
                    linewidth=1.2, zorder=3, alpha=0.8,
                    label=f'{tlbl} + γ' if lbl else '_nolegend_')
                    label=f'{tlbl} + Norm' if j == 0 else '_nolegend_')
            if ctx.tr_time_centred_q16_grid is not None:
                ax.fill_between(ctx.t_grid,
                                ctx.tr_time_centred_q16_grid,
@@ -74,22 +72,11 @@ def _render_rv_overlay(overlay_axes: list, ctx: PlotContext) -> None:
                                color='C4', alpha=0.12, zorder=2.5,
                                label='_nolegend_')

        # 4. GP + γ
        if has_gp:
            ax.plot(ctx.t_grid, ctx.norm_grid + ctx.gp_grid, '-', color='C2',
                    linewidth=1.2, zorder=4, alpha=0.8,
                    label='GP + γ' if lbl else '_nolegend_')
            if ctx.gp_centred_q16_grid is not None:
                ax.fill_between(ctx.t_grid,
                                ctx.gp_centred_q16_grid,
                                ctx.gp_centred_q84_grid,
                                color='C2', alpha=0.12, zorder=3.5,
                                label='_nolegend_')

        # Legend is now handled in render_panel_overlay to include all elements
    _set_cluster_xlims(overlay_axes, ctx)


def _render_lc_overlay(overlay_axes: list, ctx: PlotContext) -> None:
def _render_lc_overlay(overlay_axes: list, ctx: PlotContext, ax_legend=None) -> None:
    """Render per-component lines for an LC overlay panel."""
    has_trend = np.any(ctx.tr != 0)
    has_gp    = ctx.show_gp and np.any(ctx.gp != 0)
@@ -100,13 +87,12 @@ def _render_lc_overlay(overlay_axes: list, ctx: PlotContext) -> None:
    mode   = ctx.normalization_mode

    for j, ax in enumerate(overlay_axes):
        lbl = j == 0

        # All lines are now _nolegend_ on the main panels
        # 1. Baseline
        norm_baseline = build_partial_for_plot(ctx.norm_mf, z_mf, z_mf, mode, 'LC')
        ax.plot(ctx.t_mf, norm_baseline, linestyle='--', color='C0',
                linewidth=1.2, zorder=2, alpha=0.8,
                label='Baseline' if lbl else '_nolegend_')
                label='Baseline' if j == 0 else '_nolegend_')
        if ctx.norm_q16_grid is not None:
            nb_q16 = build_partial_for_plot(ctx.norm_q16_grid, z_grid, z_grid, mode, 'LC')
            nb_q84 = build_partial_for_plot(ctx.norm_q84_grid, z_grid, z_grid, mode, 'LC')
@@ -120,7 +106,7 @@ def _render_lc_overlay(overlay_axes: list, ctx: PlotContext) -> None:
                ctx.norm_mf, ctx.pm_mf, z_mf, mode, 'LC')
            ax.plot(ctx.t_mf, planet_baseline, '-', color='C1',
                    linewidth=1.5, zorder=3,
                    label='Planet + baseline' if lbl else '_nolegend_')
                    label='Planet + Baseline' if j == 0 else '_nolegend_')
            if ctx.pn_q16_grid is not None:
                ax.fill_between(ctx.t_grid, ctx.pn_q16_grid, ctx.pn_q84_grid,
                                color='C1', alpha=0.15, zorder=2.5,
@@ -132,7 +118,7 @@ def _render_lc_overlay(overlay_axes: list, ctx: PlotContext) -> None:
                ctx.norm_mf, z_mf, tr_mf_ov, mode, 'LC')
            ax.plot(ctx.t_mf, trend_baseline, '-', color='C4',
                    linewidth=1.2, zorder=3, alpha=0.8,
                    label=f'{ctx.trend_label} + baseline' if lbl else '_nolegend_')
                    label=f'{ctx.trend_label} + Baseline' if j == 0 else '_nolegend_')
            if ctx.tr_time_centred_q16_grid is not None:
                ax.fill_between(ctx.t_grid,
                                ctx.tr_time_centred_q16_grid,
@@ -140,24 +126,11 @@ def _render_lc_overlay(overlay_axes: list, ctx: PlotContext) -> None:
                                color='C4', alpha=0.12, zorder=2.5,
                                label='_nolegend_')

        # 4. GP + baseline
        if has_gp:
            gp_baseline = build_composite_for_plot(
                ctx.norm_grid, z_grid, z_grid, ctx.gp_grid, mode, 'LC')
            ax.plot(ctx.t_grid, gp_baseline, '-', color='C2',
                    linewidth=1.2, zorder=4, alpha=0.8,
                    label='GP + baseline' if lbl else '_nolegend_')
            if ctx.gp_centred_q16_grid is not None:
                ax.fill_between(ctx.t_grid,
                                ctx.gp_centred_q16_grid,
                                ctx.gp_centred_q84_grid,
                                color='C2', alpha=0.12, zorder=3.5,
                                label='_nolegend_')

    # Legend is now handled in render_panel_overlay to include all elements
    _set_cluster_xlims(overlay_axes, ctx)


def render_panel_overlay(overlay_axes_raw, ctx: PlotContext) -> None:
def render_panel_overlay(overlay_axes_raw, ctx: PlotContext, ax_legend=None) -> None:
    """Render Panel 1: data scatter + individual model-component lines.

    When ``ctx.has_non_time_trends`` is True, the data shown is the
@@ -170,6 +143,7 @@ def render_panel_overlay(overlay_axes_raw, ctx: PlotContext) -> None:
    overlay_axes_raw : list[Axes] or Axes
        Axes for the overlay row from ``create_figure_layout``.
    ctx : PlotContext
    ax_legend : Optional axis for the legend
    """
    overlay_axes = (overlay_axes_raw if isinstance(overlay_axes_raw, list)
                    else [overlay_axes_raw])
@@ -198,9 +172,9 @@ def render_panel_overlay(overlay_axes_raw, ctx: PlotContext) -> None:

    # ── Model-component lines ─────────────────────────────────────────────
    if ctx.data_type_in_mask == 'RV':
        _render_rv_overlay(overlay_axes, ctx)
        _render_rv_overlay(overlay_axes, ctx, ax_legend=ax_legend)
    else:
        _render_lc_overlay(overlay_axes, ctx)
        _render_lc_overlay(overlay_axes, ctx, ax_legend=ax_legend)

    # ── Y-label, reference bars, legend ──────────────────────────────────
    if ctx.has_non_time_trends:
@@ -215,4 +189,8 @@ def render_panel_overlay(overlay_axes_raw, ctx: PlotContext) -> None:
    add_unc_ref_bars(overlay_axes[0], ctx.x_ref, y_center,
                     ctx.med_raw, ctx.med_tot, ctx.fmt_unc,
                     ctx.is_per_inst_plot)
    overlay_axes[0].legend(loc='upper right', fontsize=7, framealpha=0.9)
    
    # Add legend to the dedicated axis using handles from the first panel
    if ax_legend is not None and len(overlay_axes) > 0:
        handles, labels = overlay_axes[0].get_legend_handles_labels()
        ax_legend.legend(handles, labels, loc='center left', fontsize=7, frameon=False)
+11 −4
Changes for PyAETNA/reporting/panel_raw.py: 11 added lines, 4 removed lines.
Original line number Diff line number Diff line
@@ -17,7 +17,7 @@ import matplotlib.pyplot as plt
from .plot_data_prep import PlotContext


def render_panel_raw(ax_raw_panel, ctx: PlotContext) -> None:
def render_panel_raw(ax_raw_panel, ctx: PlotContext, ax_legend=None) -> None:
    """Render the raw-data + non-time correction panel.

    Plots the unmodified observations together with the best-fit non-time
@@ -29,6 +29,7 @@ def render_panel_raw(ax_raw_panel, ctx: PlotContext) -> None:
    ax_raw_panel : list[Axes] or Axes or None
        Axes returned by ``create_figure_layout``.  ``None`` → no-op.
    ctx : PlotContext
    ax_legend : Optional axis for the legend
    """
    if ax_raw_panel is None:
        return
@@ -37,9 +38,9 @@ def render_panel_raw(ax_raw_panel, ctx: PlotContext) -> None:

    # Non-time correction in data space
    if ctx.data_type_in_mask == 'RV' or ctx.normalization_mode == 'additive':
        nt_line = ctx.norm + ctx.tr_nt_panel0
        nt_line = ctx.norm + (ctx.tr_nt_panel0 if ctx.tr_nt_panel0 is not None else 0.0)
    else:
        nt_line = ctx.norm * (1.0 + ctx.tr_nt_panel0)
        nt_line = ctx.norm * (1.0 + (ctx.tr_nt_panel0 if ctx.tr_nt_panel0 is not None else 0.0))

    for j, ax in enumerate(raw_axes):
        # Data points, colour-coded by visit for per-instrument plots
@@ -87,7 +88,13 @@ def render_panel_raw(ax_raw_panel, ctx: PlotContext) -> None:

    ylabel = ('Radial Velocity [m/s]' if ctx.data_type_in_mask == 'RV'
              else 'Relative Flux')
    if raw_axes:
        raw_axes[0].set_ylabel(ylabel)
    raw_axes[0].legend(loc='upper right', fontsize=7, framealpha=0.9)
        
        # Add legend to the dedicated axis using handles from the first panel
        if ax_legend is not None and len(raw_axes) > 0:
            handles, labels = raw_axes[0].get_legend_handles_labels()
            ax_legend.legend(handles, labels, loc='center left', fontsize=7, frameon=False)

    for ax in raw_axes:
        plt.setp(ax.get_xticklabels(), visible=False)
+2 −2
Changes for PyAETNA/reporting/panel_residuals.py: 2 added lines, 2 removed lines.
Original line number Diff line number Diff line
@@ -71,8 +71,8 @@ def render_panel_residuals(residual_axes_raw, ctx: PlotContext) -> None:
                     ctx.med_raw, ctx.med_tot, ctx.fmt_unc,
                     ctx.is_per_inst_plot)
    if ctx.is_per_inst_plot:
        residual_axes[0].legend(loc='lower right', fontsize=7, framealpha=0.9)
        residual_axes[0].legend(loc='best', fontsize=7, framealpha=0.9)
    else:
        residual_axes[0].legend(loc='lower right', fontsize=7, framealpha=0.9,
        residual_axes[0].legend(loc='best', fontsize=7, framealpha=0.9,
                                title='Median uncertainties')
    residual_axes[0].set_ylabel('Residuals (σ_orig error bars)')
Loading