Source code for src.tools.algorithms

"""Algorithm type definitions for EngiOpt integration.

This is the SINGLE FILE to update when adding a new algorithm type:
1. Add to AlgorithmId Literal type
2. Add to SUPPORTED_ALGORITHMS mapping with properties

Everything else (algorithm lists, registries) is derived automatically.
"""

from typing import Literal, get_args

# Type alias for algorithm IDs - single source of truth for supported algorithms
AlgorithmId = Literal["cgan_cnn_2d", "diffusion_2d_cond"]

# Supported algorithms and their properties
SUPPORTED_ALGORITHMS = {
    "cgan_cnn_2d": {
        "class": "Inverse Design",
        "dimensions": "2D",
        "conditional": True,
        "model": "GAN + CNN",
        "model_types": ["generator", "discriminator"],
    },
    "diffusion_2d_cond": {
        "class": "Inverse Design",
        "dimensions": "2D",
        "conditional": True,
        "model": "Diffusion",
        "model_types": ["model"],
    },
}

# List of supported algorithm IDs - automatically derived from AlgorithmId
ALGORITHM_IDS = list(get_args(AlgorithmId))




def canonical_algorithm(script_stem: str) -> str | None:
    """Map a candidate script's filename stem to its canonical engiopt algorithm.

    ``write_training_script`` versions candidate filenames
    (``<algo>_candidate_<timestamp>_…``), so a stem is not itself a valid engiopt
    algorithm id. The prompt instructs the agent to name candidates
    ``<algo>_candidate.py``, so the canonical id is the ``SUPPORTED_ALGORITHMS``
    entry the stem starts with. Both ``validate_training_script`` (to pick the
    contract/family) and ``generate_training_slurm`` (to pass ``--algo`` so the
    trained wandb artifact is the canonical name ``evaluate_model`` downloads)
    derive the algorithm this way, so the filename stem is the single source of
    truth across the local + cluster halves of the workflow.

    Returns ``None`` when no supported algorithm matches; callers decide whether
    to fall back to a default or surface an actionable error.
    """
    # Longest-first so a shorter id can't shadow a longer one that shares a prefix.
    for algo in sorted(SUPPORTED_ALGORITHMS, key=len, reverse=True):
        if script_stem == algo or script_stem.startswith(f"{algo}_"):
            return algo
    return None



# Export for use in other modules
__all__ = [
    "ALGORITHM_IDS",
    "SUPPORTED_ALGORITHMS",
    "AlgorithmId",
    "canonical_algorithm",
]