plot_models.py monolith split into 5 new files plus a refactored orchestrator (3c47486c) · Commits · Gaetano Scandariato / PyAETNA

PyAETNA/reporting/panel_detrended.py

0 → 100644

+119 −0

Changes for PyAETNA/reporting/panel_detrended.py: 119 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		# -- coding: utf-8 --
		"""
		panel_detrended.py
		~~~~~~~~~~~~~~~~~~
		Renderer for Panel 2: detrended data + mean-function model line.

		The data shown is ``ctx.y_detrended`` — observations with the normalisation
		baseline, polynomial/Fourier trend, and GP prediction all subtracted. The
		model line is the mean-function signal only (transit shape for LC, Keplerian
		for RV), with a 68% band from planet-parameter samples alone.

		Functions
		---------
		render_panel_detrended
		"""
		import numpy as np

		from .plot_data_prep import PlotContext
		from .plot_styling import add_unc_ref_bars


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

		Parameters
		----------
		detrended_axes_raw : list[Axes] or Axes
		ctx : PlotContext
		"""
		detrended_axes = (detrended_axes_raw if isinstance(detrended_axes_raw, list)
		else [detrended_axes_raw])

		# ── Data scatter ──────────────────────────────────────────────────────
		if ctx.is_per_inst_plot and ctx.visit_colors is not None:
		for j, ax in enumerate(detrended_axes):
		for uv, c in ctx.visit_colors.items():
		sel = ctx.vis_sorted == uv
		ax.errorbar(
		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_'),
		zorder=1)
		else:
		for j, ax in enumerate(detrended_axes):
		ax.errorbar(
		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_',
		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,
		zorder=2,
		label=('Mean function — planet only (1σ band)'
		if lbl else '_nolegend_'))
		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,
		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,
		label='_nolegend_')
		else:
		ax.plot(ctx.t_mf, 1.0 + ctx.pm_mf, '-', color='C0',
		linewidth=1.5, zorder=2,
		label=('Mean function — 1 + planet (1σ band)'
		if lbl else '_nolegend_'))
		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,
		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,
		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,
		zorder=2,
		label=('Mean function — no planet (zero RV)'
		if lbl else '_nolegend_'))
		else:
		ax.axhline(1.0, color='C0', linestyle='-', linewidth=1.5,
		zorder=2,
		label=('Mean function — no planet (flat baseline)'
		if lbl else '_nolegend_'))

		# ── X-limits per cluster ──────────────────────────────────────────────
		for j, ax in enumerate(detrended_axes):
		if ctx.is_per_inst_plot and ctx.clusters and j < len(ctx.clusters):
		c_start, c_end = ctx.clusters[j]
		buf = 0.05 * (c_end - c_start)
		ax.set_xlim(c_start - buf, c_end + buf)

		# ── Y-label, reference bars, legend ──────────────────────────────────
		ylabel = ('RV − trend − GP [m/s]' if ctx.data_type_in_mask == 'RV'
		else 'Detrended flux (trend + GP removed)')
		detrended_axes[0].set_ylabel(ylabel)

		y_center = (float(np.min(ctx.y_detrended)) +
		float(np.max(ctx.y_detrended))) / 2.0
		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)

PyAETNA/reporting/panel_overlay.py

0 → 100644

+218 −0

Changes for PyAETNA/reporting/panel_overlay.py: 218 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		# -- coding: utf-8 --
		"""
		panel_overlay.py
		~~~~~~~~~~~~~~~~
		Renderer for Panel 1: model-component overlay.

		Plots data (with non-time trends removed when applicable) together with
		individual model components (normalization, planet, trend, GP), each with
		their 68% confidence band.

		Functions
		---------
		render_panel_overlay
		"""
		import numpy as np

		from PyAETNA.model_utils import build_partial_for_plot, build_composite_for_plot
		from .plot_data_prep import PlotContext
		from .plot_styling import add_unc_ref_bars


		def _set_cluster_xlims(ax_list: list, ctx: PlotContext) -> None:
		"""Apply per-cluster x-limits when showing a per-instrument plot."""
		if not (ctx.is_per_inst_plot and ctx.clusters):
		return
		for j, ax in enumerate(ax_list):
		if j < len(ctx.clusters):
		c_start, c_end = ctx.clusters[j]
		buf = 0.05 * (c_end - c_start)
		ax.set_xlim(c_start - buf, c_end + buf)


		def _render_rv_overlay(overlay_axes: list, ctx: PlotContext) -> 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)

		# For the overlay panel use time-only trend when non-time trends exist,
		# otherwise use the full fine-grid trend.
		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_')
		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,
		label='_nolegend_')

		# 2. Planet + γ
		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_')
		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,
		label='_nolegend_')

		# 3. Trend + γ
		if has_trend:
		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_')
		if ctx.tr_time_centred_q16_grid is not None:
		ax.fill_between(ctx.t_grid,
		ctx.tr_time_centred_q16_grid,
		ctx.tr_time_centred_q84_grid,
		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_')

		_set_cluster_xlims(overlay_axes, ctx)


		def _render_lc_overlay(overlay_axes: list, ctx: PlotContext) -> 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)
		tr_mf_ov = ctx.tr_time_mf if ctx.has_non_time_trends else ctx.tr_mf

		z_grid = np.zeros_like(ctx.pm_grid)
		z_mf = np.zeros_like(ctx.pm_mf)
		mode = ctx.normalization_mode

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

		# 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_')
		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')
		ax.fill_between(ctx.t_grid, nb_q16, nb_q84,
		color='C0', alpha=0.12, zorder=1.5,
		label='_nolegend_')

		# 2. Planet + baseline
		if ctx.has_planets:
		planet_baseline = build_partial_for_plot(
		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_')
		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,
		label='_nolegend_')

		# 3. Trend + baseline
		if has_trend:
		trend_baseline = build_partial_for_plot(
		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_')
		if ctx.tr_time_centred_q16_grid is not None:
		ax.fill_between(ctx.t_grid,
		ctx.tr_time_centred_q16_grid,
		ctx.tr_time_centred_q84_grid,
		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_')

		_set_cluster_xlims(overlay_axes, ctx)


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

		When ``ctx.has_non_time_trends`` is True, the data shown is the
		non-time-detrended ``ctx.overlay_y`` and model lines use the time-only
		trend counterparts already pre-computed in ``PlotContext``. No
		save/restore of variables is required.

		Parameters
		----------
		overlay_axes_raw : list[Axes] or Axes
		Axes for the overlay row from ``create_figure_layout``.
		ctx : PlotContext
		"""
		overlay_axes = (overlay_axes_raw if isinstance(overlay_axes_raw, list)
		else [overlay_axes_raw])

		# ── Data scatter ──────────────────────────────────────────────────────
		if ctx.is_per_inst_plot and ctx.visit_colors is not None:
		for j, ax in enumerate(overlay_axes):
		for uv, c in ctx.visit_colors.items():
		sel = ctx.vis_sorted == uv
		ax.errorbar(
		ctx.t[sel], ctx.overlay_y[sel], yerr=ctx.dy[sel],
		fmt='.', markersize=3, elinewidth=0.5, alpha=0.6,
		color=c,
		label=f'Visit {uv}' if j == 0 else '_nolegend_',
		zorder=1)
		else:
		for j, ax in enumerate(overlay_axes):
		ax.errorbar(
		ctx.t, ctx.overlay_y, yerr=ctx.dy,
		fmt='.', markersize=3, elinewidth=0.5, alpha=0.6,
		color='gray', label='Data' if j == 0 else '_nolegend_',
		zorder=1)

		# X-limits for data scatter (clusters)
		_set_cluster_xlims(overlay_axes, ctx)

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

		# ── Y-label, reference bars, legend ──────────────────────────────────
		if ctx.has_non_time_trends:
		ylabel = ('RV − non-time trends [m/s]' if ctx.data_type_in_mask == 'RV'
		else 'Flux − non-time trends')
		else:
		ylabel = ('Radial Velocity [m/s]' if ctx.data_type_in_mask == 'RV'
		else 'Relative Flux')
		overlay_axes[0].set_ylabel(ylabel)

		y_center = (float(np.min(ctx.overlay_y)) + float(np.max(ctx.overlay_y))) / 2.0
		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)

PyAETNA/reporting/panel_raw.py

0 → 100644

+93 −0

Changes for PyAETNA/reporting/panel_raw.py: 93 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		# -- coding: utf-8 --
		"""
		panel_raw.py
		~~~~~~~~~~~~
		Renderer for Panel 0 (raw data + non-time correction).

		Shown only when ``ctx.show_panel0`` is True, i.e. when at least one
		polynomial/Fourier/GP component operates on a non-BJD variable.

		Functions
		---------
		render_panel_raw
		"""
		import numpy as np
		import matplotlib.pyplot as plt

		from .plot_data_prep import PlotContext


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

		Plots the unmodified observations together with the best-fit non-time
		correction line (one segment per visit) and, when samples are available,
		a 68% confidence band. X-limits are set to each cluster's extent.

		Parameters
		----------
		ax_raw_panel : list[Axes] or Axes or None
		Axes returned by ``create_figure_layout``. ``None`` → no-op.
		ctx : PlotContext
		"""
		if ax_raw_panel is None:
		return

		raw_axes = ax_raw_panel if isinstance(ax_raw_panel, list) else [ax_raw_panel]

		# 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
		else:
		nt_line = ctx.norm * (1.0 + ctx.tr_nt_panel0)

		for j, ax in enumerate(raw_axes):
		# Data points, colour-coded by visit for per-instrument plots
		if ctx.is_per_inst_plot and ctx.visit_colors is not None:
		for uv, c in ctx.visit_colors.items():
		sel = ctx.vis_sorted == uv
		ax.errorbar(
		ctx.t[sel], ctx.y[sel], yerr=ctx.dy[sel],
		fmt='.', markersize=3, elinewidth=0.5, alpha=0.6,
		color=c, label=f'Visit {uv}' if j == 0 else '_nolegend_',
		zorder=1)
		else:
		ax.errorbar(
		ctx.t, ctx.y, yerr=ctx.dy,
		fmt='.', markersize=3, elinewidth=0.5, alpha=0.6,
		color='gray', label='Data' if j == 0 else '_nolegend_',
		zorder=1)

		# Non-time correction line, drawn per visit to avoid inter-season gaps
		first_vis = True
		for uv in np.unique(ctx.vis_sorted):
		sel = ctx.vis_sorted == uv
		ax.plot(ctx.t[sel], nt_line[sel], '-', color='C1', linewidth=1.2,
		zorder=5, alpha=0.9,
		label='Non-time correction' if (first_vis and j == 0)
		else '_nolegend_')
		first_vis = False

		# 68% confidence band on the non-time correction
		if ctx.has_bands and ctx.nt_data_q16_m is not None:
		if ctx.data_type_in_mask == 'RV' or ctx.normalization_mode == 'additive':
		q16 = ctx.norm + ctx.nt_data_q16_m
		q84 = ctx.norm + ctx.nt_data_q84_m
		else:
		q16 = ctx.norm * (1.0 + ctx.nt_data_q16_m)
		q84 = ctx.norm * (1.0 + ctx.nt_data_q84_m)
		ax.fill_between(ctx.t, q16, q84, color='C1', alpha=0.15,
		zorder=4, label='_nolegend_')

		# X-limits for this cluster sub-panel
		if ctx.is_per_inst_plot and ctx.clusters and j < len(ctx.clusters):
		c_start, c_end = ctx.clusters[j]
		buf = 0.05 * (c_end - c_start)
		ax.set_xlim(c_start - buf, c_end + buf)

		ylabel = ('Radial Velocity [m/s]' if ctx.data_type_in_mask == 'RV'
		else 'Relative Flux')
		raw_axes[0].set_ylabel(ylabel)
		raw_axes[0].legend(loc='upper right', fontsize=7, framealpha=0.9)
		for ax in raw_axes:
		plt.setp(ax.get_xticklabels(), visible=False)

PyAETNA/reporting/panel_residuals.py

0 → 100644

+78 −0

Changes for PyAETNA/reporting/panel_residuals.py: 78 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		# -- coding: utf-8 --
		"""
		panel_residuals.py
		~~~~~~~~~~~~~~~~~~
		Renderer for Panel 3: residuals (data − total model).

		Plots the residuals with σ_orig error bars (raw measurement noise only,
		not jitter-inflated), adds a zero reference line, annotates with RMS, and
		draws labelled reference uncertainty bars.

		Functions
		---------
		render_panel_residuals
		"""
		import numpy as np

		from .plot_data_prep import PlotContext
		from .plot_styling import add_unc_ref_bars


		def render_panel_residuals(residual_axes_raw, ctx: PlotContext) -> None:
		"""Render Panel 3: residuals + RMS annotation + reference σ bars.

		Parameters
		----------
		residual_axes_raw : list[Axes] or Axes
		ctx : PlotContext
		"""
		residual_axes = (residual_axes_raw if isinstance(residual_axes_raw, list)
		else [residual_axes_raw])

		# ── Data scatter ──────────────────────────────────────────────────────
		if ctx.is_per_inst_plot and ctx.visit_colors is not None:
		for j, ax in enumerate(residual_axes):
		for uv, c in ctx.visit_colors.items():
		sel = ctx.vis_sorted == uv
		ax.errorbar(
		ctx.t[sel], ctx.residuals[sel], yerr=ctx.dy[sel],
		fmt='.', markersize=3, elinewidth=0.5, alpha=0.6,
		color=c,
		label=f'Visit {uv}' if j == 0 else '_nolegend_',
		zorder=1)
		else:
		for j, ax in enumerate(residual_axes):
		ax.errorbar(
		ctx.t, ctx.residuals, yerr=ctx.dy,
		fmt='.', markersize=3, elinewidth=0.5, alpha=0.6,
		color='gray', label='Data' if j == 0 else '_nolegend_',
		zorder=1)

		# ── Zero line and cluster x-limits ────────────────────────────────────
		for j, ax in enumerate(residual_axes):
		ax.axhline(0, color='k', linestyle='--', linewidth=0.5, zorder=2)
		if ctx.clusters and j < len(ctx.clusters):
		c_start, c_end = ctx.clusters[j]
		buf = 0.05 * (c_end - c_start)
		ax.set_xlim(c_start - buf, c_end + buf)

		# ── RMS annotation ────────────────────────────────────────────────────
		rms = float(np.sqrt(np.mean(ctx.residuals ** 2)))
		rms_str = (f'RMS: {rms:.2f} m/s' if ctx.data_type_in_mask == 'RV'
		else f'RMS: {rms * 1e6:.1f} ppm')
		residual_axes[0].text(
		0.02, 0.95, rms_str,
		transform=residual_axes[0].transAxes, fontsize=8,
		verticalalignment='top',
		bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))

		# ── Reference uncertainty bars and legend ─────────────────────────────
		add_unc_ref_bars(residual_axes[0], ctx.x_ref, 0.0,
		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)
		else:
		residual_axes[0].legend(loc='lower right', fontsize=7, framealpha=0.9,
		title='Median uncertainties')
		residual_axes[0].set_ylabel('Residuals (σ_orig error bars)')

PyAETNA/reporting/plot_data_prep.py

0 → 100644

+579 −0

File added.

Preview size limit exceeded, changes collapsed.