studies/adv_highest_difficulty_resnet_increasing_rollout_compensate_train_time.py

"""
Trains a ResNet to become an emulator for the 2D advection equation at high
difficulty (gamma_1=10.5) with increasing unrollment.

Compensates shorter training unrolling with additional update steps such that all
configurations should run almost equally long (running 10 seeds in parallel
should take ~45min on a modern GPU)
"""
BASE_TRAIN_DURATION = 10_000

CONFIGS = [
    {
        "scenario": "diff_adv",
        "task": "predict",
        "net": "Res;26;8;relu",  # 32'943 params, 16 receptive field per direction
        "train": f"sup;{rollout:02d}",
        "start_seed": s,
        "num_seeds": 10,
        "advection_gamma": 10.5,
        "optim_config": f"adam;{int(BASE_TRAIN_DURATION * 15 / rollout):d};warmup_cosine;0.0;1e-3;{int(BASE_TRAIN_DURATION * 15 / rollout / 5):d}",
    }
    for s in [0, 10, 20, 30, 40]
    for rollout in [1, 2, 3, 5, 10, 15]
]