Grid Search Module¶

def run_grid_search(config: GridSearchConfig) -> Path:
    """
    Run grid search over chunk contexts.

    Args:
        config: Grid search configuration

    Returns:
        Path to grid search summary CSV
    """
    from leech.constants import generate_random_seed

    # Generate random seed if not provided
    if config.seed is None:
        seed = generate_random_seed()
        logger.info(f"Generated random seed: {seed}")
        config.seed = seed
    else:
        seed = config.seed
        logger.info(f"Using provided seed: {seed}")

    # Default dwell_offsets to [0] if not provided
    dwell_offsets = config.dwell_offsets if config.dwell_offsets is not None else [0]

    # Skip dwell_offset grid for models without a feature branch
    if config.model_name not in ModelInferenceWrapper.FEATURE_MODELS:
        if dwell_offsets != [0]:
            logger.info(
                f"Model {config.model_name} has no feature branch; collapsing dwell_offsets to [0]"
            )
            dwell_offsets = [0]

    logger.info("=" * 80)
    logger.info("Starting Grid Search")
    logger.info("=" * 80)
    logger.info(f"Model: {config.model_name}")
    logger.info(f"Left contexts: {config.left_contexts}")
    logger.info(f"Right contexts: {config.right_contexts}")
    logger.info(f"Dwell offsets: {dwell_offsets}")
    logger.info(
        f"Total grid points: {len(config.left_contexts) * len(config.right_contexts) * len(dwell_offsets)}"
    )
    logger.info(f"Output directory: {config.output_dir}")
    logger.info(f"Random seed: {seed}")
    logger.info("=" * 80)

    config.output_dir.mkdir(parents=True, exist_ok=True)

    # Save seed to file
    seed_file = config.output_dir / "grid_search_seed.txt"
    with open(seed_file, "w") as f:
        f.write(f"{seed}\n")

    # Save grid search config
    config_dict = {
        "model_name": config.model_name,
        "left_contexts": config.left_contexts,
        "right_contexts": config.right_contexts,
        "dwell_offsets": dwell_offsets,
        "kmer_context": config.kmer_context,
        "epochs": config.epochs,
        "batch_size": config.batch_size,
        "learning_rate": config.learning_rate,
        "device": config.device,
        "seed": seed,
    }

    with open(config.output_dir / "grid_config.json", "w") as f:
        json.dump(config_dict, f, indent=2)

    # Pre-compute class weights from the training labels (they don't change
    # across grid points). Only the label column is needed, so read that --
    # not the corpus.
    labels_array = _training_label_column(config.train_data_path)
    logger.info(f"Read {len(labels_array)} training labels")
    unique, counts = np.unique(labels_array, return_counts=True)
    pos_weight = None
    if len(unique) == 2:
        label_counts = dict(zip(unique, counts, strict=True))
        neg_count = label_counts.get(0, 0)
        pos_count = label_counts.get(1, 0)
        if pos_count > 0:
            pos_weight = neg_count / pos_count
            logger.info(
                f"Pre-computed class weights: negative={neg_count}, "
                f"positive={pos_count}, pos_weight={pos_weight:.4f}"
            )

    # Resolve selection_metric "auto" once so every grid point uses the same
    # criterion: macro-F1 for multiclass, AUROC for binary (AUROC is threshold-
    # free and not gamed by majority-class prediction on imbalanced heads).
    resolved_metric = _resolve_selection_metric(config.selection_metric, n_classes=len(unique))
    logger.info(f"Selection metric: {config.selection_metric!r} -> {resolved_metric!r}")

    # Generate grid
    grid_points = list(
        itertools.product(config.left_contexts, config.right_contexts, dwell_offsets)
    )
    summary_path = config.output_dir / "grid_summary.csv"

    # Build grid point argument dicts
    grid_args = []
    for left, right, dwoff in grid_points:
        if len(dwell_offsets) > 1 or dwell_offsets != [0]:
            grid_output_dir = config.output_dir / f"left_{left}_right_{right}_dwoff_{dwoff}"
        else:
            grid_output_dir = config.output_dir / f"left_{left}_right_{right}"

        grid_args.append(
            {
                "train_data_path": config.train_data_path,
                "val_data_path": config.val_data_path,
                "model_name": config.model_name,
                "output_dir": grid_output_dir,
                "left_context": left,
                "right_context": right,
                "kmer_len": 2 * config.kmer_context + 1,
                "epochs": config.epochs,
                "batch_size": config.batch_size,
                "learning_rate": config.learning_rate,
                "device": config.device,
                "seed": config.seed,
                "early_stopping_patience": config.early_stopping_patience,
                "dwell_offset": dwoff,
                "pos_weight": pos_weight,
                "weight_decay": config.weight_decay,
                "max_grad_norm": config.max_grad_norm,
                "scheduler": config.scheduler,
                "scheduler_patience": config.scheduler_patience,
                "scheduler_factor": config.scheduler_factor,
                "warmup_epochs": config.warmup_epochs,
                "loss_type": config.loss_type,
                "focal_gamma": config.focal_gamma,
                "label_smoothing": config.label_smoothing,
                "mixed_precision": config.mixed_precision,
                "augment_jitter": config.augment_jitter,
                "augment_scale_min": config.augment_scale_min,
                "augment_scale_max": config.augment_scale_max,
                "augment_time_mask_bases": config.augment_time_mask_bases,
                "augment_time_mask_count": config.augment_time_mask_count,
                "augment_shift_max_bases": config.augment_shift_max_bases,
                "augment_feature_noise_scale": config.augment_feature_noise_scale,
                "num_workers": config.num_workers,
                "balance_groups": config.balance_groups,
                "oversample_minority": config.oversample_minority,
                "motif": config.motif,
                "motif_offset": config.motif_offset,
                "base_justify": config.base_justify,
                "adversarial_lambda": config.adversarial_lambda,
                "adversarial_anneal_epochs": config.adversarial_anneal_epochs,
                "confound": config.confound,
                "cl_regression": config.cl_regression,
                "cl_lambda": config.cl_lambda,
                "signal_mode": config.signal_mode,
                "selection_metric": resolved_metric,
            }
        )

    results: list[dict] = []

    if config.n_parallel > 1:
        # Parallel execution: each worker streams the corpus per grid point
        logger.info(
            f"Running {len(grid_points)} grid points with {config.n_parallel} parallel workers"
        )
        # CUDA does not support fork; use spawn to avoid hangs with multiple GPU workers
        ctx = multiprocessing.get_context("spawn" if config.device != "cpu" else "fork")
        with ctx.Pool(processes=config.n_parallel) as pool:
            for i, result in enumerate(pool.imap_unordered(_grid_point_worker, grid_args), 1):
                results.append(result)
                logger.info(f"Completed grid point {i}/{len(grid_points)}")
                save_grid_summary(results, summary_path)
    else:
        # Sequential execution. Grid points must not share a chunk list:
        # LeechDataset drops each chunk's arrays once it has tensorized them,
        # so the second point got chunks whose `signal` was None and died with
        # "'NoneType' object has no attribute 'dtype'".
        for i, args in enumerate(grid_args, 1):
            logger.info(f"\n\nGrid point {i}/{len(grid_points)}")
            result = run_grid_point(**args)
            results.append(result)
            save_grid_summary(results, summary_path)

    # Print summary with Rich tables
    console.print("\n[bold green]Grid Search Complete![/bold green]\n")

    successful_results = [r for r in results if r.get("status") == "success"]

    if not successful_results:
        logger.error(f"All {len(results)} grid points failed")
        raise RuntimeError(f"Grid search failed: all {len(results)} grid points failed")

    if successful_results:
        # Map history-key metric -> result-dict key produced by run_grid_point
        sort_key = {
            "val_acc": "best_val_acc",
            "val_f1": "best_val_f1",
            "val_auc": "best_val_auc",
        }[resolved_metric]

        # Create results table; tag the column we ranked on with [*]
        col_marker = {
            "best_val_acc": ("Val Accuracy[*]", "Val F1", "Val AUC"),
            "best_val_f1": ("Val Accuracy", "Val F1[*]", "Val AUC"),
            "best_val_auc": ("Val Accuracy", "Val F1", "Val AUC[*]"),
        }[sort_key]
        table = Table(title="Grid Search Results", show_header=True, header_style="bold magenta")
        table.add_column("Left Context", justify="right", style="cyan")
        table.add_column("Right Context", justify="right", style="cyan")
        table.add_column("Dwell Offset", justify="right", style="cyan")
        table.add_column(col_marker[0], justify="right", style="green")
        table.add_column(col_marker[1], justify="right", style="green")
        table.add_column(col_marker[2], justify="right", style="yellow")
        table.add_column("Best Epoch", justify="right", style="blue")
        table.add_column("Training Time", justify="right", style="white")

        sorted_results = sorted(successful_results, key=lambda x: x.get(sort_key, 0), reverse=True)

        for r in sorted_results[:10]:  # Show top 10
            table.add_row(
                str(r["left_context"]),
                str(r["right_context"]),
                str(r.get("dwell_offset", 0)),
                f"{r.get('best_val_acc', 0):.4f}",
                f"{r.get('best_val_f1', 0):.4f}",
                f"{r.get('best_val_auc', 0):.4f}",
                str(r.get("best_epoch", 0)),
                f"{r.get('train_time_sec', 0):.1f}s",
                style="bold" if r == sorted_results[0] else None,
            )

        console.print(table)

        # Best configuration summary
        best_result = sorted_results[0]
        summary_table = Table(
            title="Best Configuration", show_header=True, header_style="bold magenta"
        )
        summary_table.add_column("Parameter", style="cyan")
        summary_table.add_column("Value", justify="right", style="green")

        summary_table.add_row("Selection Metric", f"{config.selection_metric} -> {resolved_metric}")
        summary_table.add_row("Left Context", str(best_result["left_context"]))
        summary_table.add_row("Right Context", str(best_result["right_context"]))
        summary_table.add_row("Dwell Offset", str(best_result.get("dwell_offset", 0)))
        summary_table.add_row("Validation Accuracy", f"{best_result['best_val_acc']:.4f}")
        summary_table.add_row("Validation F1", f"{best_result.get('best_val_f1', 0):.4f}")
        summary_table.add_row("Validation AUC", f"{best_result.get('best_val_auc', 0):.4f}")
        summary_table.add_row("Best Epoch", str(best_result.get("best_epoch", 0)))
        summary_table.add_row("Model Path", str(best_result["model_path"]))

        console.print(summary_table)

        best_params_path = config.output_dir / "best_params.json"
        with open(best_params_path, "w") as f:
            json.dump(
                {
                    "left_context": best_result["left_context"],
                    "right_context": best_result["right_context"],
                    "dwell_offset": best_result.get("dwell_offset", 0),
                    "selection_metric": resolved_metric,
                },
                f,
                indent=2,
            )
        console.print(f"[bold]Best params saved to:[/bold] {best_params_path}")

    console.print(f"\n[bold]Results saved to:[/bold] {summary_path}")

    return summary_path

def parse_context_grid(
    context_grid: str | None = None,
    left_contexts: str | None = None,
    right_contexts: str | None = None,
) -> tuple[list[int], list[int]]:
    """Parse context grid strings into integer lists.

    Supports range syntax (``start:stop:step``), comma-separated lists,
    and single values.

    Args:
        context_grid: Fallback context values when left/right not provided
            (e.g., "200,500,1000" or "200:1000:200")
        left_contexts: Override left contexts, or None to use context_grid
        right_contexts: Override right contexts, or None to use context_grid

    Returns:
        Tuple of (left_contexts_list, right_contexts_list)

    Raises:
        ValueError: If context_grid is None and either left_contexts or
            right_contexts is also None

    Examples:
        >>> parse_context_grid("200,500,1000")
        ([200, 500, 1000], [200, 500, 1000])
        >>> parse_context_grid("200:1000:200")
        ([200, 400, 600, 800, 1000], [200, 400, 600, 800, 1000])
        >>> parse_context_grid("200,500", left_contexts="100,200", right_contexts="300,400")
        ([100, 200], [300, 400])
        >>> parse_context_grid(left_contexts="100,200", right_contexts="300,400")
        ([100, 200], [300, 400])
    """
    if context_grid is None:
        if left_contexts is None or right_contexts is None:
            msg = (
                "--context-grid is required when --left-contexts or "
                "--right-contexts is not provided"
            )
            raise ValueError(msg)

    left_list = parse_values(left_contexts if left_contexts is not None else context_grid)
    right_list = parse_values(right_contexts if right_contexts is not None else context_grid)

    return left_list, right_list