Commit 7222828b authored by Gaetano Scandariato's avatar Gaetano Scandariato
Browse files

refactor of group 4 started

parent 5fb868d6

PyAETNA/reporting/confidence_bands.py

+14 −19
Changes for PyAETNA/reporting/confidence_bands.py: 14 added lines, 19 removed lines.
Original line number Diff line number Diff line
@@ -145,27 +145,23 @@ def compute_data_bands( 
    planet_trend_samples  = np.array(planet_trend_samples)

    total_samples         = np.array(total_samples)



    mean_trend_median = np.percentile(mean_trend_samples, 50, axis=0)

    mean_trend_q16    = np.percentile(mean_trend_samples, 16, axis=0)

    mean_trend_q84    = np.percentile(mean_trend_samples, 84, axis=0)

    # A single np.percentile() call with a list of quantiles avoids repeating

    # the (partial-)sort of each samples array three times.

    mean_trend_q16, mean_trend_median, mean_trend_q84 = np.percentile(

        mean_trend_samples, [16, 50, 84], axis=0)



    planet_median = planet_q16 = planet_q84 = None

    planet_norm_median = planet_norm_q16 = planet_norm_q84 = None

    planet_trend_median = planet_trend_q16 = planet_trend_q84 = None

    if has_planets:

        planet_median        = np.percentile(planet_samples,       50, axis=0)

        planet_q16           = np.percentile(planet_samples,       16, axis=0)

        planet_q84           = np.percentile(planet_samples,       84, axis=0)

        planet_norm_median   = np.percentile(planet_norm_samples,  50, axis=0)

        planet_norm_q16      = np.percentile(planet_norm_samples,  16, axis=0)

        planet_norm_q84      = np.percentile(planet_norm_samples,  84, axis=0)

        planet_trend_median  = np.percentile(planet_trend_samples, 50, axis=0)

        planet_trend_q16     = np.percentile(planet_trend_samples, 16, axis=0)

        planet_trend_q84     = np.percentile(planet_trend_samples, 84, axis=0)



    total_median = np.percentile(total_samples, 50, axis=0)

    total_q16    = np.percentile(total_samples, 16, axis=0)

    total_q84    = np.percentile(total_samples, 84, axis=0)

        planet_q16, planet_median, planet_q84 = np.percentile(

            planet_samples, [16, 50, 84], axis=0)

        planet_norm_q16, planet_norm_median, planet_norm_q84 = np.percentile(

            planet_norm_samples, [16, 50, 84], axis=0)

        planet_trend_q16, planet_trend_median, planet_trend_q84 = np.percentile(

            planet_trend_samples, [16, 50, 84], axis=0)



    total_q16, total_median, total_q84 = np.percentile(total_samples, [16, 50, 84], axis=0)



    nt_data_q16 = nt_data_q50 = nt_data_q84 = None

    if has_non_time_trends and nontime_trend_data_samples:
@@ -179,9 +175,8 @@ def compute_data_bands( 
        

        for name in comp_names:

            comp_samples = np.array([s[name] for s in nontime_trend_data_samples])

            nt_data_q16[name] = np.percentile(comp_samples, 16, axis=0)

            nt_data_q50[name] = np.percentile(comp_samples, 50, axis=0)

            nt_data_q84[name] = np.percentile(comp_samples, 84, axis=0)

            nt_data_q16[name], nt_data_q50[name], nt_data_q84[name] = np.percentile(

                comp_samples, [16, 50, 84], axis=0)



    return dict(

        has_bands=True, n_samples_use=n_samples_use,

PyAETNA/reporting/model_components.py

+0 −9
Changes for PyAETNA/reporting/model_components.py: 0 added lines, 9 removed lines.
Original line number Diff line number Diff line
@@ -10,7 +10,6 @@ compute_model_components 
save_model_components_binary

"""

import numpy as np

import jax

import jax.numpy as jnp

import dill
@@ -21,14 +20,6 @@ from PyAETNA.gp import ( 
    get_gp_x_variable,

)

from PyAETNA.model_utils import build_mean_function

from PyAETNA.body_models import (

    compute_lightcurve as _compute_lightcurve_raw,

    compute_lightcurve_supersampled as _compute_lightcurve_ss,

    compute_radial_velocity as _compute_radial_velocity_raw,

    compute_rossiter_mclaughlin as _compute_rossiter_mclaughlin_raw,

    compute_transit_ttv as _compute_transit_ttv_raw,

    reparametrize_LD as _reparametrize_LD,

)



from . import logger

def _log(msg, level='INFO'):

PyAETNA/reporting/pdf_report.py

+4 −8
Changes for PyAETNA/reporting/pdf_report.py: 4 added lines, 8 removed lines.
Original line number Diff line number Diff line
@@ -2,11 +2,9 @@ 
import numpy as np

import matplotlib.pyplot as plt

from matplotlib.backends.backend_pdf import PdfPages

import os

import dill

import yaml

import traceback

import time as ttt



from . import logger

def _log(msg, level='INFO'):
@@ -15,7 +13,7 @@ def _log(msg, level='INFO'): 
from .logger import set_log_file

from .map_search import find_map_parameters

from .diagnostics import compute_mcmc_diagnostics, print_mcmc_summary_table, print_sampling_summary, compute_goodness_of_fit

from .corner_plot import plot_corner, plot_trace, plot_kde_trace, plot_chain_comparison, plot_visit_kde, _prioritize_normalization

from .corner_plot import plot_corner, plot_kde_trace, plot_chain_comparison, plot_visit_kde, _prioritize_normalization

from .plot_models import plot_models

from .plot_data_prep import prepare_plot_data

from .panel_nontime_residuals import plot_nontime_residuals
@@ -24,9 +22,7 @@ from .confidence_bands import compute_data_bands 
from .model_math import _precompute_sample_data

from .noise_plots import plot_noise_budget

from .periodogram_plots import (

    plot_residuals_histogram,

    plot_periodograms_comparison,

    plot_correlated_noise,

    plot_residuals_vs_dbasis,

)

from .phase_plots import (
@@ -84,13 +80,13 @@ def generate_report(pickle_path, output_pdf, log_file=None, config_file=None): 
    # Re-load config from YAML if provided (to pick up updated reports settings)

    if config_file is not None:

        try:

            import yaml

            with open(config_file, 'r') as f:

                config_reloaded = yaml.safe_load(f)

            if config_reloaded:

                config = config_reloaded

        except Exception:

            pass  # Use config from pickle

        except Exception as e:

            _log(f'Failed to reload config from {config_file!r}, using config from '

                 f'pickle instead: {e}', level='WARNING')

    

    # Set log file FIRST before any logging

    if log_file is not None:

PyAETNA/reporting/periodogram_plots.py

+8 −11
Changes for PyAETNA/reporting/periodogram_plots.py: 8 added lines, 11 removed lines.
Original line number Diff line number Diff line
@@ -14,17 +14,14 @@ plot_residuals_vs_dbasis 
"""

import numpy as np

import matplotlib.pyplot as plt

from matplotlib.ticker import FuncFormatter

from astropy.timeseries import LombScargle

from scipy import stats

from scipy.stats import gaussian_kde

import warnings



from . import logger

def _log(msg, level='INFO'):

    logger._log(msg, level=level)



from .model_math import _compute_planet_model_for_plotting

from .corner_plot import _get_distinct_colors
@@ -226,7 +223,8 @@ def plot_periodograms_comparison(time, y_original, y_detrended, residuals, _log, 
        try:

            fap_powers = ls.false_alarm_level(fap_levels / 100, method='bootstrap', 

                                               method_kwds={'n_bootstraps': n_bootstrap}) * rescale

        except:

        except Exception as e:

            _log(f'Failed to compute bootstrap FAP levels for {label!r}: {e}', level='WARNING')

            fap_powers = None

        

        # Find peak
@@ -265,8 +263,8 @@ def plot_periodograms_comparison(time, y_original, y_detrended, residuals, _log, 
                    # Hide period tick labels for middle and bottom rows

                    ax_lin_top.set_xticklabels([])

                    ax_lin_top.set_xlabel('')

        except:

            _log('Failed to add top-period axis to linear-frequency panel')

        except Exception as e:

            _log(f'Failed to add top-period axis to linear-frequency panel: {e}', level='WARNING')

            

        # Mark peak

        ax_lin.axvline(peak_freq, color='blue', linestyle='--', alpha=0.7, linewidth=0.8)
@@ -399,8 +397,8 @@ def plot_periodograms_comparison(time, y_original, y_detrended, residuals, _log, 
                    ax_log.set_ylim([y_min * 0.5, y_max * 1.15])

                else:

                    ax_log.set_ylim([1e-12, 1.0])

        except:

            _log('Failed to set log y-scale on log-frequency panel')

        except Exception as e:

            _log(f'Failed to set log y-scale on log-frequency panel: {e}', level='WARNING')



        # ── Overlay Noise Baselines and GP PSD on row 2, log-scale panel ──────────────

        if row == 1:
@@ -459,7 +457,7 @@ def plot_periodograms_comparison(time, y_original, y_detrended, residuals, _log, 
                if handles and labels:  # Only add if there are actual legend elements

                    all_handles.extend(handles)

                    all_labels.extend(labels)

            except:

            except Exception:

                pass  # Skip if no legend elements

    

    # Remove duplicate labels while preserving order
@@ -554,7 +552,7 @@ def plot_periodogram(time, residuals, min_period=None, max_period=None, 
        for fap, fap_power, color in zip(fap_levels, fap_powers, colors):

            ax.axhline(fap_power, color=color, linestyle=':', alpha=0.7,

                      label=f'FAP {fap*100:.1f}%')

    except:

    except Exception:

        pass

    

    ax.legend(fontsize=8, loc='upper right')
@@ -660,7 +658,6 @@ def plot_residuals_vs_dbasis(plData, residuals, plModel=None, visit_to_plot=None 
        fig: Matplotlib figure object

    """

    from scipy import stats

    import matplotlib.cm as cm

    from PyAETNA.sampling.engine.eiv import eiv_map_estimate

    

    # Get all dBasis keys (early exit if none)

PyAETNA/reporting/ttv_plots.py

+7 −29
Changes for PyAETNA/reporting/ttv_plots.py: 7 added lines, 29 removed lines.
Original line number Diff line number Diff line
@@ -12,33 +12,10 @@ plot_ttv_transit_overlay 
"""

import numpy as np

import matplotlib.pyplot as plt

import matplotlib.cm

from matplotlib.gridspec import GridSpec

from matplotlib.lines import Line2D

from matplotlib.patches import Patch

from matplotlib.ticker import FuncFormatter

from scipy import stats

from scipy.stats import gaussian_kde

import jax

import jax.numpy as jnp



from PyAETNA.gp import (

    gp_predict,

    gp_predict_per_visit,

    build_gp_kernel_for_plotting,

    get_gp_x_variable,

)

from PyAETNA.model_utils import build_mean_function, detrend_data



from PyAETNA.model_utils import detrend_data

from PyAETNA.body_models import (

    compute_lightcurve as _compute_lightcurve_raw,

    compute_lightcurve_supersampled as _compute_lightcurve_ss,

    compute_radial_velocity as _compute_radial_velocity_raw,

    compute_rossiter_mclaughlin as _compute_rossiter_mclaughlin_raw,

    compute_transit_ttv as _compute_transit_ttv_raw,

    reparametrize_LD as _reparametrize_LD,

    compute_transit_duration,

    compute_eclipse_duration,

    compute_eclipse_phase,

    reparametrize_ecc_omega,

)
@@ -46,9 +23,7 @@ from . import logger 
def _log(msg, level='INFO'):

    logger._log(msg, level=level)



from .gp_plotting import _get_gp_correlation_length, _precompute_gp_grid_context, _gp_predict_on_grid, _build_visit_time_grid, _build_lc_train_grid

from .corner_plot import _get_distinct_colors

from .model_math import _compute_planet_model_for_plotting, _apply_scoped_values

from .model_components import compute_model_components
@@ -117,7 +92,8 @@ def plot_ttv_oc(samples, param_config, planet_suffix, plData=None, timeCenter=0. 
    fixed_params   = param_config.get('fixed_params', {})



    if planet_suffix not in planet_names:

        print(f"[TTV] plot_ttv_oc: planet '{planet_suffix}' not in param_config.")

        _log(f"[TTV] plot_ttv_oc: planet '{planet_suffix}' not in param_config.",

             level='WARNING')

        return None

    planet_idx = list(planet_names).index(planet_suffix)

    p_cfg = planet_configs[planet_idx] if planet_idx < len(planet_configs) else {}
@@ -220,7 +196,9 @@ def plot_ttv_oc(samples, param_config, planet_suffix, plData=None, timeCenter=0. 
                    epoch_visit_map[ep_i] = visit_id

                    if visit_id not in visit_order:

                        visit_order.append(visit_id)

        except Exception:

        except Exception as e:

            _log(f"[TTV] plot_ttv_oc: failed to assign epochs to visits, "

                 f"per-visit O-C panels will be skipped: {e}", level='WARNING')

            epoch_visit_map = {}

            visit_order = []