Enhance plot_periodograms_comparison with improved layout and legend handling

def plot_periodograms_comparison(time, y_original, y_detrended, residuals,_log, 

                                  title='Periodogram Comparison', figsize=(12, 12),

                                  title='Periodogram Comparison', figsize=(15, 12),

                                  n_bootstrap=1000, gp_kernel=None,

                                  gp_x_variable='BJD'):

    """

        fig: matplotlib figure

    """

    from astropy.timeseries import LombScargle

    import matplotlib.gridspec as gridspec

    # Compute frequency grid

    dt = np.median(np.diff(np.sort(time)))

    fap_levels = np.array([5, 1, 0.1])  # in percent

    fap_colors = ['green', 'orange', 'red']

    # Create figure with 6 panels (3 datasets x 2 scales)

    fig, axes = plt.subplots(3, 2, figsize=figsize)

    # Create figure with specified size

    fig = plt.figure(figsize=figsize)

    # Create 3x3 grid (3 datasets x 2 scales + 1 legend column)

    # Width ratio: 0.45/0.45/0.1 for data columns / legend column

    gs = gridspec.GridSpec(3, 3, width_ratios=[0.45, 0.45, 0.1], height_ratios=[1, 1, 1], 

                           hspace=0.3, wspace=0.1)

    # Create main axes for the 3x2 data plots (6 panels) - using list of lists

    axes = []

    for row in range(3):

        row_axes = []

        for col in range(2):

            ax = plt.subplot(gs[row, col])

            row_axes.append(ax)

        axes.append(row_axes)

    # Create legend axes in the rightmost column (middle row)

    legend_ax = plt.subplot(gs[1, 2])  # Center row of the rightmost column

    # Create the figure object explicitly

    fig = plt.gcf()

    # Row labels:

    #   Row 0 — "Original data"

        peak_power = power[peak_idx]

        # Left panel: Linear frequency scale

        ax_lin = axes[row, 0]

        ax_lin = axes[row][0]

        ax_lin.plot(freqs, power, 'k-', linewidth=0.5, label='LS periodogram')

        ax_lin.set_xlabel('Frequency [d$^{-1}$]')

        ax_lin.set_ylabel('PSD [units$^2$/d$^{-1}$]')

        ax_lin.set_title(f'{label} (Linear scale)')

        # Only show x-axis labels on the bottom row

        if row == 2:

            ax_lin.set_xlabel('Frequency [d$^{-1}$]')

        else:

            # Hide x-axis labels for upper rows

            ax_lin.tick_params(axis='x', labelbottom=False)

            ax_lin.set_xlabel('')  # Clear xlabel if it exists

        # Secondary x-axis for period (linear scale) — mirror of the right-panel period axis

        # Only show period axis labels on the top row

        try:

            ax_lin_top = ax_lin.twiny()

            ax_lin_top.set_xlim(ax_lin.get_xlim())

            period_ticks = np.array([0.1, 0.5, 1, 2, 5, 10, 30, 100])

            period_ticks = period_ticks[(1/period_ticks >= f_min) & (1/period_ticks <= f_max)]

            if len(period_ticks) > 0:

                ax_lin_top.set_xticks(1.0 / period_ticks)

                if row == 0:  # Only show period labels on top row

                    ax_lin_top.set_xticklabels([f'{p:.1f}' for p in period_ticks], fontsize=7)

                    ax_lin_top.set_xlabel('Period [d]', fontsize=8)

                else:

                    # 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')

        # Mark peak

        ax_lin.axvline(peak_freq, color='blue', linestyle='--', alpha=0.7, linewidth=0.8)

        ax_lin.text(0.98, 0.95, f'Peak: {peak_freq:.3f} d$^{{-1}}$ ({peak_period:.2f} d)',

                        bbox=dict(boxstyle='round', facecolor='lightyellow',

                                  alpha=0.8))

        ax_lin.legend(fontsize=7, loc='upper right')

        #ax_lin.legend(fontsize=7)  # Removed loc='upper right' to prevent overlap

        # Set y-axis limits

        ylim = ax_lin.get_ylim()

        ax_lin.set_ylim([-ylim[1] * 0.05, ylim[1] * 1.15])

        # Secondary x-axis for period (linear scale) — mirror of the right-panel period axis

        try:

            ax_lin_top = ax_lin.twiny()

            ax_lin_top.set_xlim(ax_lin.get_xlim())

            period_ticks = np.array([0.1, 0.5, 1, 2, 5, 10, 30, 100])

            period_ticks = period_ticks[(1/period_ticks >= f_min) & (1/period_ticks <= f_max)]

            if len(period_ticks) > 0:

                ax_lin_top.set_xticks(1.0 / period_ticks)

                ax_lin_top.set_xticklabels([f'{p:.1f}' for p in period_ticks], fontsize=7)

                ax_lin_top.set_xlabel('Period [d]', fontsize=8)

        except:

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

        # Right panel: Log frequency scale

        ax_log = axes[row, 1]

        ax_log = axes[row][1]

        ax_log.plot(freqs, power, 'k-', linewidth=0.5, label='LS periodogram')

        ax_log.set_xscale('log')

        ax_log.set_xlabel('Frequency [d$^{-1}$]')

        ax_log.set_ylabel('PSD [units$^2$/d$^{-1}$]')

        ax_log.set_title(f'{label} (Log scale)')

        # Move y-axis labels and ticks to the right side

        ax_log.yaxis.tick_right()

        ax_log.yaxis.set_label_position('right')

        # Mark peak

        ax_log.axvline(peak_freq, color='blue', linestyle='--', alpha=0.7, linewidth=0.8)

        # Only show x-axis labels on the bottom row

        if row == 2:

            ax_log.set_xlabel('Frequency [d$^{-1}$]')

        else:

            # Hide x-axis labels for upper rows

            ax_log.tick_params(axis='x', labelbottom=False)

            ax_log.set_xlabel('')  # Clear xlabel if it exists

        # Secondary x-axis for period

        # Only show period axis labels on the top row

        ax_log_top = ax_log.twiny()

        ax_log_top.set_xscale('log')

        ax_log_top.set_xlim(ax_log.get_xlim())

        period_ticks = period_ticks[(1/period_ticks >= f_min) & (1/period_ticks <= f_max)]

        if len(period_ticks) > 0:

            ax_log_top.set_xticks(1.0 / period_ticks)

            if row == 0:  # Only show period labels on top row

                ax_log_top.set_xticklabels([f'{p:.1f}' for p in period_ticks], fontsize=7)

                ax_log_top.set_xlabel('Period [d]', fontsize=8)

            else:

                # Hide period tick labels for middle and bottom rows

                ax_log_top.set_xticklabels([])

                ax_log_top.set_xlabel('')

        # Mark peak

        ax_log.axvline(peak_freq, color='blue', linestyle='--', alpha=0.7, linewidth=0.8)

        # Set y-axis to log scale for easier dynamic range viewing

        try:

            ax_log.set_yscale('log')

            # Move y-axis labels and ticks to the right side

            ax_log.yaxis.tick_right()

            ax_log.yaxis.set_label_position('right')

            # Determine a positive range for plotting on log scale;

            # include GP PSD in the range if available

            all_power = power.copy()

                ax_log.axhline(fap_power, color=color, linestyle=':', alpha=0.7, linewidth=0.8,

                              label=f'FAP={fap}%')

        ax_log.legend(fontsize=7, loc='upper right')

    # Collect legend elements from all subplots

    all_handles = []

    all_labels = []

    # Get handles and labels from each subplot

    for row in range(3):

        for col in range(2):

            ax = axes[row][col]  # Fixed: use list indexing instead of tuple

            try:

                handles, labels = ax.get_legend_handles_labels()

                if handles and labels:  # Only add if there are actual legend elements

                    all_handles.extend(handles)

                    all_labels.extend(labels)

            except:

                pass  # Skip if no legend elements

    # Remove duplicate labels while preserving order

    seen = set()

    unique_handles = []

    unique_labels = []

    for h, l in zip(all_handles, all_labels):

        if l not in seen:

            seen.add(l)

            unique_handles.append(h)

            unique_labels.append(l)

    # Add legend to the center plot of the rightmost column

    if unique_handles and unique_labels:

        # Turn off the legend axes

        legend_ax.axis('off')

        # Place the legend anchored precisely to avoid overlap

        legend_ax.legend(

            unique_handles, 

            unique_labels, 

            fontsize=10, 

            loc='center left',          # Anchor point of the legend box

            bbox_to_anchor=(0.5, 0.5),  # (x, y) coordinates relative to legend_ax (0.0 is left, 1.0 is right)

            frameon=True, 

            fancybox=False, 

            shadow=False, 

            ncol=1

        )

    # Adjust layout to accommodate the legend

    fig.suptitle(title, fontsize=12, y=0.98)

    fig.tight_layout()

    return fig

- check the checkpoint_interval logic

- implement different limb darkening laws

- refactor sampling 

- rearrange priors

- better arrangement of legends and text in plots

## Contact

        configFile = sys.argv[1]

    except Exception:

        # configFile = 'tests/test_lc_only.yaml'

        # configFile = 'tests/test_rv_only.yaml'

        configFile = 'tests/test_rv_only.yaml'

        # configFile = 'tests/test_lc_rv_joint.yaml'

        # configFile = 'tests/test_lc_gp_rv.yaml'

        # configFile = 'tests/test_phase_curve.yaml'

        configFile = 'tests/test_ttv.yaml'

        # configFile = 'tests/test_ttv.yaml'

        # configFile = 'tests/test_spectral_mixture.yaml'

    # Run with the config file (data.working_dir and data.file_path are read from YAML)