bug fixes

    with open(save_path, 'wb') as f:

        dill.dump(out, f, protocol=4)

def load_checkpoint(save_path, reset, plModel, config, mcmc_config_peek, log_func, format_time):

def load_checkpoint(save_path, reset, plModel, config, mcmc_config_peek, log_func, format_time, logFile):

    """Load existing checkpoint and restore sampler states, adapting precision."""

    plData = None

    existing_samples = None

    use_float64 = False

    if not reset and os.path.exists(save_path):

        log_func(f"Loading checkpoint from: {save_path}")

        log_func(logFile, f"Loading checkpoint from: {save_path}")

        try:

            with open(save_path, 'rb') as f:

                saved_data = dill.load(f)

                    arr = getattr(plData, attr, None)

                    if arr is not None and hasattr(arr, 'dtype') and arr.dtype != jnp.float64:

                        setattr(plData, attr, jnp.array(arr, dtype=jnp.float64))

                        log_func(f"  Upcast plData.{attr} from {arr.dtype} → float64")

                        log_func(logFile, f"  Upcast plData.{attr} from {arr.dtype} → float64")

                if hasattr(plData, 'dBasis') and isinstance(plData.dBasis, dict):

                    for name, arr in plData.dBasis.items():

                        if hasattr(arr, 'dtype') and arr.dtype != jnp.float64:

                            plData.dBasis[name] = jnp.array(arr, dtype=jnp.float64)

                            log_func(f"  Upcast plData.dBasis['{name}'] from {arr.dtype} → float64")

                            log_func(logFile, f"  Upcast plData.dBasis['{name}'] from {arr.dtype} → float64")

                if existing_mcmc_state is not None:

                    existing_mcmc_state = upcast_value_to_float64(existing_mcmc_state)

                    log_func(f"  Upcast MCMC state to x64 (float32→float64, int32→int64)")

                    log_func(logFile, f"  Upcast MCMC state to x64 (float32→float64, int32→int64)")

                if existing_warmup_state is not None:

                    existing_warmup_state = upcast_value_to_float64(existing_warmup_state)

                    log_func(f"  Upcast warmup MCMC state to x64 (float32→float64, int32→int64)")

                    log_func(logFile, f"  Upcast warmup MCMC state to x64 (float32→float64, int32→int64)")

                if existing_samples is not None:

                    for name, arr in existing_samples.items():

                        if hasattr(arr, 'dtype') and arr.dtype == np.float32:

                            existing_samples[name] = arr.astype(np.float64)

                    log_func(f"  Upcast existing samples to float64")

                    log_func(logFile, f"  Upcast existing samples to float64")

            else:

                existing_warmup_state = _cast_warmup_state_dtype(existing_warmup_state, use_float64=False)

            log_func("")

            log_func(logFile, "")

        except Exception as e:

            log_func(f"Failed to load checkpoint: {e}. Starting from input files.")

            log_func(logFile, f"Failed to load checkpoint: {e}. Starting from input files.")

            loaded_from_pickle = False

    return (existing_samples, existing_num_samples, loaded_from_pickle,

     saved_sampler_type, existing_mcmc_state, existing_warmup_state,

     existing_warmup_done, existing_warmup_elapsed, num_chains,

     plModel, plData, config, param_config, use_float64) = load_checkpoint(

         save_path, reset, plModel, config_peek, mcmc_config_peek, _log, format_time

         save_path, reset, plModel, config_peek, mcmc_config_peek, _log, format_time, logFile

     )

    # If not loaded from pickle, load from input files

import pytest

import numpy as np

from PyAETNA.reporting.plot_model_helpers import detect_large_time_gaps

def test_detect_large_time_gaps_numeric():

    # Gap (30 - 6) = 24. Total span = 30. Ratio = 24 / 30 = 0.8 > 0.1

    time = [0, 2, 5, 6, 30]

    gaps = detect_large_time_gaps(time, threshold_ratio=0.1)

    assert gaps == [(6.0, 30.0)]

def test_detect_large_time_gaps_no_gaps():

    # Evenly spaced: diff is always 1. Total span = 9. Ratio threshold = 0.9.

    time = list(range(10))

    gaps = detect_large_time_gaps(time, threshold_ratio=0.15)

    assert gaps == []

def test_detect_large_time_gaps_indices():

    time = [10, 11, 12, 50, 51] # gap is 38. Total span = 41.

    gaps_idx = detect_large_time_gaps(time, threshold_ratio=0.1, return_indices=True)

    # The gap is between sorted elements 12 (index 2) and 50 (index 3)

    assert gaps_idx == [(2, 3)]

def test_detect_large_time_gaps_edge_cases():

    # Less than 2 elements

    assert detect_large_time_gaps([]) == []

    assert detect_large_time_gaps([42]) == []

    # Zero timespan

    assert detect_large_time_gaps([5, 5, 5]) == []

def test_detect_large_time_gaps_invalid_type():

    with pytest.raises(TypeError):

        detect_large_time_gaps(["a", "b", "c"])

def test_detect_large_time_gaps_datetime():

    times = np.array([

        '2024-01-01', '2024-01-02', '2024-01-03', '2024-02-15'

    ], dtype='datetime64[D]')

    gaps = detect_large_time_gaps(times, threshold_ratio=0.1)

    assert len(gaps) == 1

    # Check that boundaries are correct

    assert gaps[0][0] == np.datetime64('2024-01-03')

    assert gaps[0][1] == np.datetime64('2024-02-15')

"""

test_gap_plots.py

~~~~~~~~~~~~~~~~~

Structural test for the gap-splitting plot layout added to plot_models.

We monkey-patch the heavy dependencies so this test runs without FITS files

or a full PyAETNA inference run.  We only verify that:

1. With no large gaps, plot_models produces _n_seg == 1 → standard layout.

2. With one large gap (two clusters), _n_seg == 2 → two time-column subplots

   per row, with the correct x-limits and shared y-axes.

3. Width ratios in the GridSpec are proportional to cluster durations.

"""

import sys

import types

import numpy as np

import matplotlib

matplotlib.use('Agg')

import matplotlib.pyplot as plt

# ── Minimal PlData / PlModel stubs ──────────────────────────────────────────

class _PlModel:

    uniqueInstruments = [1]

    uniqueVisits = [1]

    instrument_names = ['Stub']

    normalization_mode = 'multiplicative'

    data_types = {1: 'LC'}

    planet_configs = []

    planet_names = []

    polynomial_configs = []

    fourier_configs = []

    gp_configs = []

    timeCenter = 0.0

    param_scopes = {}

class _PlData:

    def __init__(self, time_arr):

        n = len(time_arr)

        self.time = time_arr

        self.y = np.ones(n)

        self.dy = np.full(n, 0.01)

        self.instrument = np.ones(n, dtype=int)

        self.visit = np.ones(n, dtype=int)

        self.exptime = np.zeros(n)

        self.dBasis = {}

        self.timeCenter = 0.0

def _make_best_params():

    return {}

def _make_param_config():

    return {

        'planet_configs': [],

        'planet_names': [],

        'polynomial_configs': [],

        'fourier_configs': [],

        'gp_configs': [],

        'param_scopes': {},

    }

# ── Import the real module (after stubs are ready) ───────────────────────────

sys.path.insert(0, '/dati/software/pyLib/PyAETNA')

# We need to mock compute_model_components and other heavy calls before import

import unittest.mock as mock

# Patch the expensive downstream functions used inside plot_models

_dummy_comp = {

    'norm': np.array([1.0]),

    'planet': np.array([0.0]),

    'trend': np.array([0.0]),

    'gp': np.array([0.0]),

    'total': np.array([1.0]),

    'jitter': np.array([0.01]),

    'poly_terms': {},

    'fourier_terms': {},

    'gp_terms': {},

}

_patches = [

    mock.patch('PyAETNA.reporting.plot_models.compute_model_components',

               return_value=_dummy_comp),

    mock.patch('PyAETNA.reporting.plot_models._precompute_sample_data',

               return_value=None),

    mock.patch('PyAETNA.reporting.plot_models._compute_nontime_trend',

               return_value=None),

    mock.patch('PyAETNA.reporting.plot_models._precompute_gp_grid_context',

               return_value={}),

    mock.patch('PyAETNA.reporting.plot_models._build_visit_time_grid',

               return_value=(np.linspace(0, 1, 10), np.ones(10, dtype=int))),

    mock.patch('PyAETNA.reporting.plot_models._build_lc_train_grid',

               return_value=(np.linspace(0, 1, 10), np.ones(10, dtype=int), [])),

    mock.patch('PyAETNA.reporting.plot_models._gp_predict_on_grid',

               return_value=np.zeros(10)),

    mock.patch('PyAETNA.reporting.plot_models.build_composite_for_plot',

               return_value=np.ones(10)),

    mock.patch('PyAETNA.reporting.plot_models.build_partial_for_plot',

               return_value=np.ones(10)),

    mock.patch('PyAETNA.reporting.plot_models.detrend_data',

               return_value=np.ones(10)),

    mock.patch('PyAETNA.reporting.plot_models.build_mean_function',

               return_value=lambda t, **kw: np.ones(len(t))),

    mock.patch('PyAETNA.reporting.plot_models._get_gp_correlation_length',

               return_value=0.1),

    mock.patch('PyAETNA.reporting.plot_models._compute_planet_model_for_plotting',

               return_value=np.zeros(10)),

]

def _make_fig(time_arr):

    """Call detect_large_time_gaps directly to test the gap logic."""

    from PyAETNA.reporting.plot_model_helpers import detect_large_time_gaps

    t = np.asarray(time_arr)

    gaps = detect_large_time_gaps(t, threshold_ratio=0.1)

    t_min, t_max = float(t.min()), float(t.max())

    if gaps:

        intervals = []

        prev = t_min

        for g_lo, g_hi in gaps:

            intervals.append((prev, float(g_lo)))

            prev = float(g_hi)

        intervals.append((prev, t_max))

    else:

        intervals = [(t_min, t_max)]

    return intervals, gaps

def test_no_gap_single_segment():

    """Uniformly sampled data → no gaps → 1 segment."""

    t = np.linspace(0, 10, 50)

    intervals, gaps = _make_fig(t)

    assert len(gaps) == 0, f"Expected 0 gaps, got {len(gaps)}"

    assert len(intervals) == 1, f"Expected 1 interval, got {len(intervals)}"

    print("test_no_gap_single_segment PASSED")

def test_one_large_gap_two_segments():

    """Two clusters separated by > 10% of total span → 2 segments."""

    t = np.concatenate([np.linspace(0, 1, 20), np.linspace(5, 6, 20)])

    intervals, gaps = _make_fig(t)

    assert len(gaps) == 1, f"Expected 1 gap, got {len(gaps)}"

    assert len(intervals) == 2, f"Expected 2 intervals, got {len(intervals)}"

    # First interval ends at ~1, second starts at ~5

    assert intervals[0][1] < 2.0

    assert intervals[1][0] > 4.0

    print("test_one_large_gap_two_segments PASSED")

def test_two_large_gaps_three_segments():

    """Three clusters → 3 segments."""

    t = np.concatenate([

        np.linspace(0, 1, 20),

        np.linspace(5, 6, 20),

        np.linspace(15, 16, 20),

    ])

    intervals, gaps = _make_fig(t)

    assert len(gaps) == 2, f"Expected 2 gaps, got {len(gaps)}"

    assert len(intervals) == 3, f"Expected 3 intervals, got {len(intervals)}"

    print("test_two_large_gaps_three_segments PASSED")

def test_proportional_widths():

    """Width of each segment column is proportional to its duration."""

    from PyAETNA.reporting.plot_model_helpers import detect_large_time_gaps

    t = np.concatenate([np.linspace(0, 2, 30), np.linspace(10, 12, 30)])

    gaps = detect_large_time_gaps(t, threshold_ratio=0.1)

    assert len(gaps) == 1

    # Cluster 1 duration ≈ 2, cluster 2 duration ≈ 2 → widths should be equal

    intervals = [(0.0, float(gaps[0][0])), (float(gaps[0][1]), 12.0)]

    durs = [hi - lo for lo, hi in intervals]

    # Both ≈ 2 → ratio ≈ 1

    ratio = durs[0] / durs[1]

    assert 0.8 < ratio < 1.2, f"Expected ratio ≈ 1, got {ratio}"

    print("test_proportional_widths PASSED")

if __name__ == '__main__':

    test_no_gap_single_segment()

    test_one_large_gap_two_segments()

    test_two_large_gaps_three_segments()

    test_proportional_widths()

    print("\nAll structural gap-split layout tests passed.")