src/xp2/_run.py

"""
xp2: effect of training-data quality on CombiLatent performance.

For each percentile p in {5, 10, 15, ..., 95}:
  1. Build a dataset from the xp1 exhaustive_9_6 raw data by removing
     the top-p% best schedules (lowest makespan values).
  2. Train the surrogate (n_embd=40, n_head=4, n_layer=2, all other
     hyper-parameters identical to xp1's best config).
  3. Run Phase-2 latent optimisation (no metropolized filling).

Run from:  src/xp2/
"""

import os
import sys
import json
import shutil
import math

import numpy as np
from tqdm import tqdm

# ── make xp1 modules importable ──────────────────────────────────────────────
XP1_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../xp1"))
sys.path.insert(0, XP1_DIR)

from train import train
from recover_schedules import recover_schedules

# ── paths ─────────────────────────────────────────────────────────────────────
RAW_DATA_DIR  = os.path.abspath(os.path.join(XP1_DIR, "outputs/exhaustive_9_6"))
XP2_OUT_ROOT  = os.path.join(os.path.dirname(__file__), "outputs/exhaustive_9_6")

N_EMBD  = 40
N_HEAD  = 4
N_LAYER = 2

PERCENTILES = list(range(5, 100, 5))   # [5, 10, 15, ..., 95]


# ── dataset processing ────────────────────────────────────────────────────────

def process_dataset_pct(raw_data_dir, output_dir, top_pct, train_ratio=0.9, seed=97):
    """
    Build a processed (train/val split) dataset from the raw exhaustive data
    by removing the best `top_pct`% schedules (lowest makespan values).

    Files written to output_dir:
        schedules_train.npy, makespans_train.npy
        schedules_val.npy,   makespans_val.npy
        metadata.json, pfsp_instance.npy
        neh_makespan.npy, cds_makespan.npy, palmer_makespan.npy, min_makespan.npy
    """
    if os.path.exists(os.path.join(output_dir, "metadata.json")):
        print(f"  [process] already done – skipping {output_dir}")
        return

    os.makedirs(output_dir, exist_ok=True)

    # load raw data
    schedules = np.load(os.path.join(raw_data_dir, "schedules.npy"))
    makespans  = np.load(os.path.join(raw_data_dir, "makespans.npy"))   # (N, nb_jobs)
    with open(os.path.join(raw_data_dir, "metadata.json")) as f:
        metadata = json.load(f)

    ms = makespans[:, -1]   # final makespan per schedule

    # remove the top-pct% best (lowest) makespans
    n_total  = len(ms)
    n_remove = math.floor(top_pct / 100 * n_total)
    sorted_idx = np.argsort(ms)
    remove_set = set(sorted_idx[:n_remove].tolist())
    keep_mask  = np.array([i not in remove_set for i in range(n_total)])

    schedules = schedules[keep_mask]
    makespans  = makespans[keep_mask]
    ms         = makespans[:, -1]

    n_kept = len(schedules)
    print(f"  [process] top_pct={top_pct}%: removed {n_remove}, kept {n_kept} schedules")

    # save min makespan of remaining data
    np.save(os.path.join(output_dir, "min_makespan.npy"), ms.min())

    # stratified train/val split by makespan value
    np.random.seed(seed)
    unique_ms = np.unique(ms)
    train_idx_list, val_idx_list = [], []
    for m in unique_ms:
        group = np.where(ms == m)[0]
        np.random.shuffle(group)
        n_train = max(1, int(len(group) * train_ratio))
        if n_train >= len(group):
            train_idx_list.append(group)
        else:
            train_idx_list.append(group[:n_train])
            val_idx_list.append(group[n_train:])

    train_indices = np.concatenate(train_idx_list)
    val_indices   = np.concatenate(val_idx_list) if val_idx_list else np.array([], dtype=int)

    np.save(os.path.join(output_dir, "schedules_train.npy"), schedules[train_indices])
    np.save(os.path.join(output_dir, "makespans_train.npy"),  makespans[train_indices])
    np.save(os.path.join(output_dir, "schedules_val.npy"),   schedules[val_indices])
    np.save(os.path.join(output_dir, "makespans_val.npy"),    makespans[val_indices])

    # metadata
    split_meta = metadata.copy()
    split_meta["nb_train_samples"] = len(train_indices)
    split_meta["nb_val_samples"]   = len(val_indices)
    split_meta["data_source"]      = raw_data_dir
    split_meta["top_pct_removed"]  = top_pct
    split_meta["n_removed"]        = int(n_remove)
    split_meta["n_kept"]           = int(n_kept)
    with open(os.path.join(output_dir, "metadata.json"), "w") as f:
        json.dump(split_meta, f, indent=4)

    # copy static files from raw dir
    for fname in ["pfsp_instance.npy", "neh_makespan.npy",
                  "cds_makespan.npy", "palmer_makespan.npy"]:
        shutil.copy(os.path.join(raw_data_dir, fname),
                    os.path.join(output_dir, fname))

    print(f"  [process] done → {n_train} train / {len(val_indices)} val samples")


# ── main loop ─────────────────────────────────────────────────────────────────

for pct in (outer:=tqdm(PERCENTILES, desc="Percentiles")):
    outer.set_description(f"top_pct={pct}%")

    data_dir  = os.path.join(XP2_OUT_ROOT, f"top_pct_{pct}")
    model_dir = os.path.join(data_dir, f"train_nEmbd{N_EMBD}_nHead{N_HEAD}_nLayer{N_LAYER}")
    rs_dir    = os.path.join(model_dir, "recover_no_mf")

    # ── Stage 1: process dataset ──────────────────────────────────────────────
    process_dataset_pct(
        raw_data_dir=RAW_DATA_DIR,
        output_dir=data_dir,
        top_pct=pct,
    )

    # ── Stage 2: train surrogate ──────────────────────────────────────────────
    train(
        testing=False,
        seed=97,
        data_dir=data_dir,
        n_embd=N_EMBD,
        n_head=N_HEAD,
        n_layer=N_LAYER,
        intermediate_schedules=True,
        dropout=0.0,
        ff_width=4,
        train_batch_size=512,
        val_batch_size=256,
        nb_epochs=5,
        early_stopping_patience=15,
        checkpoint_interval_ratio=1.0,
        decay_lr=True,
        lr_partitions_ratios=[0.66],
        init_lr=1e-4,
        max_lr=1e-3,
        min_lr=5e-5,
        lr_warmup_iters_ratio=0.1,
        lr_decay_iters_ratio=0.95,
        beta1=0.9,
        beta2=0.95,
        weight_decay=10.0,
        grad_clip=1.0,
        compile=False,
        compile_mode="default",
        save_only_last_checkpoint=True,
        output_dir=model_dir,
    )

    # ── Stage 3: Phase-2 latent optimisation (no mf) ──────────────────────────
    recover_schedules(
        testing=False,
        seed=97,
        init_mode="random",
        n_optimization_steps=2000,
        epsilon=0.01,
        nb_sinkhorn_iters=40,
        decay_ls=False,
        ls_partitions_ratios=[0.66],
        init_ls=2.0,
        max_ls=10,
        min_ls=0.1,
        ls_warmup_iters_ratio=0.2,
        ls_decay_iters_ratio=0.90,
        decay_lr=True,
        lr_partitions_ratios=[0.66],
        init_lr=1e-2,
        max_lr=1e-1,
        min_lr=1e-4,
        lr_warmup_iters_ratio=0.1,
        lr_decay_iters_ratio=0.90,
        checkpoint_interval=1,
        data_dir=data_dir,
        model_dir=model_dir,
        output_dir=rs_dir,
        apply_metropolized_filling=False,
        mf_patience=10,
        mf_fifo_size=50,
        mf_penalty_strength=1.0,
        mf_min_delta=0.05,
    )