setup

#!/usr/bin/env bash

# ======================
# 0. Machine + Runtime
# ======================
#Machine="LEONARDO"   # on LEONARDO
Machine="LEONARDO"         # on HAL

current_ntasks=1      # Better to keep zero

current_cpus=16        # On LEONARDO
#current_cpus=4          # On HAL

current_workers=16      # On LEONARDO
#current_workers=4        # On HAL

debug="Flase"         # True or False

#DATA_PATH="/leonardo_work/EUHPC_D27_095/Data_LSM_1995-2000_Large"                           # Path to the dataset on LEONARDO
DATA_PATH="/net/nfs/ssd1/kardaneh/Data_LSM_1995-2000_Large"                                  # Path to the dataset on HAL
CAMS_TRAIN_DATA_PATH="/leonardo_work/EUHPC_D27_095/CAMS/RADSCHEME_data_g224_CAMS_2009-2018_sans_2014-2015.nc"      # Path to the dataset on LEONARDO
CAMS_Test_DATA_PATH="/leonardo_work/EUHPC_D27_095/CAMS/RADSCHEME_data_g224_CAMS_2014.nc"                           # Path to the dataset on LEONARDO

# ======================
# 1. Model Configuration
# ======================
type="encodertorch" # options: "lstm", "gru", "fcn", "transformer", "encodertorch", "pinn", "mlp", "mlp_residual".
hidden_size=256

# NUMBER OF LAYERS
# ----------------
# num_layers = 3
#
# Rule of thumb: num_layers = 2 to 6 for sequence length 10
# More layers = more capacity but slower training and more overfitting risk
#
# Formula: num_layers ≈ log2(seq_length) + 1 = log2(10) + 1 ≈ 3.3 + 1 = 4.3 (rounded to 4)
#           but with dropout, 2-3 is often sufficient for this problem size
#
# Recommended range: 2-4
num_layers=3

output_channel=2
seq_length=60
feature_channel=11

# EMBEDDING SIZE
# ---------------
# embed_size = 256
#
# Rule of thumb: embed_size should be between input and output dimensions
# input=121, output=120 → embed_size=256 is a good balance
#
# Formula: embed_size = 2 × max(input, output) = 2 × 121 = 242 (rounded to 256)
#
# For 256: embed_size / 64 = 4 heads (nice round number)
#
# Recommended range: 256-512 (256 is conservative, good starting point)
embed_size=256 #256

# NUMBER OF ATTENTION HEADS
# -------------------------
# nhead = 4
#
# Rule of thumb: heads = embed_size // 64
#
# Formula: heads = embed_size // d_head (where d_head is typically 64)
#
# For embed_size=256: 256 / 64 = 4 heads
#
# Must divide embed_size evenly: 256 % 4 = 0
nhead=4

# FEED-FORWARD EXPANSION FACTOR
# -----------------------------
# forward_expansion = 2
#
# Rule of thumb: forward_expansion = 2 to 4
#
# Formula: ff_dim = embed_size × forward_expansion
#
# For forward_expansion=2: ff_dim = 256 × 2 = 512
# For forward_expansion=4: ff_dim = 256 × 4 = 1024
#
# Larger expansion = more capacity but more parameters
# Recommended range: 2-4 (2 is conservative, 4 is aggressive)
forward_expansion=4

# DROPOUT RATE
# ------------
# dropout = 0.3
#
# Rule of thumb: dropout = 0.1 to 0.5
#
# For output_channel=120 (many outputs), use higher dropout to prevent overfitting
#
# Recommended range: 0.2-0.4
dropout=0.1

sblock_perc=0.1

# ======================
# 2. Training Configuration
# ======================

batch_size=4      # On LEONARDO
#batch_size=4       # On HAL

num_epochs=300
learning_rate=0.0001
loss_type="huber" # options: "mse", "mae", "huber", "smoothl1", etc
beta=0.001
beta_delta=1.0
num_workers=$current_workers
run_type="train" # options: "train", "resume_train", "inference".
seed=42

# ======================
# 3. Data Configuration
# ======================
train_data_files=$CAMS_TRAIN_DATA_PATH
test_data_files=$CAMS_Test_DATA_PATH
train_years="2009-2018"
test_year="2014"
norm="minmax"
dataset_type="CAMS"

# ======================
# 4. Output Configuration
# ======================
save_model="True"
save_checkpoint_name="model"
save_per_samples=10000
load_checkpoint_name="experiment_best_model.pth.tar"

# ======================
# 5. Tags
# ======================
lr_str="${learning_rate//./d}"
beta_str="${beta//./d}"

# ======================
# 5. Tags (Final version)
# ======================

# Format numeric values (avoid dots)
lr_str="${learning_rate//./d}"
beta_str="${beta//./d}"
dropout_str="${dropout//./d}"

# Model-specific naming
if [ "$type" == "encodertorch" ] || [ "$type" == "transformer" ]; then
    model_parts=(
        "$type"
        "e${embed_size}"
        "h${nhead}"
        "l${num_layers}"
        "fe${forward_expansion}"
        "d${dropout_str}"
    )
else
    model_parts=(
        "$type"
        "h${hidden_size}"
        "l${num_layers}"
        "d${dropout_str}"
    )
fi
model_name=$(IFS=_ ; echo "${model_parts[*]}")

tag="experiment"
main_folder="$([ "$debug" == "True" ] && echo "Debug" || echo "Prod")__${model_name}_sb_${batch_size}_nw_${num_workers}_ne_${num_epochs}_ds_${dataset_type}_tr_${train_years}_te_${test_year}"
sub_folder="nrm_${norm}_lr_${lr_str}_l_${loss_type}"
if [ "$loss_type" == "huber" ] || [ "$loss_type" == "smoothl1" ]; then
    beta_delta_str="${beta_delta//./d}"
    sub_folder="${sub_folder}_bd_${beta_delta_str}"
fi
sub_folder="${sub_folder}_beta_${beta_str}_seed_${seed}"
prefix="$tag"

# ======================
# Generate SBATCH Script Only (removed run script)
# ======================
PROJECT_ROOT="$(cd "$(dirname "$0")" && pwd)"
JOB_DIR="$PROJECT_ROOT/slurm/${main_folder}/${sub_folder}"
SLURM_IO_DIR="$PROJECT_ROOT/slurm_io/${main_folder}/${sub_folder}"
mkdir -p "$JOB_DIR"
mkdir -p "$SLURM_IO_DIR"
sbatch_script="$JOB_DIR/sbatch_${current_ntasks}_${current_cpus}_${current_workers}_rtnn_${tag}.sh"

# Create slurm_io directory if it doesn't exist


# ======================
# 6. Create SBATCH file
# ======================

cat > $sbatch_script <<EOF
#!/bin/bash
# SBATCH always at the beginning of the script!
#SBATCH --job-name=rtnn_${main_folder}_${sub_folder}_${tag}
#SBATCH --error=$SLURM_IO_DIR/rtnn_${current_ntasks}_${current_cpus}_${current_workers}_${tag}_%j.err
#SBATCH --output=$SLURM_IO_DIR/rtnn_${current_ntasks}_${current_cpus}_${current_workers}_${tag}_%j.out
#SBATCH --chdir=$(pwd)
EOF

# ======================
# 7. Machine-specific SBATCH config
# ======================

if [ "$Machine" == "LEONARDO" ]; then

cat >> $sbatch_script <<EOF
#SBATCH -A EUHPC_D36_053
#SBATCH -p boost_usr_prod
#SBATCH --time=$([ "$debug" == "True" ] && echo "00:30:00" || echo "3:00:00")
#SBATCH -N 1
#SBATCH --gpus-per-node=1
#SBATCH --mem-per-gpu=64G
EOF

# Only add ntasks-per-node if not empty
if [ -n "$current_ntasks" ]; then
    echo "#SBATCH --ntasks-per-node=$current_ntasks" >> $sbatch_script
fi

# Only add cpus-per-task if not empty
if [ -n "$current_cpus" ]; then
    echo "#SBATCH --cpus-per-task=$current_cpus" >> $sbatch_script
fi

[ "$debug" == "True" ] && echo "#SBATCH --qos=boost_qos_dbg" >> $sbatch_script

elif [ "$Machine" == "HAL" ]; then

cat >> $sbatch_script <<EOF
#SBATCH --partition=batch
#SBATCH --gpus=1
#SBATCH --mem=48G
#SBATCH --time=24:00:00
EOF

# Only add ntasks-per-node if not empty
if [ -n "$current_ntasks" ]; then
    echo "#SBATCH --ntasks-per-node=$current_ntasks" >> $sbatch_script
fi

# Only add cpus-per-task if not empty
if [ -n "$current_cpus" ]; then
    echo "#SBATCH --cpus-per-task=$current_cpus" >> $sbatch_script
fi

fi

# ======================
# 8. Environment setup
# ======================

if [ "$Machine" == "LEONARDO" ]; then

cat >> $sbatch_script <<EOF
source "$PROJECT_ROOT/.venv/bin/activate"
EOF

elif [ "$Machine" == "HAL" ]; then

cat >> $sbatch_script <<EOF

source "$PROJECT_ROOT/.venv/bin/activate"
ulimit -s unlimited
EOF

fi

# ======================
# 9. Runtime diagnostics + training
# ======================

cat >> $sbatch_script <<EOF

echo "=========================================="
echo "CONFIGURATION TEST VALUES"
echo "=========================================="
echo "ntasks-per-node: ${current_ntasks:-default (1)}"
echo "cpus-per-task: ${current_cpus:-default}"
echo "num_workers: $num_workers"
echo "=========================================="
echo ""
echo "=== Job Info ==="
echo "Job ID: \$SLURM_JOB_ID"
echo "Node: \$SLURM_NODELIST"
echo "N tasks: \$SLURM_NTASKS"
echo "CPUs per task: \${SLURM_CPUS_PER_TASK:-default}"
echo "GPUs: \$SLURM_GPUS"
echo ""

echo "=== Environment Info ==="
echo "PyTorch version: \$(python -c 'import torch; print(torch.__version__)')"
echo "CUDA available: \$(python -c 'import torch; print(torch.cuda.is_available())')"
echo "CUDA version: \$(python -c 'import torch; print(torch.version.cuda)')"
echo "GPU count: \$(python -c 'import torch; print(torch.cuda.device_count())')"

echo "=== RTnn Training ==="
echo "Version: \$(python -c 'import rtnn; print(rtnn.__version__)')"
echo "Model: $model_name"
echo "Start time: \$(date)"

rtnn \\
    --root_dir "./" \\
    --main_folder "$main_folder" \\
    --sub_folder "$sub_folder" \\
    --prefix "$prefix" \\
    --dataset_type "$dataset_type" \\
    --type "$type" \\
    --hidden_size "$hidden_size" \\
    --num_layers "$num_layers" \\
    --output_channel "$output_channel" \\
    --seq_length "$seq_length" \\
    --feature_channel "$feature_channel" \\
    --embed_size "$embed_size" \\
    --nhead "$nhead" \\
    --forward_expansion "$forward_expansion" \\
    --dropout "$dropout" \\
    --model_name "$model_name" \\
    --batch_size "$batch_size" \\
    --num_epochs "$num_epochs" \\
    --learning_rate "$learning_rate" \\
    --sblock_perc "$sblock_perc" \\
    --loss_type "$loss_type" \\
    --beta "$beta" \\
    --beta_delta "$beta_delta" \\
    --train_data_files "$train_data_files" \\
    --test_data_files "$test_data_files" \\
    --train_years "$train_years" \\
    --test_year "$test_year" \\
    --norm "$norm" \\
    --num_workers "$num_workers" \\
    --save_model "$save_model" \\
    --save_checkpoint_name "$save_checkpoint_name" \\
    --save_per_samples "$save_per_samples" \\
    --load_checkpoint_name "$load_checkpoint_name" \\
    --run_type "$run_type" \\
    --seed "$seed" \\
    --debug "$debug"

echo "End time: \$(date)"
EOF

chmod +x $sbatch_script

echo "Created: $sbatch_script"
echo "Run with: sbatch $sbatch_script"